48、模型部署与PyTorch JIT的深度解析

模型部署与PyTorch JIT的深度解析

1. 模型导出至ONNX

1.1 导出步骤

将模型的计算过程进行追踪，并将其序列化为指定名称的ONNX文件：

torch.onnx.export(seg_model, dummy_input, "seg_model.onnx")

1.2 ONNX文件的使用

生成的ONNX文件可以在运行时环境中运行，编译到边缘设备，或者上传到云服务。在安装 onnxruntime 或 onnxruntime-gpu 并将批量数据转换为NumPy数组后，可以在Python中使用该文件：

import onnxruntime
sess = onnxruntime.InferenceSession("seg_model.onnx")
input_name = sess.get_inputs()[0].name
pred_onnx, = sess.run(None, {input_name: batch})

1.3 注意事项

并非所有的TorchScript操作符都能表示为标准化的ONNX操作符。如果导出的操作是ONNX不支持的，在使用运行时环境时会出现关于未知 aten 操作符的错误。

2. PyTorch的模型导出：追踪（Tracing）

2.1 追踪的基本