Pytorch反向传播与梯度累积

一、自动求导与反向传播

如果处理的函数可微，Pytorch可以实现自动计算梯度，只需设置requires_grad=True。

用一个线性回归的例子进行说明，输入输出分别为x,y。

x = torch.tensor([35.7, 55.9, 58.2, 81.9, 56.3, 48.9,33.9, 21.8, 48.4, 60.4, 68.4])
y = torch.tensor([0.5, 14.0, 15.0, 28.0, 11.0, 8.0,3.0, -4.0, 6.0, 13.0, 21.0])
x = 0.1 * x
params=torch.tensor([1.0, 0.0], requires_grad=True)

# 模型
def model(x,w,b):
    return w*x+b

# loss
def loss_fn(y_p,y):
    return ((y_p-y)**2).mean()

训练过程  

每次迭代时，需要将梯度归零，可以使用zero_将params的梯度归零。

torch.no_grad()避免params执行自动的梯度更新，计算梯度并手动进行反向传播。

def training_loop(n_epochs, learning_rate, params, x, y):
    for epoch in range(1, n_epochs + 1):
        if params.grad is not None:  # 梯度归零
            params.grad.zero_()

        y_p = model(x, *params)
        loss = loss_fn(y_p, y)
        loss.backward()

        with torch.no_grad():  #避免自动梯度更新，手动更新参数
            params -= learning_rate * params.grad

        if epoch % 500 == 0:
            print('Epoch %d, Loss %f' % (epoch, float(loss)))

    return params

training_loop(
    n_epochs=5000,
    learning_rate=1e-2,
    params=params,  
    x=x,  
    y=y)

二、优化器实现梯度更新

在执行backward()之前执行optimizer.zero_grad()实现梯度归零。

params的值在调用step()时更新，优化器会查看params.grad，从中减去学习率乘梯度更新参数，不需要手动计算。

def training_loop_add_optim(n_epochs, optimizer,params, x, y):
    for epoch in range(1, n_epochs + 1):

        y_p = model(x, *params)
        loss = loss_fn(y_p, y)
        optimizer.zero_grad() #梯度归零
        loss.backward()
        optimizer.step() #梯度更新


        if epoch % 500 == 0:
            print('Epoch %d, Loss %f' % (epoch, float(loss)))

    return params

training_loop_add_optim(
    n_epochs=5000,
    optimizer=torch.optim.SGD([params],lr=1e-2),
    params=params,
    x=x,
    y=y)

三、梯度累积

当运算资源有限时，batch size不能设置太大，但batch size过小可能会导致loss收敛不稳定，这时可以通过梯度累积的方法解决。

梯度累积指的是每个batch的数据都计算梯度，但不是每个batch都更新梯度。对每次计算的梯度进行累加，不清空，累加到一定次数后再进行梯度更新，然后清空梯度，进入下一次循环。体现在代码上，就是n次backward()对应一次optimizer.zero_grad()和optimizer.step()

比如batch size想要设置为72，总共14400条样本，那么需要计算14400/72=200次梯度计算，200次梯度更新和200次梯度清空。

但如果计算资源有限，batch size设置为6，则可以设置累积步数accumulation_steps为12，即12个batch更新一次梯度，清空一次梯度。则更新与清空梯度的次数仍为200，但计算梯度的次数是14400/6=2400次。

def training_loop_add_optim(n_epochs, optimizer,params, x, y):
    for epoch in range(1, n_epochs + 1):
        for i, train_batch in enumerate(train_dataloader):
            y_p = model(x, *params)
            loss = loss_fn(y_p, y)
            loss.backward()
            if ((i + 1) % accumulation_steps) == 0:
                optimizer.zero_grad() #梯度清零
                optimizer.step() # 梯度累积
            i+=1
        if epoch % 500 == 0:
            print('Epoch %d, Loss %f' % (epoch, float(loss)))

    return params