在“pytorch”中跨张量广播答案

【问题标题】：broadcasting across tensors in `pytorch`在“pytorch”中跨张量广播
【发布时间】：2020-05-31 16:35:57
【问题描述】：

我使用 pytorch 作为数组处理语言（不是用于传统的深度学习目的），我想知道规范的方法是什么来做“批处理”并行。

例如，假设我想计算 3-d 张量的二维层的 svds（例如使用 torch.svd()），并且我想返回一个堆叠的元组 us，堆叠的 s，堆叠v。

据推测，通过 SIMD 并行的魔力，这应该与单层 svd（在 gpu 上）大致相同的时间完成，但是如何编程呢？

【问题讨论】：

请提供完整的minimal reproducible example 以帮助解决问题。具体来说，您可以用伪代码或类似代码替换有问题的部分，使其更易于理解。
不确定我是否理解这个问题。 Pytorch 已经支持批量数据的 SVD，您希望与现有实现有什么不同？
@jodag 我已经实现了部分 SVD（以获得 K 个最高奇异值/向量）。通过我的代码，所有的火炬功能都是可批处理的（qr、svd、mm、转置），所以一切都很好。如果你知道一个规范的手电筒功能可以做到这一点，我全神贯注。
不幸的是，目前不支持截断的 SVD（尽管有一个 open feature request 似乎已实现并等待拉取请求）。还有两个相关的问题here 和here，您可能会觉得有用。

【解决方案1】：

PyTorch 是一个高级软件库，包含大量用于高度优化的编译代码的 python 包装器。函数或运算符要么支持批处理数据，要么不支持。除了编写自己的 C/C++/CUDA 代码并使用 python 调用它之外，别无他法。

幸运的是，大多数函数都支持批处理（包括jodag 指出的torch.svd()），可以假设开发人员（或编译器）在实现中关注数据并行性。我建议你尽可能地堆叠你的张量。它通常会导致显着的加速。

请注意，批量维度始终是张量的第一个维度。 PyTorch 支持像 +, -, *, / 这样的常见运算符的广播，如文档中的 here 所示。由于可能存在歧义，您有时需要相应地重塑数据以明确您想要什么。例如，如果您想将一批标量添加到一批向量中，您需要执行以下操作：

a = torch.zeros(2, 2)
b = torch.arange(2)
a + b.view(2, 1) # or b.reshape(2, 1)
# tensor([[0., 0.],
          [1., 1.]])

【讨论】：