-
START: Literal[0, 1, 2], whether to load the model checkpoint file.(
- 0: Default, from scratch.
- 1: load checkpoint of model parameters, optimizer state, current minimum validation error,
and lr_scheduler state (if set) from
LOAD_CHK_FILE_PATH. - 2: only load model parameters from
LOAD_CHK_FILE_PATH.
)
-
DEVICE: str, the device that model and data run on. It is the same as device of torch.Tensor. Default: 'cpu'. -
VERBOSE: int, to control the output frequency and details. 0 for silence mode, 1 for brief mode. Larger number means more detailed output information. -
EPOCH: int, number of total epochs that model will train. Default: 0. -
BATCH_SIZE: int, batch size of training samples. Default: 1. -
VAL_BATCH_SIZE: int, batch size of validation samples. Default: the same asBATCH_SIZE. -
VAL_PER_STEP: int, the validation frequency. Model will validate perVAL_PER_STEPtraining step. Default: 10. -
VAL_IF_TRN_LOSS_BELOW: float, not validate until training loss is less thanVAL_IF_TRN_LOSS_BELOW. Default: inf -
ACCUMULATE_STEP: int, the step number of gradient accumulation. Default: 1. -
DEBUG_MODE: bool, whether to turn on the debug mode, which will output the range of parameters, parameter gradients and gradient differences in each layer, and turn on the NaN check. Default: False. -
CHECK_NAN: bool, whether to check NaN after inputting each batch during training. Default: False
REDIRECT: bool, whether to output training logs toOUTPUT_PATHor directly print it to the screen. Default: True.SAVE_CHK: bool, whether to save the checkpoint file of the training model. Default: False.LOAD_CHK_FILE_PATH: str, only work whenSTART== 1. The path of checkpoint file to load.STRICT_LOAD: bool, to control the parameterstrictintorch.Module.load_state_dict.CHK_SAVE_PATH: str, only work whenSAVE_CHECK== True. The directory of checkpoint file to save. Default: "./".CHK_SAVE_POSTFIX: str, the postfix of checkpoint file. Default: "". The checkpoint file name will be "best_checkpoint_CHK_SAVE_POSTFIX.pt".OUTPUT_PATH: str, only work whenREDIRECT== True. The output directory of training log. Default: "./".OUTPUT_POSTFIX: str, the postfix of training log file. Default: "Untitled". The log file name will be f"time.strftime("%Y%m%d_%H_%M_%S")_OUTPUT_POSTFIX.out".
-
LOSS: Literal["MSE", "MAE", "Huber", "CrossEntropy", "Energy_Force_Loss", "Energy_Loss"](
- "MSE": nn.MSELoss
- "MAE": nn.L1Loss
- "Hubber": nn.HuberLoss
- "CrossEntropy": nn.CrossEntropyLoss
- "Energy_Force_Loss":
$loss = coeff_E * loss_E(x, y) + coeff_F * loss_F(x, y)$ - "Energy_Loss":
$loss = coeff_E * loss_E(x, y)$ - "custom": Any, one need to set custom loss function manually by
Trainer.set_loss_fn(loss_fn, loss_config: Optional[Dict] = None)
)
-
LOSS_CONFIG: Dict, the kwargs ofLOSS.- for
LOSS== "Energy_Force_Loss":- loss_E: Literal["MAE", "MSE"], the loss function of energies.
- loss_F: Literal["MAE", "MSE"], the loss function of forces.
- coeff_E: float, coefficient of energy loss.
- coeff_F: float, coefficient of force loss.
- for
LOSS== "Energy_Loss":- loss_E: Literal["MAE", "MSE", "SmoothMAE", "Hubber"], the loss function of energies.
- for
-
METRICS: Literal["E_MAE", "F_MAE", "F_MaxE", "E_R2", "MSE", "MAE", "R2", "RMSE"], the metrics function of training and validation results.
MODEL_NAME: str, the model name.MODEL_CONFIG: Dict, the hyperparameters of model.
OPTIM: Literal["Adam", "SGD", "AdamW", "Adadelta", "Adagrad", "ASGD", "Adamax", "custom"], the model optimizer.
(
- 'Adam': th.optim.Adam,
- 'SGD': th.optim.SGD,
- 'AdamW': th.optim.AdamW,
- 'Adadelta': th.optim.Adadelta,
- 'Adagrad': th.optim.Adagrad,
- 'ASGD': th.optim.ASGD,
- 'Adamax': th.optim.Adamax,
- 'custom': Any, need to set custom optimizer manually
by
Trainer.set_optimizer(self, optimizer, optim_config: Optional[Dict])
)
OPTIM_CONFIG: Dict, the kwargs of model optimizer.GRAD_CLIP: bool, whether to use gradient clip. Default: false.GRAD_CLIP_MAX_NORM: float, only forGRAD_CLIP== True. The max norm of gradient to clip. Default: 100.LR_SCHEDULER: Literal[{'StepLR', 'ExponentialLR', 'ChainedScheduler', 'ConstantLR', 'LambdaLR', 'LinearLR', 'custom'}], the learning rate scheduler.
(
- 'StepLR': torch.optim.lr_scheduler.StepLR,
- 'ExponentialLR': torch.optim.lr_scheduler.ExponentialLR,
- 'ChainedScheduler': torch.optim.lr_scheduler.ChainedScheduler,
- 'ConstantLR': torch.optim.lr_scheduler.ConstantLR,
- 'LambdaLR': torch.optim.lr_scheduler.LambdaLR,
- 'LinearLR': torch.optim.lr_scheduler.LinearLR,
- 'custom': Any, need to set custom lr_scheduler manually
by
Trainer.set_lr_scheduler(self, lr_scheduler, lr_scheduler_config)
)
LR_SCHEDULER_CONFIG: Dict, the kwargs of lr scheduler. Default: dict().EMA: bool, whether to apply the exponential moving average strategy for training. Note that the best_checkpoint will save with EMA parameters, while checkpoint and stop_checkpoint will not. Default: False.EMA_DECAY: float, the decay coefficient of EMA. Default: 0.999.