End to End Model Execution ========================== Prelude ------- .. figure:: ../img/tensor_func_linear_relu.png Most of the MLC process can be viewed as transformation among tensor functions. The main thing we aim to answer in our following up are: - What are the possible abstractions to represent the tensor function. - What are possible transformations among the tensor functions. In the last lecture, we focus on the primitive tensor functions. In this lecture, we will talk about how to build end-to-end models. Preparations ------------ To begin with, we will import necessary dependencies and create helper functions. .. 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 Load the Dataset ~~~~~~~~~~~~~~~~ As a concrete example, we will be using a model on the fashion MNIST dataset. The following code downloads and prepares the data from ``torchvision`` in NumPy array. .. 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 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 100.0%Extracting data/FashionMNIST/raw/t10k-labels-idx1-ubyte.gz to data/FashionMNIST/raw We can plot out the image instance that we want to be able to predict. .. 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 Download Model Parameters ~~~~~~~~~~~~~~~~~~~~~~~~~ .. raw:: latex \diilbookstyleinputcell .. code:: python # Hide outputs !wget https://github.com/mlc-ai/web-data/raw/main/models/fasionmnist_mlp_params.pkl End to End Model Integration ---------------------------- In this chapter, we will use the following model as an example. This is a two-layer neural network that consists of two linear operations with relu activation. To keep things simple, we removed the final softmax layer. The output score is un-normalized, but still, the maximum value corresponds to the most likely class. .. figure:: ../img/e2e_fashionmnist_mlp_model.png Let us begin by reviewing a Numpy implementation of the model. .. 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 [[-26.194868 -34.8678 -24.12358 -20.510454 -18.109476 12.921284 -17.477413 -4.9815836 -7.2966995 -6.3170166]] [5] NumPy-MLP Prediction: Sandal The above example code shows the high-level array operations to perform the end-to-end model execution. Again from MLC’s pov, we would like to see through the details under the hood of these array computations. For the purpose of illustrating details under the hood, we will again write examples in low-level numpy: - We will use a loop instead of array functions when necessary to demonstrate the possible loop computations. - When possible, we always explicitly allocate arrays via numpy.empty and pass them around. The code block below shows a low-level numpy implementation of the same model. .. 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 Constructing an End to End IRModule in TVMScript ------------------------------------------------ With the low-level NumPy example in mind, now we are ready to introduce an MLC abstraction for the end-to-end model execution. The code block below shows a TVMScript implementation of the model. .. 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 The above code contains kinds of functions: the primitive tensor functions (``T.prim_func``) that we saw in the last lecture and a new ``R.function`` (relax function). Relax function is a new type of abstraction representing high-level neural network executions. Again it is helpful to see the TVMScript code and low-level numpy code side-by-side and check the corresponding elements, and we are going to walk through each of them in detail. Since we already learned about primitive tensor functions, we are going to focus on the high-level execution part. .. figure:: ../img/e2e_compare_to_lnumpy.png Computational Graph View ~~~~~~~~~~~~~~~~~~~~~~~~ .. figure:: ../img/e2e_computational_graph_call_tir.png It is usually helpful to use graph to visualize high-level model executions. The above figure is a graph-view of the ``main`` function: - Each of the box in the graph corresponds to computation operations. - The arrows correspond to the input-output of the intermediate tensors. We have seen this kind of visualization in earlier lectures. The graph itself can be viewed as a type of abstraction, and it is commonly known as **computational graph** in machine learning frameworks. ``call_dps_packed`` Construct ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ One thing that you may have noticed is that each step of operations in the computational graph contains an ``R.call_dps_packed`` operation. This is the operation that brings in the tensor primitive functions .. raw:: latex \diilbookstyleinputcell .. code:: python lv0 = R.call_dps_packed(linear0, (x, w0, b0), (1, 128), dtype="float32") To explain what does ``R.call_dps_packed`` mean, let us review an equivalent low-level numpy implementation of the operation, as follows: .. 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 Specifically, call_dps_packed takes in a primitive function (``prim_func``) a list of inputs. Then what it does is allocate an output tensor ``res``, then pass the inputs and the output to the ``prim_func``. After executing ``prim_func`` the result is populated in ``res``, then we can return the result. Note that ``lnumpy_call_dps_packed`` is only a reference implementation to show the meaning of ``R.call_dps_packed``. In practice, there can be different low-level ways to optimize the execution. For example, we might choose to allocate all the output memories ahead of time and then run the execution, which we will cover in future lectures. A natural question that one could ask is why do we need ``call_dps_packed`` construct? This is because our primitive tensor functions take the following calling convention. .. raw:: latex \diilbookstyleinputcell .. code:: python def low_level_prim_func(in0, in1, ..., out): # implementations This convention is called **destination passing**. The idea is that input and output are explicitly allocated outside and passed to the low-level primitive function. This style is commonly used in low-level library designs, so higher-level frameworks can handle that memory allocation decision. Note that not all tensor operations can be presented in this style (specifically, there are operations whose output shape depends on the input). Nevertheless, in common practice, it is usually helpful to write the low-level function in this style when possible. While it is possible to assemble the destination passing convention function together by explicitly allocating intermediate results and calling each function, it is hard to turn the following code into computational graph form. .. 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 We can certainly try a bit :) The above figure is one possible “failed attempt” to fit the ``lnumpy_mlp`` into a “computational graph-like” form by simply connecting function inputs to the function. We can find that it lost a few nice properties of the previous computational graphs. Specifically, a computational graph usually has the following properties - Every input edge to the box corresponds to the input to the operation. - Every outgoing edge corresponds to the output of the operations. - Each operation can be reordered arbitrarily up to the topological order of the edges. Of course, we can still generalize the graph definition by introducing the input edge and output edge, and that can complicate the possible transformations associated with the abstraction. So coming back to ``call_dps_packed``, the key insight here is that we want to hide possible allocation or explicit writing to the functions. In a more formal term, we want the function to be **pure** or **side-effect free**. A function is **pure** or **side-effect free** if: it only reads from its inputs and returns the result via its output, it will not change other parts of the program (such as incrementing a global counter). **call_dps_packed** is a way for us to hide these details of calling into low-level primitive functions and expose them into a computational graph. We can also see ``call_dps_packed`` in action in the low-level numpy as well. Now we have defined the ``lnumpy_call_dps_packed``, we can rewrite the low-level numpy execution code as: .. 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 In practice, the lowest-level implementation will have explicit memory allocations, so ``call_dps_packed`` mainly serves as a purpose for us to continue to do some high-level transformations before we generate the actual implementation. Dataflow Block ~~~~~~~~~~~~~~ Another important element in a relax function is the ``R.dataflow()`` scope annotation. .. 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) This connects back to the **computational graph** discussion we had in the last section. Recall that ideally, each computational graph operation should be side effect free. What if we still want to introduce operations that contains side effect? A dataflow block is a way for us to mark the computational graph regions of the program. Specifically, within a dataflow block, all the operations need to be side-effect free. Outside a dataflow block, the operations can contain side-effect. The program below is an example program that contains two dataflow blocks. .. 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 Most of our lectures will only deal with computational graphs (dataflow blocks). But it is good to keep the reason behind in mind. Section Checkpoint ~~~~~~~~~~~~~~~~~~ So far, we have gone through one example instance of relax program and covered most of the elements, including: - Computational graph view - ``call_dps_packed`` construct - Dataflow block. These elements should get us started in the end to end model execution and compilation. we will also cover new concepts as we encounter them in later chapters. Build and Run the Model ----------------------- In the last section, we discussed the abstraction that enables us to represent end-to-end model execution. This section introduces how to build and run an IRModule. Let us begin by reviewing the IRModule we have. .. 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.