张量综合 (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
目前，我们只提供已在 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 保存自动优化内核搜索的结果并将其保存到文件 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维度 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)) # 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。
入门
- 分步教程 帮助您快速入门，了解和使用张量综合 PyTorch 包。
- 提供 20 多个使用 TC 的各种 ML 层示例，包括
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 团队的其他成员和我们的预发布用户提供的宝贵反馈，这些反馈指导我们更好地进行了集成。