[tensorrt] Fix multi gpu issue (#1249)

Change-Id: Ib9b219b773a6fa52cac6dd6d59a2541530c91a43
deepjavalibrary · Sep 24, 2021 · 5907d50 · 5907d50
1 parent 0c311d0
commit 5907d50
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 5 deletions.
diff --git a/engines/tensorrt/src/main/native/ai_djl_tensorrt_jni_model.cc b/engines/tensorrt/src/main/native/ai_djl_tensorrt_jni_model.cc
@@ -151,15 +151,15 @@ void TrtModel::buildModel() {
 }
 
 TrtSession *TrtModel::createSession() {
-  auto *session = new TrtSession(mEngine, mParams.maxBatchSize);
-
-  CHECK(cudaSetDevice(mParams.device));
+  auto *session = new TrtSession(mEngine, mParams.device, mParams.maxBatchSize);
 
   session->init();
   return session;
 }
 
 void TrtSession::init() {
+  CHECK(cudaSetDevice(mDeviceId));
+
   mContext = mEngine->createExecutionContext();
   int bindings = mEngine->getNbBindings();
   mBufferSizes.reserve(bindings);
@@ -221,6 +221,8 @@ void TrtSession::copyOutputs() {
 }
 
 void TrtSession::predict() {
+  CHECK(cudaSetDevice(mDeviceId));
+
   copyInputs();
 
   bool status;

diff --git a/engines/tensorrt/src/main/native/ai_djl_tensorrt_jni_model.h b/engines/tensorrt/src/main/native/ai_djl_tensorrt_jni_model.h
@@ -42,8 +42,8 @@ struct ModelParams {
 
 class TrtSession {
  public:
-  explicit TrtSession(std::shared_ptr<nvinfer1::ICudaEngine> engine, int batchSize)
-      : mEngine(std::move(engine)), mBatchSize(batchSize), mContext(nullptr) {}
+  explicit TrtSession(std::shared_ptr<nvinfer1::ICudaEngine> engine, int deviceId, int batchSize)
+      : mEngine(std::move(engine)), mDeviceId(deviceId), mBatchSize(batchSize), mContext(nullptr) {}
 
   void init();
   nvinfer1::Dims getShape(const char* name);
@@ -58,6 +58,7 @@ class TrtSession {
 
  private:
   std::shared_ptr<nvinfer1::ICudaEngine> mEngine;
+  int mDeviceId;
   nvinfer1::IExecutionContext* mContext;
   int32_t mBatchSize;
   std::vector<size_t> mBufferSizes;