如何在 PyTorch 中的矩阵上实现 map 函数？答案

【问题标题】：How can I implement the map function on the matrix in PyTorch?如何在 PyTorch 中的矩阵上实现 map 函数？
【发布时间】：2021-07-25 20:51:53
【问题描述】：

在PyTorch中，我想实现一个如下图所示的计算：

也就是说，这是一个映射函数，应用于矩阵中的每个向量，每个函数都不同，例如torch.sin、torch.cos和torch.tan。如何在 PyTorch 中实现这样的计算（即 GPU 上没有 for 循环）。

【问题讨论】：

【解决方案1】：

您可以简单地遍历行：

my_funcs = [torch.sin, torch.cos, torch.tan]

for i, f in enumerate(my_funcs):
    t[i] = f(t[i])

【讨论】：

这是一个简单的代码，但由于 for 循环，无法在 PyTorch 中有效地使用 GPU。我想在 GPU 上实现 map 函数计算，而不需要 PyTorch 中的 for 循环。
@DongQiao 我不知道有任何 PyTorch 方法（通用 Python 并行化之外）以矢量化方式在张量的不同行上应用不同函数。这与您可能获得的一样有效。您是否实际测试过它是否足够快满足您的需求？
@DongQiao 注意：Python 的 map 确实在底层使用了 for 循环。
@iacob 另外，pytorch GPU 计算本质上是asynchronous。为了避免在每次循环迭代中出现潜在的同步点，您可能应该更喜欢不进行就地分配的东西，例如 t = torch.stack([f(r) for f,r in zip(my_funcs,t)])