pt转onnx

常见的模型文件包括后缀名为.pt,.pth,.pkl的模型文件，而这几种模型文件并非格式上有区别而是后缀不同而已，保存模型文件往往用的是torch.save()，后缀不同只是单纯因为每个人喜好不同而已。通常用的是pth和pt。
保存：
orch.save(model.state_dict(), mymodel.pth)#只保存模型权重参数，不保存模型结构

调用：
model = My_model(*args, **kwargs) #这里需要重新模型结构，
pthfile = r’绝对路径’
loaded_model = torch.load(pthfile, map_location=‘cpu’)
model.load_state_dict(loaded_model[‘model’])
model.eval() #不启用 BatchNormalization 和 Dropout,不改变权值

from nn.mobilenetv3 import mobilenetv3_large,mobilenetv3_large_full,mobilenetv3_small
import torch
from nn.models import DarknetWithShh
from hyp import hyp

def convert_onnx():
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    model_path = 'weights/mbv3_large_75_light_final.pt' #这是我们要转换的模型
    backone = mobilenetv3_large(width_mult=0.75)#mobilenetv3_small()  mobilenetv3_small(width_mult=0.75)  mobilenetv3_large(width_mult=0.75)
    model = DarknetWithShh(backone, hyp,light_head=True).to(device)

    model.load_state_dict(torch.load(model_path, map_location=device)['model'])

    model.to(device)
    model.eval()
    dummy_input = torch.randn(1, 3, 32, 32).to(device)#输入大小   #data type nchw
    onnx_path = 'weights/mbv3_large_75_light_final.onnx'
    torch.onnx.export(model, dummy_input, onnx_path, input_names=['input'], output_names=['output'],opset_version=11)
    print('convert retinaface to onnx finish!!!')

if __name__ == "__main__" :
    convert_onnx()

转换结果

onnx转tensorrt

有了onnx模型转化为tensorrt模型就非常简单了，其中builder是构建engine的，也就是我们需要的模型，network是网络设置，parser是解析onnx模型的工具，config是指定一些模型的设置。通过调整输入输出模型的位置以及max_batch_size的值，还有network所对应图片的shape值

#!/usr/bin/env python3

import tensorrt as trt

import sys, os
sys.path.insert(1, os.path.join(sys.path[0], ".."))


TRT_LOGGER = trt.Logger()
EXPLICIT_BATCH=1
def get_engine(onnx_file_path, engine_file_path):
    """Attempts to load a serialized engine if available, otherwise builds a new TensorRT engine and saves it."""
    def build_engine():
        """Takes an ONNX file and creates a TensorRT engine to run inference with"""
        with trt.Builder(TRT_LOGGER) as builder, builder.create_network(EXPLICIT_BATCH) as network, builder.create_builder_config() as config, trt.OnnxParser(network, TRT_LOGGER) as parser:
            builder.max_batch_size = 1
            config.max_workspace_size = 1 << 30 # 30:1GB;28:256MiB
            builder.fp16_mode=True
            # Parse model file
            if not os.path.exists(onnx_file_path):
                print('ONNX file {} not found, please run yolov3_to_onnx.py first to generate it.'.format(onnx_file_path))
                exit(0)
            print('Loading ONNX file from path {}...'.format(onnx_file_path))
            with open(onnx_file_path, 'rb') as model:
                print('Beginning ONNX file parsing')
                parser.parse(model.read())
                if not parser.parse(model.read()):
                    print ('ERROR: Failed to parse the ONNX file.')
                    for error in range(parser.num_errors):
                        print (parser.get_error(error))
                    return None
            # The actual yolov3.onnx is generated with batch size 64. Reshape input to batch size 1
            print('Completed parsing of ONNX file')
            print('Building an engine from file {}; this may take a while...'.format(onnx_file_path))
            engine = builder.build_cuda_engine(network)
            print("Completed creating Engine")
            with open(engine_file_path, "wb") as f:
                f.write(bytearray(engine.serialize()))
            return engine

    if os.path.exists(engine_file_path):
        # If a serialized engine exists, use it instead of building an engine.
        print("Reading engine from file {}".format(engine_file_path))
        with open(engine_file_path, "rb") as f, trt.Runtime(TRT_LOGGER) as runtime:
            return runtime.deserialize_cuda_engine(f.read())
    else:
        return build_engine()
 
def get_engine1(engine_path):
    # If a serialized engine exists, use it instead of building an engine.
    print("Reading engine from file {}".format(engine_path))
    with open(engine_path, "rb") as f, trt.Runtime(TRT_LOGGER) as runtime:
        return runtime.deserialize_cuda_engine(f.read())
 

if __name__ == '__main__':
    #main()
    onnx_file_path = '/home/z/Documents/4kinds_detectface_module/libfacedetection/YuFaceDetectNet_320.onnx'
    engine_file_path = "/home/z/Documents/4kinds_detectface_module/libfacedetection/YuFaceDetectNet_new_320.trt"
    get_engine(onnx_file_path, engine_file_path)
    # 可用netron查看onnx的输出数量和尺寸
    engines=get_engine1(engine_file_path)
    for binding in engines:
        size = trt.volume(engines.get_binding_shape(binding)) * 1
        dims = engines.get_binding_shape(binding)
        print('size=',size)
        print('dims=',dims)
        print('binding=',binding)
        print("input =", engines.binding_is_input(binding))
        dtype = trt.nptype(engines.get_binding_dtype(binding))

运行结果

Loading the ONNX file...
[TensorRT] WARNING: onnx2trt_utils.cpp:220: Your ONNX model has been generated with INT64 weights, while TensorRT does not natively support INT64. Attempting to cast down to INT32.
Building an engine.  This would take a while...
(Use "--verbose" or "-v" to enable verbose logging.)
Completed creating engine.
Serialized the TensorRT engine to file: /home/z/Documents/4kinds_detectface_module/libfacedetection/YuFaceDetectNet_320.trt
Reading engine from file /home/z/Documents/4kinds_detectface_module/libfacedetection/YuFaceDetectNet_new_320.trt
size= 230400
dims= (1, 3, 320, 240)
binding= input
input = True
size= 61390
dims= (1, 4385, 14)
binding= loc
input = False
size= 8770
dims= (1, 4385, 2)
binding= conf
input = False