()
| 792 | |
| 793 | |
| 794 | def get_args(): |
| 795 | parser = argparse.ArgumentParser(description='Training') |
| 796 | parser.add_argument('--load-model', action=LoadFromCheckpoint, help='Restart training using a model checkpoint') # keep first |
| 797 | parser.add_argument('--conf', '-c', type=open, action=LoadFromFile, help='Configuration yaml file') # keep second |
| 798 | |
| 799 | # training settings |
| 800 | parser.add_argument('--num-epochs', default=300, type=int, help='number of epochs') |
| 801 | parser.add_argument('--lr-warmup-steps', type=int, default=0, help='How many steps to warm-up over. Defaults to 0 for no warm-up') |
| 802 | parser.add_argument('--lr', default=1e-4, type=float, help='learning rate') |
| 803 | parser.add_argument('--lr-patience', type=int, default=10, help='Patience for lr-schedule. Patience per eval-interval of validation') |
| 804 | parser.add_argument('--lr-min', type=float, default=1e-6, help='Minimum learning rate before early stop') |
| 805 | parser.add_argument('--lr-factor', type=float, default=0.8, help='Minimum learning rate before early stop') |
| 806 | parser.add_argument('--weight-decay', type=float, default=0.0, help='Weight decay strength') |
| 807 | parser.add_argument('--early-stopping-patience', type=int, default=30, help='Stop training after this many epochs without improvement') |
| 808 | parser.add_argument('--loss-type', type=str, default='MSE', choices=['MSE', 'MAE'], help='Loss type') |
| 809 | parser.add_argument('--loss-scale-y', type=float, default=1.0, help="Scale the loss y of the target") |
| 810 | parser.add_argument('--loss-scale-dy', type=float, default=1.0, help="Scale the loss dy of the target") |
| 811 | parser.add_argument('--energy-weight', default=1.0, type=float, help='Weighting factor for energies in the loss function') |
| 812 | parser.add_argument('--force-weight', default=1.0, type=float, help='Weighting factor for forces in the loss function') |
| 813 | |
| 814 | # dataset specific |
| 815 | parser.add_argument('--dataset', default=None, type=str, choices=datasets.__all__, help='Name of the torch_geometric dataset') |
| 816 | parser.add_argument('--dataset-arg', default=None, type=str, help='Additional dataset argument') |
| 817 | parser.add_argument('--dataset-root', default=None, type=str, help='Data storage directory') |
| 818 | parser.add_argument('--derivative', default=False, action=argparse.BooleanOptionalAction, help='If true, take the derivative of the prediction w.r.t coordinates') |
| 819 | parser.add_argument('--split-mode', default=None, type=str, help='Split mode for Molecule3D dataset') |
| 820 | |
| 821 | # dataloader specific |
| 822 | parser.add_argument('--reload', type=int, default=0, help='Reload dataloaders every n epoch') |
| 823 | parser.add_argument('--batch-size', default=32, type=int, help='batch size') |
| 824 | parser.add_argument('--inference-batch-size', default=None, type=int, help='Batchsize for validation and tests.') |
| 825 | parser.add_argument('--standardize', action=argparse.BooleanOptionalAction, default=False, help='If true, multiply prediction by dataset std and add mean') |
| 826 | parser.add_argument('--splits', default=None, help='Npz with splits idx_train, idx_val, idx_test') |
| 827 | parser.add_argument('--train-size', type=number, default=950, help='Percentage/number of samples in training set (None to use all remaining samples)') |
| 828 | parser.add_argument('--val-size', type=number, default=50, help='Percentage/number of samples in validation set (None to use all remaining samples)') |
| 829 | parser.add_argument('--test-size', type=number, default=None, help='Percentage/number of samples in test set (None to use all remaining samples)') |
| 830 | parser.add_argument('--num-workers', type=int, default=4, help='Number of workers for data prefetch') |
| 831 | |
| 832 | # model architecture specific |
| 833 | parser.add_argument('--model', type=str, default='ViSNetBlock', choices=models.__all__, help='Which model to train') |
| 834 | parser.add_argument('--output-model', type=str, default='Scalar', choices=output_modules.__all__, help='The type of output model') |
| 835 | parser.add_argument('--prior-model', type=str, default=None, choices=priors.__all__, help='Which prior model to use') |
| 836 | parser.add_argument('--prior-args', type=dict, default=None, help='Additional arguments for the prior model') |
| 837 | |
| 838 | # architectural specific |
| 839 | parser.add_argument('--embedding-dimension', type=int, default=256, help='Embedding dimension') |
| 840 | parser.add_argument('--num-layers', type=int, default=6, help='Number of interaction layers in the model') |
| 841 | parser.add_argument('--num-rbf', type=int, default=64, help='Number of radial basis functions in model') |
| 842 | parser.add_argument('--activation', type=str, default='silu', choices=list(act_class_mapping.keys()), help='Activation function') |
| 843 | parser.add_argument('--rbf-type', type=str, default='expnorm', choices=list(rbf_class_mapping.keys()), help='Type of distance expansion') |
| 844 | parser.add_argument('--trainable-rbf', action=argparse.BooleanOptionalAction, default=False, help='If distance expansion functions should be trainable') |
| 845 | parser.add_argument('--attn-activation', default='silu', choices=list(act_class_mapping.keys()), help='Attention activation function') |
| 846 | parser.add_argument('--num-heads', type=int, default=8, help='Number of attention heads') |
| 847 | parser.add_argument('--cutoff', type=float, default=5.0, help='Cutoff in model') |
| 848 | parser.add_argument('--max-z', type=int, default=100, help='Maximum atomic number that fits in the embedding matrix') |
| 849 | parser.add_argument('--max-num-neighbors', type=int, default=32, help='Maximum number of neighbors to consider in the network') |
| 850 | parser.add_argument('--reduce-op', type=str, default='add', choices=['add', 'mean'], help='Reduce operation to apply to atomic predictions') |
| 851 | parser.add_argument('--lmax', type=int, default=2, help='Max order of spherical harmonics') |
no test coverage detected