From 587a3a37c57661c3a0ef710d2b309199fad632d2 Mon Sep 17 00:00:00 2001 From: Colin Wong Date: Fri, 29 Jul 2022 06:51:16 -0500 Subject: [PATCH] Dynamic batch size support for TensorRT (#8526) * Dynamic batch size support for TensorRT * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Update export.py * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fix optimization profile when batch size is 1 * Warn users if they use batch-size=1 with dynamic * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * More descriptive assertion error * Fix syntax * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * pre-commit formatting sucked * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Update export.py Co-authored-by: Colin Wong Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Glenn Jocher --- export.py | 21 +++++++++++++++------ models/common.py | 22 ++++++++++++++++------ 2 files changed, 31 insertions(+), 12 deletions(-) diff --git a/export.py b/export.py index 3629915..4846624 100644 --- a/export.py +++ b/export.py @@ -216,8 +216,9 @@ def export_coreml(model, im, file, int8, half, prefix=colorstr('CoreML:')): return None, None -def export_engine(model, im, file, train, half, simplify, workspace=4, verbose=False, prefix=colorstr('TensorRT:')): +def export_engine(model, im, file, train, half, dynamic, simplify, workspace=4, verbose=False): # YOLOv5 TensorRT export https://developer.nvidia.com/tensorrt + prefix = colorstr('TensorRT:') try: assert im.device.type != 'cpu', 'export running on CPU but must be on GPU, i.e. `python export.py --device 0`' try: @@ -230,11 +231,11 @@ def export_engine(model, im, file, train, half, simplify, workspace=4, verbose=F if trt.__version__[0] == '7': # TensorRT 7 handling https://github.com/ultralytics/yolov5/issues/6012 grid = model.model[-1].anchor_grid model.model[-1].anchor_grid = [a[..., :1, :1, :] for a in grid] - export_onnx(model, im, file, 12, train, False, simplify) # opset 12 + export_onnx(model, im, file, 12, train, dynamic, simplify) # opset 12 model.model[-1].anchor_grid = grid else: # TensorRT >= 8 check_version(trt.__version__, '8.0.0', hard=True) # require tensorrt>=8.0.0 - export_onnx(model, im, file, 13, train, False, simplify) # opset 13 + export_onnx(model, im, file, 13, train, dynamic, simplify) # opset 13 onnx = file.with_suffix('.onnx') LOGGER.info(f'\n{prefix} starting export with TensorRT {trt.__version__}...') @@ -263,6 +264,14 @@ def export_engine(model, im, file, train, half, simplify, workspace=4, verbose=F for out in outputs: LOGGER.info(f'{prefix}\toutput "{out.name}" with shape {out.shape} and dtype {out.dtype}') + if dynamic: + if im.shape[0] <= 1: + LOGGER.warning(f"{prefix}WARNING: --dynamic model requires maximum --batch-size argument") + profile = builder.create_optimization_profile() + for inp in inputs: + profile.set_shape(inp.name, (1, *im.shape[1:]), (max(1, im.shape[0] // 2), *im.shape[1:]), im.shape) + config.add_optimization_profile(profile) + LOGGER.info(f'{prefix} building FP{16 if builder.platform_has_fast_fp16 and half else 32} engine in {f}') if builder.platform_has_fast_fp16 and half: config.set_flag(trt.BuilderFlag.FP16) @@ -460,7 +469,7 @@ def run( keras=False, # use Keras optimize=False, # TorchScript: optimize for mobile int8=False, # CoreML/TF INT8 quantization - dynamic=False, # ONNX/TF: dynamic axes + dynamic=False, # ONNX/TF/TensorRT: dynamic axes simplify=False, # ONNX: simplify model opset=12, # ONNX: opset version verbose=False, # TensorRT: verbose log @@ -520,7 +529,7 @@ def run( if jit: f[0] = export_torchscript(model, im, file, optimize) if engine: # TensorRT required before ONNX - f[1] = export_engine(model, im, file, train, half, simplify, workspace, verbose) + f[1] = export_engine(model, im, file, train, half, dynamic, simplify, workspace, verbose) if onnx or xml: # OpenVINO requires ONNX f[2] = export_onnx(model, im, file, opset, train, dynamic, simplify) if xml: # OpenVINO @@ -579,7 +588,7 @@ def parse_opt(): parser.add_argument('--keras', action='store_true', help='TF: use Keras') parser.add_argument('--optimize', action='store_true', help='TorchScript: optimize for mobile') parser.add_argument('--int8', action='store_true', help='CoreML/TF INT8 quantization') - parser.add_argument('--dynamic', action='store_true', help='ONNX/TF: dynamic axes') + parser.add_argument('--dynamic', action='store_true', help='ONNX/TF/TensorRT: dynamic axes') parser.add_argument('--simplify', action='store_true', help='ONNX: simplify model') parser.add_argument('--opset', type=int, default=12, help='ONNX: opset version') parser.add_argument('--verbose', action='store_true', help='TensorRT: verbose log') diff --git a/models/common.py b/models/common.py index 5ea1c30..959c965 100644 --- a/models/common.py +++ b/models/common.py @@ -384,19 +384,24 @@ class DetectMultiBackend(nn.Module): logger = trt.Logger(trt.Logger.INFO) with open(w, 'rb') as f, trt.Runtime(logger) as runtime: model = runtime.deserialize_cuda_engine(f.read()) + context = model.create_execution_context() bindings = OrderedDict() fp16 = False # default updated below + dynamic_input = False for index in range(model.num_bindings): name = model.get_binding_name(index) dtype = trt.nptype(model.get_binding_dtype(index)) - shape = tuple(model.get_binding_shape(index)) + if model.binding_is_input(index): + if -1 in tuple(model.get_binding_shape(index)): # dynamic + dynamic_input = True + context.set_binding_shape(index, tuple(model.get_profile_shape(0, index)[2])) + if dtype == np.float16: + fp16 = True + shape = tuple(context.get_binding_shape(index)) data = torch.from_numpy(np.empty(shape, dtype=np.dtype(dtype))).to(device) bindings[name] = Binding(name, dtype, shape, data, int(data.data_ptr())) - if model.binding_is_input(index) and dtype == np.float16: - fp16 = True binding_addrs = OrderedDict((n, d.ptr) for n, d in bindings.items()) - context = model.create_execution_context() - batch_size = bindings['images'].shape[0] + batch_size = bindings['images'].shape[0] # if dynamic, this is instead max batch size elif coreml: # CoreML LOGGER.info(f'Loading {w} for CoreML inference...') import coremltools as ct @@ -466,7 +471,12 @@ class DetectMultiBackend(nn.Module): im = im.cpu().numpy() # FP32 y = self.executable_network([im])[self.output_layer] elif self.engine: # TensorRT - assert im.shape == self.bindings['images'].shape, (im.shape, self.bindings['images'].shape) + if im.shape != self.bindings['images'].shape and self.dynamic_input: + self.context.set_binding_shape(self.model.get_binding_index('images'), im.shape) # reshape if dynamic + self.bindings['images'] = self.bindings['images']._replace(shape=im.shape) + assert im.shape == self.bindings['images'].shape, ( + f"image shape {im.shape} exceeds model max shape {self.bindings['images'].shape}" if self.dynamic_input + else f"image shape {im.shape} does not match model shape {self.bindings['images'].shape}") self.binding_addrs['images'] = int(im.data_ptr()) self.context.execute_v2(list(self.binding_addrs.values())) y = self.bindings['output'].data