core.models.equiformer_v2.trainers.energy_trainer

core.models.equiformer_v2.trainers.energy_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#

EquiformerV2EnergyTrainer

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

Module Contents#

class core.models.equiformer_v2.trainers.energy_trainer.EquiformerV2EnergyTrainer(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=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) – Output property 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)

  • 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)

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

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

load_extras()#