亚马逊AWS官方博客

宣布为 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 格式方面参与合作。我们将与 FacebookMicrosoft 和深度学习社区合作,进一步开发 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。

输入图像 输出图像
 

下一步工作?

我们正在与 ONNX 合作伙伴和社区合作,以进一步开发 ONNX,增加更实用的运算符,并扩展 ONNX-MXNet 以包括导出和更大的运算符覆盖范围。我们还将与 Apache MXNet 社区合作,将 ONNX 引入 MXNet 核心 API。

想要了解更多信息?

可在此处获得示例,它源自 ONNX-MXNet GitHub 存储库

查看 ONNX,深入探讨如何对网络图片和运算符进行编码。

欢迎贡献代码!

特别感谢 dmlc/nnvm 社区,本次实施引用了其 ONNX 代码。

Facebook 博客:
https://research.fb.com/amazon-to-join-onnx-ai-format-drive-mxnet-support/

Microsoft 博客:
https://www.microsoft.com/en-us/cognitive-toolkit/blog/2017/11/framework-support-open-ai-ecosystem-grows/


作者简介

Hagay Lupesko 是 AWS 深度学习的工程经理。他专注于构建让开发人员和科学家能够构建智能应用程序的深度学习工具。在业余时间,他喜欢阅读、徒步旅行以及与家人共享天伦之乐。

Roshani Nagmote 是 AWS 深度学习的软件开发人员。她正在开发使所有人都可以利用深度学习的创新工具。在业余时间,她喜欢逗弄可爱的侄子,并且喜欢大型犬。