feat: add support of dynamic shape inference

Signed-off-by: ktro2828 <kotaro.uetake@tier4.jp>

Author: ktro2828

perception/tensorrt_mtr/config/tensorrt_mtr.param.yaml

@@ -1,14 +1,18 @@
 /**:
   ros__parameters:
-    model_path: "$(var data_path)/mtr_static.onnx"
-    precision: "fp32"
-    target_labels: ["VEHICLE", "PEDESTRIAN", "CYCLIST"]
-    num_past: 11
-    num_mode: 6
-    num_future: 80
-    max_num_polyline: 768
-    max_num_point: 20
-    point_break_distance: 1.0
-    intention_point_filepath: "$(var data_path)/intention_point.csv"
-    num_intention_point_cluster: 64
-    polyline_label_path: "$(var data_path)/polyline_label.txt"
+    model_params:
+      model_path: "$(var data_path)/mtr_static.onnx"
+      target_labels: ["VEHICLE", "PEDESTRIAN", "CYCLIST"]
+      num_past: 11
+      num_mode: 6
+      num_future: 80
+      max_num_polyline: 768
+      max_num_point: 20
+      point_break_distance: 1.0
+      intention_point_filepath: "$(var data_path)/intention_point.csv"
+      num_intention_point_cluster: 64
+      polyline_label_path: "$(var data_path)/polyline_label.txt"
+    build_params:
+      is_dynamic: false
+      precision: "FP32"
+      calibration: "MINMAX"

perception/tensorrt_mtr/include/tensorrt_mtr/builder.hpp

@@ -50,77 +50,84 @@ struct TrtDeleter
 
 template <typename T>
 using TrtUniquePtr = std::unique_ptr<T, TrtDeleter<T>>;
-using BatchConfig = std::array<int32_t, 3>;
 
-struct BuildConfig
-{
-  // type for calibration
-  std::string calib_type_str;
+// Type names of precisions.
+enum PrecisionType { FP32 = 0, FP16 = 1, INT8 = 2 };
 
-  // DLA core ID that the process uses
-  int dla_core_id;
+// Type names of calibrations.
+enum CalibrationType { ENTROPY = 0, LEGACY = 1, PERCENTILE = 2, MINMAX = 3 };
 
-  // flag for partial quantization in first layer
-  bool quantize_first_layer;  // For partial quantization
+struct BatchOptConfig
+{
+  /**
+   * @brief Construct a new OptimizationConfig for a static shape inference.
+   *
+   * @param value
+   */
+  explicit BatchOptConfig(const int32_t value) : k_min(value), k_opt(value), k_max(value) {}
+
+  /**
+   * @brief Construct a new OptimizationConfig for a dynamic shape inference.
+   *
+   * @param k_min
+   * @param k_opt
+   * @param k_max
+   */
+  BatchOptConfig(const int32_t k_min, const int32_t k_opt, const int32_t k_max)
+  : k_min(k_min), k_opt(k_opt), k_max(k_max)
+  {
+  }
 
-  // flag for partial quantization in last layer
-  bool quantize_last_layer;  // For partial quantization
+  int32_t k_min, k_opt, k_max;
+};  // struct BatchOptConfig
 
-  // flag for per-layer profiler using IProfiler
-  bool profile_per_layer;
+struct BuildConfig
+{
+  // type of precision
+  PrecisionType precision;
 
-  // clip value for implicit quantization
-  double clip_value;  // For implicit quantization
+  // type for calibration
+  CalibrationType calibration;
 
-  // Supported calibration type
-  const std::array<std::string, 4> valid_calib_type = {"Entropy", "Legacy", "Percentile", "MinMax"};
+  BatchOptConfig batch_target;
+  BatchOptConfig batch_agent;
 
   /**
    * @brief Construct a new instance with default configurations.
-   *
    */
   BuildConfig()
-  : calib_type_str("MinMax"),
-    dla_core_id(-1),
-    quantize_first_layer(false),
-    quantize_last_layer(false),
-    profile_per_layer(false),
-    clip_value(0.0)
+  : precision(PrecisionType::FP32),
+    calibration(CalibrationType::MINMAX),
+    batch_target(1, 10, 20),
+    batch_agent(1, 30, 50),
+    is_dynamic_(false)
   {
   }
 
   /**
-   * @brief Construct a new instance with custom configurations.
+   * @brief Construct a new build config.
    *
-   * @param calib_type_str The name of calibration type which must be selected from [Entropy,
-   * MinMax].
-   * @param dla_core_id DLA core ID used by the process.
-   * @param quantize_first_layer The flag whether to quantize first layer.
-   * @param quantize_last_layer The flag whether to quantize last layer.
-   * @param profile_per_layer The flag to profile per-layer in IProfiler.
-   * @param clip_value The value to be clipped in quantization implicitly.
+   * @param is_dynamic
+   * @param precision
+   * @param calibration
    */
-  explicit BuildConfig(
-    const std::string & calib_type_str, const int dla_core_id = -1,
-    const bool quantize_first_layer = false, const bool quantize_last_layer = false,
-    const bool profile_per_layer = false, const double clip_value = 0.0)
-  : calib_type_str(calib_type_str),
-    dla_core_id(dla_core_id),
-    quantize_first_layer(quantize_first_layer),
-    quantize_last_layer(quantize_last_layer),
-    profile_per_layer(profile_per_layer),
-    clip_value(clip_value)
+  BuildConfig(
+    const bool is_dynamic, const PrecisionType & precision = PrecisionType::FP32,
+    const CalibrationType & calibration = CalibrationType::MINMAX,
+    const BatchOptConfig & batch_target = BatchOptConfig(1, 10, 20),
+    const BatchOptConfig & batch_agent = BatchOptConfig(1, 30, 50))
+  : precision(precision),
+    calibration(calibration),
+    batch_target(batch_target),
+    batch_agent(batch_agent),
+    is_dynamic_(is_dynamic)
   {
-    if (
-      std::find(valid_calib_type.begin(), valid_calib_type.end(), calib_type_str) ==
-      valid_calib_type.end()) {
-      std::stringstream message;
-      message << "Invalid calibration type was specified: " << calib_type_str << std::endl
-              << "Valid value is one of: [Entropy, (Legacy | Percentile), MinMax]" << std::endl
-              << "Default calibration type will be used: MinMax" << std::endl;
-      std::cerr << message.str();
-    }
   }
+
+  bool is_dynamic() const { return is_dynamic_; }
+
+private:
+  bool is_dynamic_;
 };  // struct BuildConfig
 
 class MTRBuilder
@@ -130,15 +137,12 @@ class MTRBuilder
    * @brief Construct a new instance.
    *
    * @param model_path Path to engine or onnx file.
-   * @param precision The name of precision type.
-   * @param batch_config The configuration of min/opt/max batch.
-   * @param max_workspace_size The max workspace size.
    * @param build_config The configuration of build.
+   * @param max_workspace_size The max workspace size.
    */
   MTRBuilder(
-    const std::string & model_path, const std::string & precision,
-    const BatchConfig & batch_config = {1, 1, 1}, const size_t max_workspace_size = (1ULL << 30),
-    const BuildConfig & build_config = BuildConfig());
+    const std::string & model_path, const BuildConfig & build_config = BuildConfig(),
+    const size_t max_workspace_size = (1ULL << 63));
 
   /**
    * @brief Destroy the instance.
@@ -158,6 +162,12 @@ class MTRBuilder
    */
   bool isInitialized() const;
 
+  // Return true if the model supports dynamic shape inference.
+  bool isDynamic() const;
+
+  // Set binding dimensions for specified for dynamic shape inference.
+  bool setBindingDimensions(int index, nvinfer1::Dims dimensions);
+
   /**
    * @brief A wrapper of `nvinfer1::IExecuteContext::enqueueV2`.
    *
@@ -178,6 +188,9 @@ class MTRBuilder
    */
   bool loadEngine(const std::string & filepath);
 
+  // Create a cache path of engine file.
+  fs::path createEngineCachePath() const;
+
   /**
    * @brief Build engine from onnx file.
    *
@@ -194,8 +207,6 @@ class MTRBuilder
   TrtUniquePtr<nvinfer1::IExecutionContext> context_;
 
   fs::path model_filepath_;
-  std::string precision_;
-  BatchConfig batch_config_;
   size_t max_workspace_size_;
   std::unique_ptr<const BuildConfig> build_config_;
 

perception/tensorrt_mtr/include/tensorrt_mtr/node.hpp

@@ -73,7 +73,7 @@ class PolylineTypeMap
 public:
   explicit PolylineTypeMap(rclcpp::Node * node)
   {
-    const auto filepath = node->declare_parameter<std::string>("polyline_label_path");
+    const auto filepath = node->declare_parameter<std::string>("model_params.polyline_label_path");
     std::ifstream file(filepath);
     if (!file.is_open()) {
       RCLCPP_ERROR_STREAM(node->get_logger(), "Could not open polyline label file: " << filepath);
@@ -164,6 +164,7 @@ class MTRNode : public rclcpp::Node
 
   // MTR parameters
   std::unique_ptr<MTRConfig> config_ptr_;
+  std::unique_ptr<BuildConfig> build_config_ptr_;
   std::unique_ptr<TrtMTR> model_ptr_;
   PolylineTypeMap polyline_type_map_;
   std::shared_ptr<PolylineData> polyline_ptr_;

perception/tensorrt_mtr/include/tensorrt_mtr/trt_mtr.hpp

@@ -92,17 +92,15 @@ class TrtMTR
    * @brief Construct a new instance.
    *
    * @param model_path The path to engine or onnx file.
-   * @param precision The precision type.
    * @param config The configuration of model.
-   * @param batch_config The configuration of batch.
+   * @param build_config The configuration of build.
    * @param max_workspace_size The max size of workspace.
    * @param build_config The configuration of build.
    */
   TrtMTR(
-    const std::string & model_path, const std::string & precision,
-    const MTRConfig & config = MTRConfig(), const BatchConfig & batch_config = {1, 1, 1},
-    const size_t max_workspace_size = (1ULL << 30),
-    const BuildConfig & build_config = BuildConfig());
+    const std::string & model_path, const MTRConfig & config = MTRConfig(),
+    const BuildConfig & build_config = BuildConfig(),
+    const size_t max_workspace_size = (1ULL << 30));
 
   /**
    * @brief Execute inference.
@@ -160,20 +158,18 @@ class TrtMTR
 
   // data size
   // load from input data
-  size_t num_target_, num_agent_, time_length_;
-  size_t agent_input_dim_;
-  size_t num_polyline_, num_point_;
-  size_t polyline_input_dim_;
+  int32_t num_target_, num_agent_, num_timestamp_, num_agent_attr_;
+  int32_t num_polyline_, num_point_, num_point_dim_, num_point_attr_;
   // load from config
-  size_t max_num_polyline_, num_mode_, num_future_;
+  int32_t max_num_polyline_, num_mode_, num_future_, num_intention_point_;
 
   // source data
   cuda::unique_ptr<int[]> d_target_index_{nullptr};
   cuda::unique_ptr<int[]> d_label_index_{nullptr};
-  cuda::unique_ptr<float[]> d_timestamps_{nullptr};
+  cuda::unique_ptr<float[]> d_timestamp_{nullptr};
   cuda::unique_ptr<float[]> d_trajectory_{nullptr};
   cuda::unique_ptr<float[]> d_target_state_{nullptr};
-  cuda::unique_ptr<float[]> d_intention_points_{nullptr};
+  cuda::unique_ptr<float[]> d_intention_point_{nullptr};
   cuda::unique_ptr<float[]> d_polyline_{nullptr};
 
   // preprocessed inputs
@@ -191,6 +187,8 @@ class TrtMTR
   // outputs
   cuda::unique_ptr<float[]> d_out_score_{nullptr};
   cuda::unique_ptr<float[]> d_out_trajectory_{nullptr};
+  std::unique_ptr<float[]> h_out_score_{nullptr};
+  std::unique_ptr<float[]> h_out_trajectory_{nullptr};
 };  // class TrtMTR
 }  // namespace trt_mtr
 #endif  // TENSORRT_MTR__TRT_MTR_HPP_