PyTorch代码效率优化：嵌套循环替代方案及通用优化咨询

一、原代码逻辑梳理 #

你的输入张量形状为(batch_size, seq_len, feat_dim)（示例中是(80,5,1024)），代码核心是生成所有满足i<j的特征对，将每对特征在维度2拼接后存入temp，最后把拼接后的原输入和temp合并返回。temp的长度刚好是seq_len*(seq_len-1)/2（即5*4/2=10），与s=15-5=10对应。

二、替换嵌套循环的向量化实现 #

显式嵌套循环无法利用PyTorch的并行计算能力，尤其是GPU环境下效率极低。以下是向量化的替代方案：

def my_fun_optimized(input):
    batch_size, seq_len, feat_dim = input.shape
    # 生成所有i < j的索引对（自动排除i=j的情况）
    i_indices, j_indices = torch.triu_indices(seq_len, seq_len, offset=1)
    # 批量取出对应索引的特征张量
    i_feats = input[:, i_indices, :]  # shape: (batch_size, 10, 1024)
    j_feats = input[:, j_indices, :]  # shape: (batch_size, 10, 1024)
    # 一次性完成所有特征对的拼接
    temp = torch.cat([i_feats, j_feats], dim=-1)  # shape: (batch_size, 10, 2048)
    # 原输入的特征维度拼接
    input_ex = torch.cat([input, input], dim=-1)
    # 合并最终结果
    v = torch.cat([input_ex, temp], dim=1)
    return v

关键优化点说明： #

1. torch.triu_indices：直接生成上三角区域（offset=1表示跳过对角线）的索引对，精准匹配原代码中i<j的循环逻辑，无需手动遍历。
2. 批量索引操作：通过一次索引取出所有需要的特征张量，利用PyTorch的并行计算能力完成后续拼接，彻底消除循环开销。
3. 避免无效初始化：原代码中先创建全零temp再逐个赋值，优化后直接通过计算生成目标张量，减少内存占用和无用操作。

三、其他提升计算速度的建议 #

• 启用混合精度：如果你的模型支持，使用torch.cuda.amp模块开启自动混合精度，可大幅降低显存占用并提升计算速度，尤其适合大批次训练场景。
• JIT编译加速：用torch.jit.script或torch.jit.trace编译优化后的函数，PyTorch会优化底层计算图，进一步提升运行效率：

  optimized_jit = torch.jit.script(my_fun_optimized)
  # 后续调用optimized_jit(input)即可
  

• 固定设备一致性：确保所有输入张量和中间张量都在同一设备（如CUDA）上，避免频繁的设备间数据拷贝，原代码中temp指定了CUDA，需保证输入input也在CUDA上。
• 预计算固定索引：如果seq_len是固定值（比如示例中的5），可以提前计算好i_indices和j_indices并保存，避免每次调用函数重复生成索引。

内容的提问来源于stack exchange，提问作者Nafees Ahmed