core.models.equiformer_v2.trainers.dens_trainer

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

DenoisingPosParams

DenoisingForcesTrainer

Functions#

add_gaussian_noise_to_position(batch, std[, ...])

add_gaussian_noise_schedule_to_position(batch, ...[, ...])

denoising_pos_eval(evaluator, prediction, target, ...)

compute_atomwise_denoising_pos_and_force_hybrid_loss(...)

Module Contents#

class core.models.equiformer_v2.trainers.dens_trainer.DenoisingPosParams#
prob: float = 0.0#
fixed_noise_std: bool = False#
std: float = None#
num_steps: int = None#
std_low: float = None#
std_high: float = None#
corrupt_ratio: float = None#
all_atoms: bool = False#
denoising_pos_coefficient: float = None#
core.models.equiformer_v2.trainers.dens_trainer.add_gaussian_noise_to_position(batch, std, corrupt_ratio=None, all_atoms=False)#

  1. Update pos in batch.

  2. Add noise_vec to batch, which will serve as the target for denoising positions.

  3. Add denoising_pos_forward to switch to denoising mode during training.

  4. Add noise_mask for partially corrupted structures when corrupt_ratio is not None.

  5. If all_atoms == True, we add noise to all atoms including fixed ones.

  6. Check whether batch has md. We do not add noise to structures from MD split.

core.models.equiformer_v2.trainers.dens_trainer.add_gaussian_noise_schedule_to_position(batch, std_low, std_high, num_steps, corrupt_ratio=None, all_atoms=False)#

  1. Similar to above, update positions in batch with gaussian noise, but additionally, also save the sigmas the noise vectors are sampled from.

  2. Add noise_mask for partially corrupted structures when corrupt_ratio is not None.

  3. If all_atoms == True, we add noise to all atoms including fixed ones.

  4. Check whether batch has md. We do not add noise to structures from MD split.

core.models.equiformer_v2.trainers.dens_trainer.denoising_pos_eval(evaluator: fairchem.core.modules.evaluator.Evaluator, prediction: dict[str, torch.Tensor], target: dict[str, torch.Tensor], denoising_targets: tuple[str], prev_metrics: dict[str, torch.Tensor] | None = None, denoising_pos_forward: bool = False)#

  1. Overwrite the original Evaluator.eval() here: Open-Catalyst-Project/ocp

  2. This is to make sure we separate forces MAE and denoising positions MAE.

core.models.equiformer_v2.trainers.dens_trainer.compute_atomwise_denoising_pos_and_force_hybrid_loss(pred, target, noise_mask, force_mult, denoising_pos_mult, mask=None)#
class core.models.equiformer_v2.trainers.dens_trainer.DenoisingForcesTrainer(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: core.models.equiformer_v2.trainers.forces_trainer.EquiformerV2ForcesTrainer

  1. We add a denoising objective to the original S2EF task.

  2. The denoising objective is that we take as input atom types, node-wise forces and 3D coordinates perturbed with Gaussian noises and then output energy of the original structure (3D coordinates without any perturbation) and the node-wise noises added to the original structure.

  3. This should make models leverage more from training data and enable data augmentation for the S2EF task.

  4. We should only modify the training part.

  5. For normalizing the outputs of noise prediction, if we use fixed_noise_std = True, we use std for the normalization factor. Otherwise, we use std_high when fixed_noise_std = False.

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)

use_denoising_pos#
denoising_pos_params#
property denoising_targets#
train(disable_eval_tqdm=False)#
_compute_loss(out, batch)#
_compute_metrics(out, batch, evaluator, metrics=None)#
predict(data_loader, per_image: bool = True, results_file: str | None = None, disable_tqdm: bool = False)#