fennol.training.training

  1import os
  2import yaml
  3import sys
  4import jax
  5import io
  6import time
  7import jax.numpy as jnp
  8import numpy as np
  9import optax
 10from collections import defaultdict
 11import json
 12from copy import deepcopy
 13from pathlib import Path
 14import argparse
 15try:
 16    import torch
 17except ImportError:
 18    raise ImportError(
 19        "PyTorch is required for training models. Install the CPU version following instructions at https://pytorch.org/get-started/locally/"
 20    )
 21import random
 22from flax import traverse_util
 23import json
 24import shutil
 25import pickle
 26
 27from flax.core import freeze, unfreeze
 28from .io import (
 29    load_configuration,
 30    load_dataset,
 31    load_model,
 32    TeeLogger,
 33    copy_parameters,
 34)
 35from .utils import (
 36    get_loss_definition,
 37    get_train_step_function,
 38    get_validation_function,
 39    linear_schedule,
 40)
 41from .optimizers import get_optimizer, get_lr_schedule
 42from ..utils import deep_update, AtomicUnits as au
 43from ..utils.io import human_time_duration
 44from ..models.preprocessing import AtomPadding, check_input, convert_to_jax
 45
 46def save_restart_checkpoint(output_directory,stage,training_state,preproc_state,metrics_state):
 47    with open(output_directory + "/restart_checkpoint.pkl", "wb") as f:
 48            pickle.dump({
 49                "stage": stage,
 50                "training": training_state,
 51                "preproc_state": preproc_state,
 52                "metrics_state": metrics_state,
 53            }, f)
 54
 55def load_restart_checkpoint(output_directory):
 56    restart_checkpoint_file = output_directory + "/restart_checkpoint.pkl"
 57    if not os.path.exists(restart_checkpoint_file):
 58        raise FileNotFoundError(
 59            f"Training state file not found: {restart_checkpoint_file}"
 60        )
 61    with open(restart_checkpoint_file, "rb") as f:
 62        restart_checkpoint = pickle.load(f)
 63    return restart_checkpoint
 64
 65
 66def main():
 67    parser = argparse.ArgumentParser(prog="fennol_train")
 68    parser.add_argument("config_file", type=str)
 69    parser.add_argument("--model_file", type=str, default=None)
 70    args = parser.parse_args()
 71    config_file = args.config_file
 72    model_file = args.model_file
 73
 74    os.environ["OMP_NUM_THREADS"] = "1"
 75    sys.stdout = io.TextIOWrapper(
 76        open(sys.stdout.fileno(), "wb", 0), write_through=True
 77    )
 78
 79    restart_training = False
 80    if os.path.isdir(config_file):
 81        output_directory = Path(config_file).absolute().as_posix()
 82        restart_training = True
 83        while output_directory.endswith("/"):
 84            output_directory = output_directory[:-1]
 85        config_file = output_directory + "/config.yaml"
 86        backup_dir = output_directory + f"_backup_{time.strftime('%Y-%m-%d-%H-%M-%S')}"
 87        shutil.copytree(output_directory, backup_dir)
 88        print("Restarting training from", output_directory)
 89
 90    parameters = load_configuration(config_file)
 91
 92    ### Set the device
 93    device: str = parameters.get("device", "cpu").lower()
 94    if device == "cpu":
 95        os.environ["CUDA_VISIBLE_DEVICES"] = ""
 96    elif device.startswith("cuda") or device.startswith("gpu"):
 97        if ":" in device:
 98            num = device.split(":")[-1]
 99            os.environ["CUDA_VISIBLE_DEVICES"] = num
100        else:
101            os.environ["CUDA_VISIBLE_DEVICES"] = "0"
102        device = "gpu"
103
104    _device = jax.devices(device)[0]
105    jax.config.update("jax_default_device", _device)
106
107    if restart_training:
108        restart_checkpoint = load_restart_checkpoint(output_directory)
109    else:
110        restart_checkpoint = None
111
112    # output directory
113    if not restart_training:
114        output_directory = parameters.get("output_directory", None)
115        if output_directory is not None:
116            if "{now}" in output_directory:
117                output_directory = output_directory.replace(
118                    "{now}", time.strftime("%Y-%m-%d-%H-%M-%S")
119                )
120            output_directory = Path(output_directory).absolute()
121            if not output_directory.exists():
122                output_directory.mkdir(parents=True)
123            print("Output directory:", output_directory)
124        else:
125            output_directory = "."
126
127    output_directory = str(output_directory) + "/"
128
129    # copy config_file to output directory
130    # config_name = Path(config_file).name
131    config_ext = Path(config_file).suffix
132    with open(config_file) as f_in:
133        config_data = f_in.read()
134    with open(output_directory + "/config" + config_ext, "w") as f_out:
135        f_out.write(config_data)
136
137    # set log file
138    log_file = "train.log"  # parameters.get("log_file", None)
139    logger = TeeLogger(output_directory + log_file)
140    logger.bind_stdout()
141
142    # set matmul precision
143    enable_x64 = parameters.get("double_precision", False)
144    jax.config.update("jax_enable_x64", enable_x64)
145    fprec = "float64" if enable_x64 else "float32"
146    parameters["fprec"] = fprec
147    if enable_x64:
148        print("Double precision enabled.")
149
150    matmul_precision = parameters.get("matmul_prec", "highest").lower()
151    assert matmul_precision in [
152        "default",
153        "high",
154        "highest",
155    ], "matmul_prec must be one of 'default','high','highest'"
156    jax.config.update("jax_default_matmul_precision", matmul_precision)
157
158    # set random seed
159    rng_seed = parameters.get("rng_seed", np.random.randint(0, 2**32 - 1))
160    print(f"rng_seed: {rng_seed}")
161    rng_key = jax.random.PRNGKey(rng_seed)
162    torch.manual_seed(rng_seed)
163    np.random.seed(rng_seed)
164    random.seed(rng_seed)
165    np_rng = np.random.Generator(np.random.PCG64(rng_seed))
166
167    try:
168        if "stages" in parameters["training"]:
169            ## train in stages ##
170            params = deepcopy(parameters)
171            stages = params["training"].pop("stages")
172            assert isinstance(stages, dict), "'stages' must be a dict with named stages"
173            model_file_stage = model_file
174            print_stages_params = params["training"].get("print_stages_params", False)
175            for i, (stage, stage_params) in enumerate(stages.items()):
176                rng_key, subkey = jax.random.split(rng_key)
177                print("")
178                print(f"### STAGE {i+1}: {stage} ###")
179
180                ## remove end_event from previous stage ##
181                if i > 0 and "end_event" in params["training"]:
182                    params["training"].pop("end_event")
183
184                ## incrementally update training parameters ##
185                params = deep_update(params, {"training": stage_params})
186                if model_file_stage is not None:
187                    ## load model from previous stage ##
188                    params["model_file"] = model_file_stage
189
190                if restart_training and restart_checkpoint["stage"] != i + 1:
191                    print(f"Skipping stage {i+1} (already completed)")
192                    continue
193
194                if print_stages_params:
195                    print("stage parameters:")
196                    print(json.dumps(params, indent=2, sort_keys=False))
197
198                ## train stage ##
199                _, model_file_stage = train(
200                    (subkey, np_rng),
201                    params,
202                    stage=i + 1,
203                    output_directory=output_directory,
204                    restart_checkpoint=restart_checkpoint,
205                )
206
207                restart_training = False
208                restart_checkpoint = None
209        else:
210            ## single training stage ##
211            train(
212                (rng_key, np_rng),
213                parameters,
214                model_file=model_file,
215                output_directory=output_directory,
216                restart_checkpoint=restart_checkpoint,
217            )
218    except KeyboardInterrupt:
219        print("Training interrupted by user.")
220    finally:
221        if log_file is not None:
222            logger.unbind_stdout()
223            logger.close()
224
225
226def train(
227    rng,
228    parameters,
229    model_file=None,
230    stage=None,
231    output_directory=None,
232    restart_checkpoint=None,
233):
234    if output_directory is None:
235        output_directory = "./"
236    elif output_directory == "":
237        output_directory = "./"
238    elif not output_directory.endswith("/"):
239        output_directory += "/"
240    stage_prefix = f"_stage_{stage}" if stage is not None else ""
241
242    if isinstance(rng, tuple):
243        rng_key, np_rng = rng
244    else:
245        rng_key = rng
246        np_rng = np.random.Generator(np.random.PCG64(np.random.randint(0, 2**32 - 1)))
247
248    ### GET TRAINING PARAMETERS ###
249    training_parameters = parameters.get("training", {})
250
251    #### LOAD MODEL ####
252    if restart_checkpoint is not None:
253        model_key = None
254        model_file = output_directory + "latest_model.fnx"
255    else:
256        rng_key, model_key = jax.random.split(rng_key)
257    model = load_model(parameters, model_file, rng_key=model_key)
258
259    ### SAVE INITIAL MODEL ###
260    save_initial_model = (
261        restart_checkpoint is None 
262        and (stage is None or stage == 0)
263    )
264    if save_initial_model:
265        model.save(output_directory + "initial_model.fnx")
266
267    model_ref = None
268    if "model_ref" in training_parameters:
269        model_ref = load_model(parameters, training_parameters["model_ref"])
270        print("Reference model:", training_parameters["model_ref"])
271
272        if "ref_parameters" in training_parameters:
273            ref_parameters = training_parameters["ref_parameters"]
274            assert isinstance(
275                ref_parameters, list
276            ), "ref_parameters must be a list of str"
277            print("Reference parameters:", ref_parameters)
278            model.variables = copy_parameters(
279                model.variables, model_ref.variables, ref_parameters
280            )
281
282    ### SET FLOATING POINT PRECISION ###
283    fprec = parameters.get("fprec", "float32")
284
285    def convert_to_fprec(x):
286        if jnp.issubdtype(x.dtype, jnp.floating):
287            return x.astype(fprec)
288        return x
289
290    model.variables = jax.tree_map(convert_to_fprec, model.variables)
291
292    ### SET UP LOSS FUNCTION ###
293    loss_definition, used_keys, ref_keys = get_loss_definition(
294        training_parameters, model_energy_unit=model.energy_unit
295    )
296
297    coordinates_ref_key = training_parameters.get("coordinates_ref_key", None)
298    if coordinates_ref_key is not None:
299        compute_ref_coords = True
300        print("Reference coordinates:", coordinates_ref_key)
301    else:
302        compute_ref_coords = False
303
304    #### LOAD DATASET ####
305    dspath = training_parameters.get("dspath", None)
306    if dspath is None:
307        raise ValueError("Dataset path 'training/dspath' should be specified.")
308    batch_size = training_parameters.get("batch_size", 16)
309    batch_size_val = training_parameters.get("batch_size_val", None)
310    rename_refs = training_parameters.get("rename_refs", {})
311    training_iterator, validation_iterator = load_dataset(
312        dspath=dspath,
313        batch_size=batch_size,
314        batch_size_val=batch_size_val,
315        training_parameters=training_parameters,
316        infinite_iterator=True,
317        atom_padding=True,
318        ref_keys=ref_keys,
319        split_data_inputs=True,
320        np_rng=np_rng,
321        add_flags=["training"],
322        fprec=fprec,
323        rename_refs=rename_refs,
324    )
325
326    compute_forces = "forces" in used_keys
327    compute_virial = "virial_tensor" in used_keys or "virial" in used_keys
328    compute_stress = "stress_tensor" in used_keys or "stress" in used_keys
329    compute_pressure = "pressure" in used_keys or "pressure_tensor" in used_keys
330
331    ### get optimizer parameters ###
332    max_epochs = int(training_parameters.get("max_epochs", 2000))
333    epoch_format = len(str(max_epochs))
334    nbatch_per_epoch = training_parameters.get("nbatch_per_epoch", 200)
335    nbatch_per_validation = training_parameters.get("nbatch_per_validation", 20)
336
337    schedule, sch_state, schedule_metrics,adaptive_scheduler = get_lr_schedule(max_epochs,nbatch_per_epoch,training_parameters)
338    optimizer = get_optimizer(
339        training_parameters, model.variables, schedule(sch_state)[0]
340    )
341    opt_st = optimizer.init(model.variables)
342
343    ### exponential moving average of the parameters ###
344    ema_decay = training_parameters.get("ema_decay", -1.0)
345    if ema_decay > 0.0:
346        assert ema_decay < 1.0, "ema_decay must be in (0,1)"
347        ema = optax.ema(decay=ema_decay)
348    else:
349        ema = optax.identity()
350    ema_st = ema.init(model.variables)
351
352    #### end event ####
353    end_event = training_parameters.get("end_event", None)
354    if end_event is None or isinstance(end_event, str) and end_event.lower() == "none":
355        is_end = lambda metrics: False
356    else:
357        assert len(end_event) == 2, "end_event must be a list of two elements"
358        is_end = lambda metrics: metrics[end_event[0]] < end_event[1]
359
360    if "energy_terms" in training_parameters:
361        model.set_energy_terms(training_parameters["energy_terms"], jit=False)
362    print("energy terms:", model.energy_terms)
363
364    ### MODEL EVALUATION FUNCTION ###
365    pbc_training = training_parameters.get("pbc_training", False)
366    if compute_stress or compute_virial or compute_pressure:
367        virial_key = "virial" if "virial" in used_keys else "virial_tensor"
368        stress_key = "stress" if "stress" in used_keys else "stress_tensor"
369        pressure_key = "pressure" if "pressure" in used_keys else "pressure_tensor"
370        if compute_stress or compute_pressure:
371            assert pbc_training, "PBC must be enabled for stress or virial training"
372            print("Computing forces and stress tensor")
373            def evaluate(model, variables, data):
374                _, _, vir, output = model._energy_and_forces_and_virial(variables, data)
375                cells = output["cells"]
376                volume = jnp.abs(jnp.linalg.det(cells))
377                stress = vir / volume[:, None, None]
378                output[stress_key] = stress
379                output[virial_key] = vir
380                if pressure_key == "pressure":
381                    output[pressure_key] = -jnp.trace(stress, axis1=1, axis2=2) / 3.0
382                else:
383                    output[pressure_key] = -stress
384                return output
385
386        else:
387            print("Computing forces and virial tensor")
388            def evaluate(model, variables, data):
389                _, _, vir, output = model._energy_and_forces_and_virial(variables, data)
390                output[virial_key] = vir
391                return output
392
393    elif compute_forces:
394        print("Computing forces")
395        def evaluate(model, variables, data):
396            _, _, output = model._energy_and_forces(variables, data)
397            return output
398
399    elif model.energy_terms is not None:
400        def evaluate(model, variables, data):
401            _, output = model._total_energy(variables, data)
402            return output
403
404    else:
405        def evaluate(model, variables, data):
406            output = model.modules.apply(variables, data)
407            return output
408
409    ### TRAINING FUNCTIONS ###
410    train_step = get_train_step_function(
411        loss_definition=loss_definition,
412        model=model,
413        model_ref=model_ref,
414        compute_ref_coords=compute_ref_coords,
415        evaluate=evaluate,
416        optimizer=optimizer,
417        ema=ema,
418    )
419
420    validation = get_validation_function(
421        loss_definition=loss_definition,
422        model=model,
423        model_ref=model_ref,
424        compute_ref_coords=compute_ref_coords,
425        evaluate=evaluate,
426        return_targets=False,
427    )
428
429    #### CONFIGURE PREPROCESSING ####
430    gpu_preprocessing = training_parameters.get("gpu_preprocessing", False)
431    if gpu_preprocessing:
432        print("GPU preprocessing activated.")
433
434    minimum_image = training_parameters.get("minimum_image", False)
435    preproc_state = unfreeze(model.preproc_state)
436    layer_state = []
437    for st in preproc_state["layers_state"]:
438        stnew = unfreeze(st)
439        #     st["nblist_skin"] = nblist_skin
440        #     if nblist_stride > 1:
441        #         st["skin_stride"] = nblist_stride
442        #         st["skin_count"] = nblist_stride
443        if pbc_training:
444            stnew["minimum_image"] = minimum_image
445        if "nblist_mult_size" in training_parameters:
446            stnew["nblist_mult_size"] = training_parameters["nblist_mult_size"]
447        if "nblist_add_neigh" in training_parameters:
448            stnew["add_neigh"] = training_parameters["nblist_add_neigh"]
449        if "nblist_add_atoms" in training_parameters:
450            stnew["add_atoms"] = training_parameters["nblist_add_atoms"]
451        layer_state.append(freeze(stnew))
452
453    preproc_state["layers_state"] = tuple(layer_state)
454    preproc_state["check_input"] = False
455    model.preproc_state = freeze(preproc_state)
456
457    ### INITIALIZE METRICS ###
458    maes_prev = defaultdict(lambda: np.inf)
459    metrics_beta = training_parameters.get("metrics_ema_decay", -1.0)
460    smoothen_metrics = metrics_beta < 1.0 and metrics_beta > 0.0
461    if smoothen_metrics:
462        print("Computing smoothed metrics with beta =", metrics_beta)
463        rmses_smooth = defaultdict(lambda: 0.0)
464        maes_smooth = defaultdict(lambda: 0.0)
465        rmse_tot_smooth = 0.0
466        mae_tot_smooth = 0.0
467        nsmooth = 0
468    max_restore_count = training_parameters.get("max_restore_count", 5)    
469
470    fmetrics = open(
471        output_directory + f"metrics{stage_prefix}.traj",
472        "w" if restart_checkpoint is None else "a",
473    )
474
475    keep_all_bests = training_parameters.get("keep_all_bests", False)
476    save_model_at_epochs = training_parameters.get("save_model_at_epochs", [])
477    save_model_at_epochs= set([int(x) for x in save_model_at_epochs])
478    save_model_every = int(training_parameters.get("save_model_every_epoch", 0))
479    if save_model_every > 0:
480        save_model_every = list(range(save_model_every, max_epochs+save_model_every, save_model_every))
481        save_model_at_epochs = save_model_at_epochs.union(save_model_every)
482    save_model_at_epochs = list(save_model_at_epochs)
483    save_model_at_epochs.sort()
484
485    ### LR factor after restoring a model that has diverged
486    restore_scaling = training_parameters.get("restore_scaling", 1)
487    assert restore_scaling > 0.0, "restore_scaling must be > 0.0"
488    assert restore_scaling <= 1.0, "restore_scaling must be <= 1.0"
489    if restore_scaling != 1:
490        print("Applying LR scaling after restore:", restore_scaling)
491
492
493    metric_use_best = training_parameters.get("metric_best", "rmse_tot")  # .lower()
494    metrics_state = {}
495
496    ### INITIALIZE TRAINING STATE ###
497    if restart_checkpoint is not None:
498        training_state = deepcopy(restart_checkpoint["training"])
499        epoch_start = training_state["epoch"]
500        model.preproc_state = freeze(restart_checkpoint["preproc_state"])
501        if smoothen_metrics:
502            metrics_state = deepcopy(restart_checkpoint["metrics_state"])
503            rmses_smooth = metrics_state["rmses_smooth"]
504            maes_smooth = metrics_state["maes_smooth"]
505            rmse_tot_smooth = metrics_state["rmse_tot_smooth"]
506            mae_tot_smooth = metrics_state["mae_tot_smooth"]
507            nsmooth = metrics_state["nsmooth"]
508        print("Restored training state")
509    else:
510        epoch_start = 1
511        training_state = {
512            "rng_key": rng_key,
513            "opt_state": opt_st,
514            "ema_state": ema_st,
515            "sch_state": sch_state,
516            "variables": model.variables,
517            "variables_ema": model.variables,
518            "step": 0,
519            "restore_count": 0,
520            "restore_scale": 1,
521            "epoch": 0,
522            "best_metric": np.inf,
523        }
524        if smoothen_metrics:
525            metrics_state = {
526                "rmses_smooth": dict(rmses_smooth),
527                "maes_smooth": dict(maes_smooth),
528                "rmse_tot_smooth": rmse_tot_smooth,
529                "mae_tot_smooth": mae_tot_smooth,
530                "nsmooth": nsmooth,
531            }
532    
533    del opt_st, ema_st, sch_state, preproc_state
534    variables_save = training_state['variables']
535    variables_ema_save = training_state['variables_ema']
536
537    ### Training loop ###
538    start = time.time()
539    print("Starting training...")
540    for epoch in range(epoch_start, max_epochs+1):
541        training_state["epoch"] = epoch
542        s = time.time()
543        for _ in range(nbatch_per_epoch):
544            # fetch data
545            inputs0, data = next(training_iterator)
546
547            # preprocess data
548            inputs = model.preprocess(verbose=True,use_gpu=gpu_preprocessing,**inputs0)
549
550            rng_key, subkey = jax.random.split(rng_key)
551            inputs["rng_key"] = subkey
552            inputs["training_epoch"] = epoch
553
554            # adjust learning rate
555            current_lr, training_state["sch_state"] = schedule(training_state["sch_state"])
556            current_lr *= training_state["restore_scale"]
557            training_state["opt_state"].inner_states["trainable"].inner_state[-1].hyperparams[
558                "step_size"
559            ] = current_lr 
560
561            # train step
562            loss,training_state, _ = train_step(data,inputs,training_state)
563
564        rmses_avg = defaultdict(lambda: 0.0)
565        maes_avg = defaultdict(lambda: 0.0)
566        for _ in range(nbatch_per_validation):
567            inputs0, data = next(validation_iterator)
568
569            inputs = model.preprocess(verbose=True,use_gpu=gpu_preprocessing,**inputs0)
570
571            rng_key, subkey = jax.random.split(rng_key)
572            inputs["rng_key"] = subkey
573            inputs["training_epoch"] = epoch
574
575            rmses, maes, output_val = validation(
576                data=data,
577                inputs=inputs,
578                variables=training_state["variables_ema"],
579            )
580            for k, v in rmses.items():
581                rmses_avg[k] += v
582            for k, v in maes.items():
583                maes_avg[k] += v
584
585        jax.block_until_ready(output_val)
586
587        ### Timings ###
588        e = time.time()
589        epoch_time = e - s
590
591        elapsed_time = e - start
592        if not adaptive_scheduler:
593            remain_glob = elapsed_time * (max_epochs - epoch + 1) / epoch
594            remain_last = epoch_time * (max_epochs - epoch + 1)
595            # estimate remaining time via weighted average (put more weight on last at the beginning)
596            wremain = np.sin(0.5 * np.pi * epoch / max_epochs)
597            remaining_time = human_time_duration(
598                remain_glob * wremain + remain_last * (1 - wremain)
599            )
600        elapsed_time = human_time_duration(elapsed_time)
601        batch_time = human_time_duration(
602            epoch_time / (nbatch_per_epoch + nbatch_per_validation)
603        )
604        epoch_time = human_time_duration(epoch_time)
605
606        for k in rmses_avg.keys():
607            rmses_avg[k] /= nbatch_per_validation
608        for k in maes_avg.keys():
609            maes_avg[k] /= nbatch_per_validation
610
611        ### Print metrics ###
612        print("")
613        line = f"Epoch {epoch}, lr={current_lr:.3e}, loss = {loss:.3e}"
614        line += f", epoch time = {epoch_time}, batch time = {batch_time}"
615        line += f", elapsed time = {elapsed_time}"
616        if not adaptive_scheduler:
617            line += f", est. remaining time = {remaining_time}"
618        print(line)
619        rmse_tot = 0.0
620        mae_tot = 0.0
621        for k in rmses_avg.keys():
622            mult = loss_definition[k]["mult"]
623            loss_prms = loss_definition[k]
624            rmse_tot = rmse_tot + rmses_avg[k] * loss_prms["weight"] ** 0.5
625            mae_tot = mae_tot + maes_avg[k] * loss_prms["weight"]
626            unit = "(" + loss_prms["unit"] + ")" if "unit" in loss_prms else ""
627
628            weight_str = ""
629            if "weight_schedule" in loss_prms:
630                w = linear_schedule(epoch, *loss_prms["weight_schedule"])
631                weight_str = f"(w={w:.3f})"
632
633            if rmses_avg[k] / mult < 1.0e-2:
634                print(
635                    f"    rmse_{k}= {rmses_avg[k]/mult:10.3e} ; mae_{k}= {maes_avg[k]/mult:10.3e}   {unit}  {weight_str}"
636                )
637            else:
638                print(
639                    f"    rmse_{k}= {rmses_avg[k]/mult:10.3f} ; mae_{k}= {maes_avg[k]/mult:10.3f}   {unit}  {weight_str}"
640                )
641
642        ### CHECK IF WE SHOULD RESTORE PREVIOUS MODEL ###
643        restore = False
644        reinit = False
645        for k, mae in maes_avg.items():
646            if np.isnan(mae):
647                restore = True
648                reinit = True
649                print("NaN detected in mae")
650                # for k, v in inputs.items():
651                #     if hasattr(v,"shape"):
652                #         if np.isnan(v).any():
653                #             print(k,v)
654                # sys.exit(1)
655            if "threshold" in loss_definition[k]:
656                thr = loss_definition[k]["threshold"]
657                if mae > thr * maes_prev[k]:
658                    restore = True
659                    break
660
661        epoch_time_str = time.strftime("%Y-%m-%d-%H-%M-%S")
662        model.variables = training_state["variables_ema"]
663        if epoch in save_model_at_epochs:
664            filename = output_directory + f"model{stage_prefix}_epoch{epoch:0{epoch_format}}_{epoch_time_str}.fnx"
665            print("Scheduled model save :", filename)
666            model.save(filename)
667
668        if restore:
669            training_state["restore_count"] += 1
670            if training_state["restore_count"]  > max_restore_count:
671                if reinit:
672                    raise ValueError("Model diverged and could not be restored.")
673                else:
674                    training_state["restore_count"]  = 0
675                    print(
676                        f"{max_restore_count} unsuccessful restores, resuming training."
677                    )
678
679            training_state["variables"] = deepcopy(variables_save)
680            training_state["variables_ema"] = deepcopy(variables_ema_save)
681            model.variables = training_state["variables_ema"]
682            training_state["restore_scale"] *= restore_scaling
683            print(f"Restored previous model after divergence. Restore scale: {training_state['restore_scale']:.3f}")
684            if reinit:
685                training_state["opt_state"] = optimizer.init(model.variables)
686                training_state["ema_state"] = ema.init(model.variables)
687                print("Reinitialized optimizer after divergence.")
688            save_restart_checkpoint(output_directory,stage,training_state,model.preproc_state,metrics_state)
689            continue
690
691        training_state["restore_count"]  = 0
692
693        variables_save = deepcopy(training_state["variables"])
694        variables_ema_save = deepcopy(training_state["variables_ema"])
695        maes_prev = maes_avg
696
697        ### SAVE MODEL ###
698        model.save(output_directory + "latest_model.fnx")
699
700        # save metrics
701        step = int(training_state["step"])
702        metrics = {
703            "epoch": epoch,
704            "step": step,
705            "data_count": step * batch_size,
706            "elapsed_time": time.time() - start,
707            "lr": current_lr,
708            "loss": loss,
709        }
710        for k in rmses_avg.keys():
711            mult = loss_definition[k]["mult"]
712            metrics[f"rmse_{k}"] = rmses_avg[k] / mult
713            metrics[f"mae_{k}"] = maes_avg[k] / mult
714
715        metrics["rmse_tot"] = rmse_tot
716        metrics["mae_tot"] = mae_tot
717        if smoothen_metrics:
718            nsmooth += 1
719            for k in rmses_avg.keys():
720                mult = loss_definition[k]["mult"]
721                rmses_smooth[k] = (
722                    metrics_beta * rmses_smooth[k] + (1.0 - metrics_beta) * rmses_avg[k]
723                )
724                maes_smooth[k] = (
725                    metrics_beta * maes_smooth[k] + (1.0 - metrics_beta) * maes_avg[k]
726                )
727                metrics[f"rmse_smooth_{k}"] = (
728                    rmses_smooth[k] / (1.0 - metrics_beta**nsmooth) / mult
729                )
730                metrics[f"mae_smooth_{k}"] = (
731                    maes_smooth[k] / (1.0 - metrics_beta**nsmooth) / mult
732                )
733            rmse_tot_smooth = (
734                metrics_beta * rmse_tot_smooth + (1.0 - metrics_beta) * rmse_tot
735            )
736            mae_tot_smooth = (
737                metrics_beta * mae_tot_smooth + (1.0 - metrics_beta) * mae_tot
738            )
739            metrics["rmse_smooth_tot"] = rmse_tot_smooth / (1.0 - metrics_beta**nsmooth)
740            metrics["mae_smooth_tot"] = mae_tot_smooth / (1.0 - metrics_beta**nsmooth)
741
742            # update metrics state
743            metrics_state["rmses_smooth"] = dict(rmses_smooth)
744            metrics_state["maes_smooth"] = dict(maes_smooth)
745            metrics_state["rmse_tot_smooth"] = rmse_tot_smooth
746            metrics_state["mae_tot_smooth"] = mae_tot_smooth
747            metrics_state["nsmooth"] = nsmooth
748
749        assert (
750            metric_use_best in metrics
751        ), f"Error: metric for selectring best model '{metric_use_best}' not in metrics"
752
753        metric_for_best = metrics[metric_use_best]
754        if metric_for_best <  training_state["best_metric"]:
755            training_state["best_metric"] = metric_for_best
756            metrics["best_metric"] = training_state["best_metric"]
757            if keep_all_bests:
758                best_name = (
759                    output_directory
760                    + f"best_model{stage_prefix}_epoch{epoch:0{epoch_format}}_{epoch_time_str}.fnx"
761                )
762                model.save(best_name)
763
764            best_name = output_directory + f"best_model{stage_prefix}.fnx"
765            model.save(best_name)
766            print("New best model saved to:", best_name)
767        else:
768            metrics["best_metric"] = training_state["best_metric"]
769
770        if epoch == 1:
771            headers = [f"{i+1}:{k}" for i, k in enumerate(metrics.keys())]
772            fmetrics.write("# " + " ".join(headers) + "\n")
773        fmetrics.write(" ".join([str(metrics[k]) for k in metrics.keys()]) + "\n")
774        fmetrics.flush()
775
776        # update learning rate using current metrics
777        assert (
778            schedule_metrics in metrics
779        ), f"Error: cannot update lr, '{schedule_metrics}' not in metrics"
780        current_lr, training_state["sch_state"] = schedule(training_state["sch_state"], metrics[schedule_metrics])
781
782        # update and save training state
783        save_restart_checkpoint(output_directory,stage,training_state,model.preproc_state,metrics_state)
784
785        
786        if is_end(metrics):
787            print("Stage finished.")
788            break
789
790    end = time.time()
791
792    print(f"Training time: {human_time_duration(end-start)}")
793    print("")
794
795    fmetrics.close()
796
797    filename = output_directory + f"final_model{stage_prefix}.fnx"
798    model.save(filename)
799    print("Final model saved to:", filename)
800
801    return metrics, filename
802
803
804if __name__ == "__main__":
805    main()