Use dynamic batch size by default
Browse files- 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 |
-
|
96 |
-
|
97 |
-
|
98 |
-
|
99 |
-
|
100 |
-
|
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
|
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)
|