yolov5 pt->onnx->om yolov5模型转onnx转om模型转换

yolov5 pt->onnx->om
yolov5-6.1版本
models/yolo.py

Detect函数修改

class Detect(nn.Module):
	def forward(self, x):
	        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()
	            y = x[i].sigmoid()
	            z.append(y.view(bs, -1, self.no))
	
	        return torch.cat(z, 1)

common.py Focus修改(说是可以提升Slice算子性能,待测,可不改)

class Focus(nn.Module):
    def forward(self, x):  # x(b,c,w,h) -> y(b,4c,w/2,h/2)
        # <==== 修改内容
        if torch.onnx.is_in_onnx_export():
            a, b = x[..., ::2, :].transpose(-2, -1), x[..., 1::2, :].transpose(-2, -1)
            c = torch.cat([a[..., ::2, :], b[..., ::2, :], a[..., 1::2, :], b[..., 1::2, :]], 1).transpose(-2, -1)
            return self.conv(c)
        else:
            return self.conv(torch.cat([x[.

你可能感兴趣的:(笔记,深度学习,tensorflow,人工智能)