Tensor Comprehensions (TC) 是一款旨在降低编写高性能代码门槛的工具。它能通过简单的高级语言生成 GPU 代码，并针对特定的输入规模对代码进行自动调优。
我们强烈建议您先阅读 Tensor Comprehensions 博客文章。
如果您遇到了以下任何一种情况，TC 将会是您的得力助手。
- 您的 PyTorch 层很大且运行缓慢，您曾考虑过为其编写专门的 C++ 或 CUDA 代码，但您并不懂得如何编写 CUDA 或底层代码。
- 您编写了一个 CUDA 层，但编写、调试和优化速度花费了一周时间。您希望能在不到一小时内完成这些工作。
- 为了提升速度，您想在网络中融合多个层，例如 Conv-ReLU-BatchNorm 或 Linear-ReLU-Linear-ReLU，但实现起来非常困难。
- 您的研究涉及一些 CuDNN 和 MKL 未进行优化的特殊张量形状。例如，您在处理输入图像为 143 x 55 的 13 x 24 卷积。您尝试使用 CuDNN 运行它，但速度比预期的要慢。
- 您的代码因为频繁转置张量以适应特定的内存布局而导致运行缓慢。您希望能够轻松编写自定义代码，以高效地处理您的输入布局。
Tensor Comprehensions 在 PyTorch 中使用非常顺畅，可与 PyTorch 张量和 nn 变量实现互操作。
让我们演练一下如何在 PyTorch 中使用 TC。
1. 安装软件包
conda install -c pytorch -c tensorcomp tensor_comprehensions
目前我们仅提供在 Ubuntu 16.04 和 CentOS 7 上测试过的 Linux-64 二进制文件。
TC 依赖于 Halide、Tapir-LLVM 和 ISL 等重量级 C++ 项目。因此，我们依靠 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 张量作为输入并返回一个 PyTorch 张量。它接收输入 I、权重 W1、偏置 B1，并返回输出 O1。
4. 让我们创建一些虚拟输入张量
B, M, N = 100, 128, 100
I, 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 和多面体编译技术，汇集了数十年的前沿循环嵌套优化 (LNO) 研究成果。
TC 与 np.einsum 很接近。我们将通过示例快速学习 TC。
lang = """
def matmul(float(M,N) A, float(N,K) B) -> (output) {
output(i, j) +=! A(i, kk) * B(kk, j)
}
"""
在此示例中，我们定义了一个函数 matmul，它接收两个输入 A 和 B（形状分别为 M x N 和 N x K），并返回一个 output。output 的形状由 TC 语言自动推断（如下所述）。
让我们看看这一行
output(i, j) +=! A(i, kk) * B(kk, j)
它表示
output(i, j)表示输出是二维的。- 对于每一个位置
output(i, j)，我们将A(i, kk) * B(kk, j)进行累加 (+=)。 i被定义为A第 0 维的所有位置，即i in range(0, M)。j被定义为B第 1 维的所有位置，即j in range(0, K)。kk被推断为从0到N的所有位置。
输出的形状是根据 i 和 j 能取到的最大值（即 M 和 K）推断出来的，因此输出大小为 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 自动调优器要求提前指定所有输入的大小。例如，如果您有一个作为图像批处理的输入 I1，自动调优器需要知道 I1 的确切形状才能生成优化的内核。您不能指定：“高度在 200 到 300 之间的图像”。这在自然语言处理等序列数据中更为关键，因为每个句子的长度可能不同。
自动调优器是非参数化的，原因在于对参数化约束进行自动调优越来越困难，这是一个活跃的研究领域。因此，在第一个版本中，我们慎重决定以一种我们确信能良好运行的形式提供该工具。
作为一种变通方法，如果您知道自己关注几种特定的形状，可以使用这些多种形状来运行自动调优器。
relu = tc.define(LANG, name="relu")
batch, channels = 16, 3
tc.autotune((batch, channels, 32, 32)) # image of size 32 x 32
tc.autotune((batch, channels, 48, 48)) # image of size 48 x 48
tc.autotune((batch, channels, 64, 64)) # image of size 64 x 64
现在，自动调优器已针对这三种特定的图像大小 32x32、48x48 和 64x64 进行了调优。
缺乏循环
如果您想编写一个 RNN，很容易将其看作是一个关于时间的 for 循环。然而，TC 语言尚不支持循环。如果您确实想编写 RNN，可以编写展开的循环。
跨步张量
TC 后端尚不支持非连续张量。如果您提供的输入不是连续的，它们会在传递给 TC 后端之前被转换为连续的。
在 TC 表达式中重塑张量
您不能在 TC 中编写如下操作：torch.matmul(...).view(...).mean(...)。每当需要使用 view 来改变输入的形状时，您必须获取输出，并在 PyTorch 层面对其进行 view。
入门指南
- 教程演练，快速开始了解并使用 Tensor Comprehensions PyTorch 软件包。
- 包含 20 多个各种机器学习层的 TC 示例，包括
avgpool、maxpool、matmul、matmul（提供输出缓冲区）以及batch-matmul、convolution、strided-convolution、batchnorm、copy、cosine similarity、Linear、Linear + ReLU、group-convolutions、跨步group-convolutions、indexing、Embedding（查找表）、small-mobilenet、softmax、tensordot、transpose。 - 关于 Tensor Comprehensions 以及与 PyTorch 集成的详细文档。
交流
- Slack：关于框架集成、构建支持、协作等方面的讨论，请加入我们的 Slack 频道。
- 电子邮件：tensorcomp@fb.com
- GitHub：错误报告、功能请求、安装问题、RFC、想法等。
致谢
我们要感谢 Soumith Chintala、Edward Yang 和 Sam Gross 在使集成 API 变得优雅顺滑方面给予的巨大指导和帮助。我们还要感谢 PyTorch 团队的其他成员以及我们的预发布用户，感谢他们提供的宝贵反馈，这些反馈引导我们改进了集成效果。