YOLOV4ToTensorRT引擎 python实现记录

YOLOV4ToTensorRT引擎 python实现记录

问题:yolov4权重导出成onnx,生成tensorRT引擎失败
原因:yolov4中存网络中使用临近差值上采样在tensorRT中不支持。
解决方案:

  1. 修改yolov4网络模型中的上采样,使用下面的代码代替网络中的torch.nn.Upsample
import torch
import torch.nn as nn
import torch.nn.functional as F

class Upsample(nn.Module):
    def __init__(self, size, scale, mode, align_corners=None):
        super(Upsample, self).__init__()
        self.size = size
        self.scale = scale
        self.mode = mode
        self.align_corners = align_corners

    def forward(self, x):
        sh = torch.tensor(x.shape)
        return F.interpolate(x, size=(int(sh[2]*self.scale), int(sh[3]*self.scale)), mode=self.mode, align_corners=self.align_corners)
  1. 将权重导出成onnx时使用修改后的网络,下面创建的YoloBody网络为替换了上采样的网络
import numpy as np
from PIL import Image
import torch.onnx
from nets.yolo4_tiny import YoloBody
from utils.utils import letterbox_image

def get_classes(classes_path):
    '''loads the classes'''
    with open(classes_path) as f:
        class_names = f.readlines()
    class_names = [c.strip() for c in class_names]
    return class_names

def get_anchors(anchors_path):
    '''loads the anchors from a file'''
    with open(anchors_path) as f:
        anchors = f.readline()
    anchors = [float(x) for x in anchors.split(',')]
    return np.array(anchors).reshape([-1, 3, 2])


anchors_path = 'model_data/yolo_anchors.txt'
classes_path = 'model_data/voc_classes.txt'
model_image_size = torch.tensor([416, 416, 3])
class_names = get_classes(classes_path)
anchors = get_anchors(anchors_path)
num_classes = len(class_names)

# 创建模型
model = YoloBody(len(anchors[0]), num_classes).eval()

# # 加载权重
model_path = './Epoch1-Total_Loss3.9009-Val_Loss5.2476.pth'
state_dict = torch.load(model_path)
model.load_state_dict(state_dict)

# # 打开图片
img = Image.open(r"F:\alldata/380.jpg")
# # 图片预处理
image_shape = np.array(np.shape(img)[0:2])
crop_img = np.array(letterbox_image(img, (model_image_size[0], model_image_size[1])))
photo = np.array(crop_img, dtype=np.float32)
photo /= 255.0
photo = np.transpose(photo, (2, 0, 1))
photo = photo.astype(np.float32)
images = []
images.append(photo)
images = np.asarray(images)
images = torch.from_numpy(images)
images = images
outputs = model(images)

torch.onnx.export(model,  # model being run
                  images,  # model input (or a tuple for multiple inputs)
                  "yoloV4_tiny.onnx",  # where to save the model (can be a file or file-like object)
                  export_params=True,  # store the trained parameter weights inside the model file
                  opset_version=11,  # the ONNX version to export the model to
                  do_constant_folding=True,  # whether to execute constant folding for optimization
                  input_names=['input'],  # the model's input names
                  output_names=['output'],  # the model's output names
                  dynamic_axes={'input': {0: 'batch_size'},  # variable lenght axes
                                'output': {0: 'batch_size'}}
                  )

如果导出onnx时没有创建网络,直接加载的权重文件,则需要在训练时,使用替换了上采样的网络训练生成的权重文件导出成onnx。如果是yolov5的话就替换配置文件eg:yolov5s.yaml文件中的nn.Upsample第一步中的上采样,将带一部上采样代码放到统计目录中的common.py中即可。
在这里插入图片描述

  1. 接下来使用以下代码将onnx文件转换成tensorRT引擎文件(tensorRT使用的7.1.3.0版本)
import os
import tensorrt as trt

EXPLICIT_BATCH = []
if trt.__version__[0] >= '7':
    EXPLICIT_BATCH.append(
        1 << (int)(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))


def build_engine(onnx_file_path, engine_file_path, verbose=False):
    """Takes an ONNX file and creates a TensorRT engine."""
    TRT_LOGGER = trt.Logger(trt.Logger.VERBOSE) if verbose else trt.Logger()
    with trt.Builder(TRT_LOGGER) as builder, builder.create_network(*EXPLICIT_BATCH) as network, trt.OnnxParser(network,
                                                                                                                TRT_LOGGER) as parser:
        builder.max_workspace_size = 1 << 27  # 25#33554432#32M
        builder.max_batch_size = 4
        builder.fp16_mode = True
        # builder.strict_type_constraints = True

        # Parse model file
        print('Loading ONNX file from path {}...'.format(onnx_file_path))
        with open(onnx_file_path, 'rb') as model:
            print('Beginning ONNX file parsing')
            if not parser.parse(model.read()):
                print('ERROR: Failed to parse the ONNX file.')
                for error in range(parser.num_errors):
                    print(parser.get_error(error))
                return None
        if trt.__version__[0] >= '7':
            # The actual yolo*.onnx is generated with batch size 64.
            # Reshape input to batch size 1
            shape = list(network.get_input(0).shape)
            shape[0] = 1
            network.get_input(0).shape = shape
        print('Completed parsing of ONNX file')

        print('Building an engine; this may take a while...')
        engine = builder.build_cuda_engine(network)
        print('Completed creating engine')
        with open(engine_file_path, 'wb') as f:
            f.write(engine.serialize())
        return engine


def main():
    onnx_file_path = 'yolov4-tiny.onnx'
    if not os.path.isfile(onnx_file_path):
        raise SystemExit(
            'ERROR: file (%s) not found!  You might want to run yolo_to_onnx.py first to generate it.' % onnx_file_path)
    engine_file_path = 'yolov4-tiny.engine'
    _ = build_engine(onnx_file_path, engine_file_path, False)


if __name__ == '__main__':
    main()

  1. 后面就是用引擎文件网络的侦测过程。
©️2020 CSDN 皮肤主题: 大白 设计师:CSDN官方博客 返回首页