教程:如何使用梯度累积来克服 GPU 内存限制

1. 简介

为了训练机器学习模型,训练数据集将拆分为多个批次。该模型按批处理方式处理此数据。为了处理(评估)批处理,计算机将其加载到内存中。由于模型一次训练一个批次,因此计算机需要足够的内存来加载每个批次并完全处理它。复杂模型通常在大量详细的输入数据(如图形和音频脚本)上进行训练。数据越详细,在处理时保存数据的内存就越多。可用内存可能不足以容纳所需的批大小。当模型的内存要求超过可用内存时,它会崩溃并显示内存不足错误。

批大小越大,加载和评估其数据所需的内存就越多。解决内存约束的一种方法是减小批大小。但是,减小批量大小并不总是可取的。许多模型在更大(最多有限制)的批量大小上学习得更好、更快。

另一种方法是使用较小的批次,但不根据每个批次的反馈更新模型参数。相反,更新会累积多个批次,然后应用。在某种程度上,这模仿了使用较大批量大小的效果。这种技术称为梯度累积

所需背景

要遵循本指南,需要一些训练基于神经网络的机器学习模型的实践经验。假设您可以安装和使用 PyTorch,Torchvision 和其他软件。本指南已在采用 NVIDIA A100 GPU 的 Vultr GPU VM 上进行了测试。GPU 特定命令在其他 GPU 设备上可能有所不同。

虽然本指南确实包括基本机器学习模型的设置,但它只是为了提供一个应用梯度累积的示例,而不是解释机器学习模型本身的实现。模型的配置仅用于演示,可能不是提供最佳训练结果的最佳配置。

2. 动机

神经网络用于识别模式。神经网络接受张量形式的输入数据,并按顺序将其应用于其多层变换。模型的输出是关于输入匹配的几种模式中的哪一种的预测。例如,可以训练神经网络将动物的照片作为输入,并预测它属于哪个物种。

每个单独的变换都是一个数学运算,通常是一个线性方程。图层由大量单独的变换组成。每个变换都由模型权重参数化。这些权重(参数)确定线性方程的系数。

训练神经网络从一组随机的模型权重开始。权重会迭代更新,以使模型的输出尽可能与预期(正确)输出匹配。

在高级别上,机器学习模型训练循环中的主要步骤是:

  1. 根据(暂定)模型权重评估模型输出
  2. 计算损失(预期输出和实际输出之间的差异)
  3. 将损失反馈到模型中。损失值确定模型权重的更新量
  4. 更新模型权重

对每个批次重复这些步骤,直到模型通过整个训练数据集。训练数据集的每次传递称为一个纪元。一个典型的训练练习涉及几个时期。梯度累积技术将步骤 4 修改为:“每 N 批更新一次模型权重。

这导致保存(累积)每个批次的更新(梯度值)并继续下一批。经过N批后,应用由此累积的梯度。结果几乎与将 N 个批次的数据作为一个批次进行处理以估计梯度相同。根据硬件能够处理的批次数选择 N 的值。

3. 实现 – 手工编码算法

本节介绍如何手动实现梯度累积。本节中的示例使用 PyTorch。使用 Keras 时,过程保持不变;只有包名称不同。

此示例的目标是证明使用梯度累积的小批量大小与较大的批量大小提供类似的结果。调整准确的模型不在本指南的讨论范围之内。

3.1. PyTorch 中的基本机器学习示例

在实现梯度累积之前,请使用 PyTorch 设置一个基本的机器学习示例。假设您已经熟悉类似的示例。下一节将梯度累积应用于此示例。

注意:为了遵循代码示例,建议创建一个新的 Python 文件并将代码片段复制到其中。

导入所需模块

import torch

from torch import nn

from torch.utils.data import DataLoader

from torchvision import datasets

from torchvision.transforms import ToTensor

声明配置变量

batch_size = 64

batch_size表示每个批次的大小。此示例使用批大小 64。下一节中的示例使用较小的批大小 16 应用梯度累积。

learning_rate = 1e-3

learning_rate是一个超参数 – 它影响模型参数针对给定损失函数值的更新量。较大的学习率会使学习过程不稳定。如果学习率太小,则模型收敛时间过长。

epochs = 20

epochs是模型迭代整个训练数据集的次数。较大的批大小可能需要更多迭代。

声明变量:device

device = "cuda"

要在 GPU 上运行模型,请为值“cuda”赋值,假设为 NVIDIA GPU。模型和数据在 CPU 上实例化,需要移动到 GPU。若要在 CPU 上运行模型,请使用值“cpu”。device

导入数据集

此示例使用标准的 FashionMNIST 数据集。该数据集由服装的图片及其相应的标签组成。它被组织为一个包含 60,000 个图像的训练集和一个包含 10,000 个图像的测试集。图像缩小到 28 x 28 像素。它共有10个标签,对应10种服装。

training_data = datasets.FashionMNIST(

    root = "data",

    train = True,

    download = True,

    transform = ToTensor()

    )



test_data = datasets.FashionMNIST(

    root = "data",

    train = False,

    download = True,

    transform = ToTensor()

    )

PyTorch 的 DataLoader 实用程序模块创建了一个可迭代的结构,将数据集拆分为小批量。

loader_training_data = DataLoader(training_data, batch_size = batch_size)

loader_test_data = DataLoader(test_data, batch_size = batch_size)

定义神经网络

定义一个具有一个隐藏层的神经网络。输入层将 28 x 28 像素的图像映射到 512 个“特征”。隐藏层变换这 512 个要素,输出层将转换后的数据映射到 10 个标签。ReLU用于层内非线性激活。

class NeuralNetwork(nn.Module):

    def __init__(self):

        super(NeuralNetwork, self).__init__()

        self.flatten = nn.Flatten()

        self.linear_relu_stack = nn.Sequential(

            nn.Linear(28*28, 512),

            nn.ReLU(),

            nn.Linear(512, 512),

            nn.ReLU(),

            nn.Linear(512, 10),

        )



    def forward(self, x):

        x = self.flatten(x)

        logits = self.linear_relu_stack(x)

        return logits

将训练组件声明为变量

实例化模型(基于前面定义的神经网络)、损失函数和优化器:

model = NeuralNetwork()

# move the model to the GPU

model = model.to(device)



loss_function = nn.CrossEntropyLoss()



optimizer = torch.optim.SGD(

        model.parameters(), lr = learning_rate

        )

定义训练循环

下面的简化训练循环计算和反向传播损失,并更新每个批次的模型权重。

def training_loop(dataloader, model, loss_function, optimizer):

    size = len(dataloader.dataset)

    for batch, (X, Y_expected) in enumerate(dataloader):



        # move the data to the GPU

        X = X.to(device)

        Y_expected = Y_expected.to(device)



        Y_computed = model(X)



        # compute the loss

        loss = loss_function(Y_computed, Y_expected)



        # backpropagate the loss

        loss.backward()



        # update the model weights

        optimizer.step()



        # reset the optimizer for the next iteration

        optimizer.zero_grad()

定义测试循环

测试循环包括以下步骤:

  1. 使用纪元结束时的更新参数评估模型输出。
  2. 计算每个项目的损失,并将所有损失相加以获得累积损失。
  3. 计算正确结果的总数(其中预期输出与模型输出匹配)。
  4. 平均批次数的累积损失。这给出了该时期每批的平均损失。
  5. 正确结果数相对于数据集大小的平均值。这给出了模型的准确性。
  6. 格式化并打印结果。

下面的代码实现了这些步骤。

def test_loop(dataloader, model, loss_function):

    size = len(dataloader.dataset)

    num_batches = len(dataloader)

    cumulative_loss, correct_results = 0, 0



    with torch.no_grad():

        for X, Y_expected in dataloader:



            # move the data to the GPU

            X = X.to(device)

            Y_expected = Y_expected.to(device)



            Y_computed = model(X)

            cumulative_loss += loss_function(Y_computed, Y_expected).item()

            correct_results += (Y_computed.argmax(1) == Y_expected).type(torch.float).sum().item()



    # average out the cumulative loss and the number of correct results 

    cumulative_loss /= num_batches

    correct_results /= size

    print(f"Test Results: \n Accuracy: {(100*correct_results):>0.1f}%, Avg loss: {cumulative_loss:>8f} \n")

运行程序

在所需的纪元数上调用训练和测试循环。这将迭代训练模型,测试模型的输出,并打印出每个纪元的结果。

for t in range(epochs):

    print(f'Epoch {t+1} \n -------------------')

    training_loop(loader_training_data, model, loss_function, optimizer)

    test_loop(loader_test_data, model, loss_function)

print('Done')

3.2. 梯度累积

要使用梯度累积:

  1. 使用较小的批量大小。
  2. 为累积间隔声明一个变量 N。
  3. 通过在 N 个批次上平均损失值来规范化损失值。
  4. 修改训练循环以每 N 批调用优化器。

更新变量:batch_size

batch_size = 16

声明一个变量来保存累积间隔 N:

N = 4

更新训练循环:

def training_loop(dataloader, model, loss_function, optimizer):

    size = len(dataloader.dataset)

    for batch, (X, Y_expected) in enumerate(dataloader):



        X = X.to(device)

        Y_expected = Y_expected.to(device)



        Y_computed = model(X)



        loss = loss_function(Y_computed, Y_expected)



        # normalize the loss value

        loss = loss / N



        loss.backward()



        # call the optimizer every N steps

        if ((batch + 1) % N == 0) or (batch + 1 == len(dataloader)):

            optimizer.step()

            optimizer.zero_grad()

为实现梯度累积而添加/修改的代码行前面有注释。代码的其余部分与以前相同。

解释

loss.backward()– 反向传播函数,计算每个参数的损失(相对于每个参数的损失的偏导数),并将其累积在该参数的梯度中。对于参数:x

# pseudo-code

x.gradient += d_loss / d_x

调用根据梯度(和学习率)更新模型权重:optimizer.step()

# pseudo-code

x += -learning_rate * x.gradient

optimizer.zero_grad()在 N 次迭代结束时重置梯度 – 因此可以累积接下来 N 批的梯度。

在标准情况下,优化器和反向传播函数每批调用一次。因此,每个批次后的更新仅基于该批次的梯度。对于梯度累积,损失仍然在每个批次中反向传播(和累积),但优化器每 N 个批次调用一次。因此,模型权重根据 N 个批次中累积的梯度每 N 个批次更新一次。更新权重后,将重置梯度。

在调用反向传播函数之前,有必要通过对 N 个批次进行平均来规范化损失。因为将每批的损失反向传播到N个批次会使梯度比应有的大得多,从而导致过度校正。规范化损失可以纠正这一点。使用归一化损失累积的梯度是使用单个大批次更新模型权重时梯度的代理。

第 3.1 节中的原始示例使用批大小 64。使用梯度累积通过使用 16 的批大小并累积 N = 4 个批次的梯度来获得类似的结果。在最初的几个时期之后,有和没有梯度累积的结果应该收敛到类似的数字(结果不会相同)。

注意:要在 PyTorch 上测量模型的资源(时间和内存)消耗,请使用 PyTorch Profiler。有关探查器用法的讨论不在本指南的讨论范围之内。

4. 实施 – 使用预打包工具

梯度累积也可以使用优化器上的预打包包装器来实现。一些流行的机器学习框架(例如,PyTorch Lightning)包括对直接使用梯度累积的支持。默认情况下,Keras 和 PyTorch 不包括梯度累积支持。但是,一些机器学习库包含直接在 Keras 和 PyTorch 中实现梯度累积的模块。

本节中的代码片段仅说明如何在现有代码中使用预打包的工具。这些不是完整的代码示例。

4.1. 运行:AI 包装器

Run:AI是一家机器学习基础设施和平台公司。他们公开可用的 Python 工具库包括 Keras 和 PyTorch 的梯度累积包装器。

安装:runai

$ pip install runai

4.1.1. 凯拉斯

Keras 是一个深度学习框架,为底层 TensorFlow 模块提供 API 接口。要在 Keras 模型中使用包装器,请导入特定于 Keras 的梯度累积模块:runai

>>> import tensorflow as tf

>>> import runai.ga.keras

从 Keras 实例化优化器:

>>> optimizer = tf.keras.optimizers.Adam()

上面的行创建了一个基于 Adam 算法的优化器。

使用 Run:AI 包装器更新优化器实例:

>>> optimizer = runai.ga.keras.optimizers.Optimizer(optimizer, steps=N)

这应该输出如下内容:

Wrapping 'Adam' Keras optimizer with GA of N steps

使用新的优化器实例运行具有梯度累积的训练循环。

4.1.2. PyTorch

从以下位置导入特定于 PyTorch 的梯度累积模块:runai

>>> import runai.ga.torch

从 PyTorch 库中实例化优化器:

>>> optimizer = torch.optim.SGD(model.parameters(), lr = learning_rate)

这将创建一个基于随机梯度下降算法的优化器。

用梯度累积模块包裹:optimizer

>>> optimizer = runai.ga.torch.optim.Optimizer(optimizer, steps=N)

输出应类似于:

Wrapping 'SGD' PyTorch optimizer with GA of N steps

在训练循环中使用此优化器来利用梯度累积。

4.2. PyTorch 闪电

Lightning是一个基于 PyTorch 的框架。它带有许多预先打包的例程,无需编写样板代码。闪电还包括对梯度累积的支持。

闪电的训练器模块处理训练过程。导入培训师:

from pytorch_lightning import Trainer

使用梯度累积间隔实例化训练器对象,:accumulate_grad_batches

trainer = Trainer(accumulate_grad_batches=4)

请注意,默认值为 1。上面定义的 Trainer 实例在调用优化器之前累积 4 个批次的梯度。也可以指定每个纪元的累积间隔:accumulate_grad_batches

trainer = Trainer(accumulate_grad_batches={0: 8, 4: 2})

上面的实例为从纪元 0 开始的初始纪元累积了 8 个批次的梯度。纪元 4 开始,累积 2 批。

5. 最后

梯度累积适用于训练过程。它不适用于在内存有限的系统上运行(预训练)大型模型,如稳定扩散。

请注意,较大的批量大小并不总是对应于更好的训练 – 这取决于模型和数据的细节。可能需要一些试验和错误来确定给定用例的正确批量大小。在实践中,使用较大的批次获得良好的结果需要调整其他变量,例如 epoch 数和学习率。

该技术并非特定于 GPU 内存。它适用于正在训练模型的任何内存。如果在 CPU 上训练模型,则会将梯度累积应用于主内存 (RAM)。在实践中,大到足以保证这种技术的模型(和数据集)太大,无法在CPU上有效地训练。

赞(0)
未经允许不得转载:主机百科 » 教程:如何使用梯度累积来克服 GPU 内存限制