PyTorch自定义非线性MLP训练加速方案求助

核心瓶颈分析 #

你的代码速度慢的根源在于两处串行计算：

1. 神经元遍历循环：每层神经元通过Python循环逐个计算，完全没利用PyTorch的张量并行能力
2. 非线性函数内部递归循环：input_output_nonlinearity_torch里对每个权重维度做迭代，进一步放大串行开销

针对性加速方案 #

1. 神经元级并行：向量化层计算 #

将每层所有神经元的计算批量处理，修改input_output_nonlinearity_torch以接受批量权重，替换所有串行循环为张量操作。

关键代码调整：

首先修改非线性函数，支持批量神经元计算：

def input_output_nonlinearity_torch(x, w_batch, numpoints=100, C_total=1, readoutStrength=1):
    num_neurons = w_batch.shape[0]
    device = x.device
    dtype = torch.complex64  # 若精度允许，用complex64替代complex128可大幅提速

    # 预生成网格点并移到对应设备
    z_values = torch.linspace(1e-10, 1-1e-10, numpoints, device=device, dtype=dtype)
    w_values = torch.linspace(1e-10, 1-1e-10, numpoints, device=device, dtype=dtype)

    # 初始化批量状态张量：[num_neurons, numpoints]
    Bin = torch.zeros(num_neurons, numpoints, dtype=dtype, device=device)
    c_per_mode = C_total / w_batch.shape[1]

    # 向量化递归计算：所有神经元并行处理每个权重维度
    for i in range(w_batch.shape[1]):
        E_c_val = w_batch[:, i].unsqueeze(1)  # [num_neurons, 1]，适配广播
        E_p_val = x[i]  # 输入维度值，自动广播到所有神经元
        Bin = spinwave_recursive_calculation_torch(Bin, z_values, w_values, E_c_val, E_p_val, c_per_mode)

    # 批量计算最终输出
    output_Efield_w_z = final_readoutKernel_torch(z_values, w_values, readoutStrength, Bin, kval=1)
    output_Efield_w = torch.trapz(torch.real(output_Efield_w_z), x=z_values, dim=2)
    output_Efield = torch.trapz(torch.real(output_Efield_w), x=w_values, dim=1)
    return output_Efield

然后修改MLP的forward方法，去掉神经元循环：

def forward(self, x):
    x_squeezed = x.squeeze()  # [input_size]
    # 第一层：权重转置为[hidden_size, input_size]，批量计算所有隐藏神经元
    hidden_outputs = input_output_nonlinearity_torch(x_squeezed, self.weights1.T, C_total=C_total, readoutStrength=readoutStrength)
    hidden_outputs = F.relu(hidden_outputs)  # [hidden_size]

    # 第二层：批量计算所有输出神经元
    final_outputs = input_output_nonlinearity_torch(hidden_outputs, self.weights2.T, C_total=C_total, readoutStrength=readoutStrength)
    return final_outputs.unsqueeze(0)  # 适配[batch_size, num_classes]形状

2. GPU加速细节优化 #

• 确保所有张量（包括z_values、w_values）都移到GPU，避免CPU-GPU数据传输开销
• 使用torch.complex64替代torch.complex128（若精度满足需求），减少内存占用并提升计算速度
• 预计算固定参数：把readinKernel_torch中的ic、gamma、Np等固定值提前计算为张量，避免每次调用重复计算

3. 批量数据处理 #

当前batch_size=1，可尝试增大到8/16/32，同时修改input_output_nonlinearity_torch支持批量输入x（形状[batch_size, input_size]），进一步放大并行效率。

原理说明 #

PyTorch的GPU加速和自动微分完全依赖张量的向量化操作，Python循环会绕过所有底层优化。通过将串行逻辑转为张量操作，PyTorch可将计算任务分发到GPU多核心并行执行，同时自动维护反向传播的梯度计算链路。

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