record-----torch(.pth)模型转onnx(.onnx)模型的方法

前言

torch转onnx目的用于将torch(.pth)模型转为tensorRT(.engine)，pth->onnx->engine

code

"""
需要的环境：pytorch1.4

进行模型转换，torch->onnx
环境要求:tvm转torch模型必须在pytorch1.4的环境下进行，1.4下可以进行torch->onnx->tvm,
同时也可以进行torch->libtorch转换，但是c++中源码编译的libtorch库要和torch->libtorch转换时使用的pytorch版本一致
由于c++工程中使用的libtorch库是1.0,因此torch->libtorch的过程也要在pytorch1.0环境下进行
"""
import sys
#如果将编译的tvm添加到系统环境变量中不起作用，需要在此进行再次系统添加
# sys.path.insert(0,'/data_1/software/use_software/tvm/tvm/topi/python/')
# sys.path.insert(0,'/data_1/software/use_software/tvm/tvm/python/')
# sys.path.insert(0,'/data_1/software/use_software/tvm/tvm/nnvm/')
import torch
import torch.onnx  # 谨记：torch.onnx包 不能和　onnx包同时加载，即使同时加载了，也要onnx在torch.onnx之前进行加载。不然会报“段错误”，因此ｋｅ将torch->onnx以及onnx->tvm分开在两个py文件中
from pynvml import *     #pip install nvidia-ml-py3,显存监控
# from models_class.pose_resnet import * #关键点
from models_class.resnet_chejian_xiaolaomei import *

# from msra_resnet import * #centernet转换


# 将torch模型转为onnx模型

#转换关键点模型pth->onnx
def gen_torch_to_onnx_model_keypoints(model_path, save_onnx_path):
    # 加载torch模型，分类模型
    # 关键点模型
    model = get_pose_net('',False)  # pytorch的网络
    state_dict = torch.load(model_path)
    model.load_state_dict(state_dict)

    # 以下操作是避免pytorch转libtorch的模型显存增加
    for p in model.parameters():
        p.requires_grad = False
    model.cuda()
    model.eval()

    # 转 ONNX 模型
    example = torch.rand(32, 3, 1024, 768).cuda()   #输入图像大小,自行设定,如果使用tensorrt7.0.0多batch使用时，此处必须设定batch的大小

    input_name = ['input_1']  # 如果使用tensorrt7.0.0，必须设置输入的名字，输入可以是一个，也可以是多个,在c++中使用TRT需要知道的
    output_name = ['output_1']  # 如果使用tensorrt7.0.0，必须设置输出层的名字，输出可以是一个，也可以是多个,在c++中使用TRT需要知道的

    torch_out = torch.onnx.export(model=model,  # model being run
                                  args=(example,),  # model input (or a tuple for multiple inputs)
                                  f=save_onnx_path,
                                  input_names=input_name,
                                  output_names=output_name,
                                  export_params=True,
                                  verbose=True)  # store the trained parameter weights inside the model file   # 带参数输出)
    print("onnx convert ok")

#转换centernet检测模型pth->onnx
def gen_torch_to_onnx_model_centernet(model_path,save_onnx_path):

    #加载torch模型，分类模型

    #centernet
    num_layers = 10
    heads = {'hm': 7, 'reg': 2, 'wh': 2}
    head_conv = 64
    resnet_spec = {10: (BasicBlock, [1, 1, 1, 1]),
                   18: (BasicBlock, [2, 2, 2, 2]),
                   34: (BasicBlock, [3, 4, 6, 3]),
                   50: (Bottleneck, [3, 4, 6, 3]),
                   101: (Bottleneck, [3, 4, 23, 3]),
                   152: (Bottleneck, [3, 8, 36, 3])}
    name = 'ResNet_10'
    block_class, layers = resnet_spec[num_layers]
    model = PoseResNet(block_class, layers, heads, head_conv=head_conv)

    state_dict = torch.load(model_path)
    model.load_state_dict(state_dict['state_dict'])

    model.eval().cuda()

    # 转 ONNX 模型
    example = torch.rand(16, 3, 1120, 1120).cuda() #输入图像大小,自行设定,如果使用tensorrt7.0.0多batch使用时，此处必须设定batch的大小

    input_name = ['input_1'] #如果使用tensorrt7.0.0，必须设置输入的名字，输入可以是一个，也可以是多个
    output_name = ['output_hm','output_wh','output_reg'] #如果使用tensorrt7.0.0，必须设置输出层的名字，输出可以是一个，也可以是多个


    torch_out = torch.onnx.export(model=model,  # model being run
                                args=(example,),  # model input (or a tuple for multiple inputs)
                                f=save_onnx_path,
                                input_names=input_name,
                                output_names=output_name,
                                export_params=True,
                                verbose=True)  # store the trained parameter weights inside the model file   # 带参数输出)
    print("onnx convert ok")

#转换分类检测模型pth->onnx
def gen_torch_to_onnx_model_class(model_path,save_onnx_path):

    #加载torch模型，分类模型
    model = resnet29_cls() #pytorch的网络

    state_dict = torch.load(model_path)
    model.load_state_dict(state_dict,strict=True)

    model.eval().cuda()

    # 转 ONNX 模型
    example = torch.rand(32, 3, 400, 400).cuda() #输入图像大小,自行设定
    # example = torch.rand(2, 3, 608, 608).cuda() #输入图像大小,自行设定,如果使用tensorrt7.0.0多batch使用时，此处必须设定batch的大小

    input_name = ['input_1'] #如果使用tensorrt7.0.0，必须设置输入的名字，输入可以是一个，也可以是多个
    output_name = ['output_1'] #如果使用tensorrt7.0.0，必须设置输出层的名字，输出可以是一个，也可以是多个

    torch_out = torch.onnx.export(model=model,  # model being run
                                args=(example,),  # model input (or a tuple for multiple inputs)
                                f=save_onnx_path,
                                input_names=input_name,
                                output_names=output_name,
                                export_params=True,
                                verbose=True)  # store the trained parameter weights inside the model file   # 带参数输出)
    print("onnx convert ok")

def gen_torch_to_onnx_model(model_path,save_onnx_path):

    #加载torch模型，分类模型
    #关键点模型
    model = get_pose_net() #pytorch的网络

    #分类模型
    # model = resnet29_cls() #pytorch的网络

    #centernet
    # num_layers = 10
    # heads = {'hm': 7, 'reg': 2, 'wh': 2}
    # head_conv = 64
    # resnet_spec = {10: (BasicBlock, [1, 1, 1, 1]),
    #                18: (BasicBlock, [2, 2, 2, 2]),
    #                34: (BasicBlock, [3, 4, 6, 3]),
    #                50: (Bottleneck, [3, 4, 6, 3]),
    #                101: (Bottleneck, [3, 4, 23, 3]),
    #                152: (Bottleneck, [3, 8, 36, 3])}
    # name = 'ResNet_10'
    # block_class, layers = resnet_spec[num_layers]
    # model = PoseResNet(block_class, layers, heads, head_conv=head_conv)

    # state_dict = torch.load(model_path)
    # model.load_state_dict(state_dict['state_dict'])

    state_dict = torch.load(model_path)
    model.load_state_dict(state_dict,strict=True)

    model.eval().cuda()

    # 转 ONNX 模型
    example = torch.rand(32, 3, 200, 200).cuda() #输入图像大小,自行设定
    # example = torch.rand(2, 3, 608, 608).cuda() #输入图像大小,自行设定,如果使用tensorrt7.0.0多batch使用时，此处必须设定batch的大小

    # input_name = ['input_1'] #如果使用tensorrt7.0.0，必须设置输入的名字，输入可以是一个，也可以是多个
    # output_name = ['output_hm','output_wh','output_reg'] #如果使用tensorrt7.0.0，必须设置输出层的名字，输出可以是一个，也可以是多个


    input_name = ['input_1'] #如果使用tensorrt7.0.0，必须设置输入的名字，输入可以是一个，也可以是多个
    output_name = ['output_1'] #如果使用tensorrt7.0.0，必须设置输出层的名字，输出可以是一个，也可以是多个

    torch_out = torch.onnx.export(model=model,  # model being run
                                args=(example,),  # model input (or a tuple for multiple inputs)
                                f=save_onnx_path,
                                input_names=input_name,
                                output_names=output_name,
                                export_params=True,
                                verbose=True)  # store the trained parameter weights inside the model file   # 带参数输出)
    print("onnx convert ok")

if __name__=="__main__":
    # 显存监控
    GPU_USE=0
    nvmlInit() #初始化
    handle = nvmlDeviceGetHandleByIndex(GPU_USE) #获得指定GPU的handle
    info_begin = nvmlDeviceGetMemoryInfo(handle) #获得显存信息

    # 路径设置
    img_str = "/data_1/Working/project/torch-onnx-tvm/torch_py/data/[email protected]"
    model_path = "/data1/Working/project/weifa-all-project/video-3.0/use-tensorRT/test_centernet/models/xiang/centerNet.pth"
    onnx_save_path = "/data1/Working/project/weifa-all-project/video-3.0/use-tensorRT/test_centernet/models/xiang/centerNet_16_1120_1120.onnx"

    # 将torch模型转为onnx模型
    gen_torch_to_onnx_model_centernet(model_path,onnx_save_path)
    # gen_torch_to_onnx_model_keypoints(model_path,onnx_save_path)
    # gen_torch_to_onnx_model_class(model_path,onnx_save_path)

    # #输出显存使用mb
    info_end = nvmlDeviceGetMemoryInfo(handle)
    print("-" * 15 + "TORCH GPU MEMORY INFO" + "-" * 15)
    print("       Memory Total: " + str(info_end.total // (1024 ** 2)))
    print("       Memory Free: " + str(info_end.free // (1024 ** 2)))
    print("       Memory Used: " + str(info_end.used // (1024 ** 2) - info_begin.used // (1024 ** 2)))
    print("-" * 40)