PyTorch (可选）将模型从 PyTorch 导出到 ONNX 并使用 ONNX Runtime 运行

ONNX（Open Neural Network Exchange）作为一种开放的神经网络交换格式，为深度学习模型的跨框架部署提供了便捷的途径。ONNX Runtime 则是一个专注于性能的推理引擎，支持在多个平台和硬件上高效运行 ONNX 模型。本文将详细介绍如何将 PyTorch 模型导出为 ONNX 格式，并利用 ONNX Runtime 进行推理，使您能够在多种环境中高效地部署和运行模型。

一、环境搭建与依赖安装

在开始模型导出和推理之前，请确保已安装以下依赖库：

pip install onnx onnxruntime

同时，建议使用 PyTorch 的最新版本以确保兼容性。

二、定义并加载 PyTorch 超分辨率模型

我们将使用一个预定义的超分辨率模型，该模型通过高效的子像素卷积层提升图像分辨率。

import torch
import torch.nn as nn
import torch.nn.init as init


class SuperResolutionNet(nn.Module):
    def __init__(self, upscale_factor, inplace=False):
        super(SuperResolutionNet, self).__init__()
        self.relu = nn.ReLU(inplace=inplace)
        self.conv1 = nn.Conv2d(1, 64, (5, 5), (1, 1), (2, 2))
        self.conv2 = nn.Conv2d(64, 64, (3, 3), (1, 1), (1, 1))
        self.conv3 = nn.Conv2d(64, 32, (3, 3), (1, 1), (1, 1))
        self.conv4 = nn.Conv2d(32, upscale_factor ** 2, (3, 3), (1, 1), (1, 1))
        self.pixel_shuffle = nn.PixelShuffle(upscale_factor)
        self._initialize_weights()


    def forward(self, x):
        x = self.relu(self.conv1(x))
        x = self.relu(self.conv2(x))
        x = self.relu(self.conv3(x))
        x = self.pixel_shuffle(self.conv4(x))
        return x


    def _initialize_weights(self):
        init.orthogonal_(self.conv1.weight, init.calculate_gain('relu'))
        init.orthogonal_(self.conv2.weight, init.calculate_gain('relu'))
        init.orthogonal_(self.conv3.weight, init.calculate_gain('relu'))
        init.orthogonal_(self.conv4.weight)


## 创建超分辨率模型实例
torch_model = SuperResolutionNet(upscale_factor=3)


## 加载预训练权重
model_url = 'https://s3.amazonaws.com/pytorch/test_data/export/superres_epoch100-44c6958e.pth'
batch_size = 1
map_location = lambda storage, loc: storage
if torch.cuda.is_available():
    map_location = None
torch_model.load_state_dict(torch.utils.model_zoo.load_url(model_url, map_location=map_location))


## 将模型设置为推理模式
torch_model.eval()

三、将 PyTorch 模型导出为 ONNX 格式

使用 torch.onnx.export() 函数将模型导出为 ONNX 格式。在导出过程中，模型会被执行一次，以跟踪计算图的结构。

## 准备模型输入
x = torch.randn(batch_size, 1, 224, 224, requires_grad=True)


## 导出模型到 ONNX 格式
torch.onnx.export(torch_model,
                 x,
                 "super_resolution.onnx",
                 export_params=True,
                 opset_version=10,
                 do_constant_folding=True,
                 input_names=['input'],
                 output_names=['output'],
                 dynamic_axes={'input': {0: 'batch_size'},
                               'output': {0: 'batch_size'}})

四、验证 ONNX 模型的有效性

使用 ONNX 的 API 检查导出的模型是否有效。

import onnx


## 加载 ONNX 模型
onnx_model = onnx.load("super_resolution.onnx")


## 验证模型结构
onnx.checker.check_model(onnx_model)

五、使用 ONNX Runtime 进行推理

利用 ONNX Runtime 的 Python API 加载并运行 ONNX 模型。

import onnxruntime


## 创建推理会话
ort_session = onnxruntime.InferenceSession("super_resolution.onnx")


## 定义张量转 NumPy 数组的辅助函数
def to_numpy(tensor):
    return tensor.detach().cpu().numpy() if tensor.requires_grad else tensor.cpu().numpy()


## 计算 ONNX Runtime 输出
ort_inputs = {ort_session.get_inputs()[0].name: to_numpy(x)}
ort_outs = ort_session.run(None, ort_inputs)


## 验证 ONNX Runtime 和 PyTorch 的输出一致性
np.testing.assert_allclose(to_numpy(torch_model(x)), ort_outs[0], rtol=1e-03, atol=1e-05)
print("Exported model has been tested with ONNXRuntime, and the result looks good!")

六、示例应用：超分辨率图像处理

加载示例图像并进行预处理，使用导出的 ONNX 模型进行超分辨率转换。

from PIL import Image
import torchvision.transforms as transforms


## 加载并预处理图像
img = Image.open("cat.jpg")  # 请替换为实际图像文件路径
resize = transforms.Resize([224, 224])
img = resize(img)
img_ycbcr = img.convert('YCbCr')
img_y, img_cb, img_cr = img_ycbcr.split()
to_tensor = transforms.ToTensor()
img_y = to_tensor(img_y)
img_y.unsqueeze_(0)


## 使用 ONNX Runtime 进行推理
ort_inputs = {ort_session.get_inputs()[0].name: to_numpy(img_y)}
ort_outs = ort_session.run(None, ort_inputs)
img_out_y = ort_outs[0]


## 后处理输出图像
img_out_y = Image.fromarray(np.uint8((img_out_y[0] * 255.0).clip(0, 255)[0]), mode='L')
final_img = Image.merge("YCbCr", [
    img_out_y,
    img_cb.resize(img_out_y.size, Image.BICUBIC),
    img_cr.resize(img_out_y.size, Image.BICUBIC),
]).convert("RGB")


## 保存结果图像
final_img.save("cat_superres_with_ort.jpg")

七、总结与拓展

本教程展示了如何将 PyTorch 模型导出为 ONNX 格式，并使用 ONNX Runtime 进行推理。ONNX 为深度学习模型的跨框架部署提供了强大的支持，而 ONNX Runtime 则确保了模型在不同环境中的高效运行。通过这一过程，您可以轻松地将模型部署到各种平台，包括 Windows、Linux 和 Mac，以及利用 CPU 和 GPU 的计算能力。

未来，您可以进一步探索 ONNX Runtime 的高级功能，如模型优化、量化和部署到云端或边缘设备。编程狮将持续为您带来更多关于模型部署和优化的优质教程，助力您的项目高效落地。