use torch2.10 to export onnx, and a bit refine of trt10 code (#349)

Author: CoinCheung

tensorrt/plugins/kernels.hpp

@@ -45,7 +45,7 @@ int8_t shfl_down_sync_func(int8_t val, uint32_t delta) {
 
 template<typename scalar_t>
 __forceinline__ __device__
-scalar_t max_pair_shfl_func(scalar_t& val, int32_t& ind, const uint32_t delta) {
+void max_pair_shfl_func(scalar_t& val, int32_t& ind, const uint32_t delta) {
     scalar_t other_v = shfl_down_sync_func(val, delta);
     int32_t other_i = shfl_down_sync_func(ind, delta);
 

tensorrt/segment.cu

@@ -19,17 +19,6 @@
 #include "read_img.hpp"
 
 
-// using nvinfer1::IHostMemory;
-// using nvinfer1::IBuilder;
-// using nvinfer1::INetworkDefinition;
-// using nvinfer1::ICudaEngine;
-// using nvinfer1::IInt8Calibrator;
-// using nvinfer1::IBuilderConfig;
-// using nvinfer1::IRuntime;
-// using nvinfer1::IExecutionContext;
-// using nvinfer1::ILogger;
-// using nvinfer1::Dims;
-// using Severity = nvinfer1::ILogger::Severity;
 
 using std::string;
 using std::ios;

tensorrt/trt_dep.cu

@@ -102,42 +102,33 @@ void SemanticSegmentTrt::parse_to_engine(string onnx_pth,
     config->addOptimizationProfile(profile);
 
     config->setMemoryPoolLimit(nvinfer1::MemoryPoolType::kWORKSPACE, 1UL << 32);
-
-    if (quant == "fp16" or quant == "int8") { // fp16
-        if (builder->platformHasFastFp16() == false) {
-            cout << "fp16 is set, but platform does not support, so we ignore this\n";
-        } else {
-            config->setFlag(nvinfer1::BuilderFlag::kFP16); 
-        }
-    }
-    if (quant == "bf16") { // bf16
+    config->setBuilderOptimizationLevel(5);
+
+    if (quant == "fp16") { // fp16
+        config->setFlag(nvinfer1::BuilderFlag::kFP16); 
+    } else if (quant == "int8") { // int8
+        config->setFlag(nvinfer1::BuilderFlag::kFP16); 
+        config->setFlag(nvinfer1::BuilderFlag::kINT8); 
+    } else if (quant == "bf16") { // bf16
         config->setFlag(nvinfer1::BuilderFlag::kBF16); 
-    }
-    if (quant == "fp8") { // fp8
+    } else if (quant == "fp8") { // fp8
         config->setFlag(nvinfer1::BuilderFlag::kFP8); 
     }
 
     std::unique_ptr<IInt8Calibrator> calibrator;
     if (quant == "int8") { // int8
-        if (builder->platformHasFastInt8() == false) {
-            cout << "int8 is set, but platform does not support, so we ignore this\n";
-        } else {
-
-            int batchsize = 32;
-            int n_cal_batches = -1;
-            string cal_table_name = "calibrate_int8";
-
-            Dims indim = network->getInput(0)->getDimensions();
-            BatchStream calibrationStream(
-                    batchsize, n_cal_batches, indim,
-                    data_root, data_file);
-
-            config->setFlag(nvinfer1::BuilderFlag::kINT8); 
-
-            calibrator.reset(new Int8EntropyCalibrator2<BatchStream>(
-                calibrationStream, 0, cal_table_name.c_str(), input_name.c_str(), false));
-            config->setInt8Calibrator(calibrator.get());
-        }
+        int batchsize = 32;
+        int n_cal_batches = -1;
+        string cal_table_name = "calibrate_int8";
+
+        Dims indim = network->getInput(0)->getDimensions();
+        BatchStream calibrationStream(
+                batchsize, n_cal_batches, indim,
+                data_root, data_file);
+
+        calibrator.reset(new Int8EntropyCalibrator2<BatchStream>(
+            calibrationStream, 0, cal_table_name.c_str(), input_name.c_str(), false));
+        config->setInt8Calibrator(calibrator.get());
     }
 
     // output->setType(nvinfer1::DataType::kINT32);

tools/export_onnx.py

@@ -4,7 +4,7 @@
 sys.path.insert(0, '.')
 
 import torch
-from torch.onnx import OperatorExportTypes
+from torch.export import Dim
 
 from lib.models import model_factory
 from configs import set_cfg_from_file
@@ -37,11 +37,13 @@
 #  dummy_input = torch.randn(1, 3, 1024, 2048)
 input_names = ['input_image']
 output_names = ['preds',]
-dynamic_axes = {'input_image': {0: 'batch'}, 'preds': {0: 'batch'}}
+batch_size = Dim('batch', min=1, max=128)
+dynamic_shapes = ({0: batch_size},)
 
 torch.onnx.export(net, dummy_input, args.out_pth,
     input_names=input_names, output_names=output_names,
+    do_constant_folding=True,
     verbose=False, opset_version=18,
-    operator_export_type=OperatorExportTypes.ONNX_FALLTHROUGH,
-    dynamic_axes=dynamic_axes)
+    external_data=False,
+    dynamic_shapes=dynamic_shapes)