当前位置: 首页 > 新闻中心 > 学习PyTorch的optim模块
学习PyTorch的optim模块
时间:2024-04-22 15:05:14 点击次数:

PyTorch的optim是用于参数优化的库(可以说是花式梯度下降),optim文件夹下有12个文件,包括1个核心的父类(optimizer)、1个辅助类(lr_scheduler)以及10个常用优化算法的实现类。optim中内置的常用算法包括adadelta、adam、adagrad、adamax、asgd、lbfgs、rprop、rmsprop、sgd、sparse_adam。


Optimizer类是所有优化方法的父类,它保存参数状态并根据梯度将其更新。这里将分享Optimizer类的构造方法和有关梯度控制的两个方法。构造方法指的是_init_()方法,梯度控制方法包括大家熟悉的zero_grad()step()


1.构造方法init()

Optimizer的init函数接收两个参数:第一个是需要被优化的参数,其形式必须是Tensor或者dict;第二个是优化选项,包括学习率、衰减率等。

第一个位置通常用model.parameters()填充,如果有特殊的需求,也可以手动写一个dict来作为输入。这时只需要保证dict中有一个['params']键即可,其他键可以按照自己的要求填写。例如,如果想让大部分参数使用某个学习率(例如1e-2),同时某个参数使用另一种学习率(例如1e-3),这时就可以用手写dict的方式来进行了。

optim.SGD([               
{'params': model.base.parameters()},
{'params': model.classifier.parameters(), 'lr': 1e-3}],
    lr=1e-2, momentum=0.9)


现在我们写一个简单的卷积模型与一个具体的SGD实现,然后看看该模型的参数与SGD内部保存的参数的对比。

class TestNet(nn.Module):
    def __init__(self):
        super(TestNet, self).__init__()
        self.conv1 = nn.Conv2d(1, 5, 3)

    def forward(self, x):
        return F.relu(self.conv1(x))

首先用parameters()函数打印出卷积模型的内部参数,模型包括每个卷积层的weight和bias:

[('conv1.weight', Parameter containing:
tensor([[[[-0.0213,  0.0984, -0.2405],
[-0.0056, -0.2703, -0.3016],
[ 0.0223,  0.2061,  0.1056]]],

[[[ 0.0897,  0.0924, -0.1880],
[ 0.1791, -0.2905,  0.0672],
[-0.1899, -0.0755, -0.1943]]],

[[[ 0.1418,  0.0732, -0.0843],
[ 0.0812,  0.1916,  0.0555],
[-0.1801, -0.1689, -0.0851]]],

[[[ 0.3089,  0.0022, -0.2323],
[ 0.1079, -0.2968,  0.1566],
[ 0.1965,  0.2649,  0.1342]]],

[[[ 0.1645, -0.0703, -0.1438],
[-0.2804,  0.1042,  0.1325],
[ 0.3113, -0.0392,  0.3187]]]])), 
('conv1.bias', Parameter containing: 
tensor([ 0.2647, -0.2603,  0.1474,  0.2162,  0.0728]))]

在接收了外部输入的参数params以后,优化器会把params存在内建的self.param_groups里(这其实是一个dict)。除了保存参数的['params']键以外,Optimizer还维护着其他的优化系数,例如学习率和衰减率。经过处理的参数变成了如下的形式:

[{'params':[Parameter containing:
tensor([[[[ 0.0323, -0.0670,  0.3266],
[ 0.0829, -0.2057,  0.1977],
[ 0.1516,  0.0476, -0.1742]]],

[[[-0.0586, -0.2476,  0.1331],
[-0.1710,  0.1317,  0.1946],
[-0.2252,  0.2496,  0.3116]]],

[[[-0.2862, -0.0688, -0.0352],
[-0.2557,  0.1285, -0.1951],
[ 0.2105, -0.1723,  0.2505]]],

[[[-0.1729,  0.2929,  0.2806],
[-0.2044,  0.2730,  0.0408],
[ 0.2173,  0.2981,  0.1191]]],

[[[ 0.1528,  0.3279,  0.2510],
[ 0.2618,  0.2712,  0.0861],
[-0.1798,  0.1648,  0.2226]]]]),
Parameter containing: 
tensor([ 0.1659, -0.0458,  0.1124,  0.3329, -0.3009])],
'lr': 0.001, 
'momentum': 0, 
'dampening': 0, 
'weight_decay': 0, 
'nesterov': False}]


2.梯度控制方法zero_grad()和step()

在进行反向传播之前,必须要用zero_grad()清空梯度。具体的方法是遍历self.param_groups中全部参数,根据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_()

在反向传播backward()计算出梯度之后,就可以调用step()实现参数更新。不过在Optimizer类中,step()函数内部是空的,并且用raise NotImplementError来作为提醒。后面会根据具体的优化器来分析step()的实现思路。


lr_scheduler用于在训练过程中根据轮次灵活调控学习率。调整学习率的方法有很多种,但是其使用方法是大致相同的:用一个Schedule把原始Optimizer装饰上,然后再输入一些相关参数,然后用这个Schedule做step()。

现在介绍一个使用LambdaLR的例子。LambdaLR的构造方法是LambdaLR (optimizer,lr_lambda,last_epoch=-1),第一个函数参数是我们的原始优化器(例如SGD等),最后一个参数是最后一个轮次的index,中间的各个参数因方法而异。在使用普通的优化方法时,优化器的学习率是不可变的,LRSchedule的存在使得灵活调控成为可能。

在LambdaLR类中,lr_lambda是一个用于计算学习率的Lambda函数,具体的使用方法如下:

lambda1 = lambda epoch: epoch // 30
scheduler = LambdaLR(optimizer, lr_lambda=lambda0)
scheduler.step()


optim库中实现的算法包括Adadelta、Adagrad、Adam、基于离散张量的Adam、基于 \\infty 范式的Adam(Adamax)、Averaged SGD、L-BFGS、RMSProp、resilient BP、基于Nesterov的SGD算法。

下面以SGD算法为例,分析step()函数的实现。对于param_groups中每个group的param而言,首先要检查它的梯度p.grad,在p.grad的基础上做momentum计算,把最终的计算结果通过add_(-group['lr'], d_p)的方式给params做更新。

# SGD的step()函数
def step(self, closure=None):
    """Performs a single optimization step. """
    loss = None
    if closure is not None:
        loss = closure()

    for group in self.param_groups:
        weight_decay = group['weight_decay']
        momentum = group['momentum']
        dampening = group['dampening']
        nesterov = group['nesterov']

        for p in group['params']:
            if p.grad is None:
                continue
            d_p = p.grad.data
            if weight_decay != 0:
                d_p.add_(weight_decay, p.data)
            if momentum != 0:
                param_state = self.state[p]
                if 'momentum_buffer' not in param_state:
                    buf = param_state['momentum_buffer'] = torch.zeros_like(p.data)
                    buf.mul_(momentum).add_(d_p)
                else:
                    buf = param_state['momentum_buffer']
                    buf.mul_(momentum).add_(1 - dampening, d_p)
                if nesterov:
                    d_p = d_p.add(momentum, buf)
                else:
                    d_p = buf

            p.data.add_(-group['lr'], d_p)

    return loss

Copyright © 2012-2018 IM电竞真空泵水泵销售中心 版权所有     课ICP备985981178号

平台注册入口