open-mmlab/mmyolo

tensorrt test is wrong

Winnie202 opened this issue · 0 comments

Prerequisite

🐞 Describe the bug

run projects/easydeploy/model/backendwrapper.py when set device=0 ,the result is right, but set device=1 ,get [TRT] [E] 1: [shapeHostToDeviceRunner.cpp::execute::48] Error Code 1: Cuda Runtime (invalid resource handle)

Environment

2024-04-24 11-59-37屏幕截图

Additional information

import warnings
from collections import namedtuple
from functools import partial
from pathlib import Path
from typing import List, Optional, Union

import tensorrt as trt
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', '.trt')
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.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


def forward(self, *inputs):
    assert len(inputs) == self.num_inputs
    inputs = [inp.to(self.device) for inp in 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()
        self.context.set_binding_shape(i, tuple(contiguous_inputs[i].shape))

    # create output tensors
    outputs: List[torch.Tensor] = []

    for i in range(self.num_outputs):
        j = i + self.num_inputs
        shape = tuple(self.context.get_binding_shape(j))
        output = torch.empty(size=shape,dtype=self.outputs_info[i].dtype,device=self.device)
        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)

device = 1
model = TRTWrapper('tensorrt/model_new.trt',device)
inputs =torch.randn(2, 3, 640,640)
outputs = model(inputs)
print(outputs[0])