风火编程--模型部署相关

模型部署相关

pth模型转onnx

import torch
import torch.onnx

#Function to Convert to ONNX
def Convert_ONNX():

    # set the model to inference mode
    torch.load('weights/CodeFormer/codeformer.pth')

    # Let's create a dummy input tensor
    dummy_input = torch.randn(1, input_size, requires_grad=True)

    # Export the model
    torch.onnx.export(model,         # model being run
         dummy_input,       # model input (or a tuple for multiple inputs)
         "ImageClassifier.onnx",       # where to save the model
         export_params=True,  # store the trained parameter weights inside the model file
         opset_version=10,    # the ONNX version to export the model to
         do_constant_folding=True,  # whether to execute constant folding for optimization
         input_names = ['modelInput'],   # the model's input names
         output_names = ['modelOutput'], # the model's output names
         dynamic_axes={'modelInput' : {0 : 'batch_size'},    # variable length axes
                                'modelOutput' : {0 : 'batch_size'}})
    print(" ")
    print('Model has been converted to ONNX')

onnx模型优化

import onnx
from onnxoptimizer import optimize
from onnxsim import simplify
from onnxconverter_common import float16


def model_simplify(model_path):
    # 加载ONNX模型
    model = onnx.load(model_path)

    # 剪枝模型权重
    model, _ = simplify(model)

    # 保存剪枝后的模型
    onnx.save(model, model_path)


def model_optimize(model_path):
    # 加载ONNX模型
    model = onnx.load(model_path)

    # 优化模型
    passes = ["fuse_bn_into_conv"]
    model = optimize(model, passes)

    # 保存优化后的模型
    onnx.save(model, model_path)


def model2fp16(model_path):
    model = onnx.load(model_path)
    model_fp16 = float16.convert_float_to_float16(model)
    onnx.save(model_fp16, model_path)

onnx转tensorrt

trtexec --onnx=conv.onnx --saveEngine=flame_sim.engine --workspace=1024 --fp16

模型部署优化原理
https://zhuanlan.zhihu.com/p/663753346
onnx模型优化原理
https://zhuanlan.zhihu.com/p/625485007

tensorrt部署优化原理
https://zhuanlan.zhihu.com/p/432215219

查看模型结构
https://www.machunjie.com/dl/Visualization/index.html