什么是ONNX
- ONNX(Open Neural Network Exchange)简称ONNX
- 微软和FaceBook提出用来表示深度学习模型的开放格式。
- ONNX定义了一组和环境、平台无关的标准格式,来增强各种AI模型的可交互性。
- ONNX存储权重,模型的结构、每一层的输入输出、辅助信息。
ProtoBuf简介
- ProtoBuf 是一种轻便高效的结构化数据存储格式,可用于结构数据的序列化。
- 可用作通讯协议、数据存储等领域的语言无关、平台无关、可扩展的序列化结构数据格式。
pytorch 生成 onnx
模型定义
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.onnx
import os
class Model(torch.nn.Module):
def __init__(self):
super().__init__()
self.conv = nn.Conv2d(1, 1, 3, padding=1)
self.relu = nn.ReLU()
self.conv.weight.data.fill_(1)
self.conv.bias.data.fill_(0)
def forward(self, x):
x = self.conv(x)
x = self.relu(x)
return x
onnx导出
odel = Model()
dummy = torch.zeros(1, 1, 3, 3)
torch.onnx.export(
model,
# 这里的args,是指输入给model的参数,需要传递tuple,因此用括号
(dummy,),
# 储存的文件路径
"demo.onnx",
# 打印详细信息
verbose=False,
# 为输入和输出节点指定名称,方便后面查看或者操作
input_names=["image"],
output_names=["output"],
# 这里的opset,指,各类算子以何种方式导出,对应于symbolic_opset11
opset_version=11,
# 表示他有batch、height、width3个维度是动态的,在onnx中给其赋值为-1
# 通常,我们只设置batch为动态,其他的避免动态
dynamic_axes={
"image": {0: "batch", 2: "height", 3: "width"},
"output": {0: "batch", 2: "height", 3: "width"},
}
)
正确导出ONNX
- 对于任何用到shape、size返回值的参数时,例如:
tensor.view(tensor.size(0), -1) # ×
tensor.view(int(tensor.size(0)), -1) # √
这类操作,避免直接使用 tensor.size 的返回值,而是加上int转换。断开跟踪。
- 对于nn.Upsample或nn.functional.interpolate函数,使用scale_factor指定倍率,而不使用size参数指定大小
- 对于reshape、view操作时,-1的指定放在batch维度。其他维度可以计算出来即可。batch维度禁止指定大于-1的数字
- torch.onnx.export指定dynamic_axes参数,只指定batch维度,尽量不用其他动态
- opset_version = 11, 不要低于11
- 避免使用inplace操作, 例如: y[…,0:2] = y[…,0:2] * 2 - 0.5 不要使用批量复制
- 尽量少的出现5个维度,例如ShuffleNet Module 可以考虑合并wh避免出现5维度
- 尽量让后处理部分在onnx模型中实现,降低后处理复杂度。比如 anchor
- 简化过程的复杂度,去掉gather、shape的不必要节点。
ONNX 解析器
ONNX格式分析
- ModelProto: 包含了一些版本、生产者、一个GraphProto
- GraphProto: 包含四个repeated数组,分别是:
- node: NodeProto类型, 存放模型中所有的计算节点。
- input: ValueInfoProto类型, 存放模型的所有输入节点。
- output: valueInfoProto,存放模型的所有输出节点。
- initializer: 存放了模型的所有权重参数。
- AttributeProto: 描述节点的属性,比如Conv节点或者卷积层的属性包含group,pad
onnx.helper
- ONNX 是把一个网络的每一层或者一个算子当成节点node。
- 将node 构成一个Graph
在构建ONNX模型的时候,可以使用make_node,make_graph, make_tensor等等接口完成一个ONNX模型的构建。
yolov5 ONNX修改
1. 输入输出的宽高为动态宽高,希望修改为固定宽高,只有BatchSize是动态维度
修改方式:
# export.py
torch.onnx.export(model, im, f, verbose=False, opset_version=opset,
training=torch.onnx.TrainingMode.TRAINING if train else torch.onnx.TrainingMode.EVAL,
do_constant_folding=not train,
input_names=['images'],
output_names=['output'],
#dynamic_axes={'images': {0: 'batch', 2: 'height', 3: 'width'}, # shape(1,3,640,640)
#'output': {0: 'batch', 1: 'anchors'} # shape(1,25200,85)
#} if dynamic else None)
dynamic_axes={'images': {0: 'batch'}, # shape(1,3,640,640)
'output': {0: 'batch'} # shape(1,25200,85)
} if dynamic else None)
- 删除 除了 batch之外的动态维度
修改后的效果:
2. 修改输出
从上面可以看出,输出仍然包括三项我们推理不需要的内容。
主要是因为在yolo的forward定义中的输出:
return x if self.training else (torch.cat(z, 1), x)
这里的
# yolov5-6.0/models/yolo.py
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()
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.anchor_grid[i] = self._make_grid(nx, ny, i)
# 多余的三个sigmoid输出主要来自此处
y = x[i].sigmoid()
if self.inplace:
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] # wh
y = torch.cat((xy, wh, y[..., 4:]), -1)
z.append(y.view(bs, -1, self.no))
# 此处的 (torch.cat(z, 1), x) x并不是推理需要的,所以不需要保留,所以改为下面的这句
# return x if self.training else (torch.cat(z, 1), x)
return x if self.training else torch.cat(z, 1)
修改后的效果:
3. 删除多余的unsqueeze 和 gather
- 多余的unsqueeze 和 gather 很多是由于动态引用了 shape 属性。
bs, _, ny, nx = x[i].shape # x(bs,255,20,20) to x(bs,3,20,20,85)
修改为:
bs, _, ny, nx = map(int, x[i].shape) # x(bs,255,20,20) to x(bs,3,20,20,85)
修改前:
修改后:
4. Reshape 动态维度调整
reshape 的动态维度期望是: [-1, 1, 85]
- 由于原来并没有动态计算batch,而是动态计算最后候选框的数量。
- 这里我们改为动态计算batch, 手动计算候选框的数量
def forward(self, x):
z = [] # inference output
for i in range(self.nl):
x[i] = self.m[i](x[i]) # conv
bs, _, ny, nx = map(int, x[i].shape) # x(bs,255,20,20) to x(bs,3,20,20,85)
# 指定动态计算 batch
bs =-1
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.anchor_grid[i] = self._make_grid(nx, ny, i)
y = x[i].sigmoid()
if self.inplace:
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] # wh
y = torch.cat((xy, wh, y[..., 4:]), -1)
# z.append(y.view(bs, -1, self.no))
# bs 指定为 -1 后,动态指定第二个维度 手动计算候选框的数量
z.append(y.view(bs, self.na * ny * nx, self.no))
return x if self.training else torch.cat(z, 1)
修改后:
5. 将anchor 作为常量,不在网络中跟踪
- 目前在最后会有很多单独的分值节点,主要来源于计算anchor的部分。
- anchor在yolo中可以作为常量值,不必要跟踪,引入多余的节点。
- anchor_grid = (self.anchors[i].clone() * self.stride[i]).view(1, -1, 1, 1, 2)
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().__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
self.anchor_grid = [torch.zeros(1)] * self.nl # init anchor grid
self.register_buffer('anchors', torch.tensor(anchors).float().view(self.nl, -1, 2)) # shape(nl,na,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)
# 记录 anchors
self.anchors = anchors
def forward(self, x):
z = [] # inference output
for i in range(self.nl):
x[i] = self.m[i](x[i]) # conv
bs, _, ny, nx = map(int, x[i].shape) # x(bs,255,20,20) to x(bs,3,20,20,85)
bs =-1
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.anchor_grid[i] = self._make_grid(nx, ny, i)
# 预先生成anchor
anchor_grid = (self.anchors[i].clone() * self.stride[i]).view(1, -1, 1, 1, 2)
y = x[i].sigmoid()
if self.inplace:
y[..., 0:2] = (y[..., 0:2] * 2. - 0.5 + self.grid[i]) * self.stride[i] # xy
# 替换为生成的 anchor_grid
# y[..., 2:4] = (y[..., 2:4] * 2) ** 2 * self.anchor_grid[i] # wh
y[..., 2:4] = (y[..., 2:4] * 2) ** 2 * anchor_grid # 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] # wh
# 替换为生成的 anchor_grid
wh = (y[..., 2:4] * 2) ** 2 * anchor_grid # wh
y = torch.cat((xy, wh, y[..., 4:]), -1)
# z.append(y.view(bs, -1, self.no))
# bs 指定为 -1 后,动态指定第二个维度
z.append(y.view(bs, self.na * ny * nx, self.no))
return x if self.training else torch.cat(z, 1)
# return x if self.training else (torch.cat(z, 1), x)
def _make_grid(self, nx=20, ny=20, i=0):
d = self.anchors[i].device
yv, xv = torch.meshgrid([torch.arange(ny).to(d), torch.arange(nx).to(d)])
grid = torch.stack((xv, yv), 2).expand((1, self.na, ny, nx, 2)).float()
anchor_grid = (self.anchors[i].clone() * self.stride[i]) \
.view((1, self.na, 1, 1, 2)).expand((1, self.na, ny, nx, 2)).float()
return grid, anchor_grid
修改后
6. 多余的Identity
- 多余的Identity 这个多余的Identity 可能导致转TensorRT失败
- 这个主要是由于生成的anchor_grid.
- 我们可以使用 onnx-simplify 删除
if simplify:
try:
check_requirements(('onnx-simplifier',))
import onnxsim
print(f'{prefix} simplifying with onnx-simplifier {onnxsim.__version__}...')
model_onnx, check = onnxsim.simplify(
model_onnx,
dynamic_input_shape=dynamic,
input_shapes={'images': list(im.shape)} if dynamic else None)
assert check, 'assert check failed'
onnx.save(model_onnx, f)
except Exception as e:
print(f'{prefix} simplifier failure: {e}')
print(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
print(f"{prefix} run --dynamic ONNX model inference with: 'python detect.py --weights {f}'")
删除后
直接使用onnx-simplfy 会多出三个输出的节点。在实际使用中,可有先simpfy 一遍,然后再调整也是不错的选择。