Autotune

Franken autotune: automatic hyperparameter tuning for franken models

usage: franken.autotune [-h] [--rf-norm {leading_eig,none}]
                        [--save-every-model] [--dtype {float32,float64}]
                        [--save-fmaps] [--metrics METRICS [METRICS ...]]
                        [--best-model-selection BEST_MODEL_SELECTION [BEST_MODEL_SELECTION ...]]
                        [--global-scaling] [--jac-chunk-size JAC_CHUNK_SIZE]
                        [--run-dir RUN_DIR] [--seed SEED]
                        [--log-level {DEBUG,INFO,WARN,ERROR}]
                        [--atomic-energies ATOMIC_ENERGIES]
                        [--l2-penalty HYPERPARAMETER]
                        [--force-weight HYPERPARAMETER]
                        [--dataset-name DATASET_NAME]
                        [--train-path TRAIN_PATH] [--test-path TEST_PATH]
                        [--val-path VAL_PATH]
                        [--max-train-samples MAX_TRAIN_SAMPLES] --backbone
                        {mace,pet,sevenn} [--mace.path-or-id MACE.PATH_OR_ID]
                        [--mace.interaction-block MACE.INTERACTION_BLOCK]
                        [--pet.path-or-id PET.PATH_OR_ID]
                        [--sevenn.path-or-id SEVENN.PATH_OR_ID]
                        [--sevenn.interaction-block SEVENN.INTERACTION_BLOCK]
                        [--sevenn.extract-before-act]
                        [--sevenn.last-layer-only] --rf {gaussian,ms-gaussian}
                        [--gaussian.num-rf GAUSSIAN.NUM_RF]
                        [--gaussian.length-scale HYPERPARAMETER]
                        [--gaussian.no-use-offset]
                        [--gaussian.rng-seed GAUSSIAN.RNG_SEED]
                        [--ms-gaussian.num-rf MS_GAUSSIAN.NUM_RF]
                        [--ms-gaussian.no-use-offset]
                        [--ms-gaussian.length-scale-low MS_GAUSSIAN.LENGTH_SCALE_LOW]
                        [--ms-gaussian.length-scale-high MS_GAUSSIAN.LENGTH_SCALE_HIGH]
                        [--ms-gaussian.length-scale-num MS_GAUSSIAN.LENGTH_SCALE_NUM]
                        [--ms-gaussian.rng-seed MS_GAUSSIAN.RNG_SEED]

Named Arguments

--rf-norm

Possible choices: leading_eig, none

Normalization strategy for the covariance matrix.

--save-every-model

If true saves a checkpoint for every trial, otherwise it saves only the best model.

--dtype

Possible choices: float32, float64

Data-type for the franken solution. float64 can usually obtain slightly smaller errors while paying a small performance cost.

--save-fmaps

Whether to save training feature maps. If the dataset is small (~100 samples), setting this to True can increase the speed of hyperparameter tuning, at the cost of higher memory usage.

--metrics

Metrics to compute during evaluation.

Options include energy_MAE, forces_MAE, energy_RMSE, forces_RMSE, forces_MAE_species, forces_RMSE_species, and forces_cosim.

--best-model-selection

Metrics used to select the best model among trials. This does not affect the training loss.

--global-scaling

how to scale the GNN features, whether globally (across species) or individually per species.

--jac-chunk-size

Chunk-size for jacobian calculations. ‘auto’ attempts to set it based on available GPU memory. If you encounter out-of-memory issues, try setting this manually.

--run-dir, -p

Directory in which the hyperparameter search results will be saved

--seed

Random seed

--log-level

Possible choices: DEBUG, INFO, WARN, ERROR

Controls verbosity

--atomic-energies

Optional dictionary mapping atomic numbers to their reference energies (eV).: If provided, these energies will be subtracted from the prediction during training. Format: {atomic_number: energy_value, …}. Example: {1: -0.5, 8: -75.3}

--backbone

Possible choices: mace, pet, sevenn

The GNN backbone which will be used by franken.

--rf

Possible choices: gaussian, ms-gaussian

Choose the random-feature approximation.

Solver

--l2-penalty: The amount of regularization. Should be a small positive number.
--force-weight: Controls how much weight the forces term, as opposed to the energy term has in the loss. Should be a number between 0 and 1.

Dataset

--dataset-name

Dataset name. Can be either one of the predefined datasets, or a custom one.

--train-path

Path to the training dataset. This should be readable with ase, e.g. xyz files work.

--test-path

Path to the test dataset. When doing hyperparameter tuning this: is never used to determine the best model.

--val-path

Path to the validation dataset. Used for hyperparameter tuning if available.

--max-train-samples

Randomly subsample the training set to have at most max_train_samples points.

MACE backbone

Configure the MACE backbone. Specify --backbone=mace to enable.

--mace.path-or-id

Path to the GNN checkpoint, or an ID corresponding to a predefined checkpoint.: To list the available IDs please run the franken.backbones list command.

--mace.interaction-block

GNN layer out of which the features are extracted.

PET backbone

Configure the PET backbone. Specify --backbone=pet to enable.

--pet.path-or-id

Path to the GNN checkpoint, or an ID corresponding to a predefined checkpoint.: To list the available IDs please run the franken.backbones list command.

SevenNet backbone

Configure the MACE backbone. Specify --backbone=sevenn to enable.

--sevenn.path-or-id

Path to the GNN checkpoint, or an ID corresponding to a predefined checkpoint.: To list the available IDs please run the franken.backbones list command.

--sevenn.interaction-block

GNN layer out of which the features are extracted.

--sevenn.extract-before-act

Whether features should be extracted before or after activations.

--sevenn.last-layer-only

Whether to take only the features from the last interaction layer, or to concatenate them all.

Gaussian RFs

Configuration group for orthogonal random fourier features (ORFF), which are used to approximate a Gaussian kernel. Specify --rf=gaussian to enable this RF approximation

--gaussian.num-rf: Number of random features
--gaussian.length-scale: Gaussian length-scale sigma.
--gaussian.no-use-offset: Whether or not to use the version of Gaussian random features with a uniform offset. Set to true, increases the estimator’s efficiency.
--gaussian.rng-seed: Random seed for reproducibility of the random features

Multi-scale Gaussian RFs

Configuration group for multi-scale orthogonal random fourier features (ORFF), which are used to approximate a Gaussian kernel with multiple length-scales. Specify --rf=ms-gaussian to enable this RF approximation

--ms-gaussian.num-rf: Number of random features
--ms-gaussian.no-use-offset: Whether or not to use the version of Gaussian random features with a uniform offset. Set to true, increases the estimator’s efficiency.
--ms-gaussian.length-scale-low: Bottom of the range for the length-scale. Should be positive.
--ms-gaussian.length-scale-high: Top of the range for the length-scale. Should be positive.
--ms-gaussian.length-scale-num: Number of different length-scales to include within the specified range.
--ms-gaussian.rng-seed: Random seed for reproducibility of the random features

Franken CLI Reference

Backbones