极简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_loadweight = '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 buildonnx_dynamic = False  # ONNX export parameterdef __init__(self, nc=80, anchors=(), ch=(), inplace=True):  # detection layersuper(Detect, self).__init__()self.nc = nc  # number of classesself.no = nc + 5  # number of outputs per anchorself.nl = len(anchors)  # number of detection layersself.na = len(anchors[0]) // 2  # number of anchorsself.grid = [torch.zeros(1)] * self.nl  # init grida = 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 convself.inplace = inplace  # use in-place ops (e.g. slice assignment)def forward(self, x):# x = x.copy()  # for profilingz = []  # inference outputfor i in range(self.nl):x[i] = self.m[i](x[i])  # convbs, _, 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:  # inferenceif 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]  # xyy[..., 2:4] = (y[..., 2:4] * 2) ** 2 * self.anchor_grid[i]  # whelse:  # for YOLOv5 on AWS Inferentia https://github.com/ultralytics/yolov5/pull/2953xy = (y[..., 0:2] * 2. - 0.5 + self.grid[i]) * self.stride[i]  # xywh = (y[..., 2:4] * 2) ** 2 * self.anchor_grid[i].view(1, self.na, 1, 1, 2)  # why = 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_sizenc为类别数推理时输出为(torch.cat(z, 1), x)其中x与训练时的x相同torch.cat(z,1)为融合后的预测框'''return x if self.training else (torch.cat(z, 1), x)@staticmethoddef _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_loadweight = '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:  # inferenceif 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]  # xyy[..., 2:4] = (y[..., 2:4] * 2) ** 2 * anchor_grid[i]  # whz.append(y.reshape(1, -1, nc + 5))res = np.concatenate(z, 1)return resdef 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 resdef img_pad(img, size, pad_value=[114,114,114]):H,W,_ = img.shaper = 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 - topleft = int(lr/2)right = lr - leftpad_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)

极简yolov5转torchscript相关推荐

RepVGG：极简架构，SOTA性能，论文解读
** RepVGG:极简架构,SOTA性能,论文解读 ** 更新:RepVGG的更深版本达到了83.55%正确率!PyTorch代码和模型已经在GitHub上放出.DingXiaoH/RepVGG 2 ...
业务逻辑组件化android,AppJoint 极简 Android 组件化方案
AppJoint 极简 Android 组件化方案.仅包含 3 个注解加 1 个 API,超低学习成本,支持渐进式组件化. 开始接入在项目根目录的 build.gradle 文件中添加 AppJoi ...
Spring Boot 极简集成 Shiro
点击关注公众号,Java干货及时送达 1. 前言 Apache Shiro是一个功能强大且易于使用的Java安全框架,提供了认证,授权,加密,和会话管理. Shiro有三大核心组件: Subject: ...
7句话让Codex给我做了个小游戏，还是极简版塞尔达，一玩简直停不下来
点击上方"视学算法",选择加"星标"或"置顶" 重磅干货,第一时间送达梦晨萧箫发自凹非寺量子位 | 公众号 QbitAI 什么,7 ...
30个Python常用极简代码，拿走就用
点击上方"视学算法",选择加"星标"或"置顶" 重磅干货,第一时间送达作者丨Fatos Morina 来源丨Python 技术编辑丨极市 ...
30 段极简 Python 代码：这些小技巧你都 Get 了么？
选自 | towardsdatascienc 编译 | 机器之心学 Python 怎样才最快,当然是实战各种小项目,只有自己去想与写,才记得住规则.本文是 30 个极简任务,初学者可以尝试着自己实现 ...
《Kotlin极简教程》第三章 Kotlin基本数据类型
正式上架:<Kotlin极简教程>Official on shelves: Kotlin Programming minimalist tutorial 京东JD:https://item ...
10分钟手撸极简版ORM框架！
最近很多小伙伴对ORM框架的实现很感兴趣,不少读者在冰河的微信上问:冰河,你知道ORM框架是如何实现的吗?比如像MyBatis和Hibernte这种ORM框架,它们是如何实现的呢? 为了能够让小伙伴们 ...
负载分析及问题排查极简教程
作者 | Hollis ,来自 | Hollis 平常的工作中,在衡量服务器的性能时,经常会涉及到几个指标,load.cpu.mem.qps.rt等.每个指标都有其独特的意义,很多时候在线上出现问题时 ...

极简yolov5转torchscript

极简yolov5转torchscript–以yolov5s为例

解决方法

极简yolov5转torchscript相关推荐

最新文章

热门文章