极简yolov5转torchscript

极简yolov5转torchscript–以yolov5s为例

最近在使用yolov5转torchscript模型做推理时遇到了这样一个问题：
RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!
不是所有tensor都在同个设备上，但是检查了model的weights和input都是在一个设备上，就很疑惑，后来发现是yolov5推理时的detect
层的融框操作导致。

首先我们来看一般的torchscript模型转换操作

import torch
from models.experimental import attempt_load

weight = 'weights/yolov5s.pt'
jit_weight = weight.replace('.pt', '.jit')

x = torch.rand([1, 3, 640, 640])
device = torch.device('cuda:0')

torch_model = attempt_load(weight, map_location=device)
x = x.to(device)

torch.jit.trace(torch_model, x).save(jit_weight)

这种一般的转法，yolov5只能在你转模型时设置的device上运行（其他不复杂的模型不会出现这个问题），比如这次设置的是’cuda:0’，那模型就只能在0号gpu上运行，如果device=torch.device(‘cpu’)，那就只能在cpu上运行，之所以出现这种情况，是因为yolov5推理和训练时的输出不一样，详见models/yolo.py中的Detect类中的forword方法

class Detect(nn.Module):
    stride = None  # strides computed during build
    onnx_dynamic = False  # ONNX export parameter

    def __init__(self, nc=80, anchors=(), ch=(), inplace=True):  # detection layer
        super(Detect, self).__init__()

        self.nc = nc  # number of classes
        self.no = nc + 5  # number of outputs per anchor
        self.nl = len(anchors)  # number of detection layers
        self.na = len(anchors[0]) // 2  # number of anchors
        self.grid = [torch.zeros(1)] * self.nl  # init grid
        a = torch.tensor(anchors).float().view(self.nl, -1, 2)
        self.register_buffer('anchors', a)  # shape(nl,na,2)
        self.register_buffer('anchor_grid', a.clone().view(self.nl, 1, -1, 1, 1, 2))  # shape(nl,1,na,1,1,2)
        self.m = nn.ModuleList(nn.Conv2d(x, self.no * self.na, 1) for x in ch)  # output conv
        self.inplace = inplace  # use in-place ops (e.g. slice assignment)

    def forward(self, x):
        # x = x.copy()  # for profiling
        z = []  # inference output
        for i in range(self.nl):
            x[i] = self.m[i](x[i])  # conv
            bs, _, ny, nx = x[i].shape  # x(bs,255,20,20) to x(bs,3,20,20,85)
            x[i] = x[i].view(bs, self.na, self.no, ny, nx).permute(0, 1, 3, 4, 2).contiguous()

            if not self.training:  # inference
                if self.grid[i].shape[2:4] != x[i].shape[2:4] or self.onnx_dynamic:
                    self.grid[i] = self._make_grid(nx, ny).to(x[i].device)

                y = x[i].sigmoid()
                if self.inplace:
                # if True:
                    y[..., 0:2] = (y[..., 0:2] * 2. - 0.5 + self.grid[i]) * self.stride[i]  # xy
                    y[..., 2:4] = (y[..., 2:4] * 2) ** 2 * self.anchor_grid[i]  # wh
                else:  # for YOLOv5 on AWS Inferentia https://github.com/ultralytics/yolov5/pull/2953
                    xy = (y[..., 0:2] * 2. - 0.5 + self.grid[i]) * self.stride[i]  # xy
                    wh = (y[..., 2:4] * 2) ** 2 * self.anchor_grid[i].view(1, self.na, 1, 1, 2)  # wh
                    y = torch.cat((xy, wh, y[..., 4:]), -1)
                z.append(y.view(bs, -1, self.no))

        '''
        训练时输出为x，其为一个list，list中有三个tensor（这里以yolov5为例）
        x[0].shape = (bs,3,80,80,nc+5)
        x[1].shape = (bs,3,40,40,nc+5)
        x[2].shape = (bs,3,20,20,nc+5)
        bs为batch_size
        nc为类别数

        推理时输出为(torch.cat(z, 1), x)
        其中x与训练时的x相同
        torch.cat(z,1)为融合后的预测框
        '''
        return x if self.training else (torch.cat(z, 1), x)

    @staticmethod
    def _make_grid(nx=20, ny=20):
        yv, xv = torch.meshgrid([torch.arange(ny), torch.arange(nx)])
        return torch.stack((xv, yv), 2).view((1, 1, ny, nx, 2)).float()

forward方法中有个self.training变量，用来控制训练和推理时是否执行if not self.training:中的融框操作，
if not self.training:中的self.anchor_grid,self.grid不在模型结构中，所以在你转化时其所处的设备位置不会变化，
即你在使用torchscript模型时model.to(device)操作不会改变self.anchor_grid,self.grid的device，
它们的device仍然是模型转换时所设置的device

比如

import torch
from models.experimental import attempt_load

weight = 'weights/yolov5s.pt'
jit_weight = weight.replace('.pt', '.jit')

x = torch.rand([1, 3, 640, 640])
pre_device = torch.device('cuda:0')

torch_model = attempt_load(weight, map_location=pre_device)
x = x.to(pre_device)

torch.jit.trace(torch_model, x).save(jit_weight)

model = torch.jit.load(jit_weights)

device = torch.device('cuda:0')

model = model.to(device)

此时，model中的参数在device='cpu’中，而self.anchor_grid和self.grid都在pre_device='cuda:0’中，所以会报错：RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!

解决方法

在转换模型前将self.training设置为True

self.training = True
if not self.training:  # inference
    if self.grid[i].shape[2:4] != x[i].shape[2:4] or self.onnx_dynamic:
        self.grid[i] = self._make_grid(nx, ny).to(x[i].device)

然后用numpy重写推理时的融合大中小框的操作,这里仍然以yolov5s为例，anchor_grid为先验框，在models/yolov5s.yaml，
这里的80,40,20为最后大中小不同物体的检验格子的大小，这里以输入图片的尺度为6406403为例，不同的输入尺度可能不同

def numpy_detect(x,nc=None):
    # anchor_grid为先验眶
    anchor_grid = [10.0, 13.0, 16.0, 30.0, 33.0, 23.0, 30.0, 61.0, 62.0, 45.0, 59.0, 119.0, 116.0, 90.0, 156.0, 198.0, 373.0, 326.0]
    anchor_grid = np.array(anchor_grid).reshape(3,1,-1,1,1,2)
    stride = np.array([8, 16, 32])
    grid = [make_grid(80,80), make_grid(40,40), make_grid(20,20)]
    
    z = []
    
    for i in range(3):
        y = numpy_sigmoid(x[i])

        y[..., 0:2] = (y[..., 0:2] * 2. - 0.5 + grid[i]) * stride[i]  # xy
        y[..., 2:4] = (y[..., 2:4] * 2) ** 2 * anchor_grid[i]  # wh

        z.append(y.reshape(1, -1, nc + 5))
    
    res = np.concatenate(z, 1)

    return res

def numpy_sigmoid(x):
    return 1/(1+np.exp(-x))

def make_grid(nx=20,ny=20):
    xv,yv = np.meshgrid(np.arange(nx), np.arange(ny))
    res = np.stack((xv,yv), 2).reshape(1,1,nx,ny,2).astype(np.float32)
    return res

def img_pad(img, size, pad_value=[114,114,114]):
    H,W,_ = img.shape
    r = max(H/size[0], W/size[1])
    img_r = cv2.resize(img, (int(W/r), int(H/r)))
    tb = size[0] - img_r.shape[0]
    lr = size[1] - img_r.shape[1]
    top = int(tb/2)
    bottom = tb - top
    left = int(lr/2)
    right = lr - left
    pad_image = cv2.copyMakeBorder(img_r, top, bottom, left, right, cv2.BORDER_CONSTANT,value=pad_value)
    return pad_image

实例如下:

jit_weights = 'weights/yolov5s.jit'
model = torch.jit.load(jit_weights)

imgpath = 'test.jpg
im0 = cv2.imread(imgpath)

img = img_pad(im0, size=self.input_hw)[0]
img = img[:, :, ::-1].transpose(2, 0, 1)
img = np.ascontiguousarray(img, dtype=np.float32)

img = torch.from_numpy(img).to(self.device)
img /= 255.0
if img.ndimension() == 3:
    img = img.unsqueeze(0)

pred = model(img)

pred = [x.data.cpu().numpy() for x in pred]
pred = numpy_detect(pred, self.nc)
pred = torch.tensor(pred).to(self.device)

pred = non_max_suppression(pred, self.conf_thres, self.nms_thres)