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)

2. 将权重导出成onnx时使用修改后的网络，下面创建的YoloBody网络为替换了上采样的网络

import numpy as np
from PIL import Image
import to