张量推导 (Tensor Comprehensions, TC) 是一种降低编写高性能代码门槛的工具。它能从简单的高级语言生成 GPU 代码，并针对特定的输入大小自动调优代码。
我们强烈建议您先阅读 张量推导博客文章。
如果您遇到以下任何情况，TC 将是一个有用的工具。
- 您的 PyTorch 层很大且运行缓慢，您曾考虑为其编写专用的 C++ 或 CUDA 代码。但是您不知道如何用 CUDA 编程或编写低级代码。
- 您编写了一个 CUDA 层，但花费了一周时间来编写、调试和优化速度。您希望能在一个小时内完成。
- 您想在网络中融合多个层以提高速度，例如 Conv-ReLU-BatchNorm 或 Linear-ReLU-Linear-ReLU，但这很难理解。
- 您的研究涉及 CuDNN 和 MKL 未优化的奇怪张量形状。例如，您对 13 x 24 的输入图像进行 143 x 55 的卷积。您尝试使用 CuDNN 运行，但它比您预期的要慢。
- 您的代码因不断转置张量以适应特定的内存布局而变慢。您希望可以轻松编写能够高效操作输入布局的自定义代码。
张量推导在 PyTorch 中使用无缝，与 PyTorch 张量和 nn 变量互操作。
让我们看看如何在 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 张量作为输入并返回一个 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 中。下次您调用相同的代码行时，它会加载自动调谐结果而无需重新计算。
自动调谐器有一些超参数（就像您的 ConvNet 有学习率、层数等）。我们选择了合理的默认值，但您可以在 此处 阅读有关使用高级选项的信息。
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的 dim=0 中的所有位置都定义良好，即i in range(0, M)。j在B的 dim=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)) # 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 个包含各种机器学习层的示例，包括
avgpool、maxpool、matmul、matmul – 提供输出缓冲区和batch-matmul、convolution、strided-convolution、batchnorm、copy、cosine similarity、Linear、Linear + ReLU、group-convolutions、带步幅的group-convolutions、indexing、Embedding(查找表)、小型 mobilenet、softmax、tensordot、transpose。 - 关于张量推导及其与 PyTorch 集成的 详细文档。
联系方式
- Slack：有关框架集成、构建支持、协作等讨论，请加入我们的 Slack 频道。
- 电子邮件：tensorcomp@fb.com
- GitHub：错误报告、功能请求、安装问题、RFC、想法等。
鸣谢
我们衷心感谢 Soumith Chintala、Edward Yang 和 Sam Gross 在使集成 API 顺畅方面给予的巨大指导和帮助。我们还要感谢 PyTorch 团队的其他成员和我们的预发布用户提供的有益反馈，这些反馈指导我们改进了集成。