RuntimeError: KeyError: stress Issue

[MusannaGalib]

Hello,

I am facing the following error when trying to run the Allegro deployed .pth model in LAMMPS. Can you please suggest any possible reason for this?

[W CUDAFunctions.cpp:108] Warning: CUDA initialization: Unexpected error from cudaGetDeviceCount(). Did you run some cuda functions before calling NumCudaDevices() that might have already set an error? Error 34: CUDA driver is a stub library (function operator())
terminate called after throwing an instance of 'std::runtime_error'
  what():  The following operation failed in the TorchScript interpreter.
Traceback of TorchScript, serialized code (most recent call last):
  File "code/__torch__/nequip/nn/_rescale.py", line 31, in forward
        for _0 in range(torch.len(scale_keys)):
          field = scale_keys[_0]
          _1 = data0[field]
               ~~~~~~~~~~~~ <--- HERE
          scale_by = self.scale_by
          torch._set_item(data0, field, torch.mul(_1, scale_by))

Traceback of TorchScript, original code (most recent call last):
  File "/home/galibubc/ML_allegro/lib/python3.9/site-packages/nequip/nn/_rescale.py", line 147, in forward
            if self.has_scale:
                for field in self.scale_keys:
                    data[field] = data[field] * self.scale_by
                                  ~~~~~~~~~~~ <--- HERE
            if self.has_shift:
                for field in self.shift_keys:
RuntimeError: KeyError: stress

Yaml configuration used:

# general
root: results/alumina
run_name: alumina
seed: 123456
dataset_seed: 123456
append: true
default_dtype: float32

# -- network --
model_builders:
 - allegro.model.Allegro
 # the typical model builders from `nequip` can still be used:
 - PerSpeciesRescale
 #- ForceOutput
 - StressForceOutput
 - RescaleEnergyEtc
# - initialize_from_state 
#initial_model_state: /home/galibubc/scratch/ML_LAMMPS/ML_allergo/results_1_100_epoch/alumina-tutorial/alumina_allegro/best_model.pth



# cutoffs
r_max: 3.0
avg_num_neighbors: auto

# radial basis
BesselBasis_trainable: true
PolynomialCutoff_p: 6   

# symmetry
l_max: 1
parity: o3_full   

# Allegro layers:
num_layers: 1
env_embed_multiplicity: 8
embed_initial_edge: true

two_body_latent_mlp_latent_dimensions: [32, 64, 128]
two_body_latent_mlp_nonlinearity: silu
two_body_latent_mlp_initialization: uniform

latent_mlp_latent_dimensions: [128]
latent_mlp_nonlinearity: silu
latent_mlp_initialization: uniform
latent_resnet: true

env_embed_mlp_latent_dimensions: []
env_embed_mlp_nonlinearity: null
env_embed_mlp_initialization: uniform

# - end allegro layers -

# Final MLP to go from Allegro latent space to edge energies:
edge_eng_mlp_latent_dimensions: [32]
edge_eng_mlp_nonlinearity: null
edge_eng_mlp_initialization: uniform

# -- data --
dataset: ase                                                                   
dataset_file_name: /home/galibubc/scratch/ML_LAMMPS/ML_allergo/MLIP_data.extxyz                      # path to data set file
ase_args:
  format: extxyz
include_keys:
  - energy
#  - virial
key_mapping:
  energy: total_energy
#  virial: stress
dataset_include_keys:
   - virial

# A mapping of chemical species to type indexes is necessary if the dataset is provided with atomic numbers instead of type indexes.
#chemical_symbol_to_type:
#  Si: 0
chemical_symbols:
  - Al
  - O 

# logging
wandb: true
wandb_project: Alumina_potential_full
verbose: info
log_batch_freq: 10

# training
n_train: 10000
n_val: 1500
batch_size: 1
max_epochs: 150
learning_rate: 0.002
train_val_split: random
shuffle: true
metrics_key: validation_loss

# use an exponential moving average of the weights
use_ema: true
ema_decay: 0.99
ema_use_num_updates: true

# loss function
loss_coeffs:
  forces: 1.
  total_energy:
    - 1.
    - PerAtomMSELoss
#  stress: 1.  

# optimizer
optimizer_name: Adam
optimizer_params:
  amsgrad: false
  betas: !!python/tuple
  - 0.9
  - 0.999
  eps: 1.0e-08
  weight_decay: 0.

metrics_components:
  - - forces                               # key 
    - mae                                  # "rmse" or "mae"
  - - forces
    - rmse
  - - total_energy
    - mae    
  - - total_energy
    - mae
    - PerAtom: True                        # if true, energy is normalized by the number of atoms
 # - - stress                                # key for stress
 #   - mae                                  # "mae" or "rmse" for stress
 # - - stress
 #   - rmse

# lr scheduler, drop lr if no improvement for 50 epochs
lr_scheduler_name: ReduceLROnPlateau
lr_scheduler_patience: 50
lr_scheduler_factor: 0.5

early_stopping_lower_bounds:
  LR: 1.0e-5

early_stopping_patiences:
  validation_loss: 100

Model trained successfully :

  Train      #    Epoch      wal       LR       loss_f       loss_e         loss        f_mae       f_rmse        e_mae      e/N_mae
! Train             100 33483.959    0.002       0.0571      0.00799       0.0651       0.0575       0.0842         1.98       0.0253
! Validation        100 33483.959    0.002       0.0519       0.0063       0.0582       0.0535       0.0805          1.8        0.023
Wall time: 33483.96047454048
! Best model      100    0.058
! Stop training: max epochs
Wall time: 33484.061463536695
Cumulative wall time: 33484.061463536695

Packages used in the environment

(ML_allegro) galibubc@cedar1 ~/ML_allegro/bin $ pip list
Package                   Version
------------------------- -----------------------------
ase                       3.23.0+computecanada
e3nn                      0.5.3
mir-allegro               0.2.0
nequip                    0.5.5
torch                     1.11.0+cu113
torch_ema                 0.3+computecanada
torch-runstats            0.2.0

[MusannaGalib]

@Linux-cpp-lisp @cw-tan can you please comment on this? Thanks a lot.

[Linux-cpp-lisp]

I suggest trying again with the latest stable main branch of nequip, 0.6.1: https://github.com/mir-group/nequip/releases/tag/v0.6.1. I am not sure but suspect this may have been resolved.

[MusannaGalib]

Thanks @Linux-cpp-lisp. 0.6.1 solved the RuntimeError: KeyError: stress issue.

However when trying to the benchmark test getting the following issue. My extxyz dataset consists of frames with atom numbers ranging from 76 to 80. Didnt see this issue with 0.5.5.

!nequip-benchmark allegro/configs/tutorial.yaml --n-data 2

Using device: cuda
Loading dataset... 
    loading dataset took 0.3666s
    loaded dataset of size 12006 and sampled --n-data=2 frames
Traceback (most recent call last):
  File "/home/galibubc/ML_allegro_0.6.1/bin/nequip-benchmark", line 8, in <module>
    sys.exit(main())
  File "/home/galibubc/ML_allegro_0.6.1/lib/python3.9/site-packages/nequip/scripts/benchmark.py", line 127, in main
    raise NotImplementedError(
NotImplementedError: nequip-benchmark does not currently handle benchmarking on data frames with variable number of atoms

Packages used in the environment

import torch
import allegro  
import nequip
import wandb


# Release unused memory
torch.cuda.empty_cache()
print(torch.__version__)
print(torch.version.cuda)
print(allegro.__version__)
print(nequip.__version__)
print(wandb.__version__)


1.11.0+cu113
11.3
0.2.0
0.6.1
0.15.4

[cw-tan]

This is not a bug, data with different number of atoms is just not supported in nequip-benchmark for the current public branches. This feature is largely for equivariance testing and timing purposes., so you could just construct an xyz file with the same structure repeated multiple times to use it.