fastai 的优化器（optimizer）

发布日期：2024-04-22 来源: 网络 阅读量（）

fastai2提供了对优化器的实现，其主要的目的是实现对参数的优化，超参数的设置和参数的状态记录和统计，从而实现主流深度学习优化器，包括SGD，Adam等， 同时也包扩了对Pytorch的优化器的打包处理，和Lookhead的实现。

• fastai优化器的参数和超参数

对于参数，为了方便进行迁移学习，fastai将Pytorch的模型参数进行了分组处理，建立了一个param_lists。下面以transfomers 实现的bert 分类器模型为例，将其中的模型参数分为：

[m.bert.embeddings， m.bert.encoder， m.bert.pooler，m.classifier]

这样，就方便将其中任何一组模型参数冻结或者解冻，从而实现对参数的部分优化，比如，保持预先训练好的bert参数不变，而首先对classifier的参数进行优化（这部分其实就是一个简单的Linear模型）。

对于超参数，fastai在优化器的实现中包含了hypers列表，这个hypers列表是和param_lists中的参数组一一对应的，并在模型的优化过程中作用到每一个参数组中的单个参数上。

为了记录每一个参数的状态，即每次优化跌代中使用的超参数和其他辅助变量，同时也是为了对参数进行统计，调整，保存和调用，fastai为每个参数建立了一个状态字典，称为state，记住，是每一个参数，而不是参数组。

这样，对于每一个模型参数，p，都有自己所属的参数组，pg，都有自己的状态字典，state[p]，和对应的超参数，hyper。

在实现优化器之前，fastai 构建了一个基础类，用于实现一些基本的功能，既可以用于自己的优化器，也可以用于Pytorch优化器的封装。

• class_BaseOptimizer()

这个基础类只负责参数组的冻结与解冻和超参数的设置，并不包含参数优化的相关函数。其中的all_params()用来建立全部参数的列表，其中每个参数所在的tuple中包含了参数本身，所属参数组，状态字典和超参数。

(p,pg,self.state[p],hyper)

而所有的冻结和解冻操作就通过设置每一组参数的requiresgrad变量来实现。每一个参数的状态中可以包含一个“force_train"的辅助变量用来忽略冻结操作。

超参数的设置，比如设定的learning rate，是如何实现的呢？超参数可以通过单个数值，数值列表，或者slice的形式指定，通过set_hypers和set_hyper函数分配到每一个参数组并进一步指定个每一个参数。这样的方式给予了使用超参数的极大便利性。

在此基础上，fastai通过继承建立了优化器

• class Optimizer(_BaseOptimizer)

Optimizer实现了优化器的主要功能：梯度清零，梯度调整，以及记录参数状态。其中最重要的是用来调整梯度的step函数。

这个类的初始化函数接收模型优化需要训练的参数，超参数和Callback函数（cbs）。

模型参数可以是全部的模型参数或者是经过分组的型式，取决于是否要使用迁移学习和分步优化。

超参数以关键字变量的型式调用。比如”lr=0.01, mom=0.99“。

参数cbs是一个函数列表，其中的每个函数都有类似的调用变量，其中第一个调用变量必须是一个单个的模型参数p，同时也是一个Pytorch张量。其他变量可以是超参数或者其他的状态变量。函数格式类似于cb(p, **{**state, **hyper})的型式。所有对参数及其状态的调整通过顺序执行cbs中的Callback函数实现，从而实现SGD，Adam等等常用优化器的功能。

• 实现一个常用的优化器：带动量的SGD

下面来看一下最简单的SGD的实现：

def  sgd_step(p, lr, **kwargs):
    p.data.add_(p.grad.data, alpha=-lr)

上面的SGD仅仅调整了参数本身，有时需要对参数的梯度做出调整，比方说weight decay或者L2 regulation，下面是一个例子：

def  weight_decay(p, lr, wd, do_wd=True, **kwargs):
 "Weight decay as decaying `p` with `lr*wd`"
 if do_wd and wd!=0: p.data.mul_(1 - lr*wd)

weight_decay.defaults = dict(wd=0.)

def l2_reg(p, lr, wd, do_wd=True, **kwargs):
    "L2 regularization as adding `wd*p` to `p.grad`"
    if do_wd and wd!=0: p.grad.data.add_(p.data, alpha=wd)

l2_reg.defaults = dict(wd=0.)

更多的时候，我们需要对参数或参数的梯度进行统计，比如说计算梯度移动平均值，实现带momentum的SGD的话，可以再加上几个cb：

def average_grad(p, mom, dampening=False, grad_avg=None, **kwargs):
 "Keeps track of the avg grads of `p` in `state` with `mom`."
 if grad_avg is None: grad_avg = torch.zeros_like(p.grad.data)
    damp = 1-mom if dampening else 1.
    grad_avg.mul_(mom).add_(p.grad.data, alpha=damp)
 return {'grad_avg': grad_avg}

average_grad.defaults = dict(mom=0.9)

def momentum_step(p, lr, grad_avg, **kwargs):
    "Step for SGD with momentum with `lr`"
    p.data.add_(grad_avg, alpha=-lr)

上面的不同CB函数可以带有自己的默认数值，用来指定涉及的超参数。但是如果优化器的初始化变量中包含了相同超参数的数值的话，CB函数的默认超参数会被忽略掉。

组合以上的Callback函数，fastai可以建立如下的SGD函数，返回的结果就是一个实现了带有动量的SGD优化器。其中的cbs是一个函数的列表，

[weightdecay, averagegrad, momentum_step]

@log_args(to_return=True, but_as=Optimizer.__init__)
def SGD(params, lr, mom=0., wd=0., decouple_wd=True):
 "A `Optimizer` for SGD with `lr` and `mom` and `params`"
    cbs = [weight_decay] if decouple_wd else [l2_reg]
 if mom != 0: cbs.append(average_grad)
    cbs.append(sgd_step if mom==0 else momentum_step)
 return Optimizer(params, cbs, lr=lr, mom=mom, wd=wd)

类似的， fastai建立了其他的优化函数，包括RMSProp，Adam，RAdam，QHAdam，LARS/LARC，LAMB。

• class Lookahead(Optimizer, GetAttr):

Lookhead是用来实现”k steps forward, 1 step back", 也就是前进k步，后退1步。Lookhead可以建立在任何初始化时传入的优化器上面。还没有在模型训练中试过这个方法，看起来也是一种类似带动量的优化方式，不同的时对参数进行了移动平均，而不是对梯度。

• class OptimWrapper(_BaseOptimizer, GetAttr)

也许有时候我们会想对比一下fastai和Pytorch内建优化器的性能，但是又不愿意脱离fastai环境，又或者Pytorch增加了新的优化器，而fastai还没有跟上，这时候就可以直接利用OptimWrapper来实现对Pytorch内建优化器的包装。比如实现对AdamW的封装，这个优化器通常用于transformer类的模型。

OptimWrapper(torch.optim.AdamW(m.parameters(), betas=(0.9, 0.99), eps=1e-5, weight_decay=1e-2))

注意，由于fastai在建立训练器（learner）的时候，调用的时一个返回优化器的生成函数，类似上面提到的SGD函数，我们在训练器中使用OptimWrapper建立的封装优化器，也需要通过一个生成函数来传递给训练器。例如以下函数：

def  AdamW(params, lr, mom=0.9, sqr_mom=0.99, eps=1e-5, wd=0.01):
    params = [{"params": param} for param in params]
 return OptimWrapper(torch.optim.AdamW(params, betas=(mom, sqr_mom), eps=eps, weight_decay=wd))

这里的params=[{"params": param}for param in params]，是用来把fastai传入的参数组的列表转化成参数组的字典，因为Pytorch优化器接收的参数要么时参数列表，要么是分组的参数字典，否则会报错。

接下来，准备花点时间整理一下fastai的训练器的实现过程。Jeremy等人建立的fastai不仅仅是进行深度学习的一个快速的工具，如果深入fasiai实现的背后的源程序，可以学到很多深度学习的基础知识，而且还是学习Python本身的一个颇有乐趣的工具。