Batching input points while inference

[shasax]

Hello NEQUIP dev team,
I have a question about batching multiple input points when using a deployed NEQUIP model in C++. I saw in the pair_nequip codebase that the way it is implemented in LAMMPS is by creating one big graph of atoms. But, is there any way that could allow batching for faster inference?

Thanks in advance!

[Linux-cpp-lisp]

Hi @shasax ,

If by batching you mean running multiple frames / graphs / systems at the same time in a single call to the model (as in minibatching in training), then yes, nequip models certainly support this in inference mode.

However, it doesn't make sense to apply it for molecular dynamics in LAMMPS, since you only have the current frame to put through the model and do not know future frames until you run a time integration step.

If you want to run a large set of preexisting configurations/frames/graphs through a trained nequip model, see nequip-evaluate. This command takes a --batch-size flag that you can set to the maximum value that avoids out-of-memory errors in order to most efficiently take advantage of your GPU.

[Linux-cpp-lisp]

(You can also do batched inference calls to a model manually by loading a deployed model with nequip.scripts.deploy.load_deployed_model and passing a batched input to it.)

[shasax]

Thank you for your reply. My question is that how do I batch my input in c++ before passing it to a deployed model? Do I make independent graphs, or is there some other way of specifying different input graphs in one batch while using c++?

[Linux-cpp-lisp]

Hi @shasax ,

I see--- yes this is possible, but requires a little bit of coding. The logic for batching is all written in Python and a limited subset of it would need to be written in C++ (see https://github.com/mir-group/nequip/blob/main/nequip/utils/torch_geometric/batch.py#L40 for the full, and somewhat complex, implementation, and https://pytorch-geometric.readthedocs.io/en/latest/notes/batching.html for a description of the PyTorch Geometric batching scheme, which we use).

Luckily for you, the inputs to a typical model are pretty simple and so there's not much to do:

• Append together all the position tensors (N tensors of [N_atoms_i, 3] -> [\sum{N_atoms_i}, 3])
• Similarly with atom types
• Construct the edge index tensor so that the indexes are for the atoms in this new large batched position tensor, rather than in the individual examples (i.e. the 0th atom in the second frame has index 0 + num_atoms_first_frame)
• Stack cells (if relevant), i.e. [3, 3] -> [n_frame, 3, 3]
• Add a field batch to the input of dimension \sum{N_atoms_i}, which is a tensor of int64's giving for each atom the index of the frame to which it belongs (so something like [0, 0, 0, 0, 1, 1, 1, 2, 2, 2, 2, ...]

All of this can be done using torch::cat in C++, or just by building your input tensors this way to begin with (depending on where your data is coming from and what's easier).

[Linux-cpp-lisp]

If you write up a somewhat self-contained version of this, it would be great if you make a PR to add it to the examples/ folder for everyone's future reference 🙂

Looking forward to hear more about what you are doing with NequIP and Allegro!