80 lines
3.9 KiB
Python
80 lines
3.9 KiB
Python
import argparse
|
|
import os
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
import cv2
|
|
import torch
|
|
import torch.backends.cudnn as cudnn
|
|
import torch.nn as nn
|
|
from collections import OrderedDict, namedtuple
|
|
import numpy as np
|
|
class DetectMultiBackend(nn.Module):
|
|
# YOLOv5 MultiBackend class for python inference on various backends
|
|
def __init__(self, weights='yolov5s.pt', device=None, dnn=False, data=None):
|
|
# Usage:
|
|
# TensorRT: *.engine
|
|
|
|
super().__init__()
|
|
w = str(weights[0] if isinstance(weights, list) else weights)
|
|
stride, names = 64, [f'class{i}' for i in range(1000)] # assign defaults
|
|
engine = True
|
|
if engine: # TensorRT
|
|
print('Loading {w} for TensorRT inference...')
|
|
import tensorrt as trt # https://developer.nvidia.com/nvidia-tensorrt-download
|
|
#check_version(trt.__version__, '7.0.0', hard=True) # require tensorrt>=7.0.0
|
|
print('TRT version:',trt.__version__ )
|
|
Binding = namedtuple('Binding', ('name', 'dtype', 'shape', 'data', 'ptr'))
|
|
print(trt.Logger.INFO)
|
|
logger = trt.Logger(trt.Logger.INFO)
|
|
assert os.path.exists(w)
|
|
with open(w, 'rb') as f, trt.Runtime(logger) as runtime:
|
|
model = runtime.deserialize_cuda_engine(f.read())
|
|
bindings = OrderedDict()
|
|
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))
|
|
data = torch.from_numpy(np.empty(shape, dtype=np.dtype(dtype))).to(device)
|
|
bindings[name] = Binding(name, dtype, shape, data, int(data.data_ptr()))
|
|
binding_addrs = OrderedDict((n, d.ptr) for n, d in bindings.items())
|
|
context = model.create_execution_context()
|
|
batch_size = bindings['images'].shape[0]
|
|
self.__dict__.update(locals()) # assign all variables to self
|
|
|
|
def forward(self, im, augment=False, visualize=False, val=False):
|
|
# YOLOv5 MultiBackend inference
|
|
b, ch, h, w = im.shape # batch, channel, height, width
|
|
if self.engine: # TensorRT
|
|
assert im.shape == self.bindings['images'].shape, (im.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
|
|
|
|
y = torch.tensor(y) if isinstance(y, np.ndarray) else y
|
|
return (y, []) if val else y
|
|
|
|
def warmup(self, imgsz=(1, 3, 640, 640), half=False):
|
|
# Warmup model by running inference once
|
|
if self.engine: # warmup types
|
|
if isinstance(self.device, torch.device) and self.device.type != 'cpu': # only warmup GPU models
|
|
im = torch.zeros(*imgsz).to(self.device).type(torch.half if half else torch.float) # input image
|
|
self.forward(im) # warmup
|
|
|
|
@staticmethod
|
|
def model_type(p='path/to/model.pt'):
|
|
# Return model type from model path, i.e. path='path/to/model.onnx' -> type=onnx
|
|
from export import export_formats
|
|
suffixes = list(export_formats().Suffix) + ['.xml'] # export suffixes
|
|
check_suffix(p, suffixes) # checks
|
|
p = Path(p).name # eliminate trailing separators
|
|
pt, jit, onnx, xml, engine, coreml, saved_model, pb, tflite, edgetpu, tfjs, xml2 = (s in p for s in suffixes)
|
|
xml |= xml2 # *_openvino_model or *.xml
|
|
tflite &= not edgetpu # *.tflite
|
|
return pt, jit, onnx, xml, engine, coreml, saved_model, pb, tflite, edgetpu, tfjs
|
|
if __name__=='__main__':
|
|
device=torch.device('cuda:1')
|
|
print('###line80:' ,type(device), device )
|
|
data = torch.from_numpy(np.empty( (1,3,512,512), dtype=np.float32 )).to(device)
|
|
|