Pytorch：torch.optim详解

本篇笔记主要介绍torch.optim模块，记录学习过程

在深度学习中，我们通常会使用优化算法来调整神经网络的权重和偏差，以便模型能够更好地拟合训练数据。torch.optim是PyTorch中的一个模块，它提供了各种优化算法的实现，用于自动化地优化神经网络的参数。换句话说，torch.optim可以帮助我们让模型更好地学习，从而提高性能。

优化器 Optimizer

优化器主要是在模型训练阶段对模型可学习参数进行更新, 常用优化器有 SGD，RMSprop，Adam等
优化器初始化时传入模型的可学习参数，以及其他超参数如 lr，momentum等
在训练过程中先调用 optimizer.zero_grad() 清空梯度，再调用 loss.backward() 反向传播，最后调用optimizer.step()更新模型参数

简单使用示例如下所示：

import torch
import numpy as np
import warnings
warnings.filterwarnings('ignore') #ignore warnings
# 定义数据
x = torch.linspace(-np.pi, np.pi, 2000)
y = torch.sin(x)p = torch.tensor([1, 2, 3])
xx = x.unsqueeze(-1).pow(p)
# 定义模型
model = torch.nn.Sequential(torch.nn.Linear(3, 1),torch.nn.Flatten(0, 1))
# 定义损失函数:损失函数是一个衡量模型预测与实际值之间差距的函数
loss_fn = torch.nn.MSELoss(reduction='sum')
# 学习率:learning_rate参数表示学习率，它控制了每次参数更新的步长。
learning_rate = 1e-3
# 优化器:选择一个优化算法来优化模型的参数
optimizer = torch.optim.RMSprop(model.parameters(), lr=learning_rate)
# 迭代学习:训练模型
for t in range(1, 1001):y_pred = model(xx)loss = loss_fn(y_pred, y)if t % 100 == 0:print('No.{: 5d}, loss: {:.6f}'.format(t, loss.item()))# 梯度清零optimizer.zero_grad() # 反向传播计算梯度loss.backward() # 梯度下降法更新参数optimizer.step()

1.1 PyTorch 中的优化器

所有优化器都是继承父类 Optimizer，如下列表是 PyTorch 提供的优化器:
(具体用法和优缺点，后续更新。。。。。。。。。。。)

SGD
ASGD
Adadelta
Adagrad
Adam
AdamW
Adamax
SparseAdam
RMSprop
Rprop
LBFGS

1.2 父类Optimizer 基本原理

Optimizer 是所有优化器的父类，它主要有如下公共方法:

add_param_group(param_group): 添加模型可学习参数组
step(closure): 进行一次参数更新
zero_grad(): 清空上次迭代记录的梯度信息
state_dict(): 返回 dict 结构的参数状态
load_state_dict(state_dict): 加载 dict 结构的参数状态

上述方法一一解读：

1.2.1 初始化 Optimizer

初始化优化器只需要将模型的**可学习参数(params)和超参数(defaults)**分别传入优化器的构造函数

关键几个点为：

self.defaults = defaults
self.state = defaultdict(dict)
self.param_groups = [] 特别重要，最需要记住的属性，特别是属性的内容和其形式。
self.add_param_group方法，往self.param_groups里面放东西。

下面是Optimizer的初始化函数核心代码:

class Optimizer(object):def __init__(self, params, defaults):# 字典类型，子类传入，用于表示全部参数组的默认超参self.defaults = defaults# 判断参数是否为torch.Tensor格式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.param_groups = []param_groups = list(params)# 确保param_groups为一个字典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。

这两个分开说，先说params，最常见的就是model.parameters()，当然net.parameters()也是一样的，就是模型类的对象的变量名不同，如下所示。

optimizer = optim.SGD( net.parameters(),               # params的一种形式lr=LR, momentum=0.9)

params是个生成器，只返回各模型层的参数，没有参数名。

注意__init__方法中param_groups = list(params)，得到一个新的变量 param_groups，(这里self.param_groups是干什么的呢？？）。

param_groups = list(params)，list可以把生成器的元素都取出来，所以，很明显，param_groups就是一个Parameter类对象的列表，里面的元素是每个网络层的参数weight和bias（如果有）。

if not isinstance(param_groups[0], dict):param_groups = [{'params': param_groups}]

param_groups[0]是Parameter类，不是dict，这种形式的param_groups会被改造，将整个param_groups作为值，"params"作为键，形成一个键值对，放在字典里，然后重新赋值给param_groups。

现在我们要记得 param_groups的形式，一个列表，里面是一个字典，字典的键是"params"，值为所有网络层的参数。

for param_group in param_groups:self.add_param_group(param_group)

将param_groups中的每个元素送进self.add_param_group这个列表中。现在的param_groups里只有一个元素{“param”: [参数]}。

1.2.2 add_param_group

该方法在初始化函数中用到，主要用来向 self.param_groups添加不同分组的模型参数

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."""# 步骤1：判断传进来的参数是否是一个字典，必然是一个字典，不是字典报错。assert isinstance(param_group, dict), "param group must be a dict"# 步骤2：取出字典里的"params"的值，就是参数的列表，这是个列表，然后一系列判断，走到第3步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:# 步骤3：重新以列表的形式赋值回去param_group['params'] = list(params)# 步骤4：判断参数的列表里边的元素类型，必然是Parameter类型，也就是Tensor类型的，并且是叶子结点。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")# 利用默认参数给所有组设置统一的超参# 步骤5：将defaults这个字典里的键值对拿出来，放到现在的param_group这个字典里，这样该字典构成一个具有完整参数的字典，# 其所有键为：dict_keys(['params', 'lr', 'momentum', 'dampening', 'weight_decay', 'nesterov'])，方便step()方法调用。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)params = param_group['params']if len(params) != len(set(params)):warnings.warn("optimizer contains a parameter group with duplicate parameters; ""in future, this will cause an error; ""see github.com/pytorch/pytorch/issues/40967 for more information", stacklevel=3)# 步骤6： 判定当前字典中的参数组和之前的参数组是不是一样的。对于当前来说，self.param_groups是空的，所以直接到第7步param_set = set()for group in self.param_groups:param_set.update(set(group['params']))# 步骤7：判断param_set集合是否和param_group["params"]这个集合中具有相同元素，没有返回True，反之False。显然没有，所以7不执行。if not param_set.isdisjoint(set(param_group['params'])):raise ValueError("some parameters appear in more than one parameter group")# 步骤8：将构造完整的param_group这个字典，加到self.param_groups中去。self.param_groups.append(param_group)

将上述代码分为8个步骤：请认真看一遍代码以及注释

现在我们知道self.param_groups这个列表中具有字典，每个字典的keys为dict_keys([‘params’, ‘lr’, ‘momentum’, ‘dampening’, ‘weight_decay’, ‘nesterov’])，当然，每个键都有其对应的值。这些键值对是构建SGD实例时，传进来的参数。

params还有一种常见形式如下。

fcParamsId = list(map(id, resnet18_ft.fc.parameters()))     # 返回的是parameters的 内存地址
features_params = filter(lambda p: id(p) not in fcParamsId, resnet18_ft.parameters())optimizer = optim.SGD([{'params': features_params, 'lr': LR * 0.1},                         #  这个列表是params的另一种形式{'params': resnet18_ft.fc.parameters()}], 'lr': LR, momentum=0.9
)

我们可以对比一下 这种形式（前面讲的） 和 上述model.parameters() 这种方式的不同。

前者是包含着字典的列表，后者是生成器
进入到Optimizer类的初始化函数中，前者不变，还是包含2个字典的列表，后者先变成参数的列表，再变成只有一个字典的列表，该字典中只有一个键值对，"params"和所有参数值
在执行self.add_param_group()方法的循环里，前者需要遍历两个字典，后者只有一个字典，也就是说，最终在self.param_groups这个列表中，前者最终具有两个具有全部键值对的字典，后者只有一个
在如上self.add_param_group()方法中，步骤2处，取出前者字典中的"params"对应的值为生成器，在步骤3处，将其变成了参数的列表，后者在步骤2和步骤3处没变化，因为在初始化一开始就使用list()将model.parameters()的所有参数取出来了。
在如上self.add_param_group()方法中，标记为步骤4处没区别。在步骤5处有区别，并且该区别非常关键。这是对不同参数设置不同学习率和动量的地方。
param_group是循环中取出的字典，对于如上有两个字典的来说，就是将self.defaults里面取出来的键值对，分别放到两个字典中去。param_group.setdefault(name,
default)这句代码。如果param_group这个字典中比如第一个字典具有"lr"这个键的，保持不变，对于第二个字典里面没有"lr"这个键的，将default和name设置为新的键值对。不知道的同学可以看一下dict中setdefault这个方法的功能。
前者的步骤6和步骤7，由于存在两个字典，所以第二次需要和第一次进行对比，看看两个字典里面的"params"这个键对应的参数是否相同，如果相同，引用的就是同一块地址，设置不同的学习率等参数就没意义了，就会报错。
最后，将两个具有所有键值对，但是参数不同的两个字典，加入到self.param_groups这个列表中。

1.2.3 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... note::Unless otherwise specified, this function should not modify the``.grad`` field of the parameters."""raise NotImplementedError

子类如 SGD 需要实现 step 方法，如下所示:

	@torch.no_grad()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:with torch.enable_grad():loss = closure()#对参数进行遍历for group in self.param_groups: params_with_grad = []  #有梯度的网路参数收集列表d_p_list = [] #收集网络参数的梯度列表momentum_buffer_list = []# #以下为一些超参数的收集weight_decay = group['weight_decay']momentum = group['momentum']dampening = group['dampening']nesterov = group['nesterov']lr = group['lr']# 对而网络参数进行逐个遍历更新for p in group['params']: if p.grad is None:continued_p = p.gradif weight_decay != 0:d_p = d_p.add(p, alpha=weight_decay)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_(d_p, alpha=1 - dampening)if nesterov:d_p = d_p.add(buf, alpha=momentum)else:d_p = bufp.add_(d_p, alpha=-group['lr'])return loss

step 方法可传入闭包函数 closure，主要目的是为了实现如Conjugate Gradient和LBFGS等优化算法，这些算法需要对模型进行多次评估
Python 中闭包概念：在一个内部函数中，对外部作用域的变量进行引用(并且一般外部函数的返回值为内部函数)，那么内部函数就被认为是闭包

下面是 closure 的简单示例:

from torch.nn import CrossEntropyLossdummy_model = DummyModel().cuda()optimizer = SGD(dummy_model.parameters(), lr=1e-2, momentum=0.9, weight_decay=1e-4)
# 定义loss
loss_fn = CrossEntropyLoss()
# 定义数据
batch_size = 2
data = torch.randn(64, 3, 64, 128).cuda()  # 制造假数据shape=64 * 3 * 64 * 128
data_label = torch.randint(0, 10, size=(64,), dtype=torch.long).cuda()  # 制造假的labelfor batch_index in range(10):batch_data = data[batch_index*batch_size: batch_index*batch_size + batch_size]batch_label = data_label[batch_index*batch_size: batch_index*batch_size + batch_size]def closure():optimizer.zero_grad()  # 清空梯度output = dummy_model(batch_data)  # forwardloss = loss_fn(output, batch_label)  # 计算lossloss.backward()  # backwardprint('No.{: 2d} loss: {:.6f}'.format(batch_index, loss.item()))return lossoptimizer.step(closure=closure)  # 更新参数

1.2.4 zero_grad

在反向传播计算梯度之前对上一次迭代时记录的梯度清零，参数set_to_none 设置为 True 时会直接将参数梯度设置为 None，从而减小内存使用

但通常情况下不建议设置这个参数，因为梯度设置为 None 和 0 在 PyTorch 中处理逻辑会不一样。

def zero_grad(self, set_to_none: bool = False):r"""Sets the gradients of all optimized :class:`torch.Tensor` s to zero.Arguments:set_to_none (bool): instead of setting to zero, set the grads to None.This is will in general have lower memory footprint, and can modestly improve performance.However, it changes certain behaviors. For example:1. When the user tries to access a gradient and perform manual ops on it,a None attribute or a Tensor full of 0s will behave differently.2. If the user requests ``zero_grad(set_to_none=True)`` followed by a backward pass, ``.grad``sare guaranteed to be None for params that did not receive a gradient.3. ``torch.optim`` optimizers have a different behavior if the gradient is 0 or None(in one case it does the step with a gradient of 0 and in the other it skipsthe step altogether)."""for group in self.param_groups:for p in group['params']:if p.grad is not None:if set_to_none:p.grad = Noneelse:if p.grad.grad_fn is not None:p.grad.detach_()else:p.grad.requires_grad_(False)p.grad.zero_()

1.2.5 state_dict() 和 load_state_dict

这两个方法实现序列化和反序列化功能。

state_dict(): 将优化器管理的参数和其状态信息以 dict 形式返回
load_state_dict(state_dict): 加载之前返回的 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 order indices instead of Tensorsparam_mappings = {}start_index = 0def pack_group(group):nonlocal start_indexpacked = {k: v for k, v in group.items() if k != 'params'}param_mappings.update({id(p): i for i, p in enumerate(group['params'], start_index)if id(p) not in param_mappings})packed['params'] = [param_mappings[id(p)] for p in group['params']]start_index += len(packed['params'])return packedparam_groups = [pack_group(g) for g in self.param_groups]# Remap state to use order indices as keyspacked_state = {(param_mappings[id(k)] if isinstance(k, torch.Tensor) else k): vfor k, v in self.state.items()}return {'state': packed_state,'param_groups': param_groups,}

2.1 基类: _LRScheduler

学习率调整类主要的逻辑功能就是每个 epoch 计算参数组的学习率，更新 optimizer对应参数组中的lr值，从而应用在optimizer里可学习参数的梯度更新。

所有的学习率调整策略类的父类是torch.optim.lr_scheduler._LRScheduler，基类 _LRScheduler 定义了如下方法:

step(epoch=None): 子类公用
get_lr(): 子类需要实现
get_last_lr(): 子类公用
print_lr(is_verbose, group, lr, epoch=None): 显示 lr 调整信息
state_dict(): 子类可能会重写
load_state_dict(state_dict): 子类可能会重写

2.1.1 初始化

LR_scheduler是用于调节学习率lr的，在代码中，我们经常看到这样的一行代码

scheduler.step()

在pytorch代码中，各种类型scheduler大多基于 _LRScheduler 类

基类的初始化函数可传入两个参数,

第一是optimizer就是之前我们讲过的优化器的实例
第二个参数last_epoch是最后一次 epoch 的 index，默认值是 -1，代表初次训练模型

此时会对optimizer里的各参数组设置初始学习率 initial_lr。若last_epoch传入值大于 -1，则代表从某个 epoch 开始继续上次训练，此时要求optimizer的参数组中有initial_lr初始学习率信息。

初始化函数内部的 with_counter 函数主要是为了确保lr_scheduler.step()是在optimizer.step()之后调用的 (PyTorch=1.1 发生变化). 注意在__init__函数最后一步调用了self.step()，即_LRScheduler在初始化时已经调用过一次step()方法。

class _LRScheduler(object):def __init__(self, optimizer, last_epoch=-1, verbose=False):# Attach optimizerif not isinstance(optimizer, Optimizer):raise TypeError('{} is not an Optimizer'.format(type(optimizer).__name__))self.optimizer = optimizer# Initialize epoch and base learning ratesif last_epoch == -1:for group in optimizer.param_groups:group.setdefault('initial_lr', group['lr'])else:for i, group in enumerate(optimizer.param_groups):if 'initial_lr' not in group:raise KeyError("param 'initial_lr' is not specified ""in param_groups[{}] when resuming an optimizer".format(i))self.base_lrs = list(map(lambda group: group['initial_lr'], optimizer.param_groups))self.last_epoch = last_epoch# Following https://github.com/pytorch/pytorch/issues/20124# We would like to ensure that `lr_scheduler.step()` is called after# `optimizer.step()`def with_counter(method):if getattr(method, '_with_counter', False):# `optimizer.step()` has already been replaced, return.return method# Keep a weak reference to the optimizer instance to prevent# cyclic references.instance_ref = weakref.ref(method.__self__)# Get the unbound method for the same purpose.func = method.__func__cls = instance_ref().__class__del method@wraps(func)def wrapper(*args, **kwargs):instance = instance_ref()instance._step_count += 1wrapped = func.__get__(instance, cls)return wrapped(*args, **kwargs)# Note that the returned function here is no longer a bound method,# so attributes like `__func__` and `__self__` no longer exist.wrapper._with_counter = Truereturn wrapperself.optimizer.step = with_counter(self.optimizer.step)self.optimizer._step_count = 0self._step_count = 0self.verbose = verboseself.step()

2.1.2 step

当模型完成一个 epoch 训练时，需要调用step()方法，该方法里对last_epoch自增之后，在内部上下文管理器类里调用子类实现的get_lr()方法获得各参数组在此次 epoch 时的学习率，并更新到 optimizer的param_groups属性之中，最后记录下最后一次调整的学习率到self._last_lr，此属性将在get_last_lr()方法中返回。在这个方法中用到了上下文管理功能的内部类 _enable_get_lr_call，实例对象添加了_get_lr_called_within_step属性，这个属性可在子类中使用。此外，需要注意的是，step方法中的参数epoch已经废弃了，在使用时可以直接忽略这个参数。

def step(self, epoch=None):# Raise a warning if old pattern is detected# https://github.com/pytorch/pytorch/issues/20124if self._step_count == 1:if not hasattr(self.optimizer.step, "_with_counter"):warnings.warn("...") # 移除了警告信息# Just check if there were two first lr_scheduler.step() calls before optimizer.step()elif self.optimizer._step_count < 1:warnings.warn("...") # 移除了警告信息self._step_count += 1class _enable_get_lr_call:def __init__(self, o):self.o = odef __enter__(self):self.o._get_lr_called_within_step = Truereturn selfdef __exit__(self, type, value, traceback):self.o._get_lr_called_within_step = Falsewith _enable_get_lr_call(self):if epoch is None:self.last_epoch += 1values = self.get_lr()else:warnings.warn(EPOCH_DEPRECATION_WARNING, UserWarning)self.last_epoch = epochif hasattr(self, "_get_closed_form_lr"):values = self._get_closed_form_lr()else:values = self.get_lr()for i, data in enumerate(zip(self.optimizer.param_groups, values)):param_group, lr = dataparam_group['lr'] = lrself.print_lr(self.verbose, i, lr, epoch)self._last_lr = [group['lr'] for group in self.optimizer.param_groups]

2.1.3 get_last_lr、get_lr和print_lr

get_last_lr()方法比较简单，就是step()方法调用后，记录的最后一次 optimizer各参数组里更新后的学习率信息
get_lr() 方法是抽象方法，定义了更新学习率策略的接口，不同子类继承后会有不同的实现.其返回值是[lr1, lr2, …]结构
print_lr(is_verbose, group, lr, epoch=None)): 该方法提供了显示 lr 调整信息的功能

def get_last_lr(self):""" Return last computed learning rate by current scheduler."""return self._last_lrdef get_lr(self):# Compute learning rate using chainable form of the schedulerraise NotImplementedErrordef print_lr(self, is_verbose, group, lr, epoch=None):"""Display the current learning rate."""if is_verbose:if epoch is None:print('Adjusting learning rate'' of group {} to {:.4e}.'.format(group, lr))else:print('Epoch {:5d}: adjusting learning rate'' of group {} to {:.4e}.'.format(epoch, group, lr))

2.1.4 state_dict 和 load_state_dict

这两个方法和Optimizer里的方法功能是一样的，就是为了保存和重新加载状态信息，需要注意的是，这里不会重复记录self.optimizer属性的状态信息，因为 Optimizer 有自己实现的对应方法。

state_dict(): 以字典 dict 形式返回当前实例除 self.optimizer 之外的其他所有属性信息
load_state_dict(state_dict): 重新载入之前保存的状态信息

def state_dict(self):"""Returns the state of the scheduler as a :class:`dict`.It contains an entry for every variable in self.__dict__ whichis not the optimizer."""return {key: value for key, value in self.__dict__.items() if key != 'optimizer'}def load_state_dict(self, state_dict):"""Loads the schedulers state.Arguments:state_dict (dict): scheduler state. Should be an object returnedfrom a call to :meth:`state_dict`."""self.__dict__.update(state_dict)