此教程参考:赛道检测模型训练部署全过程讲解,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()