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
2
3
4
5
6
7
8
9
10
11
12
13
14
15

2. 将权重导出成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_na
1
2
3
4
5
6
7
8
9
10
11