【.pt模型转换为.onnx模型】模型转换 英特尔神经计算棒 树莓派

背景:树莓派要用英特尔神经计算棒,先要将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()
    
    

 结果

【.pt模型转换为.onnx模型】模型转换 英特尔神经计算棒 树莓派_第1张图片

并保存在定义位置 

【.pt模型转换为.onnx模型】模型转换 英特尔神经计算棒 树莓派_第2张图片

你可能感兴趣的:(树莓派,填坑,人工智能杂七杂八)