Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Add GPU support for ONNXRuntime #36963

Merged
merged 1 commit into from
Apr 25, 2022
Merged

Add GPU support for ONNXRuntime #36963

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 7 additions & 0 deletions PhysicsTools/ONNXRuntime/interface/ONNXRuntime.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -22,13 +22,20 @@ namespace cms::Ort {

typedef std::vector<std::vector<float>> FloatArrays;

enum class Backend {
cpu,
cuda,
};

class ONNXRuntime {
public:
ONNXRuntime(const std::string& model_path, const ::Ort::SessionOptions* session_options = nullptr);
ONNXRuntime(const ONNXRuntime&) = delete;
ONNXRuntime& operator=(const ONNXRuntime&) = delete;
~ONNXRuntime();

static ::Ort::SessionOptions defaultSessionOptions(Backend backend = Backend::cpu);

// Run inference and get outputs
// input_names: list of the names of the input nodes.
// input_values: list of input arrays for each input node. The order of `input_values` must match `input_names`.
Expand Down
19 changes: 16 additions & 3 deletions PhysicsTools/ONNXRuntime/src/ONNXRuntime.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -27,9 +27,7 @@ namespace cms::Ort {
if (session_options) {
session_ = std::make_unique<Session>(env_, model_path.c_str(), *session_options);
} else {
SessionOptions sess_opts;
sess_opts.SetIntraOpNumThreads(1);
session_ = std::make_unique<Session>(env_, model_path.c_str(), sess_opts);
session_ = std::make_unique<Session>(env_, model_path.c_str(), defaultSessionOptions());
}
AllocatorWithDefaultOptions allocator;

Expand Down Expand Up @@ -78,6 +76,17 @@ namespace cms::Ort {

ONNXRuntime::~ONNXRuntime() {}

SessionOptions ONNXRuntime::defaultSessionOptions(Backend backend) {
SessionOptions sess_opts;
sess_opts.SetIntraOpNumThreads(1);
if (backend == Backend::cuda) {
// https://www.onnxruntime.ai/docs/reference/execution-providers/CUDA-ExecutionProvider.html
OrtCUDAProviderOptions options;
sess_opts.AppendExecutionProvider_CUDA(options);
}
return sess_opts;
}

FloatArrays ONNXRuntime::run(const std::vector<std::string>& input_names,
FloatArrays& input_values,
const std::vector<std::vector<int64_t>>& input_shapes,
Expand All @@ -104,6 +113,10 @@ namespace cms::Ort {
} else {
input_dims = input_shapes[input_pos];
// rely on the given input_shapes to set the batch size
if (input_dims[0] != batch_size) {
throw cms::Exception("RuntimeError") << "The first element of `input_shapes` (" << input_dims[0]
<< ") does not match the given `batch_size` (" << batch_size << ")";
}
Comment on lines +116 to +119
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

this change is unrelated to the use of GPUs, isn't it ?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

No it's unrelated to GPU -- just a helpful error message for users.

}
auto expected_len = std::accumulate(input_dims.begin(), input_dims.end(), 1, std::multiplies<int64_t>());
if (expected_len != (int64_t)value->size()) {
Expand Down
1 change: 1 addition & 0 deletions PhysicsTools/ONNXRuntime/test/BuildFile.xml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,5 +4,6 @@
<bin name="testONNXRuntime" file="testRunner.cpp, testONNXRuntime.cc">
<use name="cppunit"/>
<use name="PhysicsTools/ONNXRuntime"/>
<use name="HeterogeneousCore/CUDAUtilities"/>
<use name="FWCore/ParameterSet"/>
</bin>
24 changes: 19 additions & 5 deletions PhysicsTools/ONNXRuntime/test/testONNXRuntime.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,26 +2,32 @@

#include "PhysicsTools/ONNXRuntime/interface/ONNXRuntime.h"
#include "FWCore/ParameterSet/interface/FileInPath.h"
#include "HeterogeneousCore/CUDAUtilities/interface/requireDevices.h"

#include <chrono>
#include <iostream>

using namespace cms::Ort;

class testONNXRuntime : public CppUnit::TestFixture {
CPPUNIT_TEST_SUITE(testONNXRuntime);
CPPUNIT_TEST(checkAll);
CPPUNIT_TEST(checkCPU);
CPPUNIT_TEST(checkGPU);
CPPUNIT_TEST_SUITE_END();

private:
void test(Backend backend);

public:
void checkAll();
void checkCPU();
void checkGPU();
};

CPPUNIT_TEST_SUITE_REGISTRATION(testONNXRuntime);

void testONNXRuntime::checkAll() {
void testONNXRuntime::test(Backend backend) {
std::string model_path = edm::FileInPath("PhysicsTools/ONNXRuntime/test/data/model.onnx").fullPath();
ONNXRuntime rt(model_path);
auto session_options = ONNXRuntime::defaultSessionOptions(backend);
ONNXRuntime rt(model_path, &session_options);
for (const unsigned batch_size : {1, 2, 4}) {
FloatArrays input_values{
std::vector<float>(batch_size * 2, 1),
Expand All @@ -35,3 +41,11 @@ void testONNXRuntime::checkAll() {
}
}
}

void testONNXRuntime::checkCPU() { test(Backend::cpu); }

void testONNXRuntime::checkGPU() {
if (cms::cudatest::testDevices()) {
test(Backend::cuda);
}
}