第十六章 MMDetection3D解析系列七_优化器封装（optimwrapper）(车道线感知)

小酒馆燃着灯

已于 2023-11-11 10:50:04 修改

阅读量312

点赞数 1

分类专栏：车道线检测文章标签： pytorch 深度学习人工智能

于 2023-10-16 18:18:03 首次发布

本文链接：https://blog.csdn.net/weixin_44302770/article/details/133863544

版权

深度学习同时被 3 个专栏收录

213 篇文章 45 订阅

订阅专栏

手写AI

47 篇文章 5 订阅

订阅专栏

车道线检测

30 篇文章 14 订阅

订阅专栏

一前言

近期参与到了手写AI的车道线检测的学习中去，以此系列笔记记录学习与思考的全过程。车道线检测系列会持续更新，力求完整精炼，引人启示。所需前期知识，可以结合手写AI进行系统的学习。

二优化器封装 vs 优化器

这里我们分别基于 Pytorch 内置的优化器和 MMEngine 的优化器封装进行单精度训练、混合精度训练和梯度累加，对比二者实现上的区别。

2.1 单精度训练

import torch
from torch.optim import SGD
import torch.nn as nn
import torch.nn.functional as F
inputs = [torch.zeros(10, 1, 1)] * 10
targets = [torch.ones(10, 1, 1)] * 10
model = nn.Linear(1, 1)
optimizer = SGD(model.parameters(), lr=0.01)
optimizer.zero_grad()

for input, target in zip(inputs, targets):
    output = model(input)
    loss = F.l1_loss(output, target)
    loss.backward()
    optimizer.step()
    optimizer.zero_grad()

from mmengine.optim import OptimWrapper
optim_wrapper = OptimWrapper(optimizer=optimizer)
for input, target in zip(inputs, targets):
    output = model(input)
    loss = F.l1_loss(output, target)
    optim_wrapper.update_params(loss)

在这里插入图片描述
优化器封装的 update_params 实现了标准的梯度计算、参数更新和梯度清零流程，可以直接用来更新模型参数。

2.2 混合精度训练

from torch.cuda.amp import autocast
model = model.cuda()
inputs = [torch.zeros(10, 1, 1, 1)] * 10
targets = [torch.ones(10, 1, 1, 1)] * 10

for input, target in zip(inputs, targets):
    with autocast():
        output = model(input.cuda())
    loss = F.l1_loss(output, target.cuda())
    loss.backward()
    optimizer.step()
    optimizer.zero_grad()

from mmengine.optim import AmpOptimWrapper
optim_wrapper = AmpOptimWrapper(optimizer=optimizer)
for input, target in zip(inputs, targets):
    with optim_wrapper.optim_context(model):
        output = model(input.cuda())
    loss = F.l1_loss(output, target.cuda())
    optim_wrapper.update_params(loss)

在这里插入图片描述
开启混合精度训练需要使用 AmpOptimWrapper，他的 optim_context 接口类似 autocast，会开启混合精度训练的上下文。除此之外他还能加速分布式训练时的梯度累加

2.3 实现混合精度训练和梯度累加

for idx, (input, target) in enumerate(zip(inputs, targets)):
    with autocast():
        output = model(input.cuda())
    loss = F.l1_loss(output, target.cuda())
    loss.backward()
    if idx % 2 == 0:
        optimizer.step()
        optimizer.zero_grad()

optim_wrapper = AmpOptimWrapper(optimizer=optimizer, accumulative_counts=2)
for input, target in zip(inputs, targets):
    with optim_wrapper.optim_context(model):
        output = model(input.cuda())
    loss = F.l1_loss(output, target.cuda())
    optim_wrapper.update_params(loss)

在这里插入图片描述
我们只需要配置 accumulative_counts 参数，并调用 update_params 接口就能实现梯度累加的功能。除此之外，分布式训练情况下，如果我们配置梯度累加的同时开启了 optim_wrapper 上下文，可以避免梯度累加阶段不必要的梯度同步。

2.4 自定义的参数更新逻辑

backward：传入损失，用于计算参数梯度。
step：同 optimizer.step，用于更新参数。
zero_grad：同 optimizer.zero_grad，用于参数的梯度。

for idx, (input, target) in enumerate(zip(inputs, targets)):
    optimizer.zero_grad()
    with optim_wrapper.optim_context(model):
        output = model(input.cuda())
    loss = F.l1_loss(output, target.cuda())
    optim_wrapper.backward(loss)
    if idx % 2 == 0:
        optim_wrapper.step()
        optim_wrapper.zero_grad()

我们同样可以为优化器封装配置梯度裁减策略：

# 基于 torch.nn.utils.clip_grad_norm_ 对梯度进行裁减
optim_wrapper = AmpOptimWrapper(
    optimizer=optimizer, clip_grad=dict(max_norm=1))

# 基于 torch.nn.utils.clip_grad_value_ 对梯度进行裁减
optim_wrapper = AmpOptimWrapper(
    optimizer=optimizer, clip_grad=dict(clip_value=0.2))

三获取学习率/动量

优化器封装提供了 get_lr 和 get_momentum 接口用于获取优化器的一个参数组的学习率：

import torch.nn as nn
from torch.optim import SGD

from mmengine.optim import OptimWrapper

model = nn.Linear(1, 1)
optimizer = SGD(model.parameters(), lr=0.01)
optim_wrapper = OptimWrapper(optimizer)

print(optimizer.param_groups[0]['lr'])  # 0.01
print(optimizer.param_groups[0]['momentum'])  # 0
print(optim_wrapper.get_lr())  # {'lr': [0.01]}
print(optim_wrapper.get_momentum())  # {'momentum': [0]}

0.01
0
{'lr': [0.01]}
{'momentum': [0]}

四导出/加载状态字典

优化器封装和优化器一样，提供了 state_dict 和 load_state_dict 接口，用于导出/加载优化器状态，对于 AmpOptimWrapper，优化器封装还会额外导出混合精度训练相关的参数：

import torch.nn as nn
from torch.optim import SGD
from mmengine.optim import OptimWrapper, AmpOptimWrapper

model = nn.Linear(1, 1)
optimizer = SGD(model.parameters(), lr=0.01)

optim_wrapper = OptimWrapper(optimizer=optimizer)
amp_optim_wrapper = AmpOptimWrapper(optimizer=optimizer)

# 导出状态字典
optim_state_dict = optim_wrapper.state_dict()
amp_optim_state_dict = amp_optim_wrapper.state_dict()

print(optim_state_dict)
print(amp_optim_state_dict)
optim_wrapper_new = OptimWrapper(optimizer=optimizer)
amp_optim_wrapper_new = AmpOptimWrapper(optimizer=optimizer)

# 加载状态字典
amp_optim_wrapper_new.load_state_dict(amp_optim_state_dict)
optim_wrapper_new.load_state_dict(optim_state_dict)

{'state': {}, 'param_groups': [{'lr': 0.01, 'momentum': 0, 'dampening': 0, 'weight_decay': 0, 'nesterov': False, 'maximize': False, 'foreach': None, 'params': [0, 1]}]}
{'state': {}, 'param_groups': [{'lr': 0.01, 'momentum': 0, 'dampening': 0, 'weight_decay': 0, 'nesterov': False, 'maximize': False, 'foreach': None, 'params': [0, 1]}], 'loss_scaler': {'scale': 65536.0, 'growth_factor': 2.0, 'backoff_factor': 0.5, 'growth_interval': 2000, '_growth_tracker': 0}}

五使用多个优化器

OptimWrapperDict 的核心功能是支持批量导出/加载所有优化器封装的状态字典；支持获取多个优化器封装的学习率、动量。如果没有 OptimWrapperDict，MMEngine 就需要在很多位置对优化器封装的类型做 if else 判断，以获取所有优化器封装的状态。

from torch.optim import SGD
import torch.nn as nn

from mmengine.optim import OptimWrapper, OptimWrapperDict

gen = nn.Linear(1, 1)
disc = nn.Linear(1, 1)
optimizer_gen = SGD(gen.parameters(), lr=0.01)
optimizer_disc = SGD(disc.parameters(), lr=0.01)

optim_wapper_gen = OptimWrapper(optimizer=optimizer_gen)
optim_wapper_disc = OptimWrapper(optimizer=optimizer_disc)
optim_dict = OptimWrapperDict(gen=optim_wapper_gen, disc=optim_wapper_disc)

print(optim_dict.get_lr())  # {'gen.lr': [0.01], 'disc.lr': [0.01]}
print(optim_dict.get_momentum())  # {'gen.momentum': [0], 'disc.momentum': [0]}

{'gen.lr': [0.01], 'disc.lr': [0.01]}
{'gen.momentum': [0], 'disc.momentum': [0]}

六在执行器中配置优化器封装

以配置一个 SGD 优化器封装为例：
优化器封装需要接受 optimizer 参数，因此我们首先需要为优化器封装配置 optimizer。MMEngine 会自动将 PyTorch 中的所有优化器都添加进 OPTIMIZERS 注册表中

optimizer = dict(type='SGD', lr=0.01, momentum=0.9, weight_decay=0.0001)
optim_wrapper = dict(type='OptimWrapper', optimizer=optimizer)

optimizer = dict(type='SGD', lr=0.01, momentum=0.9, weight_decay=0.0001)
optim_wrapper = dict(optimizer=optimizer)

optimizer = dict(type='SGD', lr=0.01, momentum=0.9, weight_decay=0.0001)
optim_wrapper = dict(type='AmpOptimWrapper', optimizer=optimizer, accumulative_counts=2)

七模型中的不同参数设置不同的超参数

PyTorch 的优化器支持对模型中的不同参数设置不同的超参数，例如对一个分类模型的骨干（backbone）和分类头（head）设置不同的学习率：

from torch.optim import SGD
import torch.nn as nn

model = nn.ModuleDict(dict(backbone=nn.Linear(1, 1), head=nn.Linear(1, 1)))
optimizer = SGD([{'params': model.backbone.parameters()},
     {'params': model.head.parameters(), 'lr': 1e-3}],
    lr=0.01,
    momentum=0.9)

八不同类型的参数设置不同的超参系数

例如，我们可以在 paramwise_cfg 中设置 norm_decay_mult=0，从而将正则化层（normalization layer）的权重（weight）和偏置（bias）的权值衰减系数（weight decay）设置为 0
具体示例如下，我们将 ToyModel 中所有正则化层（head.bn）的权重衰减系数设置为 0：

from mmengine.optim import build_optim_wrapper
from collections import OrderedDict

class ToyModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.backbone = nn.ModuleDict(
            dict(layer0=nn.Linear(1, 1), layer1=nn.Linear(1, 1)))
        self.head = nn.Sequential(
            OrderedDict(
                linear=nn.Linear(1, 1),
                bn=nn.BatchNorm1d(1)))

optim_wrapper = dict(
    optimizer=dict(type='SGD', lr=0.01, weight_decay=0.0001),
    paramwise_cfg=dict(norm_decay_mult=0))
optimizer = build_optim_wrapper(ToyModel(), optim_wrapper)

除了可以对正则化层的权重衰减进行配置外，MMEngine 的默认优化器封装构造器的 paramwise_cfg 还支持对更多不同类型的参数设置超参系数，支持的配置如下：

lr_mult：所有参数的学习率系数
decay_mult：所有参数的衰减系数
bias_lr_mult：偏置的学习率系数（不包括正则化层的偏置以及可变形卷积的 offset）
bias_decay_mult：偏置的权值衰减系数（不包括正则化层的偏置以及可变形卷积的 offset）
norm_decay_mult：正则化层权重和偏置的权值衰减系数
flat_decay_mult：一维参数的权值衰减系数
dwconv_decay_mult：Depth-wise 卷积的权值衰减系数
bypass_duplicate：是否跳过重复的参数，默认为 False
dcn_offset_lr_mult：可变形卷积（Deformable Convolution）的学习率系数

九为模型不同部分的参数设置不同的超参系数

此外，与上文 PyTorch 的示例一样，在 MMEngine 中我们也同样可以对模型中的任意模块设置不同的超参，只需要在 paramwise_cfg 中设置 custom_keys 即可。

例如我们想将 backbone.layer0 所有参数的学习率设置为 0，衰减系数设置为 0，backbone 其余子模块的学习率设置为 0.01；head 所有参数的学习率设置为 0.001，可以这样配置：

optim_wrapper = dict(
    optimizer=dict(type='SGD', lr=0.01, weight_decay=0.0001),
    paramwise_cfg=dict(
        custom_keys={
            'backbone.layer0': dict(lr_mult=0, decay_mult=0),
            'backbone': dict(lr_mult=1),
            'head': dict(lr_mult=0.1)
        }))
optimizer = build_optim_wrapper(ToyModel(), optim_wrapper)

上例中，模型的状态字典的 key 如下：

for name, val in ToyModel().named_parameters():
    print(name)

backbone.layer0.weight
backbone.layer0.bias
backbone.layer1.weight
backbone.layer1.bias
head.linear.weight
head.linear.bias
head.bn.weight
head.bn.bias

custom_keys 中每一个字段的含义如下：

‘backbone’: dict(lr_mult=1)：将名字前缀为 backbone 的参数的学习率系数设置为 1
‘backbone.layer0’: dict(lr_mult=0, decay_mult=0)：将名字前缀为 backbone.layer0 的参数学习率系数设置为 0，衰减系数设置为 0，该配置优先级比第一条高
‘head’: dict(lr_mult=0.1)：将名字前缀为 head 的参数的学习率系数设置为 0.1

十自定义优化器构造策略

例如，我们想实现一个叫做 LayerDecayOptimWrapperConstructor 的优化器封装构造器，能够对模型不同深度的层自动设置递减的学习率：

from mmengine.optim import DefaultOptimWrapperConstructor
from mmengine.registry import OPTIM_WRAPPER_CONSTRUCTORS
from mmengine.logging import print_log


@OPTIM_WRAPPER_CONSTRUCTORS.register_module(force=True)
class LayerDecayOptimWrapperConstructor(DefaultOptimWrapperConstructor):

    def __init__(self, optim_wrapper_cfg, paramwise_cfg=None):
        super().__init__(optim_wrapper_cfg, paramwise_cfg=None)
        self.decay_factor = paramwise_cfg.get('decay_factor', 0.5)

        super().__init__(optim_wrapper_cfg, paramwise_cfg)

    def add_params(self, params, module, prefix='' ,lr=None):
        if lr is None:
            lr = self.base_lr

        for name, param in module.named_parameters(recurse=False):
            param_group = dict()
            param_group['params'] = [param]
            param_group['lr'] = lr
            params.append(param_group)
            full_name = f'{prefix}.{name}' if prefix else name
            print_log(f'{full_name} : lr={lr}', logger='current')

        for name, module in module.named_children():
            chiled_prefix = f'{prefix}.{name}' if prefix else name
            self.add_params(
                params, module, chiled_prefix, lr=lr * self.decay_factor)


class ToyModel(nn.Module):

    def __init__(self) -> None:
        super().__init__()
        self.layer = nn.ModuleDict(dict(linear=nn.Linear(1, 1)))
        self.linear = nn.Linear(1, 1)


model = ToyModel()

optim_wrapper = dict(
    optimizer=dict(type='SGD', lr=0.01, weight_decay=0.0001),
    paramwise_cfg=dict(decay_factor=0.5),
    constructor='LayerDecayOptimWrapperConstructor')

optimizer = build_optim_wrapper(model, optim_wrapper)

08/23 22:20:26 - mmengine - INFO - layer.linear.weight : lr=0.0025
08/23 22:20:26 - mmengine - INFO - layer.linear.bias : lr=0.0025
08/23 22:20:26 - mmengine - INFO - linear.weight : lr=0.005
08/23 22:20:26 - mmengine - INFO - linear.bias : lr=0.005

add_params 被第一次调用时，params 参数为空列表（list），module 为模型（model）。
封装构造器函数

十一在训练过程中调整超参

实现了更通用的参数调度器，可以对学习率、动量等优化器相关的参数进行调整，并且支持多个调度器进行组合，应用更复杂的调度策略。

11.1 参数调度器的使用

我们构造一个 ExponentialLR，并且在每个 epoch 结束后调用 scheduler.step()，实现了随 epoch 指数下降的学习率调整策略。

import torch
from torch.optim import SGD
from torch.optim.lr_scheduler import ExponentialLR

model = torch.nn.Linear(1, 1)
dataset = [torch.randn((1, 1, 1)) for _ in range(20)]
optimizer = SGD(model, 0.1)
scheduler = ExponentialLR(optimizer, gamma=0.9)

for epoch in range(10):
    for data in dataset:
        optimizer.zero_grad()
        output = model(data)
        loss = 1 - output
        loss.backward()
        optimizer.step()
    scheduler.step()

在 mmengine.optim.scheduler 中，我们支持大部分 PyTorch 中的学习率调度器，例如 ExponentialLR，LinearLR，StepLR，MultiStepLR 等，使用方式也基本一致，所有支持的调度器见调度器接口文档。同时增加了对动量的调整，在类名中将 LR 替换成 Momentum 即可，例如 ExponentialMomentum，LinearMomentum。更进一步地，我们实现了通用的参数调度器 ParamScheduler，用于调整优化器的中的其他参数，包括 weight_decay 等。这个特性可以很方便地配置一些新算法中复杂的调整策略。
MMEngine 中通常不需要手动来实现训练循环以及调用 optimizer.step()，而是在执行器（Runner）中对训练流程进行自动管理，同时通过 ParamSchedulerHook 来控制参数调度器的执行。

11.2 使用单一的学习率调度器

# 基于手动构建学习率调度器的例子
from torch.optim import SGD
from mmengine.runner import Runner
from mmengine.optim.scheduler import MultiStepLR

optimizer = SGD(model.parameters(), lr=0.01, momentum=0.9)
param_scheduler = MultiStepLR(optimizer, milestones=[8, 11], gamma=0.1)

runner = Runner(
    model=model,
    optim_wrapper=dict(
        optimizer=optimizer),
    param_scheduler=param_scheduler,
    ...
    )

在这里插入图片描述
如果配合注册器和配置文件使用的话，我们可以设置配置文件中的 param_scheduler 字段来指定调度器, 执行器（Runner）会根据此字段以及执行器中的优化器自动构建学习率调度器：

# 在配置文件中设置学习率调度器字段
param_scheduler = dict(type='MultiStepLR', by_epoch=True, milestones=[8, 11], gamma=0.1)

注意这里增加了初始化参数 by_epoch，控制的是学习率调整频率，当其为 True 时表示按轮次（epoch）调整，为 False 时表示按迭代次数（iteration）调整，默认值为 True。在上面的例子中，表示按照轮次进行调整，此时其他参数的单位均为 epoch，例如 milestones 中的 [8, 11] 表示第 8 和 11 轮次结束时，学习率将会被调整为上一轮次的 0.1 倍。

当修改了学习率调整频率后，调度器中与计数相关设置的含义也会相应被改变。当 by_epoch=True 时，milestones 中的数字表示在哪些轮次进行学习率衰减，而当 by_epoch=False 时则表示在进行到第几次迭代时进行学习率衰减。下面是一个按照迭代次数进行调整的例子，在第 600 和 800 次迭代结束时，学习率将会被调整为原来的 0.1 倍。
param_scheduler = dict(type=‘MultiStepLR’, by_epoch=False, milestones=[600, 800], gamma=0.1)
在这里插入图片描述
若用户希望在配置调度器时按轮次填写参数的同时使用基于迭代的更新频率，MMEngine 的调度器也提供了自动换算的方式。用户可以调用 build_iter_from_epoch 方法，并提供每个训练轮次的迭代次数，即可构造按迭代次数更新的调度器对象：

epoch_length = len(train_dataloader)
param_scheduler = MultiStepLR.build_iter_from_epoch(optimizer, milestones=[8, 11], gamma=0.1, epoch_length=epoch_length)

如果使用配置文件构建调度器，只需要在配置中加入 convert_to_iter_based=True，执行器会自动调用 build_iter_from_epoch 将基于轮次的配置文件转换为基于迭代次数的调度器对象：

param_scheduler = dict(type='MultiStepLR', by_epoch=True, milestones=[8, 11], gamma=0.1, convert_to_iter_based=True)

为了能直观感受这两种模式的区别，我们这里再举一个例子。下面是一个按轮次更新的余弦退火（CosineAnnealing）学习率调度器，学习率仅在每个轮次结束后被修改：

param_scheduler = dict(type='CosineAnnealingLR', by_epoch=True, T_max=12)

在这里插入图片描述

param_scheduler = dict(type='CosineAnnealingLR', by_epoch=True, T_max=12, convert_to_iter_based=True)

在这里插入图片描述

11.3 组合多个学习率调度器（以学习率预热为例）

MMEngine 支持组合多个调度器一起使用，只需将配置文件中的 scheduler 字段修改为一组调度器配置的列表，SchedulerStepHook 可以自动对调度器列表进行处理。下面的例子便实现了学习率预热。

param_scheduler = [
    # 线性学习率预热调度器
    dict(type='LinearLR',
         start_factor=0.001,
         by_epoch=False,  # 按迭代更新学习率
         begin=0,
         end=50),  # 预热前 50 次迭代
    # 主学习率调度器
    dict(type='MultiStepLR',
         by_epoch=True,  # 按轮次更新学习率
         milestones=[8, 11],
         gamma=0.1)
]

在这里插入图片描述
注意这里增加了 begin 和 end 参数，这两个参数指定了调度器的生效区间。生效区间通常只在多个调度器组合时才需要去设置，使用单个调度器时可以忽略。当指定了 begin 和 end 参数时，表示该调度器只在 [begin, end) 区间内生效，其单位是由 by_epoch 参数决定。上述例子中预热阶段 LinearLR 的 by_epoch 为 False，表示该调度器只在前 50 次迭代生效，超过 50 次迭代后此调度器不再生效，由第二个调度器来控制学习率，即 MultiStepLR。在组合不同调度器时，各调度器的 by_epoch 参数不必相同。

这里再举一个例子：

param_scheduler = [
    # 在 [0, 100) 迭代时使用线性学习率
    dict(type='LinearLR',
         start_factor=0.001,
         by_epoch=False,
         begin=0,
         end=100),
    # 在 [100, 900) 迭代时使用余弦学习率
    dict(type='CosineAnnealingLR',
         T_max=800,
         by_epoch=False,
         begin=100,
         end=900)
]

在这里插入图片描述

十二调整动量

动量调度器（momentum scheduler）的使用方法和学习率调度器完全一样。同样也只需要将动量调度器的配置添加进配置文件中的 param_scheduler 字段的列表中即可。

param_scheduler = [
    # the lr scheduler
    dict(type='LinearLR', ...),
    # 动量调度器
    dict(type='LinearMomentum',
         start_factor=0.001,
         by_epoch=False,
         begin=0,
         end=1000)
]

十三通用的参数调度器

MMEngine 还提供了一组通用的参数调度器用于调度优化器的 param_groups 中的其他参数，将学习率调度器类名中的 LR 改为 Param 即可，例如 LinearParamScheduler。用户可以通过设置参数调度器的 param_name 变量来选择想要调度的参数。

param_scheduler = [
    dict(type='LinearParamScheduler',
         param_name='lr',  # 调度 `optimizer.param_groups` 中名为 'lr' 的变量
         start_factor=0.001,
         by_epoch=False,
         begin=0,
         end=1000)
]

这里设置的参数名是 lr，因此这个调度器的作用等同于直接使用学习率调度器 LinearLRScheduler。

除了动量之外，用户也可以对 optimizer.param_groups 中的其他参数名进行调度，可调度的参数取决于所使用的优化器。例如，当使用带 weight_decay 的 SGD 优化器时，可以按照以下示例对调整 weight_decay：

param_scheduler = [
    dict(type='LinearParamScheduler',
         param_name='weight_decay',  # 调度 `optimizer.param_groups` 中名为 'weight_decay' 的变量
         start_factor=0.001,
         by_epoch=False,
         begin=0,
         end=1000)
]