如何在张量的多行上并行调用多个函数

我明白你的需求——想要给张量的每一行匹配对应的复杂函数做变换，而且希望能并行执行，不想拆解那些复杂函数。你之前用vmap尝试的思路是对的，但问题出在当vmap把indices转换成批处理张量后，没法直接用它索引Python函数列表。下面给你几个可行的方案，从简单到更贴合并行需求的：

方案一：预计算所有函数结果再选择（简单高效，适合函数数量不多的场景） #

这个方法先让所有函数对整个张量做运算，然后根据indices挑选出每一行对应的结果。因为PyTorch的张量运算会自动在GPU上并行执行，所以这个方法的并行性很好，代码也很简洁：

import torch

x = torch.ones(3, 3)
factors = [lambda x: 2*x, lambda x: 3*x, lambda x: 4*x]
indices = torch.tensor([0, 1, 2])

# 让所有函数处理整个张量，得到形状为 (函数数量, 行数, 列数) 的结果
all_results = torch.stack([func(x) for func in factors], dim=0)
# 根据indices选择每一行对应的函数结果
result = all_results[indices, torch.arange(x.shape[0])]

print(result)
# 输出：
# tensor([[2., 2., 2.],
#         [3., 3., 3.],
#         [4., 4., 4.]])

如果你的实际函数是复杂的自定义变换，只要它们能接收张量输入并返回同形状张量，这个方法就完全适用，不需要做任何拆解。唯一的小缺点是如果函数数量远大于行数，会计算一些不需要的结果，但大多数场景下这个代价可以接受。

方案二：用torch.vmap结合ModuleList（适合需要严格对应行-函数，不想预计算多余结果的场景） #

如果不想预计算所有函数的结果，可以把函数包装成PyTorch模块，然后用vmap配合一个能根据索引调用对应模块的函数。这里需要把函数转换成nn.Module子类，这样PyTorch就能处理批处理索引的问题：

import torch
from torch import nn

x = torch.ones(3, 3)
indices = torch.tensor([0, 1, 2])

# 把你的复杂函数包装成Module子类
class CustomTransform(nn.Module):
    def __init__(self, factor):
        super().__init__()
        self.factor = factor
    
    def forward(self, x):
        # 这里可以替换成你的复杂变换逻辑
        return self.factor * x

# 用ModuleList管理所有变换函数
transforms = nn.ModuleList([CustomTransform(2), CustomTransform(3), CustomTransform(4)])

# 定义一个能根据索引调用对应变换的函数
def apply_transform(row, idx):
    return transforms[idx](row)

# 用vmap批处理每一行和对应的索引
result = torch.vmap(apply_transform, in_dims=(0, 0))(x, indices)

print(result)
# 输出和预期一致

这个方案的好处是只计算每一行需要的函数结果，不会浪费计算资源。而且vmap会自动处理并行逻辑，不管是CPU还是GPU都能高效执行。

方案三：用多线程/多进程并行（适合CPU上的大规模计算） #

如果你的计算主要在CPU上，而且函数非常耗时，可以用PyTorch的多线程池来并行执行：

import torch
from torch.multiprocessing import Pool

x = torch.ones(3, 3)
factors = [lambda x: 2*x, lambda x: 3*x, lambda x: 4*x]
indices = torch.tensor([0, 1, 2])

# 定义并行执行的任务
def process_row(args):
    row, idx = args
    return factors[idx](row)

# 创建线程池，并行处理每一行
with Pool(processes=3) as pool:
    results = pool.map(process_row, zip(x, indices))

# 把结果堆叠成张量
result = torch.stack(results)

print(result)

不过这个方案需要注意张量的内存共享问题，在GPU上使用多进程会比较麻烦，所以更适合CPU场景。

总结一下，如果你用GPU，方案一和方案二都是很好的选择；如果是CPU上的大规模计算，可以考虑方案三。这些方法都不用涉及PyTorch Streams，代码也比较干净。

备注：内容来源于stack exchange，提问作者gfdb