Using a model in python after training

[k50112113]

Hi,

Instead of running a deployed model through LAMMPS, I am wondering if there is an example python code to use a model after training a model.
For example, one might be able to evaluate forces and energy via the following python script:

atom_force = model.compute_force(atom_position)
energy = model.compute_energy(atom_position)

where atom_position and atom_force are (N*3) torch.Tensor and energy is a scalar

Thanks,
Shao-Chun Lee

[Linux-cpp-lisp]

Hi Shao-Chun,

Yep, you can definitely use our models from Python. Depending on your application, you can either:

1. Use our NequIPCalculator plugin for ASE (see ASE example #115 for an example). I recommend using this unless you have needs it cannot meet.
2. Use the deployed TorchScript model directly, which requires using nequip.data.AtomicData to build inputs and nequip.data.AtomicData.to_AtomicDataDict to make the final input object for the model. This is the most direct, lowest-level interface.
3. Using the LAMMPS pair style, but via LAMMPS' Python interface should be possible for larger-scale efforts.

Depending on what you're trying to do, the best solution may already be the nequip-evaluate --output command, which can do batched evaluation of the model on some provided set of input frames and write the results to an XYZ file for further postprocessing. This command handles all necessarily setup, batching, etc. without further thought.

[k50112113]

Thank you for your reply, I have successfully done inference from a trained model using NequIPCalculator with ASE module.
However, I wish I could do model inference directly from torch.Tensor, so I tried the second method. So far I am able to do inference using the second method with the following code:

import torch
import torch.nn as nn
import numpy as np
from nequip.utils import Config
from nequip.data import AtomicData, AtomicDataDict

#load data
config = Config.from_file('configs/minimal.yaml')
mydata = np.load('benchmark_data/aspirin_ccsd-test.npz')
pos = torch.tensor(mydata['R'][-1])
labeled_forces = torch.tensor(mydata['F'][-1])
z = mydata['z']
ATOMIC_NUMBERS_KEY = torch.Tensor(torch.from_numpy(z.astype(np.float32))).to(torch.int64)
ATOM_TYPE_KEY = torch.zeros_like(ATOMIC_NUMBERS_KEY)
type_unique = torch.unique(ATOMIC_NUMBERS_KEY)
for index, a_type in enumerate(type_unique): ATOM_TYPE_KEY[(ATOMIC_NUMBERS_KEY == a_type).nonzero()] = index
data = AtomicData.from_points(pos=pos, r_max=config['r_max'],
**{AtomicDataDict.ATOMIC_NUMBERS_KEY: ATOMIC_NUMBERS_KEY,
AtomicDataDict.ATOM_TYPE_KEY: ATOM_TYPE_KEY})
datadict = AtomicData.to_AtomicDataDict(data)
#load model
mymodel = torch.load("./results/aspirin/aspirin_model.pth")
mymodel.eval()
pred = mymodel(datadict)
print(pred['total_energy'])
print(pred['forces'])

However, every time I do inference, i.e., "mymodel(datadict)", the forces seems to be computed automatically. How can I "turn off" the evaluation of forces and evaluate only the energy? E.g., total_energy = mymodel.compute_total_energy(pos)

Thanks,
Shao-Chun Lee

[Linux-cpp-lisp]

If you print out the model and look at the tree, you can go in to the GradientOutput module that computes the forces and set skip = True. Something like:

model.model.skip = True

on latest main, but print(model) and check that I got the hierarchy right.

Side note: be careful with type mapping, this may not always be correct. You should instead build a TypeMapper from your config:

from nequip.data.transforms import TypeMapper
from nequip.utils import instantiate
type_mapper, _ = instantiate(TypeMapper, prefix="dataset", optional_args=config)
# ...
data = AtomicData.from_points(pos=pos, r_max=config['r_max'], **{AtomicDataDict.ATOMIC_NUMBERS_KEY: ATOMIC_NUMBERS_KEY})
data = type_mapper(data)

something like the above should do it.

[Linux-cpp-lisp]

Also very important is to use

 model, metadata = nequip.scripts.deploy.load_deployed_model(
            model_path="/path/to/whatever.pth",
            device=device,
            set_global_options="warn",
        )

and not torch.load. There are many reasons this is important and helpful, chief among them that it correctly sets certain global options to match those of the model (see ex #295 (reply in thread)). This should be called first so that those global options are already set for the rest of your code.

[k50112113]

There is no "model" attribute if I load the model by:

model, metadata = nequip.scripts.deploy.load_deployed_model(
            model_path="/path/to/whatever.pth",
            device=device,
            set_global_options="warn",
        )

but I am able to call model.model.skip if I load the model by:

model = torch. Load("/path/to/whatever.pth")

However, even if I can set model.model.skip = True, I got the following KeyError:

RuntimeError: 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/scl6/anaconda3/envs/nequip/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: forces

[Linux-cpp-lisp]

You should still print model in the nequip.scripts.deploy.load_deployed_model case to find the right submodule. Regardless, you should still use nequip.scripts.deploy.load_deployed_model.

Regarding skip... hm. Looks like a bug. We've never actually really used this feature 😆 What will definitely work is to make a copy of your original YAML file, delete the ForceOutput model builder, add load_model_state at the end of the model builders, and then set the load_model_state option to the path to the best_model.pth you originally deployed. Then you can deploy this new YAML, which will build a new model that doesn't compute forces, loading the trained weights from the old one.