fix(tensorrt_yolo): multi gpu on component container (autowarefoundat…

…ion#2297) Signed-off-by: Kaan Colak <[email protected]> Signed-off-by: Kaan Çolak <[email protected]> Signed-off-by: kminoda <[email protected]>
kminoda · Jan 6, 2023 · 604e36e · 604e36e
1 parent 35337cf
commit 604e36e
Show file tree

Hide file tree

Showing 3 changed files with 10 additions and 2 deletions.
diff --git a/perception/tensorrt_yolo/README.md b/perception/tensorrt_yolo/README.md
@@ -61,6 +61,7 @@ Jocher, G., et al. (2021). ultralytics/yolov5: v6.0 - YOLOv5n 'Nano' models, Rob
 | `calib_image_directory` | string | ""            | The directory name including calibration images for int8 inference |
 | `calib_cache_file`      | string | ""            | The calibration cache file for int8 inference                      |
 | `mode`                  | string | "FP32"        | The inference mode: "FP32", "FP16", "INT8"                         |
+| `gpu_id`                | int    | 0             | GPU device ID that runs the model                                  |
 
 ## Assumptions / Known limits
 

diff --git a/perception/tensorrt_yolo/include/tensorrt_yolo/nodelet.hpp b/perception/tensorrt_yolo/include/tensorrt_yolo/nodelet.hpp
@@ -58,6 +58,8 @@ class TensorrtYoloNodelet : public rclcpp::Node
 
   yolo::Config yolo_config_;
 
+  int gpu_device_id_;
+
   std::vector<std::string> labels_;
   std::unique_ptr<float[]> out_scores_;
   std::unique_ptr<float[]> out_boxes_;

diff --git a/perception/tensorrt_yolo/src/nodelet.cpp b/perception/tensorrt_yolo/src/nodelet.cpp
@@ -51,7 +51,7 @@ TensorrtYoloNodelet::TensorrtYoloNodelet(const rclcpp::NodeOptions & options)
   std::string calib_image_directory = declare_parameter("calib_image_directory", "");
   std::string calib_cache_file = declare_parameter("calib_cache_file", "");
   std::string mode = declare_parameter("mode", "FP32");
-  int gpu_device_id = declare_parameter("gpu_id", 0);
+  gpu_device_id_ = declare_parameter("gpu_id", 0);
   yolo_config_.num_anchors = declare_parameter("num_anchors", 3);
   auto anchors = declare_parameter(
     "anchors", std::vector<double>{
@@ -67,7 +67,7 @@ TensorrtYoloNodelet::TensorrtYoloNodelet(const rclcpp::NodeOptions & options)
   yolo_config_.use_darknet_layer = declare_parameter("use_darknet_layer", true);
   yolo_config_.ignore_thresh = declare_parameter("ignore_thresh", 0.5);
 
-  if (!yolo::set_cuda_device(gpu_device_id)) {
+  if (!yolo::set_cuda_device(gpu_device_id_)) {
     RCLCPP_ERROR(this->get_logger(), "Given GPU not exist or suitable");
   }
 
@@ -134,6 +134,11 @@ void TensorrtYoloNodelet::callback(const sensor_msgs::msg::Image::ConstSharedPtr
 
   tier4_perception_msgs::msg::DetectedObjectsWithFeature out_objects;
 
+  if (!yolo::set_cuda_device(gpu_device_id_)) {
+    RCLCPP_ERROR(this->get_logger(), "Given GPU not exist or suitable");
+    return;
+  }
+
   cv_bridge::CvImagePtr in_image_ptr;
   try {
     in_image_ptr = cv_bridge::toCvCopy(in_image_msg, sensor_msgs::image_encodings::BGR8);