将TVM集成到PyTorch

随着TVM不断展示出对深度学习执行效率的改进，很明显PyTorch将从直接利用编译器堆栈中受益。PyTorch的主要宗旨是提供无缝且强大的集成，而这不会妨碍用户。PyTorch现在具有基于TVM的官方后端torch_tvm。

用法很简单：

import torch_tvm

torch_tvm.enable()

就是这样！然后，PyTorch将尝试在其JIT编译过程中，将所有可能的算子转换为已知的Relay算子。

背景

与许多其它ML框架不同，PyTorch公开了一个执行的编程接口。这种编程风格避免了图元编程，而专注于以Python方式直接控制n维数组（张量）。该框架最初非常适合模型的试验和开发，但不适用于自动性能优化或部署。为了利用优化的编译器技术，PyTorch最近引入了一些较大的改进来解决此问题。

PyTorch 1.0引入了PyTorch IR，PyTorch专用的中间表示形式，用于类似于Relay的模型。可以通过模型跟踪，将PyTorch程序转换为IR，该跟踪记录模型或Python的子集TorchScript的执行。新的TVM后端将PyTorch的IR降低到了Relay，能够透明地提高PyTorch的性能，无需用户参与。

整合与结果

为了支持Relay，PyTorch JIT添加了两个功能：自定义转换过程和自定义子图解释器。

当torch_tvm启用时，可以转换到Relay PyTorch IR的子图Expr，旨意被标记为Relay 兼容。由于PyTorch IR并不总是包含形状信息，因此在调用之前，无法以有用的方式编译任何子图。

在用户调用期间，PyTorch JIT Runtime将确定输入形状信息，并使用新的Relay C ++构建系统，编译先前标记的子图。根据输入形状来缓存编译，以供后续运行。可以在README中找到更多详细信息。

torch_tvm建立了一个连续的基准测试系统，该系统正在监视ResNet18在CPU上的性能。对于各种ResNet型号，TVM的性能都是默认PyTorch JIT后端的两倍以上。下图详细描述了在AWS c5n.4xlarge实例上，使用16个线程实现的每秒迭代次数（越大越好）：

这些结果令人鼓舞，该项目将继续致力于在更多模型上提高CPU推理速度。

未来的工作

现在，PyTorch JIT进行了大量工作来查找其IR的纯功能子集，馈送到Relay。避免了将采样和控制流信息映射到Relay，这不是必需的。将更多的PyTorch IR映射到Relay，可能会取得性能上的胜利，这是该项目的目标。PyTorch IR在开发过程中正在迅速变化，因此必须谨慎进行。

将做更多的工作来确保PyTorch和TVM代码之间的切换是有效的。这包括统一线程模型，分配器以及减少与将输入复制到TVM相关的开销。

help文件

如果已经编写了PyTorch模型，最简单的入门方法就是使用torch.jit.trace方法

import torch_tvm

from your_model import model, inputs

torch_tvm.enable(opt_level=3)

iters = 100

warmup = 10

# Ensure your model is in eval mode and also turn off gradients.

with torch.no_grad():

  # Use tuned parameters for better performance.

  with autotvm.apply_history_best("test/autotvm_tuning.log"):

    # This is where all the compilation happens.

    trace_tvm = torch.jit.trace(model, inputs)

    # Warmup

    for _ in range(warmup):

      _ = trace_tvm(*inputs)

    # Benchmark

    start = time.time()

    for _ in range(iters):

      _ = trace_tvm(*inputs)

    tvm_time = time.time() - start

    print("Took {}s to run {} iters".format(tvm_time, iters))

注意，用于AVX2 LLVM编译的调整参数位于存储库test/文件夹中。

如果直接使用Relay，可以通过（隐式）跟踪或TorchScript直接，从PyTorch函数中提取表达式：

def add(a, b, c):

    return a + b + c

# via tracing

relay_graph = torch_tvm.to_relay(add, inputs)

@torch.jit.script

def mul(a, b, c):

    return a * b * c

# via script

relay_graph = torch_tvm.to_relay(mul, inputs)