Python实现双无穷和定义的四变量函数：求高效优化方案

你的代码运行缓慢的核心问题在于Python双重循环的固有低效、频繁的Torch/Numpy数据转换开销，以及逐元素累加未利用并行运算能力。以下是针对性的优化方案：

关键优化思路 #

1. 抛弃Python循环，改用PyTorch张量的广播机制+批量求和，利用底层C++并行计算能力；
2. 全程使用PyTorch张量运算，避免Torch与Numpy之间的来回转换，减少数据拷贝成本；
3. 预计算与输入无关的固定系数，避免重复计算。

优化后的代码 #

import torch

# 预计算求和系数（与输入x无关，仅需计算一次）
_M = torch.arange(1, 100).unsqueeze(1)  # shape: (99, 1)
_N = torch.arange(1, 100).unsqueeze(0)  # shape: (1, 99)
_COEFF = 4.0 / ((_N + _M) * torch.pi**5)

def myfunction(x):
    device = x.device
    # 将预计算系数移到输入所在设备（CPU/GPU）
    coeff = _COEFF.to(device)
    
    # 提取x的四个变量，保留batch维度
    x1 = x[:, 0:1]
    x2 = x[:, 1:2]
    x3 = x[:, 2:3]
    x4 = x[:, 3:4]
    
    # 计算各三角函数项，利用广播自动匹配维度
    term1 = torch.sin(_N * x1)  # shape: (batch_size, 99)
    term2 = torch.sin(_M**2 * x2)  # shape: (batch_size, 99)
    term3 = torch.sin(_N * torch.pi * x3)  # shape: (batch_size, 99)
    term4 = torch.cos(_M * 2 * x4)  # shape: (batch_size, 99)
    
    # 合并所有项并对m、n维度求和
    total_terms = coeff * term1.unsqueeze(1) * term2.unsqueeze(2) * term3.unsqueeze(1) * term4.unsqueeze(2)
    value = total_terms.sum(dim=(-1, -2))
    
    return value

优化细节说明 #

• 预计算系数：_COEFF在模块加载时仅计算一次，避免每次调用函数重复计算相同的固定系数；
• 张量广播：通过unsqueeze调整维度，让m、n与输入的batch维度自动匹配，实现批量并行计算；
• 设备对齐：将预计算系数移到输入x所在的设备（CPU/GPU），确保全程运算在同一设备进行，避免跨设备数据传输开销；
• 批量求和：用sum(dim=(-1,-2))一次性完成m和n维度的求和，替代原代码中10000次循环累加，效率提升数倍至数十倍。

额外优化建议 #

• 调整求和上限：如果函数收敛速度快，可测试将求和上限从100降低到50或更小，在精度可接受的前提下进一步减少计算量；
• GPU加速：若环境支持CUDA，只需将输入x移至GPU（x = x.cuda()），代码会自动利用GPU并行计算，速度提升更为显著；
• 精度权衡：若无需高精度，可将张量类型改为float16（需GPU支持），进一步减少内存占用和计算时间。

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