mirror of
https://github.com/openmm/openmm-torch.git
synced 2026-03-10 19:14:16 +09:00
b76deb4df8
* Add CUDA graph draft
* Initialize energy and force tensors in the GPU.
* Add comment on graph capture
* Catch torch exception if the model fails to capture.
* Replay graph just after construction
Finish capturing before rethrowing if an exception occurred during capture
* Add python-side test script for CUDA graphs
* Implement properties
* Update the Python bindings
* Unify the API for properties
* Pass the propery map to the constructor
* Skip graph tests if no GPU is present
* Guard CUDA graph behavior with the CUDA_GRAPH_ENABLE macro
* Check validity of the useCUDAGraphs property
* Add missing bracket to openmmtorch.i
* Fix bug in useCUDAgraph selection
* Update tests
* Add test for get/setProperty
* Update documentation with new functionality
* Add a CUDA graph test for a model that returns only energy
* Add contributors
* Reset pos grads after graph capture. Make energy and force tensors persistent.
* Add tests that execute the model many times to catch bugs related with
CUDA graph capture
* Run formatter
* Warmup model for several steps
* Include gradient reset into the graph
* Do not reset energy and force tensors before graph capture
* Remove unnecessary line
* Add tests for larger number of particles
* Remove unnecessary compilation guard now that Pytorch 1.10 is not supported
* Simplify getTensorPointer now that Pytorch 1.7 is not supported
* Change addForcesToOpenMM to addForces
* Change execute_graph to executeGraph
* Wrap graph warming up in a try/catch block
* Add correctness test for modules that only provide energy
* Revert "Add correctness test for modules that only provide energy"
This reverts commit d20f4bfa83.
* Explicit conversion to correct type in getTensorPointer
* Added a new property for TorchForce, CUDAGraphWarmupSteps.
* Clarify docs
* Document properties
* Throw if requested property does not exist
* Change getProperty(string) to getProperties()
* Add getProperties to python wrappers
* Fix formatting
* Set default properties
* Update tests
* Update some comments
---------
Co-authored-by: Raimondas Galvelis <r.galvelis@acellera.com>
94 lines
3.2 KiB
Python
94 lines
3.2 KiB
Python
__author__ = "Raul P. Pelaez"
|
|
import openmmtorch as ot
|
|
import torch
|
|
import openmm as mm
|
|
import numpy as np
|
|
import pytest
|
|
|
|
|
|
class UngraphableModule(torch.nn.Module):
|
|
def forward(self, positions):
|
|
torch.cuda.synchronize()
|
|
return (torch.sum(positions**2), -2.0 * positions)
|
|
|
|
|
|
class GraphableModule(torch.nn.Module):
|
|
def forward(self, positions):
|
|
energy = torch.einsum("ij,ij->i", positions, positions).sum()
|
|
return (energy, -2.0 * positions)
|
|
|
|
|
|
class GraphableModuleOnlyEnergy(torch.nn.Module):
|
|
def forward(self, positions):
|
|
energy = torch.einsum("ij,ij->i", positions, positions).sum()
|
|
return energy
|
|
|
|
|
|
def tryToTestForceWithModule(
|
|
ModuleType, outputsForce, useGraphs=False, warmup=10, numParticles=10
|
|
):
|
|
"""Test that the force is correctly computed for a given module type.
|
|
Warmup makes OpenMM call TorchForce execution multiple times, which might expose some bugs related to that given that with CUDA graphs the first execution is different from the rest.
|
|
"""
|
|
module = torch.jit.script(ModuleType())
|
|
torch_force = ot.TorchForce(
|
|
module, {"useCUDAGraphs": "true" if useGraphs else "false"}
|
|
)
|
|
torch_force.setOutputsForces(outputsForce)
|
|
system = mm.System()
|
|
positions = np.random.rand(numParticles, 3)
|
|
for _ in range(numParticles):
|
|
system.addParticle(1.0)
|
|
system.addForce(torch_force)
|
|
integ = mm.VerletIntegrator(1.0)
|
|
platform = mm.Platform.getPlatformByName("CUDA")
|
|
context = mm.Context(system, integ, platform)
|
|
context.setPositions(positions)
|
|
for _ in range(warmup):
|
|
state = context.getState(getEnergy=True, getForces=True)
|
|
expectedEnergy = np.sum(positions**2)
|
|
expectedForce = -2.0 * positions
|
|
energy = state.getPotentialEnergy().value_in_unit(mm.unit.kilojoules_per_mole)
|
|
force = state.getForces(asNumpy=True).value_in_unit(
|
|
mm.unit.kilojoules_per_mole / mm.unit.nanometer
|
|
)
|
|
assert np.allclose(expectedEnergy, energy)
|
|
assert np.allclose(expectedForce, force)
|
|
|
|
|
|
def testUnGraphableModelRaises():
|
|
if not torch.cuda.is_available():
|
|
pytest.skip("CUDA not available")
|
|
with pytest.raises(mm.OpenMMException):
|
|
tryToTestForceWithModule(UngraphableModule, outputsForce=True, useGraphs=True)
|
|
|
|
|
|
@pytest.mark.parametrize("numParticles", [10, 10000])
|
|
@pytest.mark.parametrize("useGraphs", [True, False])
|
|
@pytest.mark.parametrize("warmup", [1, 10])
|
|
def testGraphableModelOnlyEnergyIsCorrect(useGraphs, warmup, numParticles):
|
|
if not torch.cuda.is_available():
|
|
pytest.skip("CUDA not available")
|
|
tryToTestForceWithModule(
|
|
GraphableModuleOnlyEnergy,
|
|
outputsForce=False,
|
|
useGraphs=useGraphs,
|
|
warmup=warmup,
|
|
numParticles=numParticles,
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize("numParticles", [10, 10000])
|
|
@pytest.mark.parametrize("useGraphs", [True, False])
|
|
@pytest.mark.parametrize("warmup", [1, 10])
|
|
def testGraphableModelIsCorrect(useGraphs, warmup, numParticles):
|
|
if not torch.cuda.is_available():
|
|
pytest.skip("CUDA not available")
|
|
tryToTestForceWithModule(
|
|
GraphableModule,
|
|
outputsForce=True,
|
|
useGraphs=useGraphs,
|
|
warmup=warmup,
|
|
numParticles=numParticles,
|
|
)
|