(可选) 将模型从 PyTorch 导出到 ONNX 并使用 ONNX Runtime 运行
自 PyTorch 2.1 起，ONNX 导出器有两个版本。
torch.onnx.dynamo_export是最新的（仍在测试阶段）导出器，基于随 PyTorch 2.0 发布的 TorchDynamo 技术。
torch.onnx.export基于 TorchScript 后端，自 PyTorch 1.2.0 起已可用。
在本教程中，我们将介绍如何使用 TorchScript 的 torch.onnx.export ONNX 导出器将 PyTorch 中定义的模型转换为 ONNX 格式。
导出的模型将使用 ONNX Runtime 执行。ONNX Runtime 是一个专注于性能的 ONNX 模型引擎，可以在多种平台和硬件（Windows、Linux 和 Mac，以及 CPU 和 GPU）上高效地进行推理。正如这里所解释的，ONNX Runtime 已被证明可以显著提高多个模型的性能。
在本教程中，您需要安装 ONNX 和 ONNX Runtime。您可以通过以下方式获取 ONNX 和 ONNX Runtime 的二进制构建：
%%bash
pipinstallonnxonnxruntime
ONNX Runtime 建议为 PyTorch 使用最新的稳定运行时。
# Some standard imports
importnumpyasnp
fromtorchimport nn
importtorch.utils.model_zooasmodel_zoo
importtorch.onnx
超分辨率是一种提高图像和视频分辨率的技术，广泛应用于图像处理或视频编辑领域。在本教程中，我们将使用一个轻量级的超分辨率模型。
首先，让我们在 PyTorch 中创建一个 SuperResolution 模型。该模型使用 “Real-Time Single Image and Video Super-Resolution Using an Efficient Sub-Pixel Convolutional Neural Network” - Shi et al 中描述的高效子像素卷积层，通过上采样因子来提高图像的分辨率。该模型期望输入图像的 YCbCr 的 Y 分量，并输出超分辨率处理后的上采样 Y 分量。
该模型 直接来自 PyTorch 的示例，未作任何修改：
# Super Resolution model definition in PyTorch
importtorch.nnasnn
importtorch.nn.initasinit
classSuperResolutionNet(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()
defforward(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)
# Create the super-resolution model by using the above model definition.
torch_model = SuperResolutionNet(upscale_factor=3)
通常，您现在会训练这个模型；然而，在本教程中，我们将下载一些预训练的权重。请注意，此模型并未完全训练以达到良好的准确性，仅在此用于演示目的。
在导出模型之前，调用 torch_model.eval() 或 torch_model.train(False) 将模型切换到推理模式非常重要。这是因为像 dropout 或 batchnorm 这样的操作符在推理和训练模式下的行为不同。
# Load pretrained model weights
model_url = 'https://s3.amazonaws.com/pytorch/test_data/export/superres_epoch100-44c6958e.pth'
batch_size = 64 # just a random number
# Initialize model with the pretrained weights
map_location = lambda storage, loc: storage
if torch.cuda.is_available():
map_location = None
torch_model.load_state_dict(model_zoo.load_url(model_url, map_location=map_location))
# set the model to inference mode
torch_model.eval()
在 PyTorch 中，模型的导出可以通过追踪（tracing）或脚本化（scripting）来实现。本教程将以通过追踪导出的模型为例。要导出模型，我们调用 torch.onnx.export() 函数。这将执行模型并记录用于计算输出的操作符的轨迹。由于 export 会运行模型，我们需要提供一个输入张量 x。只要其类型和大小正确，其中的值可以是随机的。需要注意的是，除非指定为动态轴，否则输入大小将在导出的 ONNX 图中对所有输入维度进行固定。在本例中，我们导出一个输入批次大小为 1 的模型，但在 torch.onnx.export() 的 dynamic_axes 参数中将第一个维度指定为动态的。因此，导出的模型将接受大小为 [batch_size, 1, 224, 224] 的输入，其中 batch_size 可以是可变的。
要了解更多关于 PyTorch 导出接口的详细信息，请查看 torch.onnx 文档。
# Input to the model
x = torch.randn(batch_size, 1, 224, 224, requires_grad=True)
torch_out = torch_model(x)
# Export the model
torch.onnx.export(torch_model, # model being run
x, # model input (or a tuple for multiple inputs)
"super_resolution.onnx", # where to save the model (can be a file or file-like object)
export_params=True, # store the trained parameter weights inside the model file
opset_version=10, # the ONNX version to export the model to
do_constant_folding=True, # whether to execute constant folding for optimization
input_names = ['input'], # the model's input names
output_names = ['output'], # the model's output names
dynamic_axes={'input' : {0 : 'batch_size'}, # variable length axes
'output' : {0 : 'batch_size'}})
我们还计算了模型输出 torch_out，它将用于验证我们导出的模型在 ONNX Runtime 中运行时是否计算相同的值。
但在使用 ONNX Runtime 验证模型输出之前，我们将使用 ONNX API 检查 ONNX 模型。首先，onnx.load("super_resolution.onnx") 将加载保存的模型，并输出一个 onnx.ModelProto 结构（一种用于打包机器学习模型的顶级文件/容器格式。更多信息请参见 onnx.proto 文档）。然后，onnx.checker.check_model(onnx_model) 将验证模型的结构，并确认模型具有有效的模式。通过检查模型的版本、图的结构以及节点及其输入和输出，来验证 ONNX 图的有效性。
importonnx
onnx_model = onnx.load("super_resolution.onnx")
onnx.checker.check_model(onnx_model)
现在，让我们使用 ONNX Runtime 的 Python API 来计算输出。这一部分通常可以在单独的进程或另一台机器上完成，但我们将继续在同一进程中操作，以便验证 ONNX Runtime 和 PyTorch 为网络计算的值是否一致。
为了使用 ONNX Runtime 运行模型，我们需要为模型创建一个推理会话，并选择配置参数（这里我们使用默认配置）。会话创建完成后，我们使用 run() API 来评估模型。此调用的输出是一个列表，包含由 ONNX Runtime 计算的模型输出。
importonnxruntime
ort_session = onnxruntime.InferenceSession("super_resolution.onnx", providers=["CPUExecutionProvider"])
defto_numpy(tensor):
return tensor.detach().cpu().numpy() if tensor.requires_grad else tensor.cpu().numpy()
# compute ONNX Runtime output prediction
ort_inputs = {ort_session.get_inputs()[0].name: to_numpy(x)}
ort_outs = ort_session.run(None, ort_inputs)
# compare ONNX Runtime and PyTorch results
np.testing.assert_allclose(to_numpy(torch_out), ort_outs[0], rtol=1e-03, atol=1e-05)
print("Exported model has been tested with ONNXRuntime, and the result looks good!")
我们应该看到 PyTorch 和 ONNX Runtime 运行的输出在给定精度（rtol=1e-03 和 atol=1e-05）下数值上是匹配的。顺便提一下，如果它们不匹配，那么可能是 ONNX 导出器存在问题，请在这种情况下联系我们。
模型之间的时间对比
由于 ONNX 模型针对推理速度进行了优化，因此在 ONNX 模型上运行相同数据，而不是在原生 PyTorch 模型上运行，应该能带来最高 2 倍的性能提升。批量大小越大，性能提升越明显。
importtime
x = torch.randn(batch_size, 1, 224, 224, requires_grad=True)
start = time.time()
torch_out = torch_model(x)
end = time.time()
print(f"Inference of Pytorch model used {end-start} seconds")
ort_inputs = {ort_session.get_inputs()[0].name: to_numpy(x)}
start = time.time()
ort_outs = ort_session.run(None, ort_inputs)
end = time.time()
print(f"Inference of ONNX model used {end-start} seconds")
使用 ONNX Runtime 在图像上运行模型
到目前为止，我们已经从 PyTorch 导出了一个模型，并展示了如何加载它并在 ONNX Runtime 中使用一个虚拟张量作为输入来运行它。
在本教程中，我们将使用一张广泛使用的著名猫咪图片，如下图所示
首先，我们加载图像，并使用标准的 PIL Python 库对其进行预处理。请注意，这种预处理是训练/测试神经网络时处理数据的标准做法。
我们首先将图像调整为适合模型输入的大小（224x224）。然后，我们将图像分解为 Y、Cb 和 Cr 分量。这些分量分别代表灰度图像（Y）、蓝色差值（Cb）和红色差值（Cr）色度分量。由于 Y 分量对人眼更为敏感，我们关注的是这个分量，并对其进行转换。提取 Y 分量后，我们将其转换为张量，作为我们模型的输入。
fromPILimport Image
importtorchvision.transformsastransforms
img = Image.open("./_static/img/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]
此时，模型的输出是一个张量。现在，我们将处理模型的输出，从输出张量中重建最终的输出图像，并保存该图像。后处理步骤参考了PyTorch实现的超分辨率模型，具体实现在此。
img_out_y = Image.fromarray(np.uint8((img_out_y[0] * 255.0).clip(0, 255)[0]), mode='L')
# get the output image follow post-processing step from PyTorch implementation
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")
# Save the image, we will compare this with the output image from mobile device
final_img.save("./_static/img/cat_superres_with_ort.jpg")
# Save resized original image (without super-resolution)
img = transforms.Resize([img_out_y.size[0], img_out_y.size[1]])(img)
img.save("cat_resized.jpg")
以下是两张图片的对比：
低分辨率图像
超分辨率后的图像
ONNX Runtime 作为一个跨平台引擎，您可以在多个平台上运行它，并且支持 CPU 和 GPU。
ONNX Runtime 还可以部署到云端，通过 Azure 机器学习服务进行模型推理。更多信息请参阅这里。
关于 ONNX Runtime 性能的更多信息请参阅这里。
关于 ONNX Runtime 的更多信息请参阅这里。