Merge pull request #143 from peastman/derivatives

Can compute parameter derivatives

Author: RaulPPelaez

openmmapi/include/TorchForce.h

@@ -9,7 +9,7 @@
  * Biological Structures at Stanford, funded under the NIH Roadmap for        *
  * Medical Research, grant U54 GM072970. See https://simtk.org.               *
  *                                                                            *
- * Portions copyright (c) 2018-2022 Stanford University and the Authors.      *
+ * Portions copyright (c) 2018-2024 Stanford University and the Authors.      *
  * Authors: Peter Eastman                                                     *
  * Contributors: Raimondas Galvelis, Raul P. Pelaez                           *
  *                                                                            *
@@ -36,6 +36,7 @@
 #include "openmm/Force.h"
 #include <map>
 #include <string>
+#include <vector>
 #include <torch/torch.h>
 #include "internal/windowsExportTorch.h"
 
@@ -106,6 +107,11 @@ class OPENMM_EXPORT_NN TorchForce : public OpenMM::Force {
      * Get the number of global parameters that the interaction depends on.
      */
     int getNumGlobalParameters() const;
+    /**
+     * Get the number of global parameters with respect to which the derivative of the energy
+     * should be computed.
+     */
+    int getNumEnergyParameterDerivatives() const;
     /**
      * Add a new global parameter that the interaction may depend on.  The default value provided to
      * this method is the initial value of the parameter in newly created Contexts.  You can change
@@ -144,6 +150,21 @@ class OPENMM_EXPORT_NN TorchForce : public OpenMM::Force {
      * @param defaultValue   the default value of the parameter
      */
     void setGlobalParameterDefaultValue(int index, double defaultValue);
+    /**
+     * Request that this Force compute the derivative of its energy with respect to a global parameter.
+     * The parameter must have already been added with addGlobalParameter().
+     *
+     * @param name             the name of the parameter
+     */
+    void addEnergyParameterDerivative(const std::string& name);
+    /**
+     * Get the name of a global parameter with respect to which this Force should compute the
+     * derivative of the energy.
+     *
+     * @param index     the index of the parameter derivative, between 0 and getNumEnergyParameterDerivatives()
+     * @return the parameter name
+     */
+    const std::string& getEnergyParameterDerivativeName(int index) const;
     /**
      * Set a value of a property.
      *
@@ -163,6 +184,7 @@ class OPENMM_EXPORT_NN TorchForce : public OpenMM::Force {
     std::string file;
     bool usePeriodic, outputsForces;
     std::vector<GlobalParameterInfo> globalParameters;
+    std::vector<int> energyParameterDerivatives;
     torch::jit::Module module;
     std::map<std::string, std::string> properties;
     std::string emptyProperty;

openmmapi/src/TorchForce.cpp

@@ -6,7 +6,7 @@
  * Biological Structures at Stanford, funded under the NIH Roadmap for        *
  * Medical Research, grant U54 GM072970. See https://simtk.org.               *
  *                                                                            *
- * Portions copyright (c) 2018-2022 Stanford University and the Authors.      *
+ * Portions copyright (c) 2018-2024 Stanford University and the Authors.      *
  * Authors: Peter Eastman                                                     *
  * Contributors: Raimondas Galvelis, Raul P. Pelaez                           *
  *                                                                            *
@@ -112,6 +112,24 @@ void TorchForce::setGlobalParameterDefaultValue(int index, double defaultValue)
     globalParameters[index].defaultValue = defaultValue;
 }
 
+int TorchForce::getNumEnergyParameterDerivatives() const {
+    return energyParameterDerivatives.size();
+}
+
+void TorchForce::addEnergyParameterDerivative(const string& name) {
+    for (int i = 0; i < globalParameters.size(); i++)
+        if (name == globalParameters[i].name) {
+            energyParameterDerivatives.push_back(i);
+            return;
+        }
+    throw OpenMMException(string("addEnergyParameterDerivative: Unknown global parameter '"+name+"'"));
+}
+
+const string& TorchForce::getEnergyParameterDerivativeName(int index) const {
+    ASSERT_VALID_INDEX(index, energyParameterDerivatives);
+    return globalParameters[energyParameterDerivatives[index]].name;
+}
+
 void TorchForce::setProperty(const std::string& name, const std::string& value) {
     if (properties.find(name) == properties.end())
         throw OpenMMException("TorchForce: Unknown property '" + name + "'");

platforms/cuda/src/CudaTorchKernels.cpp

@@ -6,7 +6,7 @@
  * Biological Structures at Stanford, funded under the NIH Roadmap for        *
  * Medical Research, grant U54 GM072970. See https://simtk.org.               *
  *                                                                            *
- * Portions copyright (c) 2018-2022 Stanford University and the Authors.      *
+ * Portions copyright (c) 2018-2024 Stanford University and the Authors.      *
  * Authors: Peter Eastman                                                     *
  * Contributors: Raimondas Galvelis, Raul P. Pelaez                           *
  *                                                                            *
@@ -66,6 +66,10 @@ void CudaCalcTorchForceKernel::initialize(const System& system, const TorchForce
     outputsForces = force.getOutputsForces();
     for (int i = 0; i < force.getNumGlobalParameters(); i++)
         globalNames.push_back(force.getGlobalParameterName(i));
+    for (int i = 0; i < force.getNumEnergyParameterDerivatives(); i++) {
+        paramDerivs.insert(force.getEnergyParameterDerivativeName(i));
+        cu.addEnergyParameterDerivative(force.getEnergyParameterDerivativeName(i));
+    }
     int numParticles = system.getNumParticles();
 
     // Push the PyTorch context
@@ -81,6 +85,8 @@ void CudaCalcTorchForceKernel::initialize(const System& system, const TorchForce
     boxTensor = torch::empty({3, 3}, options);
     energyTensor = torch::empty({0}, options);
     forceTensor = torch::empty({0}, options);
+    for (const string& name : globalNames)
+        globalTensors[name] = torch::tensor({0}, options);
     // Pop the PyToch context
     CUcontext ctx;
     CHECK_RESULT(cuCtxPopCurrent(&ctx), "Failed to pop the CUDA context");
@@ -125,7 +131,7 @@ static void* getTensorPointer(OpenMM::CudaContext& cu, torch::Tensor& tensor) {
 /**
  * Prepare the inputs for the PyTorch model, copying positions from the OpenMM context.
  */
-std::vector<torch::jit::IValue> CudaCalcTorchForceKernel::prepareTorchInputs(ContextImpl& context) {
+void CudaCalcTorchForceKernel::prepareTorchInputs(ContextImpl& context, vector<torch::jit::IValue>& inputs, map<string, torch::Tensor>& globalTensors) {
     int numParticles = cu.getNumAtoms();
     // Get pointers to the atomic positions and simulation box
     void* posData = getTensorPointer(cu, posTensor);
@@ -145,12 +151,17 @@ std::vector<torch::jit::IValue> CudaCalcTorchForceKernel::prepareTorchInputs(Con
         CHECK_RESULT(cuCtxSynchronize(), "Failed to synchronize the CUDA context"); // Synchronize before switching to the PyTorch context
     }
     // Prepare the input of the PyTorch model
-    vector<torch::jit::IValue> inputs = {posTensor};
+    inputs = {posTensor};
     if (usePeriodic)
         inputs.push_back(boxTensor);
-    for (const string& name : globalNames)
-        inputs.push_back(torch::tensor(context.getParameter(name)));
-    return inputs;
+    for (const string& name : globalNames) {
+        // PyTorch requires us to set requires_grad to false before initializing a tensor.
+        globalTensors[name].set_requires_grad(false);
+        globalTensors[name][0] = context.getParameter(name);
+        if (paramDerivs.find(name) != paramDerivs.end())
+            globalTensors[name].set_requires_grad(true);
+        inputs.push_back(globalTensors[name]);
+    }
 }
 
 /**
@@ -173,26 +184,34 @@ void CudaCalcTorchForceKernel::addForces(torch::Tensor& forceTensor) {
 }
 
 /**
- * This function launches  the workload in a way  compatible with CUDA
- * graphs as far  as OpenMM-Torch goes.  Capturing  this function when
- * the model  is not  itself graph compatible  (due to,  for instance,
+ * This function launches the workload in a way compatible with CUDA
+ * graphs as far as OpenMM-Torch goes.  Capturing this function when
+ * the model is not itself graph compatible (due to, for instance,
  * implicit synchronizations) will result in a CUDA error.
  */
 static void executeGraph(bool outputsForces, bool includeForces, torch::jit::script::Module& module, vector<torch::jit::IValue>& inputs, torch::Tensor& posTensor, torch::Tensor& energyTensor,
-                         torch::Tensor& forceTensor) {
+                         torch::Tensor& forceTensor, map<string, torch::Tensor>& globalTensors, set<string> paramDerivs) {
+    vector<torch::Tensor> gradInputs;
+    if (!outputsForces && includeForces)
+        gradInputs.push_back(posTensor);
+    for (auto& name : paramDerivs)
+        gradInputs.push_back(globalTensors[name]);
+    auto none = torch::Tensor();
     if (outputsForces) {
         auto outputs = module.forward(inputs).toTuple();
         energyTensor = outputs->elements()[0].toTensor();
         forceTensor = outputs->elements()[1].toTensor();
+        if (gradInputs.size() > 0)
+            energyTensor.backward(none, false, false, gradInputs);
     } else {
         energyTensor = module.forward(inputs).toTensor();
         // Compute force by backpropagating the PyTorch model
+        // CUDA graph capture sometimes fails if backwards is not explicitly requested w.r.t positions
+        // See https://github.com/openmm/openmm-torch/pull/120/
+        if (gradInputs.size() > 0)
+            energyTensor.backward(none, false, false, gradInputs);
         if (includeForces) {
-            // CUDA graph capture sometimes fails if backwards is not explicitly requested w.r.t positions
-	    // See https://github.com/openmm/openmm-torch/pull/120/
-	    auto none = torch::Tensor();
-	    energyTensor.backward(none, false, false, posTensor);
-	    // This is minus the forces, we change the sign later on
+            // This is minus the forces, we change the sign later on
             forceTensor = posTensor.grad().clone();
             // Zero the gradient to avoid accumulating it
             posTensor.grad().zero_();
@@ -203,31 +222,32 @@ static void executeGraph(bool outputsForces, bool includeForces, torch::jit::scr
 double CudaCalcTorchForceKernel::execute(ContextImpl& context, bool includeForces, bool includeEnergy) {
     // Push to the PyTorch context
     CHECK_RESULT(cuCtxPushCurrent(primaryContext), "Failed to push the CUDA context");
-    auto inputs = prepareTorchInputs(context);
+    vector<torch::jit::IValue> inputs;
+    prepareTorchInputs(context, inputs, globalTensors);
     if (!useGraphs) {
-        executeGraph(outputsForces, includeForces, module, inputs, posTensor, energyTensor, forceTensor);
+        executeGraph(outputsForces, includeForces, module, inputs, posTensor, energyTensor, forceTensor, globalTensors, paramDerivs);
     } else {
         // Record graph if not already done
         bool is_graph_captured = false;
         if (graphs.find(includeForces) == graphs.end()) {
 	    //CUDA graph capture must occur in a non-default stream
             const auto stream = c10::cuda::getStreamFromPool(false, cu.getDeviceIndex());
-	    const c10::cuda::CUDAStreamGuard guard(stream);
+	        const c10::cuda::CUDAStreamGuard guard(stream);
             // Warmup the graph workload before capturing.  This first
             // run  before  capture sets  up  allocations  so that  no
             // allocations are  needed after.  Pytorch's  allocator is
             // stream  capture-aware and,  after warmup,  will provide
             // record static pointers and shapes during capture.
             try {
                 for (int i = 0; i < this->warmupSteps; i++)
-                    executeGraph(outputsForces, includeForces, module, inputs, posTensor, energyTensor, forceTensor);
+                    executeGraph(outputsForces, includeForces, module, inputs, posTensor, energyTensor, forceTensor, globalTensors, paramDerivs);
             }
             catch (std::exception& e) {
                 throw OpenMMException(string("TorchForce Failed to warmup the model before graph construction. Torch reported the following error:\n") + e.what());
             }
             graphs[includeForces].capture_begin();
             try {
-                executeGraph(outputsForces, includeForces, module, inputs, posTensor, energyTensor, forceTensor);
+                executeGraph(outputsForces, includeForces, module, inputs, posTensor, energyTensor, forceTensor, globalTensors, paramDerivs);
                 is_graph_captured = true;
                 graphs[includeForces].capture_end();
             }
@@ -237,16 +257,23 @@ double CudaCalcTorchForceKernel::execute(ContextImpl& context, bool includeForce
                 }
                 throw OpenMMException(string("TorchForce Failed to capture the model into a CUDA graph. Torch reported the following error:\n") + e.what());
             }
+            for (const string& name : paramDerivs)
+                globalTensors[name].grad().zero_();
         }
-	// Use the same stream as the OpenMM context, even if it is the default stream
+        // Use the same stream as the OpenMM context, even if it is the default stream
         const auto openmmStream = cu.getCurrentStream();
-	const auto stream = c10::cuda::getStreamFromExternal(openmmStream, cu.getDeviceIndex());
-	const c10::cuda::CUDAStreamGuard guard(stream);
+        const auto stream = c10::cuda::getStreamFromExternal(openmmStream, cu.getDeviceIndex());
+        const c10::cuda::CUDAStreamGuard guard(stream);
         graphs[includeForces].replay();
     }
     if (includeForces) {
         addForces(forceTensor);
     }
+    map<string, double>& energyParamDerivs = cu.getEnergyParamDerivWorkspace();
+    for (const string& name : paramDerivs) {
+        energyParamDerivs[name] += globalTensors[name].grad().item<double>();
+        globalTensors[name].grad().zero_();
+    }
     // Get energy
     const double energy = energyTensor.item<double>(); // This implicitly synchronizes the PyTorch context
     // Pop to the PyTorch context

platforms/cuda/src/CudaTorchKernels.h

@@ -9,7 +9,7 @@
  * Biological Structures at Stanford, funded under the NIH Roadmap for        *
  * Medical Research, grant U54 GM072970. See https://simtk.org.               *
  *                                                                            *
- * Portions copyright (c) 2018-2022 Stanford University and the Authors.      *
+ * Portions copyright (c) 2018-2024 Stanford University and the Authors.      *
  * Authors: Peter Eastman                                                     *
  * Contributors: Raimondas Galvelis, Raul P. Pelaez                           *
  *                                                                            *
@@ -37,6 +37,7 @@
 #include "openmm/cuda/CudaArray.h"
 #include <torch/version.h>
 #include <ATen/cuda/CUDAGraph.h>
+#include <set>
 
 namespace TorchPlugin {
 
@@ -71,12 +72,14 @@ class CudaCalcTorchForceKernel : public CalcTorchForceKernel {
     torch::jit::script::Module module;
     torch::Tensor posTensor, boxTensor;
     torch::Tensor energyTensor, forceTensor;
+    std::map<std::string, torch::Tensor> globalTensors;
     std::vector<std::string> globalNames;
+    std::set<std::string> paramDerivs;
     bool usePeriodic, outputsForces;
     CUfunction copyInputsKernel, addForcesKernel;
     CUcontext primaryContext;
     std::map<bool, at::cuda::CUDAGraph> graphs;
-    std::vector<torch::jit::IValue> prepareTorchInputs(OpenMM::ContextImpl& context);
+    void prepareTorchInputs(OpenMM::ContextImpl& context, std::vector<torch::jit::IValue>& inputs, std::map<std::string, torch::Tensor>& derivInputs);
     bool useGraphs;
     void addForces(torch::Tensor& forceTensor);
     int warmupSteps;

platforms/cuda/tests/TestCudaTorchForce.cpp

@@ -6,7 +6,7 @@
  * Biological Structures at Stanford, funded under the NIH Roadmap for        *
  * Medical Research, grant U54 GM072970. See https://simtk.org.               *
  *                                                                            *
- * Portions copyright (c) 2018-2022 Stanford University and the Authors.      *
+ * Portions copyright (c) 2018-2024 Stanford University and the Authors.      *
  * Authors: Peter Eastman                                                     *
  * Contributors:                                                              *
  *                                                                            *
@@ -129,7 +129,7 @@ void testPeriodicForce() {
     ASSERT_EQUAL_TOL(expectedEnergy, state.getPotentialEnergy(), 1e-5);
 }
 
-void testGlobal() {
+void testGlobal(bool useGraphs) {
     // Create a random cloud of particles.
 
     const int numParticles = 10;
@@ -143,6 +143,8 @@ void testGlobal() {
     }
     TorchForce* force = new TorchForce("tests/global.pt");
     force->addGlobalParameter("k", 2.0);
+    force->addEnergyParameterDerivative("k");
+    force->setProperty("useCUDAGraphs", useGraphs ? "true" : "false");
     system.addForce(force);
 
     // Compute the forces and energy.
@@ -151,7 +153,7 @@ void testGlobal() {
     Platform& platform = Platform::getPlatformByName("CUDA");
     Context context(system, integ, platform);
     context.setPositions(positions);
-    State state = context.getState(State::Energy | State::Forces);
+    State state = context.getState(State::Energy | State::Forces | State::ParameterDerivatives);
 
     // See if the energy is correct.  The network defines a potential of the form E(r) = k*|r|^2
 
@@ -164,15 +166,26 @@ void testGlobal() {
     }
     ASSERT_EQUAL_TOL(expectedEnergy, state.getPotentialEnergy(), 1e-5);
 
+    // Check the gradient of the energy with respect to the parameter.
+
+    double expected = 0.0;
+    for (int i = 0; i < numParticles; i++) {
+        Vec3 pos = positions[i];
+        expected += pos.dot(pos);
+    }
+    double actual = state.getEnergyParameterDerivatives().at("k");
+    ASSERT_EQUAL_TOL(expected, actual, 1e-5);
+
     // Change the global parameter and see if the forces are still correct.
 
     context.setParameter("k", 3.0);
-    state = context.getState(State::Forces);
+    state = context.getState(State::Forces | State::ParameterDerivatives);
     for (int i = 0; i < numParticles; i++) {
         Vec3 pos = positions[i];
         double r = sqrt(pos.dot(pos));
         ASSERT_EQUAL_VEC(pos*(-6.0), state.getForces()[i], 1e-5);
     }
+    ASSERT_EQUAL_TOL(expected, state.getEnergyParameterDerivatives().at("k"), 1e-5);
 }
 
 int main(int argc, char* argv[]) {
@@ -183,7 +196,8 @@ int main(int argc, char* argv[]) {
         testForce(false);
         testForce(true);
         testPeriodicForce();
-        testGlobal();
+        testGlobal(false);
+        testGlobal(true);
     }
     catch(const std::exception& e) {
         std::cout << "exception: " << e.what() << std::endl;