使用PyTorch最小化函数时参数未更新的技术问题

问题现象 #

使用PyTorch进行Lotka-Volterra模型参数估计时，执行最小化函数后模型参数始终保持初始值，未得到优化。

相关代码 #

数据获取代码 #

import pandas as pd
import torch

dataset = pd.read_csv('/kaggle/input/fish-bear/population_data.csv')
years = torch.tensor(dataset['year'], dtype = torch.float64)
fish_pop = torch.tensor(dataset['fish_hundreds'], dtype = torch.float64)
bears_pop = torch.tensor(dataset['bears_hundreds'], dtype = torch.float64)
pop = torch.cat((fish_pop.reshape((51, 1)), bears_pop.reshape((51, 1))), 1)

常微分方程求解器 #

from typing import List, Callable, Sequence, NamedTuple, Union

class _Tableau(NamedTuple):
    c: List[float]
    b: List[float]
    a: List[List[float]]


rk4_tableau = _Tableau(c=[0.0, 0.5, 0.5, 1.0],
                       b=[1 / 6., 1 / 3., 1 / 3., 1 / 6.],
                       a=[[0.0, 0.0, 0.0, 0.0], [0.5, 0.0, 0.0, 0.0],
                          [0.0, 0.5, 0.0, 0.0], [0.0, 0.0, 1.0, 0.0]])


def explicit_rk(tableau: _Tableau, fcn: Callable[..., torch.Tensor],
                y0: torch.Tensor, t: torch.Tensor,
                params: Sequence[torch.Tensor]):
    c = tableau.c
    a = tableau.a
    b = tableau.b
    s = len(c)
    nt = len(t)

    # set up the results list
    yt_lst: List[torch.Tensor] = []
    yt_lst.append(y0)
    y = yt_lst[-1]
    for i in range(nt - 1):
        t0 = t[i]
        t1 = t[i + 1]
        h = t1 - t0
        ks: List[torch.Tensor] = []
        ksum: Union[float, torch.Tensor] = 0.0
        for j in range(s):
            if j == 0:
                k = fcn(y, t0, params)
            else:
                ak: Union[float, torch.Tensor] = 0.0
                aj = a[j]
                for m in range(j):
                    ak = aj[m] * ks[m] + ak
                k = fcn(h * ak + y, t0 + c[j] * h, params)
            ks.append(k)
            ksum = ksum + b[j] * k
        y = h * ksum + y
        yt_lst.append(y)
    yt = torch.stack(yt_lst, dim=0)
    return yt

def rk4_ivp(fcn: Callable[..., torch.Tensor], y0: torch.Tensor, t: torch.Tensor,
            params: Sequence[torch.Tensor], **kwargs):
    return explicit_rk(rk4_tableau, fcn, y0, t, params)

最小化代码 #

import torch

def lotka_volterra(y, t, params):
    y1, y2 = y
    a, b, c, d = params

    return torch.tensor([a * y1 - b * y1 * y2, c * y2 * y1 - d * y2])

def loss_function(params):
    y0 = torch.tensor([fish_pop[0], bears_pop[0]], dtype = torch.float64)
    t = torch.linspace(years[0], years[-1], len(years), dtype = torch.float64)
    output = rk4_ivp(lotka_volterra, y0, t, params)
    loss = torch.sum((output - pop)**2)
    loss.requires_grad = True
    return loss

def minimize(loss_function, initial_parameters: torch.Tensor):
    list_params = []
    params = initial_parameters
    params.requires_grad = True
    optimizer = torch.optim.SGD([params], lr=0.5)

    for i in range(5):
        optimizer.zero_grad()
        loss: torch.Tensor = loss_function(params)
        loss.backward()
        optimizer.step()
        list_params.append(params.detach().clone())

    return params, list_params

starting_point = torch.nn.Parameter(torch.tensor([1.1, .4, .1, .4], dtype = torch.float64))
minimized_params, list_of_params = minimize(loss_function, starting_point)

loss_function(minimized_params), minimized_params

执行结果 #

(tensor(118.6865, dtype=torch.float64, requires_grad=True),
 Parameter containing:
 tensor([1.1000, 0.4000, 0.1000, 0.4000], dtype=torch.float64,
        requires_grad=True))

问题原因与解决方法 #

核心问题 #

1. 计算图断裂：lotka_volterra函数中用torch.tensor()包装返回值，切断了PyTorch对参数params到输出的梯度传播路径，导致梯度无法回传，优化器无法更新参数。
2. 冗余的梯度设置：loss_function中手动设置loss.requires_grad = True完全多余，loss本身由可微分运算生成，默认已具备梯度追踪能力，手动设置可能干扰梯度计算。
3. 重复的梯度属性设置：minimize函数中对已标记为Parameter的params重复设置requires_grad = True，无实际作用。

修改后的关键代码 #

修正lotka_volterra函数

直接返回运算表达式，保留计算图的梯度追踪：

def lotka_volterra(y, t, params):
    y1, y2 = y
    a, b, c, d = params
    # 直接返回运算结果，无需重新包装tensor，保留梯度链
    return torch.stack([a * y1 - b * y1 * y2, c * y2 * y1 - d * y2])

修正loss_function函数

移除多余的梯度设置代码：

def loss_function(params):
    y0 = torch.tensor([fish_pop[0], bears_pop[0]], dtype = torch.float64)
    t = torch.linspace(years[0], years[-1], len(years), dtype = torch.float64)
    output = rk4_ivp(lotka_volterra, y0, t, params)
    loss = torch.sum((output - pop)**2)
    # 移除loss.requires_grad = True
    return loss

可选优化：调整优化器参数

原SGD学习率0.5过大，易导致训练不稳定，可降低学习率或改用Adam优化器：

# 替换SGD为Adam，学习率设为0.01
optimizer = torch.optim.Adam([params], lr=0.01)

修改后效果 #

参数将能正常被优化器更新，梯度可正确回传，损失值会随迭代逐步下降。

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