| import warnings
|
| from collections import namedtuple
|
| from functools import partial
|
| from pathlib import Path
|
| from typing import List, Optional, Union
|
|
|
| import numpy as np
|
| import onnxruntime
|
|
|
| try:
|
| import tensorrt as trt
|
| except Exception:
|
| trt = None
|
| import torch
|
|
|
| warnings.filterwarnings(action='ignore', category=DeprecationWarning)
|
|
|
|
|
| class TRTWrapper(torch.nn.Module):
|
| dtype_mapping = {}
|
|
|
| def __init__(self, weight: Union[str, Path],
|
| device: Optional[torch.device]):
|
| super().__init__()
|
| weight = Path(weight) if isinstance(weight, str) else weight
|
| assert weight.exists() and weight.suffix in ('.engine', '.plan')
|
| if isinstance(device, str):
|
| device = torch.device(device)
|
| elif isinstance(device, int):
|
| device = torch.device(f'cuda:{device}')
|
| self.weight = weight
|
| self.device = device
|
| self.stream = torch.cuda.Stream(device=device)
|
| self.__update_mapping()
|
| self.__init_engine()
|
| self.__init_bindings()
|
|
|
| def __update_mapping(self):
|
| self.dtype_mapping.update({
|
| trt.bool: torch.bool,
|
| trt.int8: torch.int8,
|
| trt.int32: torch.int32,
|
| trt.float16: torch.float16,
|
| trt.float32: torch.float32
|
| })
|
|
|
| def __init_engine(self):
|
| logger = trt.Logger(trt.Logger.ERROR)
|
| self.log = partial(logger.log, trt.Logger.ERROR)
|
| trt.init_libnvinfer_plugins(logger, namespace='')
|
| self.logger = logger
|
| with trt.Runtime(logger) as runtime:
|
| model = runtime.deserialize_cuda_engine(self.weight.read_bytes())
|
|
|
| context = model.create_execution_context()
|
|
|
| names = [model.get_binding_name(i) for i in range(model.num_bindings)]
|
|
|
| num_inputs, num_outputs = 0, 0
|
|
|
| for i in range(model.num_bindings):
|
| if model.binding_is_input(i):
|
| num_inputs += 1
|
| else:
|
| num_outputs += 1
|
|
|
| self.is_dynamic = -1 in model.get_binding_shape(0)
|
|
|
| self.model = model
|
| self.context = context
|
| self.input_names = names[:num_inputs]
|
| self.output_names = names[num_inputs:]
|
| self.num_inputs = num_inputs
|
| self.num_outputs = num_outputs
|
| self.num_bindings = num_inputs + num_outputs
|
| self.bindings: List[int] = [0] * self.num_bindings
|
|
|
| def __init_bindings(self):
|
| Binding = namedtuple('Binding', ('name', 'dtype', 'shape'))
|
| inputs_info = []
|
| outputs_info = []
|
|
|
| for i, name in enumerate(self.input_names):
|
| assert self.model.get_binding_name(i) == name
|
| dtype = self.dtype_mapping[self.model.get_binding_dtype(i)]
|
| shape = tuple(self.model.get_binding_shape(i))
|
| inputs_info.append(Binding(name, dtype, shape))
|
|
|
| for i, name in enumerate(self.output_names):
|
| i += self.num_inputs
|
| assert self.model.get_binding_name(i) == name
|
| dtype = self.dtype_mapping[self.model.get_binding_dtype(i)]
|
| shape = tuple(self.model.get_binding_shape(i))
|
| outputs_info.append(Binding(name, dtype, shape))
|
| self.inputs_info = inputs_info
|
| self.outputs_info = outputs_info
|
| if not self.is_dynamic:
|
| self.output_tensor = [
|
| torch.empty(o.shape, dtype=o.dtype, device=self.device)
|
| for o in outputs_info
|
| ]
|
|
|
| def forward(self, *inputs):
|
|
|
| assert len(inputs) == self.num_inputs
|
|
|
| contiguous_inputs: List[torch.Tensor] = [
|
| i.contiguous() for i in inputs
|
| ]
|
|
|
| for i in range(self.num_inputs):
|
| self.bindings[i] = contiguous_inputs[i].data_ptr()
|
| if self.is_dynamic:
|
| self.context.set_binding_shape(
|
| i, tuple(contiguous_inputs[i].shape))
|
|
|
|
|
| outputs: List[torch.Tensor] = []
|
|
|
| for i in range(self.num_outputs):
|
| j = i + self.num_inputs
|
| if self.is_dynamic:
|
| shape = tuple(self.context.get_binding_shape(j))
|
| output = torch.empty(
|
| size=shape,
|
| dtype=self.output_dtypes[i],
|
| device=self.device)
|
|
|
| else:
|
| output = self.output_tensor[i]
|
| outputs.append(output)
|
| self.bindings[j] = output.data_ptr()
|
|
|
| self.context.execute_async_v2(self.bindings, self.stream.cuda_stream)
|
| self.stream.synchronize()
|
|
|
| return tuple(outputs)
|
|
|
|
|
| class ORTWrapper(torch.nn.Module):
|
|
|
| def __init__(self, weight: Union[str, Path],
|
| device: Optional[torch.device]):
|
| super().__init__()
|
| weight = Path(weight) if isinstance(weight, str) else weight
|
| assert weight.exists() and weight.suffix == '.onnx'
|
|
|
| if isinstance(device, str):
|
| device = torch.device(device)
|
| elif isinstance(device, int):
|
| device = torch.device(f'cuda:{device}')
|
| self.weight = weight
|
| self.device = device
|
| self.__init_session()
|
| self.__init_bindings()
|
|
|
| def __init_session(self):
|
| providers = ['CPUExecutionProvider']
|
| if 'cuda' in self.device.type:
|
| providers.insert(0, 'CUDAExecutionProvider')
|
|
|
| session = onnxruntime.InferenceSession(
|
| str(self.weight), providers=providers)
|
| self.session = session
|
|
|
| def __init_bindings(self):
|
| Binding = namedtuple('Binding', ('name', 'dtype', 'shape'))
|
| inputs_info = []
|
| outputs_info = []
|
| self.is_dynamic = False
|
| for i, tensor in enumerate(self.session.get_inputs()):
|
| if any(not isinstance(i, int) for i in tensor.shape):
|
| self.is_dynamic = True
|
| inputs_info.append(
|
| Binding(tensor.name, tensor.type, tuple(tensor.shape)))
|
|
|
| for i, tensor in enumerate(self.session.get_outputs()):
|
| outputs_info.append(
|
| Binding(tensor.name, tensor.type, tuple(tensor.shape)))
|
| self.inputs_info = inputs_info
|
| self.outputs_info = outputs_info
|
| self.num_inputs = len(inputs_info)
|
|
|
| def forward(self, *inputs):
|
|
|
| assert len(inputs) == self.num_inputs
|
|
|
| contiguous_inputs: List[np.ndarray] = [
|
| i.contiguous().cpu().numpy() for i in inputs
|
| ]
|
|
|
| if not self.is_dynamic:
|
|
|
| for i in range(self.num_inputs):
|
| assert contiguous_inputs[i].shape == self.inputs_info[i].shape
|
|
|
| outputs = self.session.run([o.name for o in self.outputs_info], {
|
| j.name: contiguous_inputs[i]
|
| for i, j in enumerate(self.inputs_info)
|
| })
|
|
|
| return tuple(torch.from_numpy(o).to(self.device) for o in outputs)
|
|
|
