JMU黑客松 - 06_export.py 将模型导出为onnx

此教程参考:赛道检测模型训练部署全过程讲解,https://developer.horizon.cc/forumDetail/185446272545810434

在此篇教程的流程和代码上作了一些补充和说明,若使用02准备的数据集,03标注,04打散,05训练,则可以直接使用06导出即可,注意需要将DATASET_NMAE修改一致。

ResNet18 残差神经网络结构(使用netron对onnx模型文件进行可视化:https://netron.app/)-

参考上文教程安装依赖和训练,可以流畅跑通,在此仅仅对原文作一点补充。-
模型如何量化为地平线的bin格式,也可参照上文,如果有转化的需求可联系我,这一步不对同学们作出要求。

若出现如下报错:

Traceback (most recent call last):
  File "export.py", line 23, in <module>
    main()
  File "export.py", line 13, in main
    torch.onnx.export(model,
  File "/home/sunrise/.local/lib/python3.8/site-packages/torch/onnx/utils.py", line 516, in export
    _export(
  File "/home/sunrise/.local/lib/python3.8/site-packages/torch/onnx/utils.py", line 1670, in _export
    proto = onnx_proto_utils._add_onnxscript_fn(
  File "/home/sunrise/.local/lib/python3.8/site-packages/torch/onnx/_internal/onnx_proto_utils.py", line 223, in _add_onnxscript_fn
    raise errors.OnnxExporterError("Module onnx is not installed!") from e
torch.onnx.errors.OnnxExporterError: Module onnx is not installed!

主要是未安装onnx模块导致,可采用以下命令安装:

pip install onnx -i https://mirrors.aliyun.com/pypi/simple

代码参考:(附件中可下载)

import torchvision
import torch

BEST_MODEL_PATH = './model_best.pth'  # 最好的训练结果

def main(args=None):
  model = torchvision.models.resnet18(pretrained=False)
  model.fc = torch.nn.Linear(512,2)
  model.load_state_dict(torch.load(BEST_MODEL_PATH, map_location='cpu'))
  device = torch.device('cpu')
  model = model.to(device)
  model.eval()
  x = torch.randn(1, 3, 224, 224, requires_grad=True)
  # torch_out = model(x)
  torch.onnx.export(model,
                    x,
                    BEST_MODEL_PATH[:-4] + ".onnx",
                    export_params=True,
                    opset_version=11,
                    do_constant_folding=True,
                    input_names=['input'],
                    output_names=['output'])

if __name__ == '__main__':
  main()

06_export.py

超哥,训练出的onnx模型能不能转为bin模型让x5的车使用

这个是没问题的,但是编译的时候要指定X5的BPU架构为bayes-e,具体请阅读X5的工具链手册

1 个赞