本文将结合源码与代码示例详细解析Optimizer的五大方法。

1. 前言

优化器主要用在模型训练阶段，用于更新模型中可学习的参数。torch.optim提供了多种优化器接口，比如Adam、RAdam、SGD、ASGD、LBFGS等，Optimizer是所有这些优化器的父类。

2. Optimizer行为解析

2.1 公共方法

Optimizer是所有优化器的父类，它主要具有以下几类公共方法：

方法名	官方注解	说明
`add_param_group`	Add a param group to the Optimizer's param_groups.	给优化器添加模型中可学习的参数组。
`step`	Performs a single optimization step (parameter update)	进行一次参数更新。
`zero_grad`	Sets the gradients of all optimized torch.Tensor to zero.	将上次记录的梯度信息置零，避免梯度累加。
`state_dict`	Returns the state of the optimizer as a dict.	以字典形式返回优化器的状态信息。
`load_state_dict`	Loads the optimizer state.	加载以字典形式存储的优化器状态信息。

2.2 行为解析

以下将结合源码 (注：小喵的torch版本为1.4.0) 与示例代码来解析Optimizer各种方法的行为。为了方便理解，小喵在相关源码中都添加了必要的注释。

2.2.1 __init___初始化

下面是Optimizer的初始化函数源码:

def __init__(self, params, defaults):""":param params: 待优化参数params，可以有两种格式，分别对应全局参数、参数组:defaults: （全局）默认超参数字典，这里的超参数主要指优化器参数，如学习率等。"""torch._C._log_api_usage_once("python.optimizer")self.defaults = defaultsif isinstance(params, torch.Tensor):raise TypeError("params argument given to the optimizer should be ""an iterable of Tensors or dicts, but got " +torch.typename(params))self.state = defaultdict(dict)self.param_groups = []param_groups = list(params)if len(param_groups) == 0:raise ValueError("optimizer got an empty parameter list")# 如果是全局参数，则转换为字典格式，并放入列表中if not isinstance(param_groups[0], dict):param_groups = [{'params': param_groups}]for param_group in param_groups:self.add_param_group(param_group)

初始化函数用于初始化优化器，需传入两个参数：

params：模型的可学习参数。
defaults：(全局)默认优化器超参，如学习率lr等。

params允许两种传入格式，其一是网络全局参数，即网络的所有可学习参数，它们共用一套优化器参数：

# 全局参数示例：
# 假设net为我们所创建的网络模型
# 直接调用net的`.parameters()`方法即可获得网络全局参数
# 网络全局参数使用相同的学习率lr=0.001进行参数更新
params = net.parameters()
optimizer = torch.optim.Adam(params, lr=0.001)

其二是参数组，每组参数可以指定自己的优化器参数，即可使用不同的优化策略：

# 参数组示例：
# 假设net为我们所创建的网络模型
# 按网络参数的参数名是否含有bert为条件，将net的所有可学习参数分为两组
# 第二组参数使用学习率lr=2e-5
params = [{"params":[p for n, p in net.named_parameters() if "bert" not in n ]},{"params":[p for n, p in net.named_parameters() if "bert" in n ],'lr': 2e-5}
]
optimizer = torch.optim.Adam(params, lr=0.001)

补充知识：named_parameters()返回关于网络层参数名字和参数，parameters()仅返回网络层参数。

2.2.2 add_param_group参数组设置

在初始化函数中，调用了add_param_group方法往Optimizer.params_groups中添加分组参数，add_param_group的源码如下：

def add_param_group(self, param_group):r"""Add a param group to the :class:`Optimizer` s `param_groups`.This can be useful when fine tuning a pre-trained network as frozen layers can be madetrainable and added to the :class:`Optimizer` as training progresses.Arguments:param_group (dict): Specifies what Tensors should be optimized along with groupspecific optimization options."""# 每组参数必须是字典格式assert isinstance(param_group, dict), "param group must be a dict"params = param_group['params']if isinstance(params, torch.Tensor):param_group['params'] = [params]elif isinstance(params, set):raise TypeError('optimizer parameters need to be organized in ordered collections, but ''the ordering of tensors in sets will change between runs. Please use a list instead.')else:param_group['params'] = list(params)# 待优化参数必须为张量for param in param_group['params']:if not isinstance(param, torch.Tensor):raise TypeError("optimizer can only optimize Tensors, ""but one of the params is " + torch.typename(param))# 待优化参数必须为叶子结点if not param.is_leaf:raise ValueError("can't optimize a non-leaf Tensor")for name, default in self.defaults.items():if default is required and name not in param_group:raise ValueError("parameter group didn't specify a value of required optimization parameter " +name)# 对该组参数没有指定的超参数，则设置为（全局）默认优化器参数中相应的值。else:param_group.setdefault(name, default)# 该组参数不允许出现在其它参数组中，即参数组之间交集为空。param_set = set()for group in self.param_groups:param_set.update(set(group['params']))if not param_set.isdisjoint(set(param_group['params'])):raise ValueError("some parameters appear in more than one parameter group")self.param_groups.append(param_group)

一方面，针对每一组参数，add_param_group会判断它们是否违背了相关条件，如不是张量、不是叶子结点、出现在了其他参数组中。

另一方面，add_param_group会将没有单独指定优化器超参的参数组的优化器参数设定（全局）默认优化器超参值。比如在之前的例子中，第一组参数的学习率lr=0.001，第二组参数的学习率lr=2e-5。

# 参数组形式：
params = [{"params":[p for n, p in net.named_parameters() if "bert" not in n ]},{"params":[p for n, p in net.named_parameters() if "bert" in n ],'lr': 2e-5}
]
optimizer = torch.optim.Adam(params, lr=0.001)

2.2.3 step更新参数

step方法作用是执行一次参数的更新。 Optimizer 定义了 step 方法接口，所有继承自它的子类都需要对step进行实现。

# Optimizer的step
def step(self, closure):r"""Performs a single optimization step (parameter update).
Arguments:closure (callable): A closure that reevaluates the model andreturns the loss. Optional for most optimizers."""raise NotImplementedError

下面，我们以SGD优化器为例，来看step方法的具体行为，相关源码如下：

# SGD的step
def step(self, closure=None):"""Performs a single optimization step.Arguments:closure (callable, optional): A closure that reevaluates the modeland returns the loss."""loss = Noneif closure is not None:loss = closure()# 针对每个参数组for group in self.param_groups:# 取出该组参数更新时所需要的优化器超参数weight_decay = group['weight_decay']momentum = group['momentum']dampening = group['dampening']nesterov = group['nesterov']# 取出该组参数中待更新的参数for p in group['params']:if p.grad is None:continue# 取出参数的梯度值d_p = p.grad.data# 加入正则项，避免过拟合if weight_decay != 0:d_p.add_(weight_decay, p.data)# 计算动量，加速梯度下降if momentum != 0:param_state = self.state[p]if 'momentum_buffer' not in param_state:buf = param_state['momentum_buffer'] = torch.clone(d_p).detach()else:buf = param_state['momentum_buffer']buf.mul_(momentum).add_(1 - dampening, d_p)if nesterov:d_p = d_p.add(momentum, buf)else:d_p = buf# 更新参数 p.data=p.data-lr*d_pp.data.add_(-group['lr'], d_p)return loss

从源码来看，针对每组参数，参数更新过程如下：

首先取出参数的梯度值；
然后为避免过拟合，利用优化器超参中的weight_decay为梯度添加了正则项；
接着为加速，利用momentum、dampening、nesterov计算动量，更新梯度；
最后，利用梯度更新参数值。

这里我们暂且不讲必包函数closure。

2.2.4 zero_grad清零梯度

zero_grad方法一般用在反向传播之前，作用是将上次反向传播时记录的梯度值清零（在Pytorch张量属性与梯度计算那些事儿（下）一文中已提到过，如果不清零梯度，梯度会叠加）。

下面是Optimizer的zero_grad方法的源码：

def zero_grad(self):r"""Clears the gradients of all optimized :class:`torch.Tensor` s."""for group in self.param_groups:for p in group['params']:if p.grad is not None:# 截断反向传播的梯度流p.grad.detach_()# 清零已存储的梯度值p.grad.zero_()

2.2.5 state_dict保存参数与状态

state_dict方法返回优化器管理的参数与状态信息，下面是Optimizer的state_dict方法的源码：

def state_dict(self):r"""Returns the state of the optimizer as a :class:`dict`.It contains two entries:* state - a dict holding current optimization state. Its contentdiffers between optimizer classes.* param_groups - a dict containing all parameter groups"""# Save ids instead of Tensorsdef pack_group(group):packed = {k: v for k, v in group.items() if k != 'params'}packed['params'] = [id(p) for p in group['params']]return packed# 保存每组参数中相关参数的地址及优化器超参param_groups = [pack_group(g) for g in self.param_groups]# Remap state to use ids as keys# 将state中所有参数张量替换为张量对象地址保存packed_state = {(id(k) if isinstance(k, torch.Tensor) else k): vfor k, v in self.state.items()}return {'state': packed_state,'param_groups': param_groups,}

从源码可知，当我们调用state_dict方法时会获得一个字典形式的返回值，它包含两项内容：

state (dict) ：保存了优化器状态信息，键为网络参数的地址，值为在更新网络参数的过程中计算出的与该参数相关的缓存变量信息，如momentum_buffer。
param_groups(list) ：列表中每一项是一个字典，存储了相应的一组参数的参数地址及该组参数相关的优化器参数。

下面我们沿用之前在Pytorch张量属性与梯度计算那些事儿（下）中给出的示例来让大家对state_dict方法对返回值有更直接和清晰的认识：

#!/usr/local/bin/python3
# -*-coding:utf-8 -*-import torch
import torch.nn as nn
import jsonclass DemoNet(nn.Module):def __init__(self):super(DemoNet, self).__init__()self.fc = nn.Linear(3, 3)self.fc1 = nn.Linear(3, 1)def forward(self, x):x1 = self.fc(x)y = self.fc1(x1)z = torch.sigmoid(y)return zif __name__ == "__main__":# 模拟输入和标签x = torch.randn((100, 3))# 每个样本标签为1或0y = torch.empty(100).random_(2)# 网络初始化demoNet = DemoNet()# 定义损失loss_func = nn.BCELoss()# 定义优化器optimizer=optim.SGD(demoNet.parameters(), lr=1e-3)# 清空梯度optimizer.zero_grad()# 前向传播，进行预测y_hat = demoNet(x)# 反向传播，计算梯度y_hat = torch.squeeze(y_hat, dim=-1)loss = loss_func(y_hat, y)loss.backward()optimizer.step()print(json.dumps(optimizer.state_dict(),indent=4))

运行之后的结果如下，state中是空的，param_groups则保存了参数组中参数的地址及优化器参数：

{"state": {},"param_groups": [{"lr": 0.001,"momentum": 0,"dampening": 0,"weight_decay": 0,"nesterov": false,"params": [4821053656,4821053728,4821053800,4821053872]}]
}

更改下优化器与打印语句（因为张量是不能被序列化为json对象，会报Object of type Tensor is not JSON serializable错误）

if __name__ == "__main__":# 模拟输入和标签x = torch.randn((100, 3))# 每个样本标签为1或0y = torch.empty(100).random_(2)# 网络初始化demoNet = DemoNet()# 定义损失loss_func = nn.BCELoss()# 定义优化器# optimizer=optim.SGD(demoNet.parameters(), lr=1e-3)optimizer = torch.optim.SGD(demoNet.parameters(),momentum=0.8,weight_decay=1e-5,lr=1e-3)# 清空梯度optimizer.zero_grad()# 前向传播，进行预测y_hat = demoNet(x)# 反向传播，计算梯度y_hat = torch.squeeze(y_hat, dim=-1)loss = loss_func(y_hat, y)loss.backward()optimizer.step()# print(json.dumps(optimizer.state_dict(),indent=4))print(optimizer.state_dict())

这下，state中保存了参数组中参数的地址及参数相应的缓存变量momentum_buffer的信息。

{'state': {4727062392: {'momentum_buffer': tensor([[ 0.0028,  0.0005,  0.0047],[ 0.0145,  0.0027,  0.0243],[-0.0184, -0.0034, -0.0308]])}, 4727062464: {'momentum_buffer': tensor([ 0.0016,  0.0083, -0.0106])}, 4727083080: {'momentum_buffer': tensor([[-0.0516, -0.0293,  0.0673]])}, 4727083152: {'momentum_buffer': tensor([-0.0328])}},
'param_groups': [{'lr': 0.001, 'momentum': 0.8, 'dampening': 0, 'weight_decay': 1e-05, 'nesterov': False, 'params': [4727062392, 4727062464, 4727083080, 4727083152]}]
}

2.2.6 load_state_dict加载参数与状态

相应地，load_state_dict则是用于加载所保存的优化器管理的参数与状态信息，其源码如下：

def load_state_dict(self, state_dict):r"""Loads the optimizer state.Arguments:state_dict (dict): optimizer state. Should be an object returnedfrom a call to :meth:`state_dict`."""# deepcopy, to be consistent with module APIstate_dict = deepcopy(state_dict)# Validate the state_dictgroups = self.param_groupssaved_groups = state_dict['param_groups']# 判断当前优化器中的待优化参数与加载的待优化参数在size上是否一致if len(groups) != len(saved_groups):raise ValueError("loaded state dict has a different number of ""parameter groups")param_lens = (len(g['params']) for g in groups)saved_lens = (len(g['params']) for g in saved_groups)if any(p_len != s_len for p_len, s_len in zip(param_lens, saved_lens)):raise ValueError("loaded state dict contains a parameter group ""that doesn't match the size of optimizer's group")# Update the state# 建立参数对象旧地址与参数对象的映射关系id_map = {old_id: p for old_id, p inzip(chain(*(g['params'] for g in saved_groups)),chain(*(g['params'] for g in groups)))}def cast(param, value):r"""Make a deep copy of value, casting all tensors to device of param."""if isinstance(value, torch.Tensor):# Floating-point types are a bit special here. They are the only ones# that are assumed to always match the type of params.if param.is_floating_point():value = value.to(param.dtype)value = value.to(param.device)return valueelif isinstance(value, dict):return {k: cast(param, v) for k, v in value.items()}elif isinstance(value, container_abcs.Iterable):return type(value)(cast(param, v) for v in value)else:return value# Copy state assigned to params (and cast tensors to appropriate types).# State that is not assigned to params is copied as is (needed for# backward compatibility).state = defaultdict(dict)for k, v in state_dict['state'].items():# 参数对象的旧地址if k in id_map:# 参数对象param = id_map[k]# 转换参数的数据类型与devicestate[param] = cast(param, v)else:state[k] = v# Update parameter groups, setting their 'params' value# 以保存好的参数值来更新模型参数def update_group(group, new_group):new_group['params'] = group['params']return new_groupparam_groups = [update_group(g, ng) for g, ng in zip(groups, saved_groups)]self.__setstate__({'state': state, 'param_groups': param_groups})

这里还是通过示例来让大家对load_state_dict方法有一个更直接的印象。

沿用刚才那个例子，在进行完一次反向传播和参数更新完以后，将优化器管理的参数与状态信息保存起来：

# 保存优化器信息
if __name__ == "__main__":save_path="optimizer.pt"# 模拟输入和标签x = torch.randn((100, 3))# 每个样本标签为1或0y = torch.empty(100).random_(2)# 网络初始化demoNet = DemoNet()# 定义损失loss_func = nn.BCELoss()# 定义优化器optimizer = torch.optim.SGD(demoNet.parameters(),momentum=0.8,weight_decay=1e-5,lr=1e-3)# 清空梯度optimizer.zero_grad()# 前向传播，进行预测y_hat = demoNet(x)# 反向传播，计算梯度y_hat = torch.squeeze(y_hat, dim=-1)loss = loss_func(y_hat, y)loss.backward()optimizer.step()print(optimizer.state_dict())# 保存参数与状态信息torch.save(optimizer.state_dict(), save_path)

然后再修改代码重新运行，利用load_state_dict将其加载进来：

# 加载保存的优化器信息
if __name__ == "__main__":save_path="optimizer.pt"# 模拟输入和标签x = torch.randn((100, 3))# 每个样本标签为1或0y = torch.empty(100).random_(2)# 网络初始化demoNet = DemoNet()# 定义损失loss_func = nn.BCELoss()# 定义优化器optimizer = torch.optim.SGD(demoNet.parameters(),momentum=0.8,weight_decay=1e-5,lr=1e-3)# 将保存好的参数与状态信息加载到内存loaded_stated = torch.load(save_path)# 加载的结果print("loaded: \n", loaded_stated)# 将加载到内存的参数与状态信息加载到优化器optimizer.load_state_dict(loaded_stated)# 再次调用state_dict查看优化器参数与状态信息print("state_dict: \n",optimizer.state_dict())

这里，我们将保存时的参数与状态信息、加载到内存的参数与状态信息、从内存加载到优化器的参数与状态信息一并展示。

可以发现一个之前在2.2.5小节中已经出现但还没被提及的现象：参数的id会发生变化。

# 保存时打印结果
{'state': {4531957696: {'momentum_buffer': tensor([[ 0.0215,  0.0111, -0.0120],[-0.0610, -0.0314,  0.0342],[-0.0103, -0.0053,  0.0058]])}, 4531978312: {'momentum_buffer': tensor([ 0.0061, -0.0174, -0.0029])}, 4531978384: {'momentum_buffer': tensor([[-0.0171,  0.0468, -0.0662]])}, 4531978456: {'momentum_buffer': tensor([-0.0339])}},
'param_groups': [{'lr': 0.001, 'momentum': 0.8, 'dampening': 0, 'weight_decay': 1e-05, 'nesterov': False, 'params': [4531957696, 4531978312, 4531978384, 4531978456]}]
}# 加载到内存
loaded:{'state': {4531957696: {'momentum_buffer': tensor([[ 0.0215,  0.0111, -0.0120],[-0.0610, -0.0314,  0.0342],[-0.0103, -0.0053,  0.0058]])}, 4531978312: {'momentum_buffer': tensor([ 0.0061, -0.0174, -0.0029])}, 4531978384: {'momentum_buffer': tensor([[-0.0171,  0.0468, -0.0662]])}, 4531978456: {'momentum_buffer': tensor([-0.0339])}}, 'param_groups': [{'lr': 0.001, 'momentum': 0.8, 'dampening': 0, 'weight_decay': 1e-05, 'nesterov': False, 'params': [4531957696, 4531978312, 4531978384, 4531978456]}]
}# 加载到优化器
state_dict：{'state': {4683133000: {'momentum_buffer': tensor([[ 0.0215,  0.0111, -0.0120],[-0.0610, -0.0314,  0.0342],[-0.0103, -0.0053,  0.0058]])}, 4683133072: {'momentum_buffer': tensor([ 0.0061, -0.0174, -0.0029])}, 4683133144: {'momentum_buffer': tensor([[-0.0171,  0.0468, -0.0662]])}, 4683133216: {'momentum_buffer': tensor([-0.0339])}}, 'param_groups': [{'lr': 0.001, 'momentum': 0.8, 'dampening': 0, 'weight_decay': 1e-05, 'nesterov': False, 'params': [4683133000, 4683133072, 4683133144, 4683133216]}]
}

id() 函数返回的是相应对象在其生命周期内位于内存中的地址。比如下面这个例子中，我们相当于两次运行了相同的代码，变量a的值是一样的，但是python为a分配的内存地址是不同的。

➜  ~ python3
Python 3.7.3 (default, Mar 27 2019, 09:23:39)
[Clang 10.0.0 (clang-1000.11.45.5)] on darwin
Type "help", "copyright", "credits" or "license" for more information.
>>> a=1
>>> id(a)
4504543392
>>> exit()
➜  ~ python3
Python 3.7.3 (default, Mar 27 2019, 09:23:39)
[Clang 10.0.0 (clang-1000.11.45.5)] on darwin
Type "help", "copyright", "credits" or "license" for more information.
>>> a=1
>>> id(a)
4376338592
>>>

好了，我们回到之前保存和加载优化器的例子中来。

当我们调用torch.load将保存的参数与状态信息加载进内存的时候，参数的id是没有变化的；但，实际上当我们重新运行代码走到网络初始化这步时，网络的参数已经被分配了新的内存地址，即参数有新的id了。

然后，在使用 load_state_dict 将内存中加载好的优化器信息置入模型当中时，源码中id_map建立了参数对象旧地址与参数对象的映射关系，依据这个映射关系，优化器会更新它的state。所以，当我们调用state_dict后，所看到的参数id就是新的。

# 总结

本文结合源码和代码示例解析了优化器Optimizer的相关行为。在本文知识点的基础上，下一次，我们将一起学习与深度学习模型训练相关且需要掌握的两点内容：分组优化、继续训练。

参考资料

[1] torch.optim: https://pytorch.org/docs/stable/optim.html

一文详解Pytorch中的优化器Optimizer相关推荐

python中squeeze函数_详解pytorch中squeeze()和unsqueeze()函数介绍
squeeze的用法主要就是对数据的维度进行压缩或者解压. 先看torch.squeeze() 这个函数主要对数据的维度进行压缩,去掉维数为1的的维度,比如是一行或者一列这种,一个一行三列(1,3)的 ...
一文详解编程中的随机数
一文详解编程中的随机数随机数的类型真随机数生成器 TRNG - True Random Number Generator 伪随机数生成器 PRNG - Pseudo Random Number G ...
Pytorch框架之优化器 Optimizer
Pytorch框架之优化器 Optimizer 基本用法优化器主要是在模型训练阶段对模型可学习参数进行更新, 常用优化器有 SGD,RMSprop,Adam等优化器初始化时传入传入模型的可学习参数 ...
详解Pytorch中的requires_grad、叶子节点与非叶子节点、with torch.no_grad()、model.eval()、model.train()、BatchNorm层
requires_grad requires_grad意为是否需要计算梯度使用backward()函数反向传播计算梯度时,并不是计算所有tensor的梯度,只有满足下面条件的tensor的梯度才会被 ...
一文详解pytorch的“动态图”与“自动微分”技术
前言众所周知,Pytorch是一个非常流行且深受好评的深度学习训练框架.这与它的两大特性"动态图"."自动微分"有非常大的关系."动态图" ...
深度盘点：一文详解数据分析中100个常用指标和术语
大家好,有个朋友是金融行业产品经理,最近在对已有的站内用户做分层与标签分类,需要对用户进行聚类分析.一般从事数据分析行业的朋友对这类词并不陌生,但是像市场运营人员就会把这类些名词概念搞混,导致结果不准 ...
详解PyTorch中的ModuleList和Sequential
点击上方"视学算法",选择加"星标"或"置顶" 重磅干货,第一时间送达作者丨小占同学@知乎(已授权) 来源丨https://zhuanla ...
详解pytorch中的常见的Tensor数据类型以及类型转换
文章目录概览 Tensor的构建补充类型转换附录概览本文主要讲pytorch中的常见的Tensor数据类型,例如:float32,float64,int32,int64.构造他们分别使用如 ...
tensor torch 构造_详解Pytorch中的网络构造
背景在PyTroch框架中,如果要自定义一个Net(网络,或者model,在本文中,model和Net拥有同样的意思),通常需要继承自nn.Module然后实现自己的layer.比如,在下面的示例中 ...

一文详解Pytorch中的优化器Optimizer