ONNX（Open Neural Network Exchange）作為一種開(kāi)放的神經(jīng)網(wǎng)絡(luò)交換格式，為深度學(xué)習(xí)模型的跨框架部署提供了便捷的途徑。ONNX Runtime 則是一個(gè)專(zhuān)注于性能的推理引擎，支持在多個(gè)平臺(tái)和硬件上高效運(yùn)行 ONNX 模型。本文將詳細(xì)介紹如何將 PyTorch 模型導(dǎo)出為 ONNX 格式，并利用 ONNX Runtime 進(jìn)行推理，使您能夠在多種環(huán)境中高效地部署和運(yùn)行模型。

一、環(huán)境搭建與依賴(lài)安裝

在開(kāi)始模型導(dǎo)出和推理之前，請(qǐng)確保已安裝以下依賴(lài)庫(kù)：

pip install onnx onnxruntime

同時(shí)，建議使用 PyTorch 的最新版本以確保兼容性。

二、定義并加載 PyTorch 超分辨率模型

我們將使用一個(gè)預(yù)定義的超分辨率模型，該模型通過(guò)高效的子像素卷積層提升圖像分辨率。

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)


## 創(chuàng)建超分辨率模型實(shí)例
torch_model = SuperResolutionNet(upscale_factor=3)


## 加載預(yù)訓(xùn)練權(quán)重
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))


## 將模型設(shè)置為推理模式
torch_model.eval()

三、將 PyTorch 模型導(dǎo)出為 ONNX 格式

使用 torch.onnx.export() 函數(shù)將模型導(dǎo)出為 ONNX 格式。在導(dǎo)出過(guò)程中，模型會(huì)被執(zhí)行一次，以跟蹤計(jì)算圖的結(jié)構(gòu)。

## 準(zhǔn)備模型輸入
x = torch.randn(batch_size, 1, 224, 224, requires_grad=True)


## 導(dǎo)出模型到 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'}})

四、驗(yàn)證 ONNX 模型的有效性

使用 ONNX 的 API 檢查導(dǎo)出的模型是否有效。

import onnx


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


## 驗(yàn)證模型結(jié)構(gòu)
onnx.checker.check_model(onnx_model)

五、使用 ONNX Runtime 進(jìn)行推理

利用 ONNX Runtime 的 Python API 加載并運(yùn)行 ONNX 模型。

import onnxruntime


## 創(chuàng)建推理會(huì)話(huà)
ort_session = onnxruntime.InferenceSession("super_resolution.onnx")


## 定義張量轉(zhuǎn) NumPy 數(shù)組的輔助函數(shù)
def to_numpy(tensor):
    return tensor.detach().cpu().numpy() if tensor.requires_grad else tensor.cpu().numpy()


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


## 驗(yàn)證 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!")

六、示例應(yīng)用：超分辨率圖像處理

加載示例圖像并進(jìn)行預(yù)處理，使用導(dǎo)出的 ONNX 模型進(jìn)行超分辨率轉(zhuǎn)換。

from PIL import Image
import torchvision.transforms as transforms


## 加載并預(yù)處理圖像
img = Image.open("cat.jpg")  # 請(qǐng)?zhí)鎿Q為實(shí)際圖像文件路徑
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 進(jìn)行推理
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")


## 保存結(jié)果圖像
final_img.save("cat_superres_with_ort.jpg")

七、總結(jié)與拓展

本教程展示了如何將 PyTorch 模型導(dǎo)出為 ONNX 格式，并使用 ONNX Runtime 進(jìn)行推理。ONNX 為深度學(xué)習(xí)模型的跨框架部署提供了強(qiáng)大的支持，而 ONNX Runtime 則確保了模型在不同環(huán)境中的高效運(yùn)行。通過(guò)這一過(guò)程，您可以輕松地將模型部署到各種平臺(tái)，包括 Windows、Linux 和 Mac，以及利用 CPU 和 GPU 的計(jì)算能力。

未來(lái)，您可以進(jìn)一步探索 ONNX Runtime 的高級(jí)功能，如模型優(yōu)化、量化和部署到云端或邊緣設(shè)備。編程獅將持續(xù)為您帶來(lái)更多關(guān)于模型部署和優(yōu)化的優(yōu)質(zhì)教程，助力您的項(xiàng)目高效落地。