optimizer.state.get()函数介绍

torch.optim.Adam optimizer:

optimizer.param_groups

每个优化器实例都有一个名为 param_groups 的属性，它是一个包含一个或多个参数组的列表。每个参数组本身是一个字典，可以包含以下键：

• 'params': 一个包含实际参数（torch.nn.Parameter 对象）的列表。
• 'lr': 学习率。
• 'weight_decay': 权重衰减系数。
• 其他与该参数组相关的超参数，如动量 (momentum)、β值 (betas) 等。

示例

假设我们有一个简单的模型，并使用 Adam 优化器对其进行优化：

import torch
from torch.optim import Adam

# 定义一个简单的线性模型
model = torch.nn.Linear(2, 1)

# 创建一个Adam优化器，并为不同层设置不同的学习率
optimizer = Adam([
    {'params': model.weight, 'lr': 0.01},
    {'params': model.bias, 'lr': 0.001}
])

# 查看 optimizer.param_groups
print(optimizer.param_groups)

输出如下所示：

[{'params': [Parameter containing:
tensor([[-0.5061, -0.1177]], requires_grad=True)], 
'lr': 0.01, 'betas': (0.9, 0.999), 'eps': 1e-08, 'weight_decay': 0, 'amsgrad': False, 'maximize': False, 'foreach': None, 'capturable': False}, 
{'params': [Parameter containing:
tensor([0.4852], requires_grad=True)], 
'lr': 0.001, 'betas': (0.9, 0.999), 'eps': 1e-08, 'weight_decay': 0, 'amsgrad': False, 'maximize': False, 'foreach': None, 'capturable': False}]

获取“第一个参数”

optimizer.param_groups[0]['params'][0]，即第一个参数组中的第一个参数。例如：

first_param_group = optimizer.param_groups[0]
first_param = first_param_group['params'][0]

print(first_param)  # 输出: tensor([[...]])

在这个例子中，first_param 就是 model.weight，因为它是在定义优化器时被放在第一个参数组中的第一个参数。

应用到 optimizer.state.get()

当我们使用 optimizer.state.get() 方法时，传递给它的参数就是上述提到的某个具体的参数对象。例如：

state = optimizer.state.get(first_param, None)

if state is None:
    print("Parameter has no state yet.")
else:
    print("Parameter state:", state)

这段代码尝试从优化器状态字典中获取 first_param 的状态信息。如果该参数还没有参与任何优化步骤（例如，尚未调用过 optimizer.step()），那么 state 将会是 None。
如果调用 optimizer.step() 执行一次优化步骤后，优化器会更新模型参数，并根据需要更新其内部状态。此时，如果调用 optimizer.state.get(param, None)，会得到与该参数相关的一组状态信息，而不是 None。具体的状态信息取决于使用的优化器类型。

以 Adam 优化器为例
假设我们使用的是 Adam 优化器，它为每个参数维护以下状态：

exp_avg: 一阶矩估计（指数加权移动平均值）。
exp_avg_sq 二阶矩估计（指数加权移动平均值）。
step: 参数更新的次数。

案例：简单的线性模型和 Adam 优化器：

import torch
from torch.optim import Adam

# 定义一个简单的线性模型
model = torch.nn.Linear(10, 1)

# 创建一个Adam优化器
optimizer = Adam(model.parameters(), lr=0.01)

# 获取第一个参数的状态（优化前）
# iter() 用于获取一个可迭代对象的迭代器。
# next() 用于从迭代器中获取下一个元素。
first_param = next(iter(optimizer.param_groups[0]['params']))
state_before = optimizer.state.get(first_param, None)
print("State before optimization:", state_before)  # 输出: None

# 模拟一次前向传播和反向传播
input_data = torch.randn(1, 10, requires_grad=True)
output = model(input_data)
loss = output.sum()
first_param_group = optimizer.param_groups[0]
first_param = first_param_group['params'][0]

loss.backward()

# 执行一次优化步骤
optimizer.step()

# 获取第一个参数的状态（优化后）
state_after = optimizer.state.get(first_param, None)
print("State after optimization:", state_after)
# Adam 优化器输出：
# 'exp_avg': 一阶矩估计（指数加权移动平均值）。
# 'exp_avg_sq': 二阶矩估计（指数加权移动平均值）。
# 'step': 参数更新的次数。