教程 | 如何通过PyTorch上手Tensor Comprehensions？

Posted 2021-05-02 机器之心

选自pytorch

机器之心编译

参与：乾树、黄小天

Tensor Comprehensions 是一个降低高性能代码编写门槛的工具，可以将高级语言代码直接生成 GPU 代码并针对特定输入自动调整代码。本文作者将展示 TC 如何在 PyTorch 中无缝使用，并与 PyTorch Tensors 和 nn Variables 进行互操作。

Tensor Comprehensions (TC) 是一个降低高性能代码编写门槛的工具，可以将高级语言代码直接生成 GPU 代码并针对特定输入自动调整代码。本文强烈推荐你首先阅读 Tensor Comprehensions 发布博文（详见 ）。

如果你面临如下任一场景，TC 将会帮到你。

• 你的 Pytorch 层又大又慢，你打算为此写一段 C++ 或者 CUDA 的优化代码，但是你又不擅长。
  

• 你写好了 CUDA 代码，但是你还要花费大量时间去优化。你希望可以在最短时间内搞定这些。

• 为了加速，你想在网络中融合多个层，例如 Conv-ReLU-BatchNorm 或者 Linear-ReLU-Linear-ReLU。但是这很难理解。

• 你的研究涉及 CuDNN 和 MKL 未能优化的不寻常的张量。例如，你要使用 13 x 24 的卷积核对 143 x 55 的输入图像进行卷积。你试着用 CUDNN 跑，并且发现它慢的超乎想象。

• 通过调整 Tensors 以适应特定的内存布局，你的代码会变慢。你希望编写在你的输入布局上高效运行的自定义代码很容易。

TC 可以在 PyTorch 中无缝使用，与 PyTorch Tensors 和 nn Variables 进行互操作。

教程现在开始！

1. 安装包

conda install -c pytorch -c tensorcomp tensor_comprehensions

此次我们只提供在 Ubuntu 16.04 和 CentOS7 上测试过的 Linux-64 二进制包。TC 依赖于大型 C++ 项目，如 Halide、Tapir-LLVM 及 ISL。因此，我们使用 Anaconda 来可靠地分配这些依赖。因此 TC 不提供 PyPI 安装包。

2. 导入 Python 包

import tensor_comprehensions as tc

3. 定义 TC 表达式并创建一个 Python 函数

lang = """
def fcrelu(float(B,M) I, float(N,M) W1, float(N) B1) -> (O1) {
    O1(b, n) +=! I(b, m) * W1(n, m)
    O1(b, n) = O1(b, n) + B1(n)
    O1(b, n) = fmax(O1(b, n), 0)
}
"""fcrelu = tc.define(lang, name="fcrelu")

fcrelu 以 PyTorch Tensors 为输入并返回一个 PyTorch Tensor。输入 inputI、权重 W1、偏置项 B1 并返回一个 O1。

4. 创建一些虚拟输入张量

B, M, N = 100, 128, 100I, W1, B1 = torch.randn(B, M).cuda(), torch.randn(N, M).cuda(), torch.randn(N).cuda()

5. 自动调节函数以适用你的输入大小

fcrelu.autotune(I, W1, B1, cache="fcrelu_100_128_100.tc")

自动调节器是你最好的朋友。你通常在调用 tc 函数前都会先自动调节它。当自动调节开始运行，则会显示当前的最佳性能。如果你满足于当前结果或时间耗完了，请按 Ctrl + C 终止调节进程。

cache 保留自动自动调节的核搜索结果并存到 fcrelu_100_128_100.tc 文件中。下一次再调用这句代码时，它就会自动加载已调节的结果而无需再计算。

自动调节器有几个超参数 (就像你的卷积网络有学习率，层数等等）。我们预设了一些缺省项，但是你可以在这里了解如何使用高级选项。

6. 调用函数并传参，返回你的结果

out = fcrelu(I, W1, B1)

现在，我们来看看怎么写 TC 表达式。

TC 语言快速入门

TC 符号关注层的数学本质，将性能优化交给使用 Halide 和 polyhedral 编译技术的后端代码，这些技术累积了数十年来最先进的 Loop Nest Optimization（LNO）研究成果。

TC 和 np.einsum 很像。我们用一个实例快速展开学习。

lang = """
def matmul(float(M,N) A, float(N,K) B) -> (output) {
  output(i, j) +=! A(i, kk) * B(kk, j)
}
"""

我们定义了这个实例中的一个函数，它接受两个输入 M x N 的 A 和 N x K 的Ｂ并返回一个结果。output 的尺寸由 TC 语言自动推导得出 (下面展开讨论）。看看这行代码：

output(i, j) +=! A(i, kk) * B(kk, j)

意思是:

• output(i, j) 表示输出是 2 维的

• 对于每个坐标 output(i, j), 加上 (+=) A(i, kk) * B(kk, j)

• i 是 A 中 dim=0 的所有坐标, 即，i 在 range(0, M) 里

• j 是 B dim = 1 中的所有坐标，即，j 在 range(0, K) 里

• kk 是由 0 到 N 推断出的所有坐标

output 的尺寸由 i 和 j 可取的最大值推理而来，也就是 M 和 K，所以 output 的尺寸就是 M x K。

! 符号用 0.0 初始化，相当于：

output(i, j) = 0
output(i, j) += A(i, kk) * B(kk, j)

标量输入和范围约束：实现 AvgPool2d

"""

def avgpool(float(B, C, H, W) input) -> (output) {{
    output(b, c, h, w) += input(b, c, h * {sH} + kh, w * {sW} + kw) where kh in 0:{kH}, kw in 0:{kW}
}}

"""avgpool = tc.define(LANG, name="avgpool", constants={"sH":1, "sW":1, "kH":2, "kW":2})

这里 where 关键字可以选取取值范围来进行运算。0:{kH} 类似与 Python 中的 range(kH）。

注意：传入标量的语法在下一版本中可能会发生变化。

torch.nn 层

我们对 TC 的基本 PyTorch 集成进行了封装，以便通过定义正向和反向 TC 表达式并有 Variable 输入/输出来轻松地将 TC 集成到更大的 torch.nn 模块。这里有一个用 TC 定义卷积操作的实例。

你可能忽略的基础知识 (这正是我们在做的）

自动调节变量长度序列

TC 自动调节器要求预先自动调节好所有的输入尺寸。例如，如果你有一批输入图像 I1，调节器想要知道 I1 的确切尺寸以便生成优化的卷积核。你不能指定：高度介于 200 到 300 之间的图像。这对于如 NLP 之类的序列数据更为重要，NLP 每个语句可以具有不同的长度。

自动调节器不参数化是因为自动调整参数约束越来越难，并且这是一项热门研究。由于这是首次发布，我们决定主动提供工具，并保证其良好工作。

作为解决方法，如果你事先知道有一些要用到的特定尺寸，则可以输入它们运行自动调节器。

relu = tc.define(LANG, name="relu")batch, channels = 16, 3tc.autotune((batch, channels, 32, 32)) # image of size 32 x 32tc.autotune((batch, channels, 48, 48)) # image of size 48 x 48tc.autotune((batch, channels, 64, 64)) # image of size 64 x 64

当前的自动调节器自动调节三个特定的图片尺寸： 32x32、48x48 和 64x64。

缺少循环

如果你想写一个 RNN，很容易把它视作相对于时间的 for 循环。但是，TC 语言目前还没有循环结构。如果你真的想写 RNN，可以试着把它展开。

Strided-Tensors

TC 后台暂不支持不连续的张量。如果你输入的张量不连续，TC 在传递到后台代码前会作连续化处理。

使用 TC 表达式进行张量变形

在 TC 中你不能这么写操作：torch.matmul(...).view(...).mean(...)。任何时候都需要一个 view 来改变输入张量的尺寸，你必须获取到输出张量，在 PyTorch 级别上 view 它。

如何开始

• 《Walk through Tutorial》会帮助你理解并上手入门 TC PyTorch 包。
  

• 超过 20 个实例含有 TC 写的不同的 ML 层，包括 avgpool、maxpool、matmul、matmul - give output buffers、batch-matmul、convolution、strided-convolution、batchnorm、copy、cosine similarity、Linear、Linear + ReLU、group-convolutions、strided group-convolutions、indexing、Embedding (lookup table)、small-mobilenet、softmax、tensordot、transpose。

• 有关 TC 的详细文档以及与 PyTorch 的整合。

原文链接：http://pytorch.org/2018/03/05/tensor-comprehensions.html

✄------------------------------------------------

加入机器之心（全职记者/实习生）：[e-mail]

投稿或寻求报道：[e-mail]

广告&商务合作：[e-mail]