NNPOps
The goal of this project is to promote the use of neural network potentials (NNPs) by providing highly optimized, open source implementations of bottleneck operations that appear in popular potentials. These are the core design principles.
-
Each operation is entirely self contained, consisting of only a few source files that can easily be incorporated into any code that needs to use it.
-
Each operation has a simple, clearly documented API allowing it to be easily used in any context.
-
We provide both CPU (pure C++) and CUDA implementations of all operations.
-
The CUDA implementations are highly optimized. The CPU implementations are written in a generally efficient way, but no particular efforts have been made to tune them for optimum performance.
-
This code is designed for inference (running simulations), not training (creating new potential functions). It computes gradients with respect to particle positions, not model parameters.
Installation
Install with Conda
A conda package can be installed from the conda-forge channel:
conda install -c conda-forge nnpopsIf you don't have conda, we recommend installing Miniconda.
Build from source
Prerequisites
- CUDA Toolkit (https://developer.nvidia.com/cuda-downloads)
- Miniconda (https://docs.conda.io/en/latest/miniconda.html#linux-installers)
Build & install
- Get the source code
$ git clone https://github.com/openmm/NNPOps.git- Set
CUDA_HOME
$ export CUDA_HOME=/usr/local/cuda-11.2- Crate and activate a Conda environment
$ cd NNPOps
$ conda env create -n nnpops -f environment.yml
$ conda activate nnpops- Configure, build, and install
$ mkdir build && cd build
$ cmake .. \
-DTorch_DIR=$CONDA_PREFIX/lib/python3.9/site-packages/torch/share/cmake/Torch \
-DCMAKE_INSTALL_PREFIX=$CONDA_PREFIX
$ make install- Run the tests
$ ctest --verboseUsage
Accelerated TorchANI operations:
Example
import mdtraj
import torch
import torchani
from NNPOps.SpeciesConverter import TorchANISpeciesConverter
from NNPOps.SymmetryFunctions import TorchANISymmetryFunctions
from NNPOps.BatchedNN import TorchANIBatchedNN
from NNPOps.EnergyShifter import TorchANIEnergyShifter
from NNPOps import OptimizedTorchANI
device = torch.device('cuda')
# Load a molecule
molecule = mdtraj.load('molecule.mol2')
species = torch.tensor([[atom.element.atomic_number for atom in molecule.top.atoms]], device=device)
positions = torch.tensor(molecule.xyz * 10, dtype=torch.float32, requires_grad=True, device=device)
# Construct ANI-2x and replace its operations with the optimized ones
nnp = torchani.models.ANI2x(periodic_table_index=True).to(device)
nnp.species_converter = TorchANISpeciesConverter(nnp.species_converter, species).to(device)
nnp.aev_computer = TorchANISymmetryFunctions(nnp.species_converter, nnp.aev_computer, species).to(device)
nnp.neural_networks = TorchANIBatchedNN(nnp.species_converter, nnp.neural_networks, species).to(device)
nnp.energy_shifter = TorchANIEnergyShifter(nnp.species_converter, nnp.energy_shifter, species).to(device)
# Compute energy and forces
energy = nnp((species, positions)).energies
energy.backward()
forces = -positions.grad.clone()
print(energy, forces)
# Alternatively, all the optimizations can be applied with OptimizedTorchANI
nnp2 = torchani.models.ANI2x(periodic_table_index=True).to(device)
nnp2 = OptimizedTorchANI(nnp2, species).to(device)
# Compute energy and forces again
energy = nnp2((species, positions)).energies
positions.grad.zero_()
energy.backward()
forces = -positions.grad.clone()
print(energy, forces)