• 问答
  • 技术
  • 实践
  • 资源
PyTorch 源码解读 | torch.optim:优化算法接口详解

OpenMMLab 商汤学术

商汤学术

原文链接:

https://zhuanlan.zhihu.com/p/346205754

0/前言

本篇笔记主要介绍 torch.optim 模块,主要包含模型训练的优化器 Optimizer , 学习率调整策略 LRScheduler以及 SWA 相关优化策略. 本文中涉及的源码以 torch==1.7.0 为准。

本文主要目录结构

1. 优化器 Optimizer

2. 学习率调节器 lr_scheduler

3. 随机参数平均 swa_utils

4. 参考资料

1/优化器 Optimizer

1.0 基本用法

● 优化器主要是在模型训练阶段对模型可学习参数进行更新, 常用优化器有 SGD,RMSprop,Adam 等。

● 优化器初始化时传入传入模型的可学习参数,以及其他超参数如lr , momentum 等。

● 在训练过程中先调用 optimizer.zero_grad() 清空梯度,再调用 loss.backward() 反向传播,最后调用 optimizer.step()更新模型参数。

简单使用示例如下所示:

image.png

image.png

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)分别传入优化器的构造函数,下面是 Optimizer 的初始化函数核心代码:

image.png

1.2.2 add_param_group

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

image.png

利用 add_param_group 函数功能,可以对模型不同的可学习参数组设定不同的超参数,初始化优化器可传入元素是 dict 的 list,每个 dict 中的 key 是 params 或者其他超参数的名字如 lr ,下面是一个实用的例子,对模型的 fc 层参数设置不同的学习率:

image.png

1.2.3 step

此方法主要完成一次模型参数的更新。

● 基类 Optimizer 定义了 step 方法接口,如下所示:

image.png

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

image.png

● step 方法可传入闭包函数 closure,主要目的是为了实现如Conjugate Gradient和LBFGS等优化算法,这些算法需要对模型进行多次评估

● Python 中闭包概念:在一个内部函数中,对外部作用域的变量进行引用(并且一般外部函数的返回值为内部函数),那么内部函数就被认为是闭包

下面是 closure 的简单示例:

image.png

1.2.4 zero_grad

● 在反向传播计算梯度之前对上一次迭代时记录的梯度清零,参数 set_to_none 设置为 True 时会直接将参数梯度设置为 None,从而减小内存使用, 但通常情况下不建议设置这个参数,因为梯度设置为 None 和 0 在 PyTorch 中处理逻辑会不一样。

image.png

1.2.5 state_dict() 和 load_state_dict

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

● state_dict(): 将优化器管理的参数和其状态信息以 dict 形式返回

● load_state_dict(state_dict): 加载之前返回的 dict,更新参数和其状态

● 两个方法可用来实现模型训练中断后继续训练功能:

image.png

1.3 常见优化器简介

1.3.1 SGD(params, lr, momentum=0, dampening=0, weight_decay=0, nesterov=False)

实现带 momentum 和 dampening 的 SGD,公式如下:

image.png

1.3.2 Adagrad(params, lr=0.01, lr_decay=0, weight_decay=0, initial_accumulator_value=0, eps=1e-10)

自适应学习率,考虑历史所有梯度信息, 公式如下: 

image.png

1.3.3 RMSprop(params, lr=0.01, alpha=0.99, eps=1e-08, weight_decay=0, momentum=0, centered=False)

image.png

1.3.4 Adam(params, lr=1e-3, betas=(0.9, 0.999), eps=1e-8, weight_decay=0, amsgrad=False)

image.png

2/学习率调节器 lr_scheduler

有了优化器,还需要根据 epoch 来调整学习率,lr_schedluer提供了在训练模型时学习率的调整策略。

目前 PyTorch 提供了如下学习率调整策略:

● StepLR: 等间隔调整策略

● MultiStepLR: 多阶段调整策略

● ExponentialLR: 指数衰减调整策略

● ReduceLROnPlateau: 自适应调整策略

● CyclicLR: 循环调整策略

● OneCycleLR: 单循环调整策略

● CosineAnnealingLR: 余弦退火调整策略

● CosineAnnealingWarmRestarts: 带预热启动的余弦退火调整策略

● LambdaLR: 自定义函数调整策略

● MultiplicativeLR: 乘法调整策略

学习率调整策略可粗略分为以下三大类:

有序调整策略:

● StepLR

● MultiStepLR

● ExponentialLR

● CyclicLR

● OneCycleLR

● CosineAnnealingLR

● CosineAnnealingWarmRestarts

自适应调整策略:

● ReduceLROnPlateau

自定义调整策略:

● LambdaLR

● MultiplicativeLR 

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 初始化

基类的初始化函数可传入两个参数, 第一是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() 方法。

image.png
image.png

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 已经废弃了,在使用时可以直接忽略这个参数。

image.png
image.png

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 调整信息的功能

image.png

2.1.4 state_dict 和 load_state_dict

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

● state_dict(): 以字典 dict 形式返回当前实例除self.optimizer 之外的其他所有属性信息

● load_state_dict(state_dict): 重新载入之前保存的状态信息

image.png

2.2 学习率调整策略示例

2.2.1 StepLR(optimizer, step_size, gamma=0.1, last_epoch=-1, verbose=False)

StepLR 是根据 epoch 的等间隔学习率调整策略,实现了get_lr() 方法。初始化函数须传入优化器,epoch 间隔step_size,gamma 是学习率的衰减系数,默认是 0.1。

image.png

image.png

image.png

2.2.2 MultiStepLR(optimizer, milestones, gamma=0.1, last_epoch=-1, verbose=False)

多阶段学习率调整策略,参数 milestones 是包含多个学习率调整点列表

image.png

image.png

2.2.3 MultiplicativeLR(optimizer, lr_lambda, last_epoch=-1, verbose=False)

乘法调整策略实现了学习率的衰减系数gamma 可变,即在每个调整节点,可对各参数组的学习率乘上一个不同的衰减率 gamma ,初始化函数中 lr_lambda 参数可以是一个 lambda 函数,也可是 lambda 函数列表,每个 lambda 函数输入是 epoch,输出是 gamma 。

image.png

image.png

2.2.4 LambdaLR(optimizer, lr_lambda, last_epoch=-1, verbose=False)

该策略可传入自定义的 lambda 函数,lambda 函数参数为 epoch ,返回值为学习率。

image.png

image.png

2.2.5 ExponentialLR(optimizer, gamma, last_epoch=-1, verbose=False)

指数衰减学习率调整策略

image.png

2.2.6 CosineAnnealingLR(optimizer, T_max, eta_min=0, last_epoch=-1, verbose=False)

余弦退火调整策略,T_max 是最大迭代次数,eta_min 是最小学习率值,其公式如下, eta_max 为初始学习率,T_cur 是自重新启动后的 epoch 数 

image.png

2.2.7 CosineAnnealingWarmRestarts(optimizer, T_0, T_mult=1, eta_min=0, last_epoch=-1, verbose=False)

在 SGDR(Stochastic Gradient Descent with Warm Restarts)中提出:

● T_0 : 第一次启动时的迭代数

● T_mult : 启动后,改变周期 T 的因子

● eta_min : 学习率下限

image.png

2.2.8 CyclicLR(optimizer, base_lr, max_lr, step_size_up=2000, step_size_down=None, mode='triangular', ...)

类似三角波形状的学习率调整策略,以下是几个重要初始化参数:

●  base_lr: 基准学习率,也是最小的学习率

●  max_lr: 学习率上限

●  step_size_up: 一个周期里上升阶段 epoch 数

● step_size_down: 一个周期里下降阶段 epoch 数

image.png

2.2.9 OneCycleLR(optimizer, max_lr, total_steps=None, epochs=None, steps_per_epoch=None, pct_start=0.3,...)

只有 1 次循环的学习率调整策略

● max_lr ‍‍‍‍‍‍: float/list, 学习率调整的上限

● total_steps : int 循环中的总步数

image.png

2.2.10 ReduceLROnPlateau(optimizer, mode='min', factor=0.1, patience=10, threshold=0.0001, threshold_mode='rel', ...)

自适应学习率调整策略,比如只有当 loss 在几个 epoch 里都不发生下降时,才调整学习率。注意在调用时,需要在其 step() 方法中传入对应的参考变量,例如:  scheduler.step(val_loss) 

●  mode : 评价模型训练质量的模式, 传入值为 min或 max

●  factor : 学习率衰减因子, 类似 gamma

●  patience : 控制何时调整学习率

示例:

image.png

3/swa_utils里SWA相关类和函数

该模块中只有 2 个类和一个函数:

●  AveragedModel: 实现 SWA 算法的权重平均模型

●  SWALR: 与 AverageModel 配合使用的学习率调整策略

●  update_bn: 更新模型中的 bn

3.0 SWA 简介

随机权重平均(SWA)是一种优化算法,在SWA 论文的结果证明,取 SGD 轨迹的多点简单平均值,以一个周期或者不变的学习率,会比传统训练有更好的泛化效果。论文的结果同样了证明了,随机权重平均 (SWA) 可以找到更广的最优值域。

3.1 AveragedModel

●  该类实现 SWA 算法的权重平均模型,初始化时传入模型 model 和参数平均化函数avg_fn ,然后在初始化函数中对model 的参数进行深拷贝, 注册模型计数器。

●  在 update_parameters(self, model) 方法中再次传入模型后,根据参数 avg_fn 对模型参数进行平均后更新 swa 模型参数。

image.png

3.2 update_bn

该函数主要是通过传入的某个训练时刻的模型model 和dataloader ,来允许 swa 模型计算和更新 bn

image.png

Example:

image.png

3.3 SWALR

SWALR 类继承_LRScheduler 基类,实现了供 swa 模型的学习率调整策略

在此就只放出其使用示例:

image.png

4/参考资料

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

[2]https://github.com/pytorch/pytorch/

[3]https://blog.csdn.net/shanglianlm/article/details/85143614

[4]https://blog.csdn.net/lis\_12/article/details/53521554

[5]https://pymotw.com/3/weakref/index.html

  • 1
  • 0
  • 1018
收藏
暂无评论
Find me
大咖

一个大的公司

  • 21,319

    关注
  • 275

    获赞
  • 54

    精选文章
近期动态
  • 哈工大深圳研究生院CV汪,请原谅我这一生放纵不羁爱CV~
文章专栏
  • Awsome-Github 资源列表