Allow extra inputs to model (#5)

* Allow extra inputs to model

* Serialization of inputs

* Python interface for inputs

* Documentation for inputs

* Fixed compilation error on Python 3.12 and earlier

Author: peastman

README.md

@@ -159,6 +159,54 @@ on the `Context`.
 context.setParameter("k", 5.0)
 ```
 
+## Extra Inputs
+
+You also can specify extra inputs that should be passed to the model.  Unlike global parameters,
+which are always scalars, extra inputs can be tensors of any size and shape.  On the other hand,
+their values are fixed at Context creation time, and can only be changed by reinitializing the
+Context.
+
+This example is similar to the one above, but `k` is now a vector containing a different force
+constant for every particle.
+
+```python
+class ForceWithInput(torch.nn.Module):
+    def forward(self, positions, k):
+        positions.grad = None
+        r2 = torch.sum(positions*positions, dim=1)
+        energy = torch.sum(k*r2)
+        energy.backward()
+        forces = -positions.grad
+        return energy, forces
+
+torch.onnx.export(model=ForceWithInput(),
+                  args=(torch.ones(1, 3, requires_grad=True), torch.ones(1)),
+                  f="ForceWithInput.onnx",
+                  input_names=["positions", "k"],
+                  output_names=["energy", "forces"],
+                  dynamic_axes={"positions":[0], "forces":[0], "k":[0]})
+```
+
+Notice that we included `k[0]` in `dynamic_axes` when exporting the model.  This allows its length
+to be variable, so we can use the model for systems with any number of particles.
+
+Here is how we create the OnnxForce.
+
+```python
+import openmmonnx
+force = OnnxForce("ForceWithInput.onnx")
+force.addInput(openmmonnx.FloatInput("k", k, [len(k)]))
+```
+
+The three arguments to the FloatInput constructor are the name of the input (matching the name we
+specified in `input_names` when exporting the model), a list or array containing the values, and
+the shape of the tensor.  In this case the tensor has one dimension, so the shape argument contains
+only a single value.  Higher dimensional tensors are also allowed.  In that case, the second
+argument should contain the values in flattened order.
+
+In addition to FloatInput, which specifies a tensor of 32 bit floating point values, there is also
+an IntegerInput class, which specifies a tensor of 32 bit integer values.
+
 ## Execution Providers
 
 ONNX Runtime supports a variety of backends that can be used to compute the neural network.  They

openmmapi/include/OnnxForce.h

@@ -46,6 +46,9 @@ namespace OnnxPlugin {
 
 class OPENMM_EXPORT_ONNX OnnxForce : public OpenMM::Force {
 public:
+    class Input;
+    class IntegerInput;
+    class FloatInput;
     /**
      * This is an enumeration of ONNX execution providers.
      */
@@ -89,6 +92,7 @@ class OPENMM_EXPORT_ONNX OnnxForce : public OpenMM::Force {
      * @param properties optional map of properties
      */
     OnnxForce(const std::vector<uint8_t>& model, const std::map<std::string, std::string>& properties={});
+    ~OnnxForce();
     /**
      * Get the binary representation of the model in ONNX format.
      */
@@ -161,6 +165,32 @@ class OPENMM_EXPORT_ONNX OnnxForce : public OpenMM::Force {
      * @param defaultValue   the default value of the parameter
      */
     void setGlobalParameterDefaultValue(int index, double defaultValue);
+    /**
+     * Get the number of extra tensors to pass to the model.
+     */
+    int getNumInputs() const;
+    /**
+     * Add an extra tensor that should be passed to the model.  The Input object should have
+     * been created on the heap with the "new" operator.  The OnnxForce takes over ownership
+     * of it, and deletes it when the OnnxForce itself is deleted.
+     *
+     * @param input       the tensor to pass in.  This should be the appropriate subclass for
+     *                    the type of value.
+     * @return the index of the input that was added
+     */
+    int addInput(Input* input);
+    /**
+     * Get a const reference to an extra input that is passed to the model.
+     *
+     * @param index     the index of the input to return
+     */
+    const Input& getInput(int index) const;
+    /**
+     * Get a writable reference to an extra input that is passed to the model.
+     *
+     * @param index     the index of the input to return
+     */
+    Input& getInput(int index);
     /**
      * Set the value of a property.
      *
@@ -183,9 +213,118 @@ class OPENMM_EXPORT_ONNX OnnxForce : public OpenMM::Force {
     ExecutionProvider provider;
     bool periodic;
     std::vector<GlobalParameterInfo> globalParameters;
+    std::vector<Input*> inputs;
     std::map<std::string, std::string> properties;
 };
 
+/**
+ * An Input defines a tensor that should be passed to the model.  This is an abstract class.  Subclasses
+ * define particular types of tensors.
+ */
+class OnnxForce::Input {
+public:
+    virtual ~Input() {
+    }
+    /**
+     * Get the name of the input.
+     */
+    const std::string& getName() const {
+        return name;
+    }
+    /**
+     * Get the shape of the tensor.
+     */
+    const std::vector<int>& getShape() const {
+        return shape;
+    }
+    /**
+     * Set the shape of the tensor.
+     */
+    void setShape(const std::vector<int>& shape) {
+        this->shape = shape;
+    }
+protected:
+    Input() {
+    }
+    Input(const std::string& name, const std::vector<int>& shape) : name(name), shape(shape) {
+    }
+private:
+    std::string name;
+    std::vector<int> shape;
+};
+
+/**
+ * A tensor containing integer values that should be passed to the model.
+ */
+class OnnxForce::IntegerInput : public Input {
+public:
+    /**
+     * Create an IntegerInput.
+     *
+     * @param name      the name of the input
+     * @param values    the values contained by the tensor, in flattened order
+     * @param shape     the shape of the tensor
+     */
+    IntegerInput(const std::string& name, const std::vector<int>& values, const std::vector<int>& shape) : Input(name, shape), values(values) {
+    }
+    /**
+     * Get a const reference to the values contained in the tensor, in flattened order.
+     */
+    const std::vector<int>& getValues() const {
+        return values;
+    }
+    /**
+     * Get a writable reference to the values contained in the tensor, in flattened order.
+     */
+    std::vector<int>& getValues() {
+        return values;
+    }
+    /**
+     * Set the values contained in the tensor, in flattened order.
+     */
+    void setValues(const std::vector<int>& values) {
+        this->values = values;
+    }
+private:
+    std::vector<int> values;
+};
+
+/**
+ * A tensor containing float values that should be passed to the model.
+ */
+class OnnxForce::FloatInput : public Input {
+public:
+    /**
+     * Create an FloatInput.
+     *
+     * @param name      the name of the input
+     * @param values    the values contained by the tensor, in flattened order
+     * @param shape     the shape of the tensor
+     */
+    FloatInput(const std::string& name, const std::vector<float>& values, const std::vector<int>& shape) : Input(name, shape), values(values) {
+    }
+    /**
+     * Get a const reference to the values contained in the tensor, in flattened order.
+     */
+    const std::vector<float>& getValues() const {
+        return values;
+    }
+    /**
+     * Get a writable reference to the values contained in the tensor, in flattened order.
+     */
+    std::vector<float>& getValues() {
+        return values;
+    }
+    /**
+     * Set the values contained in the tensor, in flattened order.
+     */
+    void setValues(const std::vector<float>& values) {
+        this->values = values;
+    }
+private:
+    std::vector<float> values;
+};
+
 /**
  * This is an internal class used to record information about a global parameter.
  * @private

openmmapi/include/internal/OnnxForceImpl.h

@@ -56,12 +56,15 @@ class OPENMM_EXPORT_ONNX OnnxForceImpl : public OpenMM::CustomCPPForceImpl {
     double computeForce(OpenMM::ContextImpl& context, const std::vector<OpenMM::Vec3>& positions, std::vector<OpenMM::Vec3>& forces);
 private:
     const OnnxForce& owner;
+    void validateInput(const std::string& name, const std::vector<int>& shape, int size);
     Ort::Env env;
     Ort::Session session;
     std::vector<Ort::Value> inputTensors, outputTensors;
     std::vector<const char*> inputNames;
     std::vector<int> particleIndices;
     std::vector<float> positionVec, paramVec;
+    std::vector<OnnxForce::IntegerInput> integerInputs;
+    std::vector<OnnxForce::FloatInput> floatInputs;
     float boxVectors[9];
 };
 

openmmapi/src/OnnxForce.cpp

@@ -52,6 +52,11 @@ OnnxForce::OnnxForce(const std::vector<uint8_t>& model, const map<string, string
     initProperties(properties);
 }
 
+OnnxForce::~OnnxForce() {
+    for (Input* input : inputs)
+        delete input;
+}
+
 void OnnxForce::initProperties(const std::map<std::string, std::string>& properties) {
     const std::map<std::string, std::string> defaultProperties = {{"UseGraphs", "false"}, {"DeviceIndex", "0"}};
     this->properties = defaultProperties;
@@ -123,6 +128,25 @@ void OnnxForce::setGlobalParameterDefaultValue(int index, double defaultValue) {
     globalParameters[index].defaultValue = defaultValue;
 }
 
+int OnnxForce::getNumInputs() const {
+    return inputs.size();
+}
+
+int OnnxForce::addInput(OnnxForce::Input* input) {
+    inputs.push_back(input);
+    return inputs.size() - 1;
+}
+
+const OnnxForce::Input& OnnxForce::getInput(int index) const {
+    ASSERT_VALID_INDEX(index, inputs);
+    return *inputs[index];
+}
+
+OnnxForce::Input& OnnxForce::getInput(int index) {
+    ASSERT_VALID_INDEX(index, inputs);
+    return *inputs[index];
+}
+
 void OnnxForce::setProperty(const string& name, const string& value) {
     if (properties.find(name) == properties.end())
         throw OpenMMException("OnnxForce: Unknown property '" + name + "'");

openmmapi/src/OnnxForceImpl.cpp

@@ -32,6 +32,7 @@
 #include "internal/OnnxForceImpl.h"
 #include "openmm/OpenMMException.h"
 #include "openmm/internal/ContextImpl.h"
+#include <sstream>
 
 using namespace OnnxPlugin;
 using namespace OpenMM;
@@ -125,6 +126,46 @@ void OnnxForceImpl::initialize(ContextImpl& context) {
         inputTensors.emplace_back(Value::CreateTensor<float>(memoryInfo, &paramVec[i], 1, paramShape, 1));
         inputNames.push_back(owner.getGlobalParameterName(i).c_str());
     }
+
+    // Process extra inputs.
+
+    for (int i = 0; i < owner.getNumInputs(); i++) {
+        const OnnxForce::IntegerInput* integerInput = dynamic_cast<const OnnxForce::IntegerInput*>(&owner.getInput(i));
+        if (integerInput != nullptr) {
+            validateInput(integerInput->getName(), integerInput->getShape(), integerInput->getValues().size());
+            integerInputs.push_back(OnnxForce::IntegerInput(integerInput->getName(), integerInput->getValues(), integerInput->getShape()));
+        }
+        const OnnxForce::FloatInput* floatInput = dynamic_cast<const OnnxForce::FloatInput*>(&owner.getInput(i));
+        if (floatInput != nullptr) {
+            validateInput(floatInput->getName(), floatInput->getShape(), floatInput->getValues().size());
+            floatInputs.push_back(OnnxForce::FloatInput(floatInput->getName(), floatInput->getValues(), floatInput->getShape()));
+        }
+    }
+    for (OnnxForce::IntegerInput& input : integerInputs) {
+        vector<int64_t> shape;
+        for (int i : input.getShape())
+            shape.push_back(i);
+        inputTensors.emplace_back(Value::CreateTensor<int>(memoryInfo, input.getValues().data(), input.getValues().size(), shape.data(), shape.size()));
+        inputNames.push_back(input.getName().c_str());
+    }
+    for (OnnxForce::FloatInput& input : floatInputs) {
+        vector<int64_t> shape;
+        for (int i : input.getShape())
+            shape.push_back(i);
+        inputTensors.emplace_back(Value::CreateTensor<float>(memoryInfo, input.getValues().data(), input.getValues().size(), shape.data(), shape.size()));
+        inputNames.push_back(input.getName().c_str());
+    }
+}
+
+void OnnxForceImpl::validateInput(const string& name, const vector<int>& shape, int size) {
+    int expected = 1;
+    for (int i : shape)
+        expected *= i;
+    if (expected != size) {
+        stringstream message;
+        message<<"Incorrect length for input '"<<name<<"'.  Expected "<<expected<<" elements, found "<<size<<".";
+        throw OpenMMException(message.str());
+    }
 }
 
 map<string, double> OnnxForceImpl::getDefaultParameters() {