前言
本文我们解读PyTorch中的优化器源码，并且简单实现一个SGD。
Optimizer
在PyTorch中，所有优化器诸如SGD、Adam都是继承自基类Optimizer，优化器的使用方法为：
for input, target in dataset:
optimizer.zero_grad() # 清空梯度
output = model(input) # 计算输出
loss = loss_fn(output, target) # 计算loss
loss.backward() # 计算梯度
optimizer.step() # 参数更新由于绝大部分深度学习的优化算法都是基于梯度的，所以Optimizer要做的事情最重要的只有两件：参数更新以及清空梯度。当然为了更好的拓展，Optimizer中还有一些其他接口，接下来我们结合源码进行解读。
__init__：初始化
下面是Optimizer的初始化函数源码:
def __init__(self, params, defaults):
"""
:param params: 待优化参数params，可以有两种格式，分别对应全局参数、参数组
:defaults: （全局）默认超参数字典，这里的超参数主要指优化器参数，如学习率等。
"""
torch._C._log_api_usage_once("python.optimizer")
self.defaults = defaults
if 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允许两种传入格式，其一是网络全局参数，即网络的所有可学习参数，它们共用一套优化器参数：
# 全局参数示例：
# 假设model为我们所创建的网络模型
# 直接调用net的`.parameters()`方法即可获得网络全局参数
# 网络全局参数使用相同的学习率lr=0.001进行参数更新
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)其二是参数组，每组参数可以指定自己的优化器参数，即可使用不同的优化策略。我们可以定制参数从而玩出一些花样，例如经典的weight decay方案：
no_decay = ['bias', 'LayerNorm.weight']
optimizer_grouped_parameters = [
{'params': [p for n, p in model.named_parameters() if not any(nd in n for nd in no_decay)],
'weight_decay': args.weight_decay},
{'params': [p for n, p in model.named_parameters() if any(nd in n for nd in no_decay)],
'weight_decay': 0.0}
]还有打比赛时常用的层间差分学习率trick，出自：How to Fine-Tune BERT for Text Classification？：
def get_parameters(model, model_init_lr, multiplier, classifier_lr):
parameters = []
lr = model_init_lr
for layer in range(24,-1,-1):
layer_params = {
'params': [p for n,p in model.named_parameters() if f'encoder.layer.{layer}.' in n],
'lr': lr
}
parameters.append(layer_params)
lr *= multiplier
classifier_params = {
'params': [p for n,p in model.named_parameters() if 'layer_norm' in n or 'linear' in n
or 'pooling' in n],
'lr': classifier_lr
}
parameters.append(classifier_params)
return parameters
parameters=get_parameters(model,2e-5,0.95, 1e-4)
optimizer=AdamW(parameters)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 made
trainable and added to the :class:`Optimizer` as training progresses.
Arguments:
param_group (dict): Specifies what Tensors should be optimized along with group
specific 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)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 and
returns the loss. Optional for most optimizers.
"""
raise NotImplementedErrorzero_grad：清空清零梯度
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_() 除了这些比较重要的方法之外，还有返回优化器管理的参数与状态信息的state_dict方法；加载所保存的优化器管理的参数与状态信息的load_state_dict等等，实现也很简单，这里就不详细说了。
写一个简单的SGD
了解了基类的源码之后，我们可以自己来写优化器了。事实上只需要重写几个比较重要的方法就好。PyTorch中已经内置了很多优化器，SGD自然也有，但是耦合了比较多的操作，我们现在来解耦一个清爽版的SGD。模型的话选择一个最简单的二元标量函数：
$$ f(x_1,x_2;\theta_1,\theta_2)=(\sin^2 x_1+\sin^2 x_2)\cdot(\theta_1^2+\theta_2^2) $$
模型用nn.Module实现如下：
class MyModel(nn.Module):
def __init__(self, limits=5):
super(MyModel, self).__init__()
self.params = Parameter(torch.FloatTensor([0.9 * limits, 0.9 * limits]))
self.register_parameter('params', self.params)
def forward(self, input):
return (torch.sin(input[0]) ** 2 + torch.sin(input[1]) ** 2) * (self.params[0] ** 2 + self.params[1] ** 2)实现一个SGD，我们只需要重写__init__和step两个方法，具体地：
from torch.optim.optimizer import Optimizer, required
class MySGD(Optimizer):
def __init__(self, params, lr=required):
if lr <= 0.0:
raise ValueError("Invalid learning rate: {}".format(lr))
defaults = dict(lr=lr)
super(MySGD, self).__init__(params, defaults)
@torch.no_grad()
def step(self):
for group in self.param_groups:
params_with_grad = []
d_p_list = []
lr = group['lr']
for p in group['params']:
if p.grad is not None:
params_with_grad.append(p)
d_p_list.append(p.grad)
for i, param in enumerate(params_with_grad):
d_p = d_p_list[i]
param.add_(d_p, alpha=-lr)我们再测试一下是否能正确更新模型参数：
inputs = torch.FloatTensor((np.random.rand(2) - 0.5) * 10)
model = MyModel()
optimizer = MySGD(model.parameters(), lr=1.0)
loss = model(inputs)
loss.backward() # 计算梯度
print(model.params.data) # tensor([4.5000, 4.5000])
print(model.params.grad) # tensor([5.6980, 5.6980])
optimizer.step() # 参数更新
print(model.params.data) # tensor([-1.1980, -1.1980])
optimizer.zero_grad() # 清空梯度
print(model.params.grad) # tensor([0., 0.])显然是没有任何问题的。
