宣布为 Apache MXNet 推出 ONNX 支持

今天，AWS 宣布推出 ONNX-MXNet，它是一种用于将 Open Neural Network Exchange (ONNX) 深度学习模型导入到 Apache MXNet 的开源 Python 程序包。MXNet 是功能齐全且可扩展的深度学习框架，可以跨 Python、Scala 和 R 等多种热门语言提供 API。通过 MXNet 的 ONNX 格式支持，开发人员可以使用 PyTorch、Microsoft Cognitive Toolkit 或 Caffe2 等其他框架构建和训练模型，然后将这些模型导入 MXNet 中运行，从而使用 MXNet 高度优化且可扩展的引擎进行推理。

我们还很激动地告诉大家，AWS 将在 ONNX 格式方面参与合作。我们将与 Facebook、Microsoft 和深度学习社区合作，进一步开发 ONNX，让深度学习从业人员都可以访问和使用它。

什么是 ONNX？

ONNX 是一种用于对深度学习模型进行编码的开源格式。ONNX 定义神经网络计算图的格式以及图中使用的大量运算符的格式。随着越来越多的框架和硬件供应商支持 ONNX，从事深度学习的开发人员可以轻松地在框架间移动，选择最适合当前任务的框架。

快速入门

我们将介绍如何使用 ONNX-MXNet 将 ONNX 模型导入 MXNet，以及如何使用导入的模型进行推理，从 MXNet 优化的执行引擎中获益。

步骤 1：安装

首先，根据 ONNX 存储库相关说明安装 ONNX。

然后，安装 ONNX-MXNet 程序包：

$ pip install onnx-mxnet

步骤 2：准备要导入的 ONNX 模型

在此示例中，我们将演示导入 Super Resolution 模型，以增加图像的空间分辨率。此模型使用 PyTorch 构建和训练，并且使用 PyTorch 的 ONNX 导出 API 导出到 ONNX。有关模型设计的更多详情，请参阅 PyTorch 示例。

将 Super Resolution ONNX 模型下载到您的工作目录：

$ wget https://s3.amazonaws.com/onnx-mxnet/examples/super_resolution.onnx

步骤 3：将 ONNX 模型导入 MXNet

现在我们已经准备好 ONNX 模型文件，接下来让我们使用 ONNX-MXNet 导入 API 将其导入 MXNet。在 Python 外壳中运行以下代码：

import onnx_mxnet
sym, params = onnx_mxnet.import_model('super_resolution.onnx')

此代码在 Python 运行时中创建了两个实例： sym – 模型的符号图和 params – 模型的权重。现在已完成导入 ONNX 模型，我们生成了一个标准的 MXNet 模型。

步骤 4：准备输入进行推理

下一步，我们将准备输入图像用于推理。以下步骤下载示例图像，然后将其调整为模型期望的输入形状，最后将其转换为 numpy 数组。

从外壳控制台中，将示例输入图像下载到工作目录：

$ wget https://s3.amazonaws.com/onnx-mxnet/examples/super_res_input.jpg

然后安装 Pillow (Python Imaging Library)，我们便可加载和预处理输入图像：

$ pip install Pillow

下一步，从 Python 外壳运行代码，准备 MXNet NDArray 格式的图像：

import numpy as np
import mxnet as mx
from PIL import Image
img = Image.open("super_res_input.jpg").resize((224, 224))
img_ycbcr = img.convert("YCbCr")
img_y, img_cb, img_cr = img_ycbcr.split()
test_image = mx.nd.array(np.array(img_y)[np.newaxis, np.newaxis, :, :])

步骤 5：创建 MXNet 模块

我们将使用 MXNet 模块 API 创建和绑定模块并分配加载权重。
请注意，ONNX-MXNet 导入 API 向输入层分配了名称“input_0”，我们在初始化和绑定模块时将使用该名称。

mod = mx.mod.Module(symbol=sym, data_names=['input_0'], label_names=None)
mod.bind(for_training=False, data_shapes=[('input_0',test_image.shape)])
mod.set_params(arg_params=params, aux_params=None)

步骤 6：运行推理

现在我们已经加载和绑定了 MXNet 模块并分配了训练权重，我们已准备好运行推理。我们将准备单个输入批处理，并通过网络前馈：

from collections import namedtuple
Batch = namedtuple('Batch', ['data'])
mod.forward(Batch([test_image]))
output = mod.get_outputs()[0][0][0]

步骤 7：检查结果

现在，我们来看看对 Super Resolution 图像运行推理后收到的结果：

img_out_y = Image.fromarray(np.uint8((output.asnumpy().clip(0, 255)), mode='L'))
result_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")
result_img.save("super_res_output.jpg")

以下是输入图像和生成的输出图像。如您所见，该模型能够将图像空间分辨率从 256 × 256 增加到 672 × 672。