Relay 神经网络推理

Relay 神经网络推理#

import numpy as np
from PIL import Image
import tvm
from tvm import relay

使用 PIL 读取图像，mxnet 获取预训练的神经网络，以及在 TVM 中的 relay 模块 [Roesch et al., 2019] 转换和优化神经网络。

Relay 是 TVM 中表示神经网络的高级中间表示（intermediate representation，简称 IR）。

获得预训练模型#

预训练模型是指在数据集上训练好参数的神经网络。在这里，通过从 MXNet 的模型动物园 [Roesch et al., 2019] 指定 pretrained=True 来下载和加载 ResNet-18 模型。如果你想了解该模型，可以参考 Chapter 7.6 in D2L。

参见

MXNet model zoo 可以找到更多信息。或者参考 GluonCV 和 GluonNLP 使用更多的计算机视觉和自然语言模型。

from mxnet.gluon.model_zoo.vision import get_model

model_name = 'resnet18_v2'
model = get_model(model_name, pretrained=True)
len(model.features), model.output

(13, Dense(512 -> 1000, linear))

加载的模型在 Imagenet 1K 数据集上训练，该数据集包含 1000 个类中大约 100 万张自然物体图像。模型分为两部分，主体部分 model.features 包含 13 个块，输出层是 dense 层，有 1000 个输出。

下面的代码块为 Imagenet 数据集中的每个类加载文本标签。

from d2py.download import get_github_content

labels = get_github_content("xinetzone", "meta-data", "vision/imagenet1k_labels.txt")
labels = labels.split("\n")

数据预处理示例#

读取样本图像并调整其大小，即 224 像素的宽度和高度，这是训练的神经网络的尺寸。

from d2py.download import iter_github_bytes

# 获取图片
image_byte = next(iter_github_bytes("xinetzone", "meta-data", "vision/images"))
with Image.open(image_byte) as im:
    image = np.array(im.resize((224, 224)))

根据 动物园模型页面。图像像素在每个颜色通道上进行归一化，数据布局为 (batch, RGB channels, height, width)。下面的函数对输入图像进行变换，使其满足要求。

def image_preprocessing(image):
    mean_rgb = [123.68, 116.779, 103.939]
    std_rgb = [58.393, 57.12, 57.375]
    image = image - np.array([mean_rgb])
    image /= np.array([std_rgb])
    image = image.transpose((2, 0, 1))
    image = image[np.newaxis, :]
    return image.astype('float32')

x = image_preprocessing(image)
x.shape

(1, 3, 224, 224)

编译预训练模型#

为了编译模型，使用 from_mxnet 方法导入 MXNet 模型并变换为 Relay IR。在该方法为模型提供输入数据形状。一些神经网络可能需要稍后确定的数据形状的某些维数。然而，在 ResNet 模型中，数据形状是固定的，这使得编译器更容易实现高性能。推荐固定的数据形状。在后面的章节中，只涉及动态数据形状（即在运行时确定的某些维度）。

input_name = 'data'
relay_mod, relay_params = relay.frontend.from_mxnet(model, {input_name: x.shape})
type(relay_mod), type(relay_params)

(tvm.ir.module.IRModule, dict)

from_mxnet() 方法将返回 program relay_mod，它是 relay 模块，以及 relay_params 参数字典，它将字符串键映射到 TVM ndarray。

查看每个参数：

type(relay_params['resnetv20_dense0_weight'])

tvm.runtime.ndarray.NDArray

接下来，将模块 lower 到一些可以被 llvm 后端使用的低级 IR。LLVM 定义了被多种编程语言采用的 IR。然后，LLVM 编译器能够将生成的程序编译成 CPU 的机器码。

此外，将优化级别设置为级别 3。您可能会收到警告消息，并不是每个算子都得到了很好的优化，现在可以忽略它。

target = 'llvm'
# 将模型与标准优化一起构建成 TVM 库
with tvm.transform.PassContext(opt_level=3):
    lib = relay.build(relay_mod, target, params=relay_params)

lib

One or more operators have not been tuned. Please tune your model for better performance. Use DEBUG logging level to see more details.

<tvm.relay.backend.executor_factory.GraphExecutorFactoryModule at 0x7f68f2e05ea0>

编译模块 lib 有：

  • params：映射参数名到权重的字典。

  • graph_json：将被 graph compiler 部署成 JSON 格式输出的 json graph。

  • function_metadata：字符串到 FunctionInfo 的 Map。这保存了映射函数名称到它们的信息。graph 中可以包含指向 libmod 中 PackedFunc 名称的算子 (tvm_op)。

  • ir_mod：构建的 IR 模块。

  • executor：Executor 的内部表示。

  • libmod_name：模块名称。

type(lib.graph_json), type(lib.params)

(str, dict)

len(lib.params), len(relay_params)

(43, 98)

可以将 lib.graph_json 中间转化为 Python 字典：

graph_bunch = eval(lib.graph_json)
type(graph_bunch)

dict

lib.libmod_name # 模块名称

'default'

获取模块所对应的函数：

func = lib[lib.libmod_name]
func

<tvm.runtime.packed_func.PackedFunc at 0x7f68f571f140>

func 是在 TVM 中使用的 PackedFunc 对象。

参见

更多信息可参考：全局函数。

从 TVM 库中创建 TVM graph 运行时模块。包含已编译算子机器码的库，带有可以从目标构建的设备上下文（ctx）。这里的设备是 CPU，由 llvm 指定。

from tvm.contrib.graph_executor import GraphModule

ctx = tvm.device(target, 0)
module = GraphModule(func(ctx))
module

<tvm.contrib.graph_executor.GraphModule at 0x7f6cc413dcc0>

推理#

借由创建的运行时模块来运行模型推理，即神经网络的前向传播。使用 set_input 加载参数，并通过输入数据运行工作负载（workload）。

# dtype = "float32"
module.set_input(input_name, x)
module.run()

也可以直接使用 run 加载参数：

module.run(**{input_name:x})

由于此网络只有单个输出层，可以通过 get_output(0) 得到 (1, 1000) 形状矩阵。最终输出长度为 1000 的 NumPy 向量。

tvm_output = module.get_output(0).numpy()
tvm_output.shape

(1, 1000)

该向量包含每个类的预测置信度得分（confidence score）。注意，预训练的模型没有 softmax 算子，所以这些得分没有映射到概率 (0,1) 中。现在可以找到两个最大的分数并报告它们的标签。

scores = tvm_output[0]
a = np.argsort(scores)[-1:-5:-1]
labels[a[0]], labels[a[1]]

('tiger cat', 'Egyptian cat')

保存已编译的库#

可以保存 relay.build 的输出到磁盘以便以后重用它们。下面的代码块保存了 json 字符串、库和参数。

from d2py.utils.file import mkdir

mkdir("outputs") # 创建目录
!rm -rf outputs/resnet18*
graph_fn, mod_fn, params_fn = ['outputs/'+model_name+ext for ext in ('.json','.tar','.params')]
lib.export_library(mod_fn)
with open(graph_fn, 'w') as f:
    f.write(lib.graph_json)
with open(params_fn, 'wb') as f:
    f.write(relay.save_param_dict(lib.params))

!ls -alht outputs/resnet18*

-rw-rw-r-- 1 ai ai 45M 9月  23 14:03 outputs/resnet18_v2.params
-rw-rw-r-- 1 ai ai 36K 9月  23 14:03 outputs/resnet18_v2.json
-rw-rw-r-- 1 ai ai 42M 9月  23 14:03 outputs/resnet18_v2.tar

加载已保存的模块。

with open(graph_fn) as fp:
    loaded_graph = fp.read()

loaded_mod = tvm.runtime.load_module(mod_fn)

with open(params_fn, "rb") as fp:
    loaded_params = fp.read()

可以使用 create() 加载运行时模块：

from tvm.contrib.graph_executor import create

loaded_rt = create(loaded_graph, loaded_mod, ctx)

也可以像前面一样构造运行时模块：

loaded_rt = GraphModule(loaded_mod["default"](ctx))

验证结果：

import numpy as np

loaded_rt.load_params(loaded_params)
loaded_rt.run(data=tvm.nd.array(x))
loaded_scores = loaded_rt.get_output(0).numpy()[0]
np.testing.assert_allclose(loaded_scores, scores)

小结

  • 可以利用 TVM 的 relay 将神经网络转换并编译成模块以进行模型推理。

  • 可以将编译后的模块保存到磁盘中，以方便将来的部署。