ONNX模型部署

pytorch模型转换为onnx

输入model为训练好的模型，x为符合其输入形状的随机输入
其中 'input' 和 'output' 可以更改为自定义的名称

x = torch.randn(1, 256, 1024)
def trans_onnx(model, x):
    # print('model device:', next(model.parameters()).device)
    # print('input device:', x.device)
    with torch.no_grad():
        torch.onnx.export(
            model,
            x,
            f='model.onnx',
            input_names=['input'],
            output_names=['output'],
            opset_version=15
            # export_params=False
            # use_external_data_format=True
        )
    print('finish!')

onnxruntime推理

CPU推理

首先安装CPU版本推理库

pip install onnxruntime

GPU推理

首先安装GPU版本推理库，需要卸载CPU版本

pip uninstall onnxruntime
pip install onnxruntime-gpu

注意：onnxruntime、cuda、cudnn版本需要对应，详见 Requirements1

我使用的是：pytorch 1.12.0 + onnxruntime-gpu 1.14.1 + cudatoolkit 11.3.1 + cudnn 8.2.1

可以使用如下代码检验当前设备是否可用GPU推理，无论安装的是CPU还是GPU版本库，名字都由onnxruntime导入

import onnxruntime
print(onnxruntime.get_device())
ort_session = onnxruntime.InferenceSession("model.onnx", providers=['CUDAExecutionProvider'])
print(ort_session.get_providers())

# 输出如下
# GPU
# ['CUDAExecutionProvider', 'CPUExecutionProvider']

推理

ort_session = onnxruntime.InferenceSession('model.onnx', providers=['CUDAExecutionProvider'])
# onnx runtime 输入
img = img_transforms(image).view(1, 256, 1024).numpy()
ort_inputs = {'input': img}
# onnx runtime 输出
ort_output = ort_session.run(['output'], ort_inputs)[0]
print(ort_output)

References

• 文档
• browser-onnx
• onnxruntime-github
• OpenMMLab教程