facebook

/

OMAT24

like 67

Follow

AI at Meta 5.7k

Python version
3.11.8

fairchem-core version
1.9.0

pytorch version
2.6.9

cuda version
12.4

Operating system version
No response

Minimal example

from fairchem.core.models.model_registry import model_name_to_local_file
from fairchem.core.common.relaxation.ase_utils import OCPCalculator

calculator = OCPCalculator(checkpoint_path="./esen_30m_oam.pt",cpu=False)

Current behavior

Traceback (most recent call last):
File "", line 1, in
File "/scratch/user/u.sa119259/envs/mlff/lib/python3.11/site-packages/fairchem/core/common/relaxation/ase_utils.py", line 200, in init
checkpoint = torch.load(checkpoint_path, map_location=torch.device("cpu"))
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/scratch/user/u.sa119259/envs/mlff/lib/python3.11/site-packages/torch/serialization.py", line 1425, in load
with _open_file_like(f, "rb") as opened_file:
^^^^^^^^^^^^^^^^^^^^^^^^
File "/scratch/user/u.sa119259/envs/mlff/lib/python3.11/site-packages/torch/serialization.py", line 751, in _open_file_like
return _open_file(name_or_buffer, mode)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/scratch/user/u.sa119259/envs/mlff/lib/python3.11/site-packages/torch/serialization.py", line 732, in init
super().init(open(name, mode))
^^^^^^^^^^^^^^^^
FileNotFoundError: [Errno 2] No such file or directory: 'esen_30M_oam.pt'

calculator = OCPCalculator(checkpoint_path="./esen_30m_oam.pt",cpu=False)
Traceback (most recent call last):
File "", line 1, in
File "/scratch/user/u.sa119259/envs/mlff/lib/python3.11/site-packages/fairchem/core/common/relaxation/ase_utils.py", line 236, in init
self.trainer = registry.get_trainer_class(config["trainer"])(
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/scratch/user/u.sa119259/envs/mlff/lib/python3.11/site-packages/fairchem/core/common/registry.py", line 302, in get_trainer_class
return cls.get_class(name, "trainer_name_mapping")
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/scratch/user/u.sa119259/envs/mlff/lib/python3.11/site-packages/fairchem/core/common/registry.py", line 273, in get_class
raise cls.__import_error(name, mapping_name)
RuntimeError: Failed to find the trainer 'mlip_trainer'. You may either use a trainer from the registry (one of 'base', 'forces', 'energy', 'ocp', 'equiformerv2_forces', 'equiformerv2_dens' or 'equiformerv2_energy') or provide the full import path to the trainer (e.g., 'fairchem.core.models.equiformer_v2.trainers.energy_trainer.EquiformerV2EnergyTrainer').

Expected Behavior
it should have been loaded, because it works for equiformerV2. Does ESEN require a different way to load?

Relevant files to reproduce this bug
No response

If I use,

calculator = OCPCalculator(checkpoint_path="./esen_30m_oam.pt",cpu=False.trainer="ocp")

It shows

INFO:root:amp: false
cmd:
  checkpoint_dir: /ztank/scratch/user/u.sa119259/fuse/bench/es-phase-5/checkpoints/2025-04-17-14-41-52
  commit: core:None,experimental:NA
  identifier: ''
  logs_dir: /ztank/scratch/user/u.sa119259/fuse/bench/es-phase-5/logs/wandb/2025-04-17-14-41-52
  print_every: 100
  results_dir: /ztank/scratch/user/u.sa119259/fuse/bench/es-phase-5/results/2025-04-17-14-41-52
  seed: null
  timestamp_id: 2025-04-17-14-41-52
version: 1.9.0
dataset:
  a2g_args:
    r_energy: true
    r_forces: true
    r_stress: true
  format: ase_db
  metadata_path: ''
  transforms:
    element_references: ''
    normalizer:
      energy:
        mean: 0.0
        stdev: 0.6274358
      forces:
        mean: 0.0
        stdev: 0.6274358
      stress:
        mean: 0.0
        stdev: 0.6274358
evaluation_metrics:
  metrics:
    energy:
    - mae
    - per_atom_mae
    forces:
    - mae
    - cosine_similarity
    stress:
    - mae
  primary_metric: forces_mae
gp_gpus: null
gpus: 0
logger: wandb
loss_functions:
-energy:
    coefficient: 20
    fn: per_atom_mae
- forces:
    coefficient: 20
    fn: l2mae
- stress:
    coefficient: 5
    fn: mae
model:
  backbone:
    act_type: gate
    cutoff: 6.0
    direct_forces: false
    distance_function: gaussian
    edge_channels: 128
    hidden_channels: 128
    lmax: 3
    max_neighbors: 300
    max_num_elements: 100
    mlp_type: spectral
    mmax: 2
    model: esen_backbone
    norm_type: rms_norm_sh
    num_distance_basis: 64
    num_layers: 10
    otf_graph: true
    regress_forces: true
    regress_stress: true
    sphere_channels: 128
    use_envelope: true
    use_pbc: true
    use_pbc_single: true
  heads:
    mptrj:
module: esen_mlp_efs_head
  name: hydra
  otf_graph: true
  pass_through_head_outputs: true
optim:
  batch_size: 4
  clip_grad_norm: 100
  ema_decay: 0.999
  eval_batch_size: 4
  eval_every: 5000
  load_balancing: atoms
  lr_initial: 0.0004
  max_epochs: 1
  num_workers: 4
  optimizer: AdamW
  optimizer_params:
    weight_decay: 0.001
  scheduler: LambdaLR
  scheduler_params:
    epochs: 85271
    lambda_type: cosine
    lr: 0.0004
    lr_min_factor: 0.1
    warmup_epochs: 8527
    warmup_factor: 0.2
outputs:
  energy:
    level: system
    property: energy
  forces:
    eval_on_free_atoms: true
    level: atom
    property: forces
    train_on_free_atoms: true
  stress:
    level: system
property: stress
relax_dataset: {}
slurm: {}
task: {}
test_dataset: {}
trainer: ocp
val_dataset: {}

INFO:root:Loading model: hydra
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/scratch/user/u.sa119259/envs/mlff/lib/python3.11/site-packages/fairchem/core/common/relaxation/ase_utils.py", line 236, in __init__
    self.trainer = registry.get_trainer_class(config["trainer"])(
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/scratch/user/u.sa119259/envs/mlff/lib/python3.11/site-packages/fairchem/core/trainers/ocp_trainer.py", line 109, in __init__
    super().__init__(
  File "/scratch/user/u.sa119259/envs/mlff/lib/python3.11/site-packages/fairchem/core/trainers/base_trainer.py", line 220, in __init__
    self.load(inference_only)
  File "/scratch/user/u.sa119259/envs/mlff/lib/python3.11/site-packages/fairchem/core/trainers/base_trainer.py", line 243, in load
    self.load_model()
  File "/scratch/user/u.sa119259/envs/mlff/lib/python3.11/site-packages/fairchem/core/trainers/base_trainer.py", line 561, in load_model
    self.model = registry.get_model_class(model_name)(
                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/scratch/user/u.sa119259/envs/mlff/lib/python3.11/site-packages/fairchem/core/models/base.py", line 276, in __init__
    self.backbone: BackboneInterface = registry.get_model_class(
                                       ^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/scratch/user/u.sa119259/envs/mlff/lib/python3.11/site-packages/fairchem/core/common/registry.py", line 294, in get_model_class
    return cls.get_class(name, "model_name_mapping")
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/scratch/user/u.sa119259/envs/mlff/lib/python3.11/site-packages/fairchem/core/common/registry.py", line 273, in get_class
    raise cls.__import_error(name, mapping_name)
RuntimeError: Failed to find the model 'esen_backbone'. You may either use a model from the registry (one of 'hydra', 'schnet', 'dimenetplusplus_energy_and_force_head', 'dimenetplusplus', 'dimenetplusplus_backbone', 'scn', 'equiformer_v2', 'rank2_symmetric_head', 'equiformerV2_scalar_head', 'equiformerV2_vector_head', 'equiformer_v2_force_head', 'equiformer_v2_energy_head', 'equiformer_v2_backbone', 'equiformer_v2_dens_backbone', 'eqV2_DeNS_scalar_head', 'eqV2_DeNS_vector_head', 'dens_rank2_symmetric_head', 'gp_gemnet_t', 'gemnet_oc', 'gemnet_oc_backbone', 'gemnet_oc_energy_and_grad_force_head', 'gemnet_oc_force_head', 'gemnet_t', 'gemnet_t_backbone', 'gemnet_t_energy_and_grad_force_head', 'gemnet_t_force_head', 'painn', 'painn_backbone', 'painn_energy_head', 'painn_force_head', 'escn', 'escn_backbone', 'escn_energy_head', 'escn_force_head' or 'escn_export') or provide the full import path to the model (e.g., 'fairchem.core.models.escn.escn_exportable.eSCN').