First PyTorch tests for TorchScript inference CPU/CUDA

PhysicsTools/PyTorch/BuildFile.xml

@@ -0,0 +1,8 @@
+<use name="FWCore/Framework"/>
+<use name="FWCore/MessageLogger"/>
+<use name="FWCore/Utilities"/>
+<use name="FWCore/ServiceRegistry"/>
+<use name="pytorch"/>
+<export>
+  <lib name="1"/>
+</export>

PhysicsTools/PyTorch/test/BuildFile.xml

@@ -0,0 +1,33 @@
+<use name="python_tools"/>
+
+<bin name="testTorchSimpleDnn" file="testRunner.cc,testTorchSimpleDnn.cc">
+  <use name="pytorch"/>
+  <use name="boost_filesystem"/>
+  <use name="cppunit"/>
+</bin>
+
+<iftool name="cuda">
+  <bin name="testTorchSimpleDnnCUDA" file="testRunner.cc,testTorchSimpleDnnCUDA.cc">
+    <use name="boost_filesystem"/>
+    <use name="catch2"/>
+    <use name="cppunit"/>
+    <use name="cuda"/>
+    <use name="pytorch"/>
+    <use name="FWCore/ParameterSet"/>
+    <use name="FWCore/ParameterSetReader"/>
+    <use name="FWCore/PluginManager"/>
+    <use name="FWCore/ServiceRegistry"/>
+    <use name="FWCore/Utilities"/>
+    <use name="HeterogeneousCore/CUDAServices"/>
+  </bin>
+</iftool>
+
+
+<bin name="testTorch" file="testTorch.cc">
+  <use name="pytorch"/>
+</bin>
+
+<bin name="testTorchTimeSeries" file="time_serie_prediction.cpp">
+  <use name="pytorch"/>
+  <use name="pytorch-cuda"/>
+</bin>

PhysicsTools/PyTorch/test/create_simple_dnn.py

@@ -0,0 +1,31 @@
+import sys
+import os
+import torch
+
+# prepare the datadir
+if len(sys.argv) >= 2:
+    datadir = sys.argv[1]
+else:
+    thisdir = os.path.dirname(os.path.abspath(__file__))
+    datadir = os.path.join(os.path.dirname(thisdir), "bin", "data")
+
+os.makedirs(datadir, exist_ok=True)
+
+class MyModule(torch.nn.Module):
+    def __init__(self, N, M):
+        super(MyModule, self).__init__()
+        self.weight = torch.nn.Parameter(torch.ones(N, M))
+        self.bias = torch.nn.Parameter(torch.ones(N))
+
+    def forward(self, input):
+          return torch.sum(torch.nn.functional.elu(self.weight.mv(input) + self.bias))
+
+
+module = MyModule(10, 10)
+x = torch.ones(10)
+
+tm = torch.jit.trace(module.eval(), x)
+
+tm.save(f"{datadir}/simple_dnn.pt")
+
+print("simple_dnn.pt created successfully!")

PhysicsTools/PyTorch/test/testBase.h

@@ -0,0 +1,68 @@
+/*
+ * Base class for tests.
+ *
+ */
+
+#ifndef PHYSICSTOOLS_PYTORCH_TEST_TESTBASE_H
+#define PHYSICSTOOLS_PYTORCH_TEST_TESTBASE_H
+
+#include <boost/filesystem.hpp>
+#include <filesystem>
+#include <cppunit/extensions/HelperMacros.h>
+#include <stdexcept>
+
+class testBasePyTorch : public CppUnit::TestFixture {
+public:
+  std::string dataPath_;
+
+  void setUp();
+  void tearDown();
+  std::string cmsswPath(std::string path);
+
+  virtual void test() = 0;
+
+  virtual std::string pyScript() const = 0;
+};
+
+void testBasePyTorch::setUp() {
+  dataPath_ =
+      cmsswPath("/test/" + std::string(std::getenv("SCRAM_ARCH")) + "/" + boost::filesystem::unique_path().string());
+
+  // create the graph using apptainer
+  std::string testPath = cmsswPath("/src/PhysicsTools/PyTorch/test");
+  std::string cmd = "apptainer exec -B " + cmsswPath("") +
+                    "  /cvmfs/unpacked.cern.ch/registry.hub.docker.com/cmsml/cmsml:3.11  python " + testPath + "/" +
+                    pyScript() + " " + dataPath_;
+  std::cout << "cmd: " << cmd << std::endl;
+  std::array<char, 128> buffer;
+  std::string result;
+  std::shared_ptr<FILE> pipe(popen(cmd.c_str(), "r"), pclose);
+  if (!pipe) {
+    throw std::runtime_error("Failed to run apptainer to prepare the PyTorch test model: " + cmd);
+  }
+  while (!feof(pipe.get())) {
+    if (fgets(buffer.data(), 128, pipe.get()) != NULL) {
+      result += buffer.data();
+    }
+  }
+  std::cout << std::endl << result << std::endl;
+}
+
+void testBasePyTorch::tearDown() {
+  if (std::filesystem::exists(dataPath_)) {
+    std::filesystem::remove_all(dataPath_);
+  }
+}
+
+std::string testBasePyTorch::cmsswPath(std::string path) {
+  if (path.size() > 0 && path.substr(0, 1) != "/") {
+    path = "/" + path;
+  }
+
+  std::string base = std::string(std::getenv("CMSSW_BASE"));
+  std::string releaseBase = std::string(std::getenv("CMSSW_RELEASE_BASE"));
+
+  return (std::filesystem::exists(base.c_str()) ? base : releaseBase) + path;
+}
+
+#endif  // PHYSICSTOOLS_PYTORCH_TEST_TESTBASE_H

PhysicsTools/PyTorch/test/testBaseCUDA.h

@@ -0,0 +1,78 @@
+/*
+ * Base class for tests.
+ *
+ */
+
+#ifndef PHYSICSTOOLS_PYTORCH_TEST_TESTBASECUDA_H
+#define PHYSICSTOOLS_PYTORCH_TEST_TESTBASECUDA_H
+
+#include <boost/filesystem.hpp>
+#include <filesystem>
+#include <cppunit/extensions/HelperMacros.h>
+#include <stdexcept>
+
+#include "FWCore/ParameterSet/interface/ParameterSet.h"
+#include "FWCore/ParameterSet/interface/ConfigurationDescriptions.h"
+#include "FWCore/ParameterSetReader/interface/ParameterSetReader.h"
+#include "FWCore/PluginManager/interface/PluginManager.h"
+#include "FWCore/PluginManager/interface/standard.h"
+#include "FWCore/ServiceRegistry/interface/Service.h"
+#include "FWCore/ServiceRegistry/interface/ServiceRegistry.h"
+#include "FWCore/ServiceRegistry/interface/ServiceToken.h"
+#include "FWCore/Utilities/interface/Exception.h"
+#include "FWCore/Utilities/interface/ResourceInformation.h"
+
+class testBasePyTorchCUDA : public CppUnit::TestFixture {
+public:
+  std::string dataPath_;
+
+  void setUp();
+  void tearDown();
+  std::string cmsswPath(std::string path);
+
+  virtual std::string pyScript() const = 0;
+
+  virtual void test() = 0;
+};
+
+void testBasePyTorchCUDA::setUp() {
+  dataPath_ =
+      cmsswPath("/test/" + std::string(std::getenv("SCRAM_ARCH")) + "/" + boost::filesystem::unique_path().string());
+
+  // create the graph using apptainer
+  std::string testPath = cmsswPath("/src/PhysicsTools/PyTorch/test");
+  std::string cmd = "apptainer exec -B " + cmsswPath("") +
+                    "  /cvmfs/unpacked.cern.ch/registry.hub.docker.com/cmsml/cmsml:3.11  python " + testPath + "/" +
+                    pyScript() + " " + dataPath_;
+  std::cout << "cmd: " << cmd << std::endl;
+  std::array<char, 128> buffer;
+  std::string result;
+  std::shared_ptr<FILE> pipe(popen(cmd.c_str(), "r"), pclose);
+  if (!pipe) {
+    throw std::runtime_error("Failed to run apptainer to prepare the PyTorch test model: " + cmd);
+  }
+  while (!feof(pipe.get())) {
+    if (fgets(buffer.data(), 128, pipe.get()) != NULL) {
+      result += buffer.data();
+    }
+  }
+  std::cout << std::endl << result << std::endl;
+}
+void testBasePyTorchCUDA::tearDown() {
+  if (std::filesystem::exists(dataPath_)) {
+    std::filesystem::remove_all(dataPath_);
+  }
+}
+
+std::string testBasePyTorchCUDA::cmsswPath(std::string path) {
+  if (path.size() > 0 && path.substr(0, 1) != "/") {
+    path = "/" + path;
+  }
+
+  std::string base = std::string(std::getenv("CMSSW_BASE"));
+  std::string releaseBase = std::string(std::getenv("CMSSW_RELEASE_BASE"));
+
+  return (std::filesystem::exists(base.c_str()) ? base : releaseBase) + path;
+}
+
+#endif  // PHYSICSTOOLS_PYTORCH_TEST_TESTBASE_H

PhysicsTools/PyTorch/test/testRunner.cc

@@ -0,0 +1 @@
+#include <Utilities/Testing/interface/CppUnit_testdriver.icpp>

PhysicsTools/PyTorch/test/testTorchSimpleDnn.cc

@@ -0,0 +1,42 @@
+#include <torch/script.h>
+#include "testBase.h"
+#include <iostream>
+#include <memory>
+#include <vector>
+
+class testSimpleDNN : public testBasePyTorch {
+  CPPUNIT_TEST_SUITE(testSimpleDNN);
+  CPPUNIT_TEST(test);
+  CPPUNIT_TEST_SUITE_END();
+
+public:
+  std::string pyScript() const override;
+  void test() override;
+};
+
+CPPUNIT_TEST_SUITE_REGISTRATION(testSimpleDNN);
+
+std::string testSimpleDNN::pyScript() const { return "create_simple_dnn.py"; }
+
+void testSimpleDNN::test() {
+  std::string model_path = dataPath_ + "/simple_dnn.pt";
+  torch::Device device(torch::kCPU);
+  torch::jit::script::Module module;
+  try {
+    // Deserialize the ScriptModule from a file using torch::jit::load().
+    module = torch::jit::load(model_path);
+    module.to(device);
+  } catch (const c10::Error& e) {
+    std::cerr << "error loading the model\n" << e.what() << std::endl;
+    CPPUNIT_ASSERT(false);
+  }
+  // Create a vector of inputs.
+  std::vector<torch::jit::IValue> inputs;
+  inputs.push_back(torch::ones(10, device));
+
+  // Execute the model and turn its output into a tensor.
+  at::Tensor output = module.forward(inputs).toTensor();
+  std::cout << "output: " << output << '\n';
+  CPPUNIT_ASSERT(output.item<float_t>() == 110.);
+  std::cout << "ok\n";
+}

PhysicsTools/PyTorch/test/testTorchSimpleDnnCUDA.cc

@@ -0,0 +1,64 @@
+#include <torch/script.h>
+#include "testBaseCUDA.h"
+#include <iostream>
+#include <memory>
+#include <vector>
+#include "HeterogeneousCore/CUDAServices/interface/CUDAInterface.h"
+
+class testSimpleDNNCUDA : public testBasePyTorchCUDA {
+  CPPUNIT_TEST_SUITE(testSimpleDNNCUDA);
+  CPPUNIT_TEST(test);
+  CPPUNIT_TEST_SUITE_END();
+
+public:
+  std::string pyScript() const override;
+  void test() override;
+};
+
+CPPUNIT_TEST_SUITE_REGISTRATION(testSimpleDNNCUDA);
+
+std::string testSimpleDNNCUDA::pyScript() const { return "create_simple_dnn.py"; }
+
+void testSimpleDNNCUDA::test() {
+  std::vector<edm::ParameterSet> psets;
+  edm::ServiceToken serviceToken = edm::ServiceRegistry::createSet(psets);
+  edm::ServiceRegistry::Operate operate(serviceToken);
+
+  // Setup the CUDA Service
+  edmplugin::PluginManager::configure(edmplugin::standard::config());
+
+  std::string const config = R"_(import FWCore.ParameterSet.Config as cms
+process = cms.Process('Test')
+process.add_(cms.Service('ResourceInformationService'))
+process.add_(cms.Service('CUDAService'))
+)_";
+  std::unique_ptr<edm::ParameterSet> params;
+  edm::makeParameterSets(config, params);
+  edm::ServiceToken tempToken(edm::ServiceRegistry::createServicesFromConfig(std::move(params)));
+  edm::ServiceRegistry::Operate operate2(tempToken);
+  edm::Service<CUDAInterface> cuda;
+  std::cout << "CUDA service enabled: " << cuda->enabled() << std::endl;
+
+  std::cout << "Testing CUDA backend" << std::endl;
+
+  std::string model_path = dataPath_ + "/simple_dnn.pt";
+  torch::Device device(torch::kCUDA);
+  torch::jit::script::Module module;
+  try {
+    // Deserialize the ScriptModule from a file using torch::jit::load().
+    module = torch::jit::load(model_path);
+    module.to(device);
+  } catch (const c10::Error& e) {
+    std::cerr << "error loading the model\n" << e.what() << std::endl;
+    CPPUNIT_ASSERT(false);
+  }
+  // Create a vector of inputs.
+  std::vector<torch::jit::IValue> inputs;
+  inputs.push_back(torch::ones(10, device));
+
+  // Execute the model and turn its output into a tensor.
+  at::Tensor output = module.forward(inputs).toTensor();
+  std::cout << "output: " << output << '\n';
+  CPPUNIT_ASSERT(output.item<float_t>() == 110.);
+  std::cout << "ok\n";
+}

PhysicsTools/PythonAnalysis/test/time_serie_prediction.cpp → PhysicsTools/PyTorch/test/time_serie_prediction.cpp

@@ -26,7 +26,7 @@ int main(int /*argc*/, char* /*argv*/[]) {
     std::cout << pair.key() << ": " << pair.value() << std::endl;
   }
 
-  std::cout << net.forward(torch::ones({2, 4})) << std::endl;
+  std::cout << net.forward(torch::ones({2, 4}).to(device)) << std::endl;
 
   return 0;
 }

PhysicsTools/PythonAnalysis/test/BuildFile.xml

@@ -130,11 +130,3 @@
 <test name="run-ipython" command="ipython -h"/>
 <test name="run-pylint" command="pylint -h"/>
 
-<bin name="testTorch" file="testTorch.cc">
-  <use name="pytorch"/>
-</bin>
-
-<bin name="testTorchTimeSeries" file="time_serie_prediction.cpp">
-  <use name="pytorch"/>
-  <use name="pytorch-cuda"/>
-</bin>