pesi
/

Luigi commited on
Commit
09ccc6e
1 Parent(s): 92676db

Use dynamic batch size by default

Browse files
Files changed (1) hide show
  1. onnx_to_engine.py +18 -8
onnx_to_engine.py CHANGED
@@ -91,14 +91,24 @@ def main(onnx_path, engine_path, batch_size):
91
  # You don't have to use it with Polygraphy loaders if you don't want to.
92
  calibrator = Calibrator(data_loader=calib_data_from_video(), cache=f"{onnx_path}-calib.cache")
93
 
 
94
 
95
- profiles = [
96
- # The low-latency case. For best performance, min == opt == max.
97
- Profile().add("input",
98
- min=(1, 3, MODEL_INPUT_SIZE[0], MODEL_INPUT_SIZE[1]),
99
- opt=(4, 3, MODEL_INPUT_SIZE[0], MODEL_INPUT_SIZE[1]),
100
- max=(9, 3, MODEL_INPUT_SIZE[0], MODEL_INPUT_SIZE[1])),
101
- ]
 
 
 
 
 
 
 
 
 
102
 
103
  # We must enable int8 mode in addition to providing the calibrator.
104
  build_engine = EngineFromNetwork(
@@ -126,7 +136,7 @@ if __name__ == "__main__":
126
  parser.add_argument("video_path", type=str, help="The path to the video file used to calibrate int8 engine")
127
  parser.add_argument("onnx_path", type=str, help="The path to the input ONNX model file")
128
  parser.add_argument("engine_path", type=str, help="The path to the exported TensorRT Engine model file")
129
- parser.add_argument("--batch_size", type=int, default=1, help="Input batch size")
130
  args = parser.parse_args()
131
  VIDEO_PATH = args.video_path
132
  MODEL_INPUT_SIZE=(416,416) if 'rtmo-t' in args.onnx_path else (640,640)
 
91
  # You don't have to use it with Polygraphy loaders if you don't want to.
92
  calibrator = Calibrator(data_loader=calib_data_from_video(), cache=f"{onnx_path}-calib.cache")
93
 
94
+ if batch_size < 1: # dynamic batch size
95
 
96
+ profiles = [
97
+ # The low-latency case. For best performance, min == opt == max.
98
+ Profile().add("input",
99
+ min=(1, 3, MODEL_INPUT_SIZE[0], MODEL_INPUT_SIZE[1]),
100
+ opt=(4, 3, MODEL_INPUT_SIZE[0], MODEL_INPUT_SIZE[1]),
101
+ max=(9, 3, MODEL_INPUT_SIZE[0], MODEL_INPUT_SIZE[1])),
102
+ ]
103
+
104
+ else: # fixed
105
+ profiles = [
106
+ # The low-latency case. For best performance, min == opt == max.
107
+ Profile().add("input",
108
+ min=(batch_size, 3, MODEL_INPUT_SIZE[0], MODEL_INPUT_SIZE[1]),
109
+ opt=(batch_size, 3, MODEL_INPUT_SIZE[0], MODEL_INPUT_SIZE[1]),
110
+ max=(batch_size, 3, MODEL_INPUT_SIZE[0], MODEL_INPUT_SIZE[1])),
111
+ ]
112
 
113
  # We must enable int8 mode in addition to providing the calibrator.
114
  build_engine = EngineFromNetwork(
 
136
  parser.add_argument("video_path", type=str, help="The path to the video file used to calibrate int8 engine")
137
  parser.add_argument("onnx_path", type=str, help="The path to the input ONNX model file")
138
  parser.add_argument("engine_path", type=str, help="The path to the exported TensorRT Engine model file")
139
+ parser.add_argument("--batch_size", type=int, default=-1, help="Input batch size (not specified if dynamic)")
140
  args = parser.parse_args()
141
  VIDEO_PATH = args.video_path
142
  MODEL_INPUT_SIZE=(416,416) if 'rtmo-t' in args.onnx_path else (640,640)