本文介绍如何利用autograd对sir微分方程模型进行可微分参数估计，通过替换scipy的odeint为autograd封装版本、规范参数传递方式，使损失函数可自动求导，从而实现端到端的梯度下降优化。
本文介绍如何利用autograd对sir微分方程模型进行可微分参数估计，通过替换scipy的odeint为autograd封装版本、规范参数传递方式，使损失函数可自动求导，从而实现端到端的梯度下降优化。
在传染病建模中，SIR模型（Susceptible–Infected–Recovered）是基础但极具代表性的常微分方程系统。当面对真实观测数据时，反向估计传播率（β）和康复率（γ）是典型的数据拟合任务。本文面向具备NumPy基础但尚未接触深度学习框架的本科生，采用轻量级自动微分库 autograd 实现可微分ODE拟合——无需JAX或PyTorch，仅需修改两处关键细节即可打通梯度流。
✅ 正确实现的关键要点
-
必须使用 autograd 封装的
odeint
原生scipy.integrate.odeint不支持自动微分。需显式导入：from autograd.scipy.integrate import odeint
该版本已注册梯度规则，使整个ODE求解过程可被autograd追踪。
-
参数必须以
tuple形式传入odeint，且禁用关键字参数
autograd 当前不支持带关键字参数（如args=(beta, gamma)）调用封装后的odeint（issue #681）。正确写法为：params_tuple = tuple(params) # e.g., [0.4, 0.07] → (0.4, 0.07) sol = odeint(sir, Y0, t, params_tuple) # 全部位置参数，无 keywords
损失函数需保持纯函数式结构
避免副作用（如全局变量、plt绘图）、确保所有计算均通过autograd.numpy（即np）完成。
✅ 完整可运行示例
import autograd
import autograd.numpy as np
import matplotlib.pyplot as plt
from autograd.scipy.integrate import odeint
from autograd.builtins import tuple
from autograd import grad
# SIR ODE system — uses autograd.numpy for differentiability
def sir(y, t, beta, gamma):
S, I, R = y
dS_dt = -beta * S * I
dI_dt = beta * S * I - gamma * I
dR_dt = gamma * I
return np.array([dS_dt, dI_dt, dR_dt])
# Loss function: L2 norm between observed and simulated trajectories
def loss(params, Y0, t, y_obs):
beta, gamma = params
params_tuple = tuple(params)
sol = odeint(sir, Y0, t, params_tuple) # ✅ autograd-aware ODE solver
return np.linalg.norm(y_obs - sol) # scalar loss
# Generate synthetic noisy observations
np.random.seed(42)
Y0 = np.array([0.95, 0.05, 0.0])
t = np.linspace(0, 30, 101)
true_beta, true_gamma = 0.5, 1/14
sol_true = odeint(sir, Y0, t, tuple([true_beta, true_gamma]))
y_obs = sol_true + np.random.normal(0, 0.05, size=sol_true.shape)
# Gradient computation & gradient descent
loss_grad = grad(loss, argnum=0)
params = np.array([0.3, 0.05]) # initial guess
learning_rate = 0.01
n_iterations = 200
for i in range(n_iterations):
grads = loss_grad(params, Y0, t, y_obs)
params -= learning_rate * grads
if i % 50 == 0:
print(f"Iter {i}: β={params[0]:.4f}, γ={params[1]:.4f}, loss={loss(params, Y0, t, y_obs):.4f}")
print(f"\n✅ Final estimate: β ≈ {params[0]:.4f}, γ ≈ {params[1]:.4f}")
print(f"True values: β = {true_beta:.4f}, γ = {true_gamma:.4f}")⚠️ 注意事项与常见陷阱
-
不要混用
numpy和autograd.numpy：务必全程使用import autograd.numpy as np，否则grad()无法追踪。 -
odeint输入顺序固定：odeint(func, y0, t, args)——args必须是tuple，不可为list或dict。 -
避免
solve_ivp：autograd当前未提供其封装版本，solve_ivp无法求导。 -
梯度爆炸风险：SIR对初值敏感，建议对
params做简单裁剪（如np.clip(params, 1e-4, 5)），或使用自适应学习率（如Adam可自行基于autograd实现）。 -
性能提示：
autograd.odeint比原生scipy.odeint略慢，但对教学级问题（百量级时间点）完全可接受。
通过以上改造，你将获得一个真正端到端可微分的SIR拟合流程：从ODE求解、残差计算到梯度更新全部由 autograd 自动完成。这不仅为参数估计提供了理论严谨性，也为后续拓展至SEIR、空间SIR或多目标损失（如加权I曲线拟合）打下坚实基础。
