端到端模型执行 ============== 前言 ---- .. figure:: ../img/tensor_func_linear_relu.png 大多数 MLC 过程可以看作是张量函数之间的转换。我们打算在上一节课中回答的主要问题是: - 什么是可能的张量函数抽象表达形式? - 什么是可能的张量函数变换? 在上一课中,我们主要关注点是元张量函数。在本次讲座中,我们将讨论如何构建端到端模型。 准备工作 -------- 首先,我们将导入必要的依赖项。 .. raw:: latex \diilbookstyleinputcell .. code:: python import IPython import numpy as np import tvm from tvm import relax from tvm.ir.module import IRModule from tvm.script import relax as R from tvm.script import tir as T 加载数据集 ~~~~~~~~~~ 作为一个具体的例子,我们将在 fashion MNIST 数据集上使用一个模型。 以下代码从 ``torchvision`` 下载并准备数据并转换成NumPy数组。 .. raw:: latex \diilbookstyleinputcell .. code:: python import torch import torchvision test_data = torchvision.datasets.FashionMNIST( root="data", train=False, download=True, transform=torchvision.transforms.ToTensor() ) test_loader = torch.utils.data.DataLoader(test_data, batch_size=1, shuffle=True) class_names = ['T-shirt/top', 'Trouser', 'Pullover', 'Dress', 'Coat', 'Sandal', 'Shirt', 'Sneaker', 'Bag', 'Ankle boot'] img, label = next(iter(test_loader)) img = img.reshape(1, 28, 28).numpy() .. raw:: latex \diilbookstyleoutputcell .. parsed-literal:: :class: output Downloading http://fashion-mnist.s3-website.eu-central-1.amazonaws.com/train-images-idx3-ubyte.gz Downloading http://fashion-mnist.s3-website.eu-central-1.amazonaws.com/train-images-idx3-ubyte.gz to data/FashionMNIST/raw/train-images-idx3-ubyte.gz 100.0% Extracting data/FashionMNIST/raw/train-images-idx3-ubyte.gz to data/FashionMNIST/raw Downloading http://fashion-mnist.s3-website.eu-central-1.amazonaws.com/train-labels-idx1-ubyte.gz Downloading http://fashion-mnist.s3-website.eu-central-1.amazonaws.com/train-labels-idx1-ubyte.gz to data/FashionMNIST/raw/train-labels-idx1-ubyte.gz 100.0% Extracting data/FashionMNIST/raw/train-labels-idx1-ubyte.gz to data/FashionMNIST/raw Downloading http://fashion-mnist.s3-website.eu-central-1.amazonaws.com/t10k-images-idx3-ubyte.gz Downloading http://fashion-mnist.s3-website.eu-central-1.amazonaws.com/t10k-images-idx3-ubyte.gz to data/FashionMNIST/raw/t10k-images-idx3-ubyte.gz 100.0% Extracting data/FashionMNIST/raw/t10k-images-idx3-ubyte.gz to data/FashionMNIST/raw Downloading http://fashion-mnist.s3-website.eu-central-1.amazonaws.com/t10k-labels-idx1-ubyte.gz 100.0%Downloading http://fashion-mnist.s3-website.eu-central-1.amazonaws.com/t10k-labels-idx1-ubyte.gz to data/FashionMNIST/raw/t10k-labels-idx1-ubyte.gz Extracting data/FashionMNIST/raw/t10k-labels-idx1-ubyte.gz to data/FashionMNIST/raw 我们可以绘制出我们希望能够预测的图像。 .. raw:: latex \diilbookstyleinputcell .. code:: python import matplotlib.pyplot as plt plt.figure() plt.imshow(img[0]) plt.colorbar() plt.grid(False) plt.show() print("Class:", class_names[label[0]]) .. figure:: output_index_e758e2_5_0.png .. raw:: latex \diilbookstyleoutputcell .. parsed-literal:: :class: output Class: Sandal 下载模型参数 ~~~~~~~~~~~~ .. raw:: latex \diilbookstyleinputcell .. code:: python # Hide outputs !wget https://github.com/mlc-ai/web-data/raw/main/models/fasionmnist_mlp_params.pkl 端到端模型整合 -------------- 在本章中,我们将使用以下模型作为示例。这是一个两层神经网络,由两个全连接层和一个有 relu 激活层组成。为了简化问题,我们删除了最终的 softmax 层。输出分数是未标准化的,但最大值仍然对应于最可能的类别。 .. figure:: ../img/e2e_fashionmnist_mlp_model.png 让我们从模型的 Numpy 实现开始。 .. raw:: latex \diilbookstyleinputcell .. code:: python def numpy_mlp(data, w0, b0, w1, b1): lv0 = data @ w0.T + b0 lv1 = np.maximum(lv0, 0) lv2 = lv1 @ w1.T + b1 return lv2 .. raw:: latex \diilbookstyleinputcell .. code:: python import pickle as pkl mlp_params = pkl.load(open("fasionmnist_mlp_params.pkl", "rb")) res = numpy_mlp(img.reshape(1, 784), mlp_params["w0"], mlp_params["b0"], mlp_params["w1"], mlp_params["b1"]) print(res) pred_kind = res.argmax(axis=1) print(pred_kind) print("NumPy-MLP Prediction:", class_names[pred_kind[0]]) .. raw:: latex \diilbookstyleoutputcell .. parsed-literal:: :class: output [[-23.49592 -24.32377 -17.936213 -15.301368 -18.97627 22.026953 -25.760815 -19.639427 -3.5585783 -10.392135 ]] [5] NumPy-MLP Prediction: Sandal 上面的示例代码显示了利用高层NumPy数组操作,执行端到端模型执行的过程。 回到 MLC 的视角中,我们希望了解这些数组计算的细节。 为了说明底层细节,我们将再次用底层 NumPy 编写示例: - 我们将在使用循环而不是数组函数来演示循环计算。 - 我们总是通过 ``np.empty`` 显式分配数组并传递它们。(不包含数组初始化) 下面的代码展示了同一模型的底层 NumPy 实现: .. raw:: latex \diilbookstyleinputcell .. code:: python def lnumpy_linear0(X: np.ndarray, W: np.ndarray, B: np.ndarray, Z: np.ndarray): Y = np.empty((1, 128), dtype="float32") for i in range(1): for j in range(128): for k in range(784): if k == 0: Y[i, j] = 0 Y[i, j] = Y[i, j] + X[i, k] * W[j, k] for i in range(1): for j in range(128): Z[i, j] = Y[i, j] + B[j] def lnumpy_relu0(X: np.ndarray, Y: np.ndarray): for i in range(1): for j in range(128): Y[i, j] = np.maximum(X[i, j], 0) def lnumpy_linear1(X: np.ndarray, W: np.ndarray, B: np.ndarray, Z: np.ndarray): Y = np.empty((1, 10), dtype="float32") for i in range(1): for j in range(10): for k in range(128): if k == 0: Y[i, j] = 0 Y[i, j] = Y[i, j] + X[i, k] * W[j, k] for i in range(1): for j in range(10): Z[i, j] = Y[i, j] + B[j] def lnumpy_mlp(data, w0, b0, w1, b1): lv0 = np.empty((1, 128), dtype="float32") lnumpy_linear0(data, w0, b0, lv0) lv1 = np.empty((1, 128), dtype="float32") lnumpy_relu0(lv0, lv1) out = np.empty((1, 10), dtype="float32") lnumpy_linear1(lv1, w1, b1, out) return out result =lnumpy_mlp( img.reshape(1, 784), mlp_params["w0"], mlp_params["b0"], mlp_params["w1"], mlp_params["b1"]) pred_kind = result.argmax(axis=1) print("Low-level Numpy MLP Prediction:", class_names[pred_kind[0]]) .. raw:: latex \diilbookstyleoutputcell .. parsed-literal:: :class: output Low-level Numpy MLP Prediction: Sandal 在 TVMScript 中构建端到端 IRModule ---------------------------------- 考虑到低级 NumPy 示例,现在我们准备为端到端模型执行引入 MLC 抽象。 下面的代码块显示了模型的 TVMScript 实现。 有了底级 NumPy 代码作为参考,现在我们准备利用 MLC 抽象来实现端到端模型运行。下面的代码展示了模型的 TVMScript 实现。 .. raw:: latex \diilbookstyleinputcell .. code:: python @tvm.script.ir_module class MyModule: @T.prim_func def relu0(x: T.handle, y: T.handle): n = T.int64() X = T.match_buffer(x, (1, n), "float32") Y = T.match_buffer(y, (1, n), "float32") for i, j in T.grid(1, n): with T.block("Y"): vi, vj = T.axis.remap("SS", [i, j]) Y[vi, vj] = T.max(X[vi, vj], T.float32(0)) @T.prim_func def linear0(x: T.handle, w: T.handle, b: T.handle, z: T.handle): m, n, k = T.int64(), T.int64(), T.int64() X = T.match_buffer(x, (1, m), "float32") W = T.match_buffer(w, (n, m), "float32") B = T.match_buffer(b, (n, ), "float32") Z = T.match_buffer(z, (1, n), "float32") Y = T.alloc_buffer((1, n), "float32") for i, j, k in T.grid(1, n, m): with T.block("Y"): vi, vj, vk = T.axis.remap("SSR", [i, j, k]) with T.init(): Y[vi, vj] = T.float32(0) Y[vi, vj] = Y[vi, vj] + X[vi, vk] * W[vj, vk] for i, j in T.grid(1, n): with T.block("Z"): vi, vj = T.axis.remap("SS", [i, j]) Z[vi, vj] = Y[vi, vj] + B[vj] @R.function def main(x: R.Tensor((1, "m"), "float32"), w0: R.Tensor(("n", "m"), "float32"), b0: R.Tensor(("n", ), "float32"), w1: R.Tensor(("k", "n"), "float32"), b1: R.Tensor(("k", ), "float32")): m, n, k = T.int64(), T.int64(), T.int64() with R.dataflow(): lv0 = R.call_dps_packed("linear0", (x, w0, b0), R.Tensor((1, n), "float32")) lv1 = R.call_dps_packed("relu0", (lv0, ), R.Tensor((1, n), "float32")) out = R.call_dps_packed("linear0", (lv1, w1, b1), R.Tensor((1, k), "float32")) R.output(out) return out 上面的代码包含各种函数:我们在上一课中看到的元张量函数(\ ``T.prim_func``\ )和一个新的\ ``R.function``\ (Relax 函数)。 Relax 函数是一种表示上层神经网络执行的全新抽象。 同样,对比查看 TVMScript 代码和底级 NumPy 代码会对我们理解代码有很大帮助。我们这里将详细介绍它们中的每一个元素。由于我们已经了解了元张量函数,因此我们将专注于高层的执行部分。 .. figure:: ../img/e2e_compare_to_lnumpy.png 计算图 ~~~~~~ .. figure:: ../img/e2e_computational_graph_call_tir.png 使用图 (graph) 来可视化高层模型执行通常很有帮助。 上图是 ``main`` 函数的计算图视图: - 图中的每个框都对应于计算操作。 - 箭头对应于中间张量的输入输出。 我们在之前的课程中已经看到了这种可视化。 图本身可以看作是一种抽象,在机器学习框架中通常称为\ **计算图 (computational graph)**\ 。 ``call_dps_packed`` ~~~~~~~~~~~~~~~~~~~ 您可能已经注意到的一件事是,计算图中的每个操作步骤都包含一个\ ``R.call_dps_packed``\ 操作。 这是引入元张量函数的过程: .. raw:: latex \diilbookstyleinputcell .. code:: python lv0 = R.call_dps_packed("linear0", (x, w0, b0), R.Tensor((1, n), dtype="float32")) 为了解释 ``R.call_dps_packed`` 的含义,让我们回顾一下操作的等效底层 NumPy 实现,如下所示: .. raw:: latex \diilbookstyleinputcell .. code:: python def lnumpy_call_dps_packed(prim_func, inputs, shape, dtype): res = np.empty(shape, dtype=dtype) prim_func(*inputs, res) return res 具体来说,\ ``call_dps_packed`` 接受一个元函数 (``prim_func``) 的输入列表,并分配一个输出张量\ ``res``\ ,然后将输入和输出传递给\ ``prim_func``\ 。 执行 ``prim_func`` 后,结果会填充到 ``res`` 中,然后我们可以返回结果。 请注意,\ ``lnumpy_call_dps_packed`` 只是一个参考实现,以显示 ``R.call_dps_packed`` 的含义。 在实际应用中,可以有不同的底层方法来优化执行。 例如,我们可能会选择提前分配所有输出内存,然后运行,我们将在以后的课程中介绍。 一个很自然的问题:为什么我们需要 ``call_dps_packed``\ ? 这是因为我们的元张量函数采用以下调用约定: .. raw:: latex \diilbookstyleinputcell .. code:: python def low_level_prim_func(in0, in1, ..., out): # implementations 此约定称为\ **目标传递 (destination passing)**\ 。 这个想法是输入和输出在外部显式分配并传递给底层元函数。 这种风格通常用于底层库设计,因此高层框架可以处理内存分配决策。 请注意,并非所有张量操作都可以通过这种方式呈现(比如,有些操作的输出形状取决于输入)。 然而,在通常的实践中,如果可能的话,以这种风格编写底层函数通常是有帮助的。 虽然可以通过显式分配中间结果并调用每个函数来将目标传递的函数组装在一起,但很难将以下代码转换为计算图形式。 .. raw:: latex \diilbookstyleinputcell .. code:: python def lnumpy_mlp(data, w0, b0, w1, b1): lv0 = np.empty((1, 128), dtype="float32") lnumpy_linear0(data, w0, b0, lv0) lv1 = np.empty((1, 128), dtype="float32") lnumpy_relu0(lv0, lv1) out = np.empty((1, 10), dtype="float32") lnumpy_linear1(lv1, w1, b1, out) return out .. figure:: ../img/e2e_computational_graph_numpy.png 我们当然可以尝试一下 :) 上图是通过简单地将函数输入连接到函数来将\ ``lnumpy_mlp``\ 拟合成“类计算图”形式的一种可能的“失败尝试”。 我们可以发现它失去了之前计算图的一些特性。 具体来说,计算图通常具有以下性质: - 框的每个输入边对应于操作的输入; - 每个出边对应于操作的输出; - 每个操作可以任意重新排序,直到边缘的拓扑顺序。 当然,我们仍然可以通过引入输入边和输出边来概括图定义,这会使抽象的定义和变换复杂化。 所以回到\ ``call_dps_packed``\ ,这里的关键思想是我们想要隐藏可能的分配或对函数的显式写入。 用更正式的术语来说,我们希望函数是 **pure** 或 **side-effect free**\ 。(译者注:“pure”和“side-effect”是 PL 中的术语,译者不确定中文的准确名称,故不进行翻译。欢迎社区中的专业人士参与完善) 如果一个函数只从其输入中读取并通过其输出返回结果,它不会改变程序的其他部分(例如递增全局计数器),那么它是\ **pure**\ 或\ **side-effect free**\ 的。 **call_dps_packed** 使我们能够隐藏调用低层元函数细节,并将它们应用到计算图中。 我们还可以在底层 NumPy 中看到 ``call_dps_packed`` 的作用。 现在我们已经定义了 ``lnumpy_call_dps_packed``\ ,我们可以将底层 NumPy 代码重写为: .. raw:: latex \diilbookstyleinputcell .. code:: python def lnumpy_mlp_with_call_dps_packed(data, w0, b0, w1, b1): lv0 = lnumpy_call_dps_packed(lnumpy_linear0, (data, w0, b0), (1, 128), dtype="float32") lv1 = lnumpy_call_dps_packed(lnumpy_relu0, (lv0, ), (1, 128), dtype="float32") out = lnumpy_call_dps_packed(lnumpy_linear1, (lv1, w1, b1), (1, 10), dtype="float32") return out result = lnumpy_mlp_with_call_dps_packed( img.reshape(1, 784), mlp_params["w0"], mlp_params["b0"], mlp_params["w1"], mlp_params["b1"]) pred_kind = np.argmax(result, axis=1) print("Low-level Numpy with CallTIR Prediction:", class_names[pred_kind[0]]) .. raw:: latex \diilbookstyleoutputcell .. parsed-literal:: :class: output Low-level Numpy with CallTIR Prediction: Sandal 实际上,最底层的实现会有显式的内存分配,所以\ ``call_dps_packed``\ 主要是为了让我们在生成实际实现之前继续做一些高层的转换。 Dataflow Block ~~~~~~~~~~~~~~ Relax 函数中的另一个重要元素是 ``R.dataflow()`` 范围标注: .. raw:: latex \diilbookstyleinputcell .. code:: python with R.dataflow(): lv0 = R.call_dps_packed("linear0", (x, w0, b0), R.Tensor((1, n), "float32")) lv1 = R.call_dps_packed("relu0", (lv0, ), R.Tensor((1, n), "float32")) out = R.call_dps_packed("linear0", (lv1, w1, b1), R.Tensor((1, k), "float32")) R.output(out) 这又回到了我们在上一节中讨论的\ **计算图**\ 。 回想一下,在理想情况下,每个计算图操作都应该没有side-effect。 如果我们仍然想引入包含side-effect的操作怎么办? dataflow block是我们标记程序计算图区域的一种方式。 具体来说,在dataflow block中,所有操作都需要side-effect free。 在dataflow block之外,操作可能包含side-effect。 下面的程序是一个包含两个dataflow block的示例程序。 .. raw:: latex \diilbookstyleinputcell .. code:: python @R.function def main(x: R.Tensor((1, "m"), "float32"), w0: R.Tensor(("n", "m"), "float32"), b0: R.Tensor(("n", ), "float32"), w1: R.Tensor(("k", "n"), "float32"), b1: R.Tensor(("k", ), "float32")): m, n, k = T.int64(), T.int64(), T.int64() with R.dataflow(): lv0 = R.call_dps_packed("linear0", (x, w0, b0), R.Tensor((1, n), "float32")) gv0 = R.call_dps_packed("relu0", (lv0, ), R.Tensor((1, n), "float32")) R.output(gv0) with R.dataflow(): out = R.call_dps_packed("linear0", (gv0, w1, b1), R.Tensor((1, k), "float32")) R.output(out) return out 我们的大部分课程只会处理计算图(dataflow blocks),但最好能够理解背后的原因。 章节小结 ~~~~~~~~ 到目前为止,我们已经完成了一个 Relax 程序的示例,并涵盖了大部分元素,包括: - 计算图 - ``call_dps_packed`` - Dataflow block 这些元素应该让我们能够开始端到端的模型执行和编译。 当我们在后面的章节中遇到新概念时,我们还将介绍它们。 构建并运行模型 -------------- 在上一节中,我们讨论了端到端模型执行的抽象。 本节介绍如何构建和运行 IRModule。 让我们首先回顾一下我们现有的 IRModule。 .. raw:: latex \diilbookstyleinputcell .. code:: python IPython.display.Code(MyModule.script(), language="python") .. raw:: html
# from tvm.script import ir as I
# from tvm.script import tir as T
# from tvm.script import relax as R
@I.ir_module
class Module:
@T.prim_func
def linear0(x: T.handle, w: T.handle, b: T.handle, z: T.handle):
m = T.int64()
X = T.match_buffer(x, (1, m))
n = T.int64()
W = T.match_buffer(w, (n, m))
B = T.match_buffer(b, (n,))
Z = T.match_buffer(z, (1, n))
# with T.block("root"):
Y = T.alloc_buffer((1, n))
for i, j, k in T.grid(1, n, m):
with T.block("Y"):
vi, vj, vk = T.axis.remap("SSR", [i, j, k])
T.reads(X[vi, vk], W[vj, vk])
T.writes(Y[vi, vj])
with T.init():
Y[vi, vj] = T.float32(0)
Y[vi, vj] = Y[vi, vj] + X[vi, vk] * W[vj, vk]
for i, j in T.grid(1, n):
with T.block("Z"):
vi, vj = T.axis.remap("SS", [i, j])
T.reads(Y[vi, vj], B[vj])
T.writes(Z[vi, vj])
Z[vi, vj] = Y[vi, vj] + B[vj]
@T.prim_func
def relu0(x: T.handle, y: T.handle):
n = T.int64()
X = T.match_buffer(x, (1, n))
Y = T.match_buffer(y, (1, n))
# with T.block("root"):
for i, j in T.grid(1, n):
with T.block("Y"):
vi, vj = T.axis.remap("SS", [i, j])
T.reads(X[vi, vj])
T.writes(Y[vi, vj])
Y[vi, vj] = T.max(X[vi, vj], T.float32(0))
@R.function
def main(x: R.Tensor((1, "m"), dtype="float32"), w0: R.Tensor(("n", "m"), dtype="float32"), b0: R.Tensor(("n",), dtype="float32"), w1: R.Tensor(("k", "n"), dtype="float32"), b1: R.Tensor(("k",), dtype="float32")) -> R.Tensor((1, "k"), dtype="float32"):
k = T.int64()
m = T.int64()
n = T.int64()
with R.dataflow():
lv0 = R.call_dps_packed("linear0", (x, w0, b0), out_sinfo=R.Tensor((1, n), dtype="float32"))
lv1 = R.call_dps_packed("relu0", (lv0,), out_sinfo=R.Tensor((1, n), dtype="float32"))
out = R.call_dps_packed("linear0", (lv1, w1, b1), out_sinfo=R.Tensor((1, k), dtype="float32"))
R.output(out)
return out
# from tvm.script import ir as I
# from tvm.script import tir as T
# from tvm.script import relax as R
@I.ir_module
class Module:
@T.prim_func
def linear0(x: T.handle, w: T.handle, b: T.handle, z: T.handle):
m = T.int64()
X = T.match_buffer(x, (1, m))
n = T.int64()
W = T.match_buffer(w, (n, m))
B = T.match_buffer(b, (n,))
Z = T.match_buffer(z, (1, n))
# with T.block("root"):
Y = T.alloc_buffer((1, n))
for i, j, k in T.grid(1, n, m):
with T.block("Y"):
vi, vj, vk = T.axis.remap("SSR", [i, j, k])
T.reads(X[vi, vk], W[vj, vk])
T.writes(Y[vi, vj])
with T.init():
Y[vi, vj] = T.float32(0)
Y[vi, vj] = Y[vi, vj] + X[vi, vk] * W[vj, vk]
for i, j in T.grid(1, n):
with T.block("Z"):
vi, vj = T.axis.remap("SS", [i, j])
T.reads(Y[vi, vj], B[vj])
T.writes(Z[vi, vj])
Z[vi, vj] = Y[vi, vj] + B[vj]
@R.function
def main(x: R.Tensor((1, 784), dtype="float32")) -> R.Tensor((1, 10), dtype="float32"):
with R.dataflow():
lv0 = R.call_dps_packed("linear0", (x, metadata["relax.expr.Constant"][0], metadata["relax.expr.Constant"][1]), out_sinfo=R.Tensor((1, 128), dtype="float32"))
lv1 = R.call_dps_packed("env.relu", (lv0,), out_sinfo=R.Tensor((1, 128), dtype="float32"))
out = R.call_dps_packed("env.linear", (lv1, metadata["relax.expr.Constant"][2], metadata["relax.expr.Constant"][3]), out_sinfo=R.Tensor((1, 10), dtype="float32"))
R.output(out)
return out
# Metadata omitted. Use show_meta=True in script() method to show it.