PyTorch中自定义多对一非线性函数的反向传播与训练方法咨询

我写了一个输入输出非线性函数的代码，它接收输入列表X和权重W，生成单个非线性输出。我想把它当作「神经元」，通过反向传播来训练（理想情况可以把多个这类神经元串联起来，目前先搞定单个神经元的情况就行）。

之前我问过多对一非线性函数的训练可行性，得到的答复说可以用autograd实现反向传播。以下是我的代码：

import numpy as np
from scipy import integrate, special
from scipy.constants import epsilon_0

z_values =  np.linspace(1e-10, 1-1e-10, 100)

def readinKernel(wdummy, z, Ec, Ep, kval=1):
    return (Ec * kval * special.jv(0, Ec * kval * np.sqrt(np.outer(z, (1 - wdummy))))* Ep / np.cosh(np.arctanh(wdummy)))

def steep_sigmoid(x, k=50):    
    return 1.0 / (1.0 + np.exp(-k*x))

def readoutKernel(zdummy, z, B_in, Ec, kval=1):
    return (1 / np.sqrt(np.maximum(1e-10, np.subtract.outer(z, zdummy))) * 
            special.jv(1, 2 * Ec * kval * 
                       np.sqrt(np.maximum(1e-10, np.subtract.outer(z, zdummy)))) *
            Ec * kval *
            steep_sigmoid(np.subtract.outer(z, zdummy), 50) * 
            np.repeat(B_in, len(zdummy)).reshape(len(B_in), len(zdummy)))

def spinwave_recursive_calculation(B_in, z_values, Ec, Ep):

    wdummy_values = np.linspace(1e-10, 1-1e-10, 100)
    zdummy_values = np.linspace(1e-10, 1-1e-10, 100)

    readin_values = readinKernel(wdummy_values, z_values, Ec, Ep)
    readout_values = readoutKernel(zdummy_values, z_values, B_in, Ec)

    readin_integrals = np.trapz(readin_values, wdummy_values, axis=1)
    readout_integrals = np.trapz(readout_values, zdummy_values, axis=1)

    spinwave = readin_integrals - readout_integrals + B_in
    return spinwave

def input_output_nonlinearity(x, w):

    Bin = np.zeros(len(z_values))
    BoutMatrix = np.tile(Bin, [len(w),1])

    for i in range(len(w)):
        E_c_val = w[i]
        E_p_val = x[i]
        # Origionally this was a recursive funciton that returned a single output array, but now I have implemented it as a for loop that returns a matrix of outputs
        # This was to try to avoid a potential bug in pytorch which was recommended by a user DerekG. (I'm not sure my implementation fixes this issue though)
        # Bout = spinwave_recursive_calculation(Bin, z_values, E_c_val, E_p_val)
        BoutMatrix[i, :] = spinwave_recursive_calculation(Bin, z_values, E_c_val, E_p_val)
        Bin = BoutMatrix[i, :]
    output = print(np.sum(np.abs(Bout)))
    return output

x = np.array([1, 0, 0])
w = np.array([1, .5, 1])

input_output_nonlinearity(x, w)

请问当前代码能不能直接实现反向传播？有哪些可行的训练思路可以参考？

一、当前代码能否直接实现反向传播？ #

不行，当前代码完全基于NumPy和SciPy编写，这两个库不支持自动微分（autograd），没办法直接计算权重W的梯度，自然没法做反向传播。具体问题点包括：

• NumPy/SciPy的所有运算都是静态数值计算，不会记录计算图，反向传播需要的梯度信息根本没法追踪
• 代码里用了np.trapz（数值积分）、special.jv（贝塞尔函数）这类SciPy函数，这些函数没有对应的自动微分实现
• 最后返回的是print的结果（output = print(np.sum(np.abs(Bout)))），这会导致函数返回None，连输出值都没法正确获取，更别说计算梯度了

二、可行的训练思路 #

1. 改用PyTorch/TensorFlow实现自动微分 #

把现有代码迁移到PyTorch（推荐，因为对科学计算的支持更灵活）或TensorFlow，替换所有NumPy/SciPy的运算为框架自带的API：

• 用torch.linspace替代np.linspace，torch.trapz替代np.trapz
• PyTorch有torch.special.jv实现贝塞尔函数，直接替换scipy.special.jv
• 把权重W定义成torch.nn.Parameter，输入X定义成torch.Tensor，这样框架会自动构建计算图，调用loss.backward()就能自动计算梯度
• 修复代码里的bug：output = print(np.sum(np.abs(Bout)))要改成output = torch.sum(torch.abs(BoutMatrix[-1]))，确保返回有效的张量输出

2. 手动实现反向传播（不推荐，复杂度高） #

如果不想用深度学习框架，就得手动推导每个运算的梯度，然后用链式法则把梯度从输出传到权重W：

• 先推导spinwave_recursive_calculation对Ec（即w[i]）的偏导数，再推导input_output_nonlinearity对每个w[i]的梯度
• 数值积分的梯度可以用积分交换法则推导，贝塞尔函数的梯度可以用其递推公式计算
• 这种方法工作量极大，容易出错，只适合简单场景，你的代码里有多个嵌套运算和特殊函数，手动推导几乎不现实

3. 用JAX实现自动微分 #

JAX是针对科学计算的自动微分框架，兼容大部分NumPy/SciPy的API，改动成本比PyTorch小：

• 把import numpy as np改成import jax.numpy as np，from scipy import special改成from jax.scipy import special
• 用jax.grad装饰器直接获取input_output_nonlinearity对W的梯度
• JAX支持自动微分数值积分、特殊函数，不需要大量修改原有代码结构

4. 单个神经元的训练流程（以PyTorch为例） #

• 定义数据集：准备输入X和对应的目标输出y
• 初始化权重W为nn.Parameter(torch.tensor([1., 0.5, 1.]))
• 定义损失函数（比如MSE损失：loss_fn = torch.nn.MSELoss()）
• 定义优化器（比如SGD或Adam：optimizer = torch.optim.Adam([W], lr=1e-3)）
• 训练循环：
  1. 前向传播：计算模型输出y_pred = input_output_nonlinearity(X, W)
  2. 计算损失：loss = loss_fn(y_pred, y)
  3. 反向传播：optimizer.zero_grad(); loss.backward()
  4. 更新权重：optimizer.step()

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