注意

本教程可以与 Google Colab 互动使用！您也可以点击这里在本地运行 Jupyter notebook。

快速入门

本教程面向 Apache TVM 的新手。通过一个简单的例子展示如何使用 Apache TVM 编译一个简单的神经网络。

概述 

Apache TVM 是一个机器学习编译框架，遵循Python 优先开发和通用部署的原则。它接收预训练的机器学习模型，编译并生成可部署的模块，这些模块可以嵌入并在任何地方运行。Apache TVM 还支持自定义优化流程，以引入新的优化、库、代码生成等。

Apache TVM 可以帮助您：

优化 ML 工作负载的性能，组合库和代码生成。
部署 ML 工作负载到各种新的环境，包括新的运行时和新的硬件。
通过快速自定义库调度、引入自定义运算符和代码生成，持续改进和自定义 Python 中的 ML 部署管道。

总体流程 

接下来，我们将展示使用 Apache TVM 编译神经网络模型的总体流程，展示如何优化、部署和运行模型。总体流程如下图所示

../../_static/downloads/tvm_overall_flow.svg

总体流程包括以下步骤

构建或导入模型：构建神经网络模型或从其他框架（例如 PyTorch、ONNX）导入预训练模型，并创建 TVM IRModule，其中包含编译所需的所有信息，包括用于计算图的高级 Relax 函数，以及用于张量程序的低级 TensorIR 函数。
执行可组合的优化：执行一系列优化转换，例如图优化、张量程序优化和库调度。
构建和通用部署：将优化的模型构建为可部署的模块，以用于通用运行时，并在不同的设备（例如 CPU、GPU 或其他加速器）上执行它。

构建或导入模型 

在我们开始之前，让我们首先构建一个神经网络模型。在本教程中，为了简化，我们将使用 TVM Relax 前端直接在本脚本中定义一个两层 MLP 网络，它类似于 PyTorch 的 API。

import tvm
from tvm import relax
from tvm.relax.frontend import nn


class MLPModel(nn.Module):
    def __init__(self):
        super(MLPModel, self).__init__()
        self.fc1 = nn.Linear(784, 256)
        self.relu1 = nn.ReLU()
        self.fc2 = nn.Linear(256, 10)

    def forward(self, x):
        x = self.fc1(x)
        x = self.relu1(x)
        x = self.fc2(x)
        return x

然后我们可以将模型导出到 TVM IRModule，它是 TVM 中的中心中间表示。

mod, param_spec = MLPModel().export_tvm(
    spec={"forward": {"x": nn.spec.Tensor((1, 784), "float32")}}
)
mod.show()

# from tvm.script import ir as I
# from tvm.script import relax as R

@I.ir_module
class Module:
    @R.function
    def forward(x: R.Tensor((1, 784), dtype="float32"), fc1_weight: R.Tensor((256, 784), dtype="float32"), fc1_bias: R.Tensor((256,), dtype="float32"), fc2_weight: R.Tensor((10, 256), dtype="float32"), fc2_bias: R.Tensor((10,), dtype="float32")) -> R.Tensor((1, 10), dtype="float32"):
        R.func_attr({"num_input": 1})
        with R.dataflow():
            permute_dims: R.Tensor((784, 256), dtype="float32") = R.permute_dims(fc1_weight, axes=None)
            matmul: R.Tensor((1, 256), dtype="float32") = R.matmul(x, permute_dims, out_dtype="void")
            add: R.Tensor((1, 256), dtype="float32") = R.add(matmul, fc1_bias)
            relu: R.Tensor((1, 256), dtype="float32") = R.nn.relu(add)
            permute_dims1: R.Tensor((256, 10), dtype="float32") = R.permute_dims(fc2_weight, axes=None)
            matmul1: R.Tensor((1, 10), dtype="float32") = R.matmul(relu, permute_dims1, out_dtype="void")
            add1: R.Tensor((1, 10), dtype="float32") = R.add(matmul1, fc2_bias)
            gv: R.Tensor((1, 10), dtype="float32") = add1
            R.output(gv)
        return gv

执行优化转换 

Apache TVM 利用 pipeline 来转换和优化程序。pipeline 封装了一系列转换，旨在实现两个目标（在同一级别）：

模型优化：例如算子融合、布局重写。
张量程序优化：将算子映射到低级实现（库或代码生成）。

注意

这两个是目标，而不是 pipeline 的阶段。这两个优化在同一级别执行，或者在两个阶段分别执行。

注意

在本教程中，我们仅演示总体流程，通过利用 zero 优化 pipeline，而不是针对任何特定目标进行优化。

mod = relax.get_pipeline("zero")(mod)

构建和通用部署 

优化之后，我们可以将模型构建为可部署的模块，并在不同的设备上运行它。

import numpy as np

target = tvm.target.Target("llvm")
ex = tvm.compile(mod, target)
device = tvm.cpu()
vm = relax.VirtualMachine(ex, device)
data = np.random.rand(1, 784).astype("float32")
tvm_data = tvm.nd.array(data, device=device)
params = [np.random.rand(*param.shape).astype("float32") for _, param in param_spec]
params = [tvm.nd.array(param, device=device) for param in params]
print(vm["forward"](tvm_data, *params).numpy())

[[23886.062 25407.682 26255.523 24531.725 26409.904 24607.594 24389.584
  25355.768 23537.834 24064.062]]

我们的目标是以任何感兴趣的语言将机器学习引入应用程序，并具有最少的运行时支持。

IRModule 中的每个函数都成为运行时中可运行的函数。例如，在 LLM 案例中，我们可以直接调用 prefill 和 decode 函数。
```
prefill_logits = vm["prefill"](inputs, weight, kv_cache)
decoded_logits = vm["decode"](inputs, weight, kv_cache)
```

TVM 运行时带有原生数据结构，例如 NDArray，还可以与现有生态系统进行零拷贝交换（与 PyTorch 进行 DLPack 交换）

# Convert PyTorch tensor to TVM NDArray
x_tvm = tvm.nd.from_dlpack(x_torch.to_dlpack())
# Convert TVM NDArray to PyTorch tensor
x_torch = torch.from_dlpack(x_tvm.to_dlpack())

TVM 运行时可以在非 Python 环境中工作，因此它适用于移动设备等设置

// C++ snippet
runtime::Module vm = ex.GetFunction("load_executable")();
vm.GetFunction("init")(...);
NDArray out = vm.GetFunction("prefill")(data, weight, kv_cache);

// Java snippet
Module vm = ex.getFunction("load_executable").invoke();
vm.getFunction("init").pushArg(...).invoke;
NDArray out = vm.getFunction("prefill").pushArg(data).pushArg(weight).pushArg(kv_cache).invoke();

下一步阅读 

本教程演示了使用 Apache TVM 编译神经网络模型的总体流程。有关更高级或特定的主题，请参阅以下教程

由 Sphinx-Gallery 生成的图库

快速入门

概述

总体流程

构建或导入模型

执行优化转换

构建和通用部署

下一步阅读