大BatchSize下多Torch函数高效执行的实现方案问询

首先纠正原代码中的笔误：循环里的y[1,:]应该是y[i,:]，否则所有行都会被最后一次循环覆盖。

针对批量较大时循环效率低的问题，核心思路是避免Python层面的循环，利用PyTorch的向量化操作和底层并行计算能力，具体实现如下：

import torch

# 示例：function_list为固定的element-wise函数集合
function_list = [torch.sin, torch.exp, torch.tanh]
# function_choice为batchsize长度的索引数组，指定每行对应使用的函数
function_choice = [0,1,2,0,1]
batchsize = len(function_choice)
dimension = 3  # 示例特征维度

def optimized_weird_function(x):
    # x shape: [1, dimension]
    # 一次性计算所有候选函数对x的输出，得到形状 [num_funcs, dimension] 的张量
    all_func_outputs = torch.stack([func(x) for func in function_list], dim=0)
    # 将function_choice转为PyTorch长整型张量，用于索引
    choice_indices = torch.tensor(function_choice, dtype=torch.long)
    # 根据索引直接选取对应行，自动扩展为 [batchsize, dimension]
    y = all_func_outputs[choice_indices]
    return y

# 测试示例
x = torch.randn(1, dimension)
result = optimized_weird_function(x)
print(result.shape)  # 输出: torch.Size([5, 3])

性能提升原因 #

原代码的Python循环会逐次调用函数，无法利用PyTorch的向量化加速；优化后的代码先批量计算所有函数的输出，再通过张量索引一次性完成选取，所有操作都在PyTorch的底层C++/CUDA执行路径中，能充分利用CPU/GPU的并行计算能力——即使在CPU环境下，向量化操作的效率也远高于Python循环，GPU环境下性能提升会更显著。

特殊场景适配 #

如果你的function_list确实是每个batch元素对应一个独立函数（即长度等于batchsize），可以使用PyTorch 1.10+支持的torch.vmap来批量映射函数：

def optimized_weird_function(x):
    # 将x扩展为 [batchsize, dimension]，每个元素都是原x的复制
    x_batch = x.repeat(batchsize, 1)
    # 使用vmap批量应用每个函数到对应位置的x
    y = torch.vmap(lambda func, x: func(x))(function_list, x_batch)
    return y

这种方式同样规避了Python循环，利用PyTorch的自动微分和并行机制完成高效计算。

内容的提问来源于stack exchange，提问作者喵喵露