PyTorch模型部署流程(ONNX Runtime)

模型部署指让训练好的深度学习模型在特定环境中运行的过程。模型部署会面临的难题：

运行模型所需的环境难以配置。深度学习模型通常是由一些框架编写，比如 PyTorch、TensorFlow。由于框架规模、依赖环境限制，框架不适合在手机、开发板等生产环境中安装。

深度学习模型的结构通常比较大，需要大量算力才能满足实时运行需求，运行效率需要优化。

因为这些难题的存在，模型部署不能靠简单的环境配置与安装完成。目前模型部署有一条流行的流水线：

为了让模型最终能够部署到某一环境上，可以使用任意一种深度学习框架来定义网络结构，并通过训练确定网络中的参数。之后，模型的结构和参数会被转换成一种只描述网络结构的中间表示，一些针对网络结构的优化会在中间表示上进行。最后，用面向硬件的高性能编程框架（如 CUDA，OpenCL）编写，能高效执行深度学习网络中算子的推理引擎会把中间表示转换成特定的文件格式，并在对应硬件平台上高效运行模型。

创建PyTorch模型

a.配置环境

# 创建预安装 Python 3.7 的名叫 deploy 虚拟环境
conda create -n deploy python=3.7 -y
# 进入虚拟环境
conda activate deploy
# 安装 Gpu 版本的 PyTorch
# 从官网中选择合适的配置并复制下载路径---https://pytorch.org/get-started/locally/
cconda install pytorch torchvision cudatoolkit=11.3 -c pytorch
# 安装 ONNX Runtime, ONNX, OpenCV
pip install onnxruntime onnx opencv-python

b.创建PyTorch模型

import os

import cv2
import numpy as np
import requests
import torch
import torch.onnx
from torch import nn

class SuperResolutionNet(nn.Module):
    def __init__(self, upscale_factor):
        super().__init__()
        self.upscale_factor = upscale_factor
        self.img_upsampler = nn.Upsample(
            scale_factor=self.upscale_factor,
            mode='bicubic',
            align_corners=False)

        self.conv1 = nn.Conv2d(3,64,kernel_size=9,padding=4)
        self.conv2 = nn.Conv2d(64,32,kernel_size=1,padding=0)
        self.conv3 = nn.Conv2d(32,3,kernel_size=5,padding=2)

        self.relu = nn.ReLU()

   def forward(self, x):
        x = self.img_upsampler(x)
        out = self.relu(self.conv1(x))
        out = self.relu(self.conv2(out))
        out = self.conv3(out)
        return out

# Download checkpoint and test image
urls = ['https://download.openmmlab.com/mmediting/restorers/srcnn/srcnn_x4k915_1x16_1000k_div2k_20200608-4186f232.pth',
        'https://raw.githubusercontent.com/open-mmlab/mmediting/master/tests/data/face/000001.png']
names = ['srcnn.pth', 'face.png']
for url, name in zip(urls, names):
    if not os.path.exists(name):
        open(name, 'wb').write(requests.get(url).content)

def init_torch_model():
    torch_model = SuperResolutionNet(upscale_factor=3)

    state_dict = torch.load('srcnn.pth')['state_dict']

    # Adapt the checkpoint
    for old_key in list(state_dict.keys()):
        new_key = '.'.join(old_key.split('.')[1:])
        state_dict[new_key] = state_dict.pop(old_key)

    torch_model.load_state_dict(state_dict)
    torch_model.eval()
    return torch_model

model = init_torch_model()
input_img = cv2.imread('face.png').astype(np.float32)

# HWC to NCHW
input_img = np.transpose(input_img, [2, 0, 1])
input_img = np.expand_dims(input_img, 0)

# Inference
torch_output = model(torch.from_numpy(input_img)).detach().numpy()

# NCHW to HWC
torch_output = np.squeeze(torch_output, 0)
torch_output = np.clip(torch_output, 0, 255)
torch_output = np.transpose(torch_output, [1, 2, 0]).astype(np.uint8)

# Show image
cv2.imwrite("face_torch.png", torch_output)

代码创建一个经典的超分辨率网络 SRCNN。SRCNN 先把图像上采样到对应分辨率，再用 3 个卷积层处理图像。如果脚本正常运行，一幅超分辨率的人脸照片会保存在 face_torch.png 中。

在 PyTorch 模型测试正确后，我们来正式开始部署这个模型。我们下一步的任务是把 PyTorch 模型转换成用中间表示 ONNX 描述的模型。

c.将PyTorch模型转换成用ONNX描述的模型

ONNX（Open Neural Network Exchange）是 Facebook 和微软在 2017 年共同发布的，用于标准描述计算图的一种格式。目前，在数家机构的共同维护下，ONNX 已经对接了多种深度学习框架和多种推理引擎。因此，ONNX 被当成了深度学习框架到推理引擎的桥梁，就像编译器的中间语言一样。由于各框架兼容性不一，通常只用 ONNX 表示更容易部署的静态图。

# PyTorch 的模型转换成 ONNX 格式的模型
x = torch.randn(1, 3, 256, 256)

with torch.no_grad():
    torch.onnx.export(
        model,
        x,
        "srcnn.onnx",
        opset_version=11,
        input_names=['input'],
        output_names=['output'])

其中torch.onnx.export 是 PyTorch 自带的把模型转换成 ONNX 格式的函数。前三个参数分别是要转换的模型、模型的任意一组输入、导出的 ONNX 文件的文件名。

从 PyTorch 的模型到 ONNX 的模型，PyTorch提供了一种叫做追踪（trace）的模型转换方法：给定一组输入，再实际执行一遍模型，即把这组输入对应的计算图记录下来，保存为 ONNX 格式。export 函数用的就是追踪导出方法，需要给任意一组输入，让模型跑起来。测试图片是三通道256x256 大小的，这里也构造一个同样形状的随机张量。

opset_version 表示 ONNX 算子集的版本。input_names, output_names 是输入、输出 tensor 的名称。代码运行成功，目录下会新增一个 srcnn.onnx 的 ONNX 模型文件

# 验证模型文件是否正确，直接加在前面的代码后面就行
import onnx

onnx_model = onnx.load("srcnn.onnx")
try:
    onnx.checker.check_model(onnx_model)
except Exception:
  print("Model incorrect")
else:
    print("Model correct")

d.在ONNX Runtime上运行模型和推理

推理引擎-ONNX Runtime 是由微软维护的一个跨平台机器学习推理加速器。ONNX Runtime 是直接对接ONNX的，即ONNX Runtime可以直接读取并运行.onnx文件，而不需要再把.onnx格式的文件转换成其他格式的文件。也就是说，对于 PyTorch - ONNX - ONNX Runtime 这条部署流水线，只要在目标设备中得到 .onnx 文件，并在 ONNX Runtime 上运行模型，模型部署就算大功告成了。

ONNX Runtime 提供了 Python 接口。

# ONNX Runtime完成模型推理，还是在之前脚本后添加代码
import onnxruntime

ort_session = onnxruntime.InferenceSession("srcnn.onnx")
ort_inputs = {'input': input_img}
ort_output = ort_session.run(['output'], ort_inputs)[0]

ort_output = np.squeeze(ort_output, 0)
ort_output = np.clip(ort_output, 0, 255)
ort_output = np.transpose(ort_output, [1, 2, 0]).astype(np.uint8)
cv2.imwrite("face_ort.png", ort_output)

onnxruntime.InferenceSession 用于获取一个 ONNX Runtime 推理器，其参数是用于推理的 ONNX 模型文件。推理器的 run 方法用于模型推理，其第一个参数为输出张量名的列表，第二个参数为输入值的字典。其中输入值字典的 key 为张量名，value 为 numpy 类型的张量值。输入输出张量的名称需要和 torch.onnx.export 中设置的输入输出名对应。

如果代码正常运行的话，另一幅超分辨率照片会保存在 face_ort.png 中。这幅图片和刚刚得到的 face_torch.png 是一样的。即ONNX Runtime 成功运行了 SRCNN 模型，完成模型部署。

以后再想实现超分辨率的操作，只需要提供一个srcnn.onnx文件，并帮助用户配置好 ONNX Runtime 的 Python 环境，用几行代码就可以运行模型。或者可以利用 ONNX Runtime 编译出一个可以直接执行模型的应用程序。只需给用户提供 ONNX 模型文件，并让用户在应用程序选择要执行的 ONNX 模型文件名就可以运行模型了。

PyTorch模型部署流程(ONNX Runtime)

猜你喜欢