core.models.equiformer_v2.trainers.forces_trainer

core.models.equiformer_v2.trainers.forces_trainer#

Copyright (c) Meta, Inc. and its affiliates.

This source code is licensed under the MIT license found in the LICENSE file in the root directory of this source tree.

Classes#

EquiformerV2ForcesTrainer

Trainer class for the Structure to Energy & Force (S2EF) and Initial State to

Module Contents#

class core.models.equiformer_v2.trainers.forces_trainer.EquiformerV2ForcesTrainer(task: dict[str, str | Any], model: dict[str, Any], outputs: dict[str, str | int], dataset: dict[str, str | float], optimizer: dict[str, str | float], loss_functions: dict[str, str | float], evaluation_metrics: dict[str, str], identifier: str, local_rank: int, timestamp_id: str | None = None, run_dir: str | None = None, is_debug: bool = False, print_every: int = 100, seed: int | None = None, logger: str = 'wandb', amp: bool = False, cpu: bool = False, name: str = 'ocp', slurm: dict | None = None, gp_gpus: int | None = None, inference_only: bool = False)#

Bases: fairchem.core.trainers.OCPTrainer

Trainer class for the Structure to Energy & Force (S2EF) and Initial State to Relaxed State (IS2RS) tasks.

Note

Examples of configurations for task, model, dataset and optimizer can be found in configs/ocp_s2ef and configs/ocp_is2rs.

Parameters:

  • task (dict) – Task configuration.

  • model (dict) – Model configuration.

  • outputs (dict) – Dictionary of model output configuration.

  • dataset (dict) – Dataset configuration. The dataset needs to be a SinglePointLMDB dataset.

  • optimizer (dict) – Optimizer configuration.

  • loss_functions (dict) – Loss function configuration.

  • evaluation_metrics (dict) – Evaluation metrics configuration.

  • identifier (str) – Experiment identifier that is appended to log directory.

  • run_dir (str, optional) – Path to the run directory where logs are to be saved. (default: None)

  • timestamp_id (str, optional) – timestamp identifier.

  • run_dir – Run directory used to save checkpoints and results.

  • is_debug (bool, optional) – Run in debug mode. (default: False)

  • print_every (int, optional) – Frequency of printing logs. (default: 100)

  • seed (int, optional) – Random number seed. (default: None)

  • logger (str, optional) – Type of logger to be used. (default: wandb)

  • local_rank (int, optional) – Local rank of the process, only applicable for distributed training. (default: 0)

  • amp (bool, optional) – Run using automatic mixed precision. (default: False)

  • cpu (bool) – If True will run on CPU. Default is False, will attempt to use cuda.

  • name (str) – Trainer name.

  • slurm (dict) – Slurm configuration. Currently just for keeping track. (default: {})

  • gp_gpus (int, optional) – Number of graph parallel GPUs.

  • inference_only (bool) – If true trainer will be loaded for inference only. (ie datasets, optimizer, schedular, etc, will not be instantiated)

load_extras() → None#