if __name__ == '__main__':
parser = argparse.ArgumentParser()
parser.add_argument('--weights', type=str, default='runs/train/exp2/weights/best.pt', help='weights path')
parser.add_argument('--img-size', nargs='+', type=int, default=[640, 640], help='image size') # height, width
parser.add_argument('--batch-size', type=int, default=1, help='batch size')
parser.add_argument('--dynamic', action='store_true', help='dynamic ONNX axes')
parser.add_argument('--dynamic-batch', action='store_true', help='dynamic batch onnx for tensorrt and onnx-runtime')
parser.add_argument('--grid', action='store_true', help='export Detect() layer grid')
parser.add_argument('--end2end', action='store_true', help='export end2end onnx')
parser.add_argument('--max-wh', type=int, default=None, help='None for tensorrt nms, int value for onnx-runtime nms')
parser.add_argument('--topk-all', type=int, default=100, help='topk objects for every images')
parser.add_argument('--iou-thres', type=float, default=0.45, help='iou threshold for NMS')
parser.add_argument('--conf-thres', type=float, default=0.55, help='conf threshold for NMS')
parser.add_argument('--device', default='0', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
parser.add_argument('--simplify', action='store_true', help='simplify onnx model')
parser.add_argument('--include-nms', action='store_true', help='export end2end onnx')
parser.add_argument('--fp16', action='store_true', help='CoreML FP16 half-precision export')
parser.add_argument('--int8', action='store_true', help='CoreML INT8 quantization')
opt = parser.parse_args()
opt.img_size *= 2 if len(opt.img_size) == 1 else 1 # expand
opt.dynamic = opt.dynamic and not opt.end2end
opt.dynamic = False if opt.dynamic_batch else opt.dynamic
print(opt)
set_logging()
t = time.time()
# Load PyTorch model
device = select_device(opt.device)
model = attempt_load(opt.weights, map_location=device) # load FP32 model
labels = model.names
# Checks
gs = int(max(model.stride)) # grid size (max stride)
opt.img_size = [check_img_size(x, gs) for x in opt.img_size] # verify img_size are gs-multiples
# Input
img = torch.zeros(opt.batch_size, 3, *opt.img_size).to(device) # image size(1,3,320,192) iDetection
# Update model
for k, m in model.named_modules():
m._non_persistent_buffers_set = set() # pytorch 1.6.0 compatibility
if isinstance(m, models.common.Conv): # assign export-friendly activations
if isinstance(m.act, nn.Hardswish):
m.act = Hardswish()
elif isinstance(m.act, nn.SiLU):
m.act = SiLU()
# elif isinstance(m, models.yolo.Detect):
# m.forward = m.forward_export # assign forward (optional)
model.model[-1].export = not opt.grid # set Detect() layer grid export
y = model(img) # dry run
if opt.include_nms:
model.model[-1].include_nms = True
y = None
# TorchScript export
try:
print('\nStarting TorchScript export with torch %s...' % torch.__version__)
f = opt.weights.replace('.pt', '.torchscript.pt') # filename
ts = torch.jit.trace(model, img, strict=False)
ts.save(f)
print('TorchScript export success, saved as %s' % f)
except Exception as e:
print('TorchScript export failure: %s' % e)
-
定义了命令行参数的解析器,并解析命令行参数。这些参数包括模型权重路径、图像尺寸、批量大小、设备选择等。
-
根据参数选择的设备,使用select_device函数选择运行模型的设备。
-
使用attempt_load函数加载指定路径的模型权重,并返回一个加载的FP32精度模型。
-
从加载的模型中获取类别标签。
-
通过计算模型的最大步长(stride),确定图像尺寸是否是步长的倍数并进行验证。
-
创建一个shape为(批量大小,3,*图像尺寸)的全零张量img,用于模型推理的输入。
-
更新模型的处理方式。对于模型中的Conv层,将其激活函数替换为与TorchScript兼容的Hardswish或SiLU。最后一层的Detect层是否导出网格根据参数opt.grid来确定。
-
进行一次模型的前向传播,用于进行一次“干跑”。
-
根据参数opt.include_nms,决定是否将NMS操作包含在导出的模型中。
-
尝试使用TorchScript转换为TorchScript模型。创建一个文件路径f,将模型以TorchScript格式保存在该路径下。如果成功保存,则打印成功信息;如果失败,则打印失败信息。
这段代码的作用是根据给定的参数加载PyTorch模型,将其转换为TorchScript模型,并保存在指定路径下。转换的模型可以在没有Python环境的设备上进行运行和部署。
try:
import onnx
print('\nStarting ONNX export with onnx %s...' % onnx.__version__)
f = opt.weights.replace('.pt', '.onnx') # filename
model.eval()
output_names = ['classes', 'boxes'] if y is None else ['output']
dynamic_axes = None
if opt.dynamic:
dynamic_axes = {'images': {0: 'batch', 2: 'height', 3: 'width'}, # size(1,3,640,640)
'output': {0: 'batch', 2: 'y', 3: 'x'}}
if opt.dynamic_batch:
opt.batch_size = 'batch'
dynamic_axes = {
'images': {
0: 'batch',
}, }
if opt.end2end and opt.max_wh is None:
output_axes = {
'num_dets': {0: 'batch'},
'det_boxes': {0: 'batch'},
'det_scores': {0: 'batch'},
'det_classes': {0: 'batch'},
}
else:
output_axes = {
'output': {0: 'batch'},
}
dynamic_axes.update(output_axes)
if opt.grid:
if opt.end2end:
print('\nStarting export end2end onnx model for %s...' % 'TensorRT' if opt.max_wh is None else 'onnxruntime')
model = End2End(model,opt.topk_all,opt.iou_thres,opt.conf_thres,opt.max_wh,device,len(labels))
if opt.end2end and opt.max_wh is None:
output_names = ['num_dets', 'det_boxes', 'det_scores', 'det_classes']
shapes = [opt.batch_size, 1, opt.batch_size, opt.topk_all, 4,
opt.batch_size, opt.topk_all, opt.batch_size, opt.topk_all]
else:
output_names = ['output']
else:
model.model[-1].concat = True
torch.onnx.export(model, img, f, verbose=False, opset_version=12, input_names=['images'],
output_names=output_names,
dynamic_axes=dynamic_axes)
# Checks
onnx_model = onnx.load(f) # load onnx model
onnx.checker.check_model(onnx_model) # check onnx model
if opt.end2end and opt.max_wh is None:
for i in onnx_model.graph.output:
for j in i.type.tensor_type.shape.dim:
j.dim_param = str(shapes.pop(0))
首先,它导入了onnx
库,并打印出所使用的onnx
版本。然后,它根据输入的权重文件路径生成导出的ONNX文件的路径。
接下来,它将模型设为评估模式,并确定输出的名称。如果y
参数为空,则输出的名称为['classes', 'boxes'];否则,输出的名称为['output']。以及动态轴设置为None
。
如果opt.dynamic
为True
,则根据模型输入和输出的维度确定动态轴的设置。例如,如果输入的维度为(1, 3, 640, 640),那么动态轴的设置为 {'images': {0: 'batch', 2: 'height', 3: 'width'}, 'output': {0: 'batch', 2: 'y', 3: 'x'}}
。
如果opt.dynamic_batch
为True
,则将opt.batch_size
设置为'batch'
,并根据opt.end2end
和opt.max_wh
的值确定输出轴的设置。如果opt.end2end
为True
且opt.max_wh
为None
,则输出轴的设置为{'num_dets': {0: 'batch'}, 'det_boxes': {0: 'batch'}, 'det_scores': {0: 'batch'}, 'det_classes': {0: 'batch'}}
。否则,输出轴的设置为{'output': {0: 'batch'}}
。
如果opt.grid
为True
,则根据opt.end2end
的值进行下一步操作。如果opt.end2end
为True
,并且opt.max_wh
为None
,则将模型替换为End2End
对象,并更新输出的名称和形状。输出的名称为['num_dets', 'det_boxes', 'det_scores', 'det_classes']
;形状为[opt.batch_size, 1, opt.batch_size, opt.topk_all, 4, opt.batch_size, opt.topk_all, opt.batch_size, opt.topk_all]
。否则,输出的名称为['output']
。
最后,使用torch.onnx.export
函数将PyTorch模型导出为ONNX模型。img
是用于模型推断的输入数据,f
是导出的ONNX文件的路径。其中,input_names
和output_names
分别表示输入和输出的名称,dynamic_axes
表示需要设置为动态轴的维度。
接下来,代码加载导出的ONNX模型,并使用onnx.checker.check_model
函数对模型进行检查。最后,如果opt.end2end
为True
且opt.max_wh
为None
,则为每个输出轴更新维度。
end2end
是一个布尔值参数,表示是否进行端到端的导出操作。具体来说,如果end2end
为True
,则会执行一些与端到端推理相关的操作。
在目标检测任务中,通常的做法是先使用目标检测模型进行物体检测,然后将检测到的物体作为输入传递给后续的任务,如分类、跟踪等。而端到端的导出操作则是指将整个目标检测模型以及后续的任务一起导出为一个统一的模型。
在代码中,当end2end
为True
时,会将模型替换为End2End
对象,并且将输出的名称和形状相应地更新。这意味着,导出的ONNX模型将包含整个端到端的任务,而不仅仅是目标检测模型。
因此,end2end
的意思是在导出ONNX模型时,是否将整个端到端的任务一起导出,而不仅仅是目标检测模型。