onnx

CoinCheung · Sep 26, 2024 · 5c8ced3 · 5c8ced3
1 parent 2a882dd
commit 5c8ced3
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 2 deletions.
diff --git a/tensorrt/trt_dep.cpp b/tensorrt/trt_dep.cpp
@@ -95,7 +95,7 @@ void SemanticSegmentTrt::parse_to_engine(string onnx_pth,
     profile->setDimensions(input->getName(), OptProfileSelector::kMAX, dmax);
     config->addOptimizationProfile(profile);
 
-    if (quant == "fp16") { // fp16
+    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 {
@@ -275,6 +275,7 @@ void SemanticSegmentTrt::test_speed_fps() {
 
     cout << "\ntest with cropsize of (" << iH << ", " << iW << "), "
         << "and batch size of " << batchsize << " ...\n";
+    context->executeV2(buffs.data()); // run one batch ahead
     auto start = std::chrono::steady_clock::now();
     const int n_loops{2000};
     for (int i{0}; i < n_loops; ++i) {

diff --git a/tools/export_onnx.py b/tools/export_onnx.py
@@ -35,8 +35,10 @@
 dummy_input = torch.randn(1, 3, 1024, 2048)
 input_names = ['input_image']
 output_names = ['preds',]
+dynamic_axes = {'input_image': {0: 'batch'}, 'preds': {0: 'batch'}}
 
 torch.onnx.export(net, dummy_input, args.out_pth,
     input_names=input_names, output_names=output_names,
-    verbose=False, opset_version=11, )
+    verbose=False, opset_version=18,
+    dynamic_axes=dynamic_axes)