48、模型部署与PyTorch JIT的深入探索

模型部署与PyTorch JIT的深入探索

1. 模型导出至ONNX

在模型部署过程中，将模型计算过程序列化到ONNX文件是常见的操作。可以使用以下代码将模型转换为ONNX格式：

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

生成的ONNX文件可以在运行时环境中运行，编译到边缘设备，或上传到云服务。安装 onnxruntime 或 onnxruntime-gpu 后，可在Python中使用该文件，以下是具体示例：

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

需要注意的是，并非所有TorchScript运算符都能表示为标准化的ONNX运算符。若导出ONNX不支持的操作，使用运行时会出现未知 aten 运算符的错误。

2. PyTorch的模型追踪导出

当不需要考虑互操作性，但需要摆脱Python全局解释器锁（GIL）或导出网络时，可以使用PyTorch自己的表示方法——TorchScript图。创建Torch