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 和 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 张量作为输入并返回一个 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 会保存自动调优的 kernel 搜索结果，并将其保存到文件 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，它接受形状为 M x N 和 N x K 的两个输入 A 和 B，并返回一个输出。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在range(0, M)范围内j被明确定义为B在维度 1 中的所有位置，即j在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 的确切形状才能生成优化的 kernel。您不能指定：高度在 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 个使用 TC 实现各种 ML 层的示例，包括
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：错误报告、功能请求、安装问题、RFC、想法等。
致谢
我们要感谢 Soumith Chintala、Edward Yang 和 Sam Gross 在使集成 API 变得良好流畅方面给予的巨大指导和帮助。我们还要感谢 PyTorch 团队的其他成员和我们的预发布用户提供的有益反馈，这些反馈指导我们改进了集成。
