Spaces:
Runtime error
Runtime error
| from typing import Tuple, Dict,List | |
| import tensorrt as trt | |
| import numpy as np | |
| import torch | |
| import time | |
| import os | |
| from collections import OrderedDict, namedtuple | |
| class TRT_Base(): | |
| def __init__(self, | |
| input_shape: Tuple[int, int, int], | |
| model_path: str, | |
| device: str='0'): | |
| """ Tensor RT base class for inference. | |
| Args: | |
| input_shape (Tuple[int, int, int]): image size (3, H, W) | |
| model_path (str): path to the model.trt | |
| device (str, optional): CUDA device. Defaults to '0'. | |
| """ | |
| self.input_shape = input_shape | |
| self.model_path = model_path | |
| self.device = self.select_device(device) | |
| self.init_model() | |
| def select_device(self, device: str)->torch.device: | |
| """ Select device to be used for inference. | |
| Args: | |
| param device: 'cpu' or '0' or '0,1,2,3' | |
| Return: | |
| torch.device | |
| """ | |
| cpu = device.lower() == "cpu" | |
| if cpu: | |
| os.environ['CUDA_VISIBLE_DEVICES'] = '-1' | |
| return torch.device("cpu") | |
| else: | |
| assert torch.cuda.is_available(), f'CUDA unavailable, invalid device {device} requested' | |
| torch.cuda.set_device(int(device)) | |
| return torch.device("cuda") | |
| def init_model(self): | |
| """ Initialize TensorRT engine and context.""" | |
| logger = trt.Logger(trt.Logger.INFO) | |
| trt.init_libnvinfer_plugins(logger, namespace="") | |
| with open(self.model_path, 'rb') as f, trt.Runtime(logger) as runtime: | |
| engine = runtime.deserialize_cuda_engine(f.read()) | |
| context = engine.create_execution_context() | |
| self.model = { | |
| "engine": engine, | |
| "context": context | |
| } | |
| bindings, binding_addrs = self.get_bindings(input_shape=self.input_shape) | |
| input_names = [binding_name for binding_name in binding_addrs.keys() if (self.model["engine"].binding_is_input(binding_name))] | |
| for _ in range(10): | |
| for name in input_names: | |
| binding_addrs[name] = int(torch.randn(bindings[name].shape).to(self.device).data_ptr()) | |
| context.execute_v2(list(binding_addrs.values())) | |
| self.model.update({ | |
| 'binding_addrs': binding_addrs, | |
| 'bindings': bindings, | |
| 'rt_shapes': self.input_shape | |
| }) | |
| def get_bindings(self, input_shape: Tuple[int, int, int]): | |
| """ Get bindings and binding addresses for TensorRT engine. | |
| Args: | |
| input_shape (Tuple[int, int, int]): image size (3, H, W) | |
| """ | |
| self.model["context"].set_binding_shape(0, input_shape) | |
| bindings = OrderedDict() | |
| Binding = namedtuple('Binding', ('name', 'dtype', 'shape', 'data', 'ptr')) | |
| for index in range(self.model["engine"].num_bindings): | |
| name = self.model["engine"].get_binding_name(index) | |
| dtype = trt.nptype(self.model["engine"].get_binding_dtype(index)) | |
| shape = tuple(self.model["context"].get_binding_shape(index)) | |
| data = torch.from_numpy(np.empty(shape, dtype=np.dtype(dtype))).to(self.device) | |
| bindings[name] = Binding(name, dtype, shape, data, int(data.data_ptr())) | |
| binding_addrs = OrderedDict((n, d.ptr) for n, d in bindings.items()) | |
| return bindings, binding_addrs | |
| def change_runtime_dimension(self, input_shape: Tuple[int, int, int]): | |
| """ Support inference with Dynamic shape. | |
| Args: | |
| input_shape (Tuple[int, int, int]): image size (3, H, W) | |
| """ | |
| if (input_shape == self.model["rt_shapes"]): return | |
| bindings, binding_addrs = self.get_bindings(input_shape) | |
| self.model['binding_addrs'] = binding_addrs | |
| self.model['bindings'] = bindings | |
| self.model['rt_shapes'] = input_shape | |