自定义PyTorch优化器必须继承torch.optim.Optimizer，正确调用super().__init__(params, defaults)，并在step()中按param_groups遍历参数、检查p.grad、用self.state[p]管理状态，确保设备/dtype一致且原地更新。
PyTorch里自定义优化器必须继承 torch.optim.Optimizer
直接写个函数或类而不继承 torch.optim.Optimizer，PyTorch不会识别为合法优化器——调用 step() 时会报 AttributeError: 'MyOptim' object has no attribute 'param_groups'。这个基类不仅提供 param_groups 结构，还封装了 zero_grad()、状态管理（state 字典）、参数遍历逻辑等底层契约。
正确做法是：显式继承 torch.optim.Optimizer，并在 __init__ 中调用 super().__init__(params, defaults)；defaults 是一个字典，用于统一设置超参默认值（如 {'lr': 1e-3, 'beta': 0.9}），后续每个 param_group 可覆盖它。
常见错误包括：
- 漏掉
super().__init__()→ 缺失self.param_groups和self.state，后续所有操作崩溃 - 在
step()中直接修改p.data而不检查p.grad是否为None→ 遇到无梯度的参数（如冻结层）时触发RuntimeError: grad can be None only when foreach=True - 把状态更新写在
__init__里（比如预分配momentum_buffer）→ 每个参数组的参数数量不同，必须在step()中按需初始化
实现带动量的SGD变体：手动更新 p.data 并维护 self.state
PyTorch不自动帮你存中间变量（如动量缓存），所有状态必须显式存在 self.state[p] 中。例如实现带Nesterov修正的SGD，核心逻辑是：
立即学习“Python免费学习笔记（深入）”；
让 macOS 自动化像说话一样简单 —— 用代码控制你的 Mac，无需学习 AppleScript 系统要求 macOS 14+ (Sonoma) Python 3.9+ 需要开启：辅助功能 + 屏幕录制权限
for group in self.param_groups:
for p in group['params']:
if p.grad is None:
continue
grad = p.grad.data
state = self.state[p]
# 第一次访问时初始化 buffer
if len(state) == 0:
state['momentum_buffer'] = torch.zeros_like(p.data)
buf = state['momentum_buffer']
buf.mul_(group['momentum']).add_(grad)
# Nesterov：先加动量再减 lr * grad
p.data.add_(buf, alpha=-group['lr'])
注意三点：
-
torch.zeros_like(p.data)必须和参数同设备、同dtype，否则在GPU模型上会报Expected all tensors to be on the same device - 别用
p.data = ...赋值（会切断计算图引用），始终用p.data.add_()或p.data.copy_()原地操作 - 如果想支持
foreach=True（批量参数更新加速），得重写整个step()逻辑并调用torch._foreach_*系列函数，普通自定义优化器建议保持foreach=False（默认）
集成到训练循环前，务必验证 param_groups 行为是否符合预期
PyTorch允许按层设置不同学习率，例如：optimizer = MyOptim([{'params': model.layer1.parameters(), 'lr': 1e-4}, {'params': model.layer2.parameters(), 'lr': 1e-3}])。你的自定义优化器必须能正确处理这种结构——self.param_groups 是列表，每个元素是字典，含 'params'（参数列表）和其他键（如 'lr'）。常见陷阱：
- 在
step()中硬编码group['lr'] = 1e-3→ 覆盖用户传入的分组学习率 - 没对
group['params']做isinstance(..., torch.Tensor)过滤 → 遇到nn.ParameterList或None参数时报错 - 把超参（如
beta）只从第一个group读取 → 后续分组的自定义超参被忽略
验证方法很简单：构造多组参数，打印 optimizer.param_groups[0]['lr'] 和 optimizer.param_groups[1]['lr']，确认它们和初始化时一致。
调试时优先检查 self.state 的键和生命周期
self.state 是一个嵌套字典：self.state[p] 对应单个参数的状态，p 是 torch.nn.Parameter 对象本身（不是 id(p) 或名字）。这意味着：
- 不能用
str(p)或p.name当 key —— PyTorch 内部靠对象身份（identity）索引 - 如果参数被替换（如
model.conv.weight = nn.Parameter(...)），旧p的状态不会自动迁移，导致“状态丢失” - 调用
optimizer.zero_grad(set_to_none=True)不影响self.state，但set_to_none=False（默认）会把p.grad置零，不影响你存的momentum_buffer
最容易被忽略的是：训练中途新增参数（如动态插入层）后，这些新参数不在任何 param_groups 中，也不会被 step() 访问，更不会自动获得 self.state 条目——必须手动调用 add_param_group() 并初始化对应状态。
