张量推导式 (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 运行它，但速度比你希望的要慢。
-
你的代码因不断转置张量以适应特定的内存布a局而变慢。你希望可以轻松编写能够在你输入布局上高效运行的自定义代码。
Tensor Comprehensions 在 PyTorch 中可以无缝使用，与 PyTorch Tensors 和 nn Variables 交互操作。
让我们快速过一遍在 PyTorch 中使用 TC 的方法。
1. 安装包
conda install -c pytorch -c tensorcomp tensor_comprehensions
目前我们仅提供 Linux-64 位二进制文件，这些文件已在 Ubuntu 16.04 和 CentOS7 上测试过。
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。它接收输入 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")
自动调优器 (autotuner) 是你最好的朋友。通常，在没有先对其进行自动调优的情况下，你不应使用 tc 函数。
当自动调优运行时，会显示当前的最佳性能。如果你对当前结果满意或者时间不够，可以通过按 Ctrl+C 来停止调优过程。
cache 会保存自动调优内核搜索的结果，并将其保存到文件 fcrelu_100_128_100.tc 中。下次你调用同一行代码时，它会加载自动调优的结果而无需重新计算。
自动调优器有一些超参数（就像你的 ConvNet 有学习率、层数等）。我们选择了合理的默认值，但你可以在这里阅读关于使用高级选项的信息。
6. 使用输入调用该函数，以获取结果
out = fcrelu(I, W1, B1)
现在，让我们看看如何编写 TC 表达式。
TC 语言快速入门
TC 表示法侧重于层的数学本质，将性能考量留给其后端代码，后端代码使用了 Halide 和多面体编译技术，这些技术积累了数十年前沿的循环嵌套优化 (Loop Nest Optimization, 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 之间的图像。这在序列数据（如 NLP）中更为关键，其中每个句子的长度可能不同。
自动调优器是非参数化的，原因在于自动调优参数化约束越来越难，这是一个活跃的研究领域。因此，在第一个版本中，我们审慎决定以我们知道它能良好工作的形式将工具提供给你。
作为一种变通方法，如果你知道你有一些特定的感兴趣的形状，你可以使用这些多个形状来运行自动调优器。
relu = tc.define(LANG, name="relu")
batch, channels = 16, 3
tc.autotune((batch, channels, 32, 32)) # 32 x 32 大小的图像
tc.autotune((batch, channels, 48, 48)) # 48 x 48 大小的图像
tc.autotune((batch, channels, 64, 64)) # 64 x 64 大小的图像
现在，自动调优器针对这三种特定的图像尺寸 32x32、48x48 和 64x64 进行了调优。
缺少循环
如果你想编写一个 RNN，很容易将其视为一个关于时间的 for 循环。然而，TC 语言目前还没有循环。如果你确实想编写 RNN，你可以编写展开的循环。
跨步张量 (Strided-Tensors)
TC 后端目前还不支持非连续张量 (non-contiguous Tensors)。如果你给出的输入不是连续的，它们会在传递给 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、stridedgroup-convolutions、indexing、Embedding(查找表)、small-mobilenet、softmax、tensordot、transpose - 关于 Tensor Comprehensions 及其与 PyTorch 集成的详细文档。
交流沟通
- Slack：关于框架集成、构建支持、协作等的讨论，请加入我们的 slack 频道。
- 电子邮件：tensorcomp@fb.com
- GitHub：报告 bug、功能请求、安装问题、RFC、想法等。
致谢
我们要感谢 Soumith Chintala、Edward Yang 和 Sam Gross 在使集成 API 友好和顺畅方面给予的巨大指导和帮助。我们也要感谢 PyTorch 团队的其他成员以及我们的预发布用户，他们的有益反馈指导我们改进了集成。
