背景:树莓派要用英特尔神经计算棒,先要将pytorch训练生成的.pt文件转换为ONNX,再转换为IR,本来直接在树莓派上的原生树莓派系统尝试,明明已经显示输出了模型详细信息,可文件夹中没有出现.onnx文件,于是直接拿windows系统新建有pytorch的环境来尝试,尝试成功。
转换代码
注意点:要根据你的代码进行修改,修改最初的包等
import torch
from models.with_mobilenet import PoseEstimationWithMobileNet
from modules.load_state import load_state
from action_detect.net import NetV2
def convert_onnx():
print('start!!!')
device = 'cuda' if torch.cuda.is_available() else 'cpu'
#model_path = '/home/pi/xg_openpose_fall_detect-master/action_detect/checkPoint/action.pt' #这是我们要转换的模型
#backone = mobilenetv3_large(width_mult=0.75)#mobilenetv3_small() mobilenetv3_small(width_mult=0.75) mobilenetv3_large(width_mult=0.75)
model = NetV2().to(device)
checkpoint = torch.load(r'E:/xg_openpose_fall_detect-master/action_detect/checkPoint/action.pt', map_location='cpu')
model.load_state_dict(checkpoint)
#model.load_state_dict(torch.load(model_path, map_location=device)['model'])
model.to(device)
model.eval()
dummy_input = torch.randn(1, 16384).to(device)#输入大小 #data type nchw
#onnx_path = '/home/pi/xg_openpose_fall_detect-master/action_detect/checkPoint/action.onnx'
onnx_path = 'E:/xg_openpose_fall_detect-master/action_detect/checkPoint/action.onnx'
print("----- pt模型导出为onnx模型 -----")
output_name = "action.onnx"
torch.onnx.export(model, dummy_input,onnx_path,export_params=True, input_names=['input'], output_names=['output'])
print('finish!!!')
if __name__ == "__main__" :
convert_onnx()
结果
并保存在定义位置