端到端优化模型

本教程演示了如何使用 Apache TVM 优化机器学习模型。我们将使用 PyTorch 的预训练 ResNet-18 模型，并使用 TVM 的 Relax API 对其进行端到端优化。请注意，默认的端到端优化可能不适用于复杂模型。

准备工作

首先，我们准备模型和输入信息。我们使用 PyTorch 的预训练 ResNet-18 模型。

import os
import numpy as np
import torch
from torch.export import export
from torchvision.models.resnet import ResNet18_Weights, resnet18

torch_model = resnet18(weights=ResNet18_Weights.DEFAULT).eval()

Downloading: "https://download.pytorch.org/models/resnet18-f37072fd.pth" to /workspace/.cache/torch/hub/checkpoints/resnet18-f37072fd.pth

  0%|          | 0.00/44.7M [00:00<?, ?B/s]
 36%|███▋      | 16.2M/44.7M [00:00<00:00, 170MB/s]
 87%|████████▋ | 38.8M/44.7M [00:00<00:00, 208MB/s]
100%|██████████| 44.7M/44.7M [00:00<00:00, 206MB/s]

回顾整体流程

整体流程包括以下步骤

• 构建或导入模型：构建神经网络模型或从其他框架（例如 PyTorch、ONNX）导入预训练模型，并创建 TVM IRModule，其中包含编译所需的所有信息，包括用于计算图的高级 Relax 函数，以及用于张量程序的低级 TensorIR 函数。

• 执行可组合的优化：执行一系列优化转换，例如图优化、张量程序优化和库调度。

• 构建和通用部署：将优化后的模型构建为可部署到通用运行时的模块，并在不同的设备（例如 CPU、GPU 或其他加速器）上执行它。

将模型转换为 IRModule

下一步，我们使用 Relax 前端为 PyTorch 将模型转换为 IRModule，以便进一步优化。

import tvm
from tvm import relax
from tvm.relax.frontend.torch import from_exported_program

# Give an example argument to torch.export
example_args = (torch.randn(1, 3, 224, 224, dtype=torch.float32),)

# Skip running in CI environment
IS_IN_CI = os.getenv("CI", "") == "true"

if not IS_IN_CI:
    # Convert the model to IRModule
    with torch.no_grad():
        exported_program = export(torch_model, example_args)
        mod = from_exported_program(exported_program, keep_params_as_input=True)

    mod, params = relax.frontend.detach_params(mod)
    mod.show()

IRModule 优化

Apache TVM Unity 提供了一种灵活的方式来优化 IRModule。围绕 IRModule 优化的所有内容都可以与现有的 pipeline 组合。请注意，每个转换都可以通过 tvm.ir.transform.Sequential 组合成一个优化 pipeline。

在本教程中，我们专注于通过自动调优对模型进行端到端优化。我们利用 MetaSchedule 调优模型并将调优日志存储到数据库中。我们还将数据库应用于模型以获得最佳性能。

TOTAL_TRIALS = 8000  # Change to 20000 for better performance if needed
target = tvm.target.Target("nvidia/geforce-rtx-3090-ti")  # Change to your target device
work_dir = "tuning_logs"

if not IS_IN_CI:
    mod = relax.get_pipeline("static_shape_tuning", target=target, total_trials=TOTAL_TRIALS)(mod)

    # Only show the main function
    mod["main"].show()

构建和部署

最后，我们构建优化后的模型并将其部署到目标设备。我们在 CI 环境中跳过此步骤。

if not IS_IN_CI:
    ex = tvm.compile(mod, target="cuda")
    dev = tvm.device("cuda", 0)
    vm = relax.VirtualMachine(ex, dev)
    # Need to allocate data and params on GPU device
    gpu_data = tvm.nd.array(np.random.rand(1, 3, 224, 224).astype("float32"), dev)
    gpu_params = [tvm.nd.array(p, dev) for p in params["main"]]
    gpu_out = vm["main"](gpu_data, *gpu_params).numpy()

    print(gpu_out.shape)