太炸裂了!分享9个让 PyTorch 模型训练飞快的技巧!

节前,我们组织了一场算法岗技术&面试讨论会,邀请了一些互联网大厂朋友、今年参加社招和校招面试的同学。

针对大模型技术趋势、大模型落地项目经验分享、新手如何入门算法岗、该如何准备面试攻略、面试常考点等热门话题进行了深入的讨论。

汇总合集:

《PyTorch 实战宝典》重磅发布!

《大模型面试宝典》(2024版) 发布!


也许你仍然在使用32位精度进行计算,或者甚至只是在单个GPU上进行训练。

然而,随着科技的进步,我们已经有了更好的选择。使用更高精度的计算,如16位浮点数或混合精度,可以提高训练速度并减少内存消耗。同时,利用多个GPU进行并行训练,可以大大加快训练过程。

笔者在这里总结了提升Pytorch模型训练速度的9个技巧,与大家分享~

PyTorch-Lightning 是建立在PyTorch之上的一个封装,它提供了自动化训练的功能,同时允许开发者完全控制关键的模型组件。

这里以MNIST定义LightningModel并使用Trainer来训练模型为例。

#导入PyTorch-Lightning库中的Trainer类,用于管理训练过程
from pytorch_lightning import Trainer
#创建LightningModule实例,作为要训练的模型
model = LightningModule()
#创建Trainer实例,用于配置和管理训练过程
trainer = Trainer()
#开始训练
trainer.fit(model)

trainer.fit()方法将根据LightningModule中定义的训练逻辑来执行训练步骤,并自动处理数据加载、优化器配置、学习率调度等细节。

使用DataLoaders

使用DataLoaders来加载数据是获得训练速度提升的最简单方法之一。通过保存h5py或numpy文件以加速数据加载的时代已经过去了,而现在可以 使用PyTorch的DataLoader来轻松加载图像数据(对于NLP数据,请参考TorchText库)。

在PyTorch-Lightning中,不需要显式地编写训练循环,只需要定义好DataLoaders和Trainer,PyTorch-Lightning会在需要的时候自动调用它们。

下面是一个加载MNIST数据集并使用DataLoader进行训练的示例代码:

from torch.utils.data import DataLoader
from torchvision.datasets import MNIST

dataset = MNIST(root=self.hparams.data_root, train=train, download=True)
loader = DataLoader(dataset, batch_size=32, shuffle=True)

for batch in loader:
    x, y = batch
    model.training_step(x, y)
    # 其他训练逻辑

在这个示例中,首先创建了一个MNIST数据集的实例,然后使用DataLoader将其封装成一个可迭代的数据加载器。在训练循环中,可以遍历DataLoader,每次获取一个batch的数据,并将其传递给模型的training_step()方法进行训练。

通过使用DataLoaders,可以更高效地加载和处理大量的训练数据,从而提高训练速度。此外,DataLoader还支持数据的随机打乱(shuffle)、批量大小(batch_size)等参数设置,可以根据实际需求进行调整。

DataLoaders中设置num_workers参数

在DataLoaders中,可以通过设置num_workers参数来允许批量并行加载数据,从而加速训练过程。下面是一个示例代码:

# 慢的方式
loader = DataLoader(dataset, batch_size=32, shuffle=True)

# 快的方式(使用10个workers)
loader = DataLoader(dataset, batch_size=32, shuffle=True, num_workers=10)

在第一个示例中,创建了一个DataLoader,没有指定num_workers参数,这意味着数据加载将在主进程中进行,而不会并行化。

在第二个示例中,通过将num_workers设置为10,启用了批量并行加载。这意味着数据加载将在10个worker进程中进行,并行地加载多个batch,从而加速数据加载的过程。

通过适当设置num_workers参数,可以根据系统的硬件和资源情况,选择合适的worker数量来提高数据加载的效率。然而,需要注意的是,并不是worker越多越好,过多的worker可能会导致资源竞争和性能下降。

Batch size

在进行下一个优化步骤之前,增加批量大小(batch size)到CPU-RAM或GPU-RAM允许的最大范围是一个重要的优化策略。

增加批量大小可以带来以下好处:

  • 更高效地利用计算资源,尤其是GPU的并行计算能力。

  • 减少数据加载和传输的次数,提高训练速度。

  • 可以获得更稳定的梯度估计,有助于模型收敛。

然而,增加批量大小也会带来一些挑战:

  • 内存占用增加:较大的批量大小需要更多的内存空间来存储数据和梯度。

  • 学习率调整:增加批量大小后,通常需要相应地增加学习率,以保持相似的收敛行为。

因此,在增加批量大小之前,需要确保你的硬件和资源可以支持更大的批量大小,并相应地调整学习率。

梯度累积

梯度累积(Gradient Accumulation)是一种在计算资源有限的情况下,模拟较大批量大小的技术。通过多次执行前向传播、反向传播和优化步骤,将梯度累积起来,以获得与较大批量大小相同的效果。

下面是一个使用梯度累积的示例代码:

# 清除上一步的梯度
optimizer.zero_grad()

# 16次梯度累积步骤
scaled_loss = 0
for accumulated_step_i in range(16):
    out = model.forward()
    loss = some_loss(out, y)
    loss.backward()
    scaled_loss += loss.item()

# 更新权重
optimizer.step()

# 损失值现在按累积批次数量进行缩放
actual_loss = scaled_loss / 16

在这个示例中,通过循环执行16个梯度累积步骤,每个步骤进行前向传播、计算损失、反向传播和梯度累积。然后调用optimizer.step()来更新权重。

在PyTorch-Lightning中,只需要设置accumulate_grad_batches参数来指定梯度累积的次数。例如:

trainer = Trainer(accumulate_grad_batches=16)
trainer.fit(model)

保留的计算图

在记录损失值时,为了避免撑爆内存,只存储损失的数值而不是整个计算图。可以使用.item()方法来获取损失的数值。

# 方式1
losses.append(loss)

# 方式2
losses.append(loss.item())

在方式1中,损失值loss会保留整个计算图的副本,这会占用大量的内存空间。而方式2中,使用loss.item()来获取损失的数值,并将其存储到列表中,这样就只保留了数值,而不会占用过多的内存。

PyTorch-Lightning会非常小心地确保不会保留计算图的副本,尽量减少内存的占用。因此,在使用PyTorch-Lightning时,可以放心地使用.item()方法来获取损失的数值,而不必担心内存问题。

单个GPU训练

完成上述步骤之后,即可开始在GPU上进行训练。GPU上进行训练可以利用多个GPU核心之间的并行计算,从而加速训练过程。

在进行GPU训练时,需要做两件事情:

  1. 将模型移动到GPU上;

  2. 在每次数据通过时将数据放到GPU上。

下面是在PyTorch中进行GPU训练的示例代码:

# 将模型放到GPU上
model.cuda()
# 将数据放到GPU上
x = x.cuda()
# 在GPU上运行
model(x)

如果使用PyTorch-Lightning,几乎不需要做任何额外的工作,只需要设置Trainer的gpus参数来指定要使用的GPU数量。

# 指定训练的gpu id
trainer = Trainer(gpus=[0])
trainer.fit(model)

在进行GPU训练时,需注意限制CPU和GPU之间的数据传输次数。尽量避免频繁地在CPU和GPU之间复制数据。

此外,还要注意调用强制GPU同步的操作,如清空内存缓存torch.cuda.empty_cache()。这样的操作会阻塞所有GPU,直到它们都完成同步。

然而,如果使用PyTorch-Lightning,则通常不需要担心这些问题。PyTorch-Lightning会小心地处理GPU同步和内存管理,以确保高效的训练过程。

使用16-bit精度

使用16-bit精度是一种惊人的技术,可以将内存占用减半。大多数模型通常使用32位精度进行训练,但是研究表明,使用16位精度的模型也可以表现得很好。混合精度则意味着在某些部分使用16位精度,但将权重等内容保持在32位精度。

要在PyTorch中使用16位精度,可以安装NVIDIA的apex库,并对模型进行如下更改:

# 在模型和优化器上启用16位精度
model, optimizers = amp.initialize(model, optimizers, opt_level='O2')

# 当进行.backward()时,让amp处理以便它可以对损失进行缩放
with amp.scale_loss(loss, optimizer) as scaled_loss:
    scaled_loss.backward()

apex库会处理大部分工作,包括梯度缩放,以防止梯度爆炸或接近零。

在PyTorch-Lightning中,启用16位精度不需要修改模型的任何内容,也不需要执行上述代码。只需在创建Trainer时设置precision=16即可。

trainer = Trainer(amp_level='O2', use_amp=False)
trainer.fit(model)

通过这样的设置,PyTorch-Lightning将自动启用16位精度,并根据需要进行梯度缩放和其他必要的操作。

移动到多个GPUs中

在多个GPU上进行训练有几种方法可以选择。以下是其中三种常见的方法:

分批次训练(Batch Splitting)

分批次训练是指将模型复制到每个GPU上,并将每个GPU中的一部分批次数据进行训练。

# 在每个GPU上复制模型,并将批次的四分之一分配给每个GPU
model = DataParallel(model, devices=[0, 1, 2, 3])

# out有4个输出(每个GPU一个)
out = model(x.cuda(0))

在PyTorch-Lightning中,只需要增加gpus参数来指定使用的GPU数量,其他的无需更改。

trainer = Trainer(gpus=[0, 1, 2, 3])
trainer.fit(model)

模型分布训练(Model Parallelism)

模型可能太大无法完全放入内存中。例如,带有编码器和解码器的序列到序列模型可能需要占用大量内存。在这种情况下,可以将编码器和解码器放在不同的GPU上进行训练。

# 将编码器放在GPU 0上,将解码器放在GPU 1上
encoder_rnn.cuda(0)
decoder_rnn.cuda(1)

# 在GPU 0上运行输入数据通过编码器
encoder_out = encoder_rnn(x.cuda(0))

# 在GPU 1上运行输出通过解码器
out = decoder_rnn(encoder_out.cuda(1))

# 将输出数据移回GPU 0上
out = out.cuda(0)

在PyTorch-Lightning中,不需要指定任何GPU,只需将模型的模块放在正确的GPU上即可。

class MyModule(LightningModule):
    def __init__():
        self.encoder = RNN(...)
        self.decoder = RNN(...)
    def forward(x):
        self.encoder.cuda(0)
        self.decoder.cuda(1)
        out = self.encoder(x)
        out = self.decoder(out.cuda(1))

model = MyModule()
trainer = Trainer()
trainer.fit(model)

混合使用(Hybrid Approach)

# 更改这些行
self.encoder = RNN(...)
self.decoder = RNN(...)

# 更改为
# 现在每个RNN都基于不同的GPU设备
self.encoder = DataParallel(self.encoder, devices=[0, 1, 2, 3])
self.decoder = DataParallel(self.encoder, devices=[4, 5, 6, 7])

# 在forward中...
out = self.encoder(x.cuda(0))

# 注意输入数据放在设备列表中的第一个设备上
out = self.decoder(out.cuda(4))  # <--- 这里的4

在使用多个GPU进行训练时,需要考虑以下注意事项:

  • 如果模型已经在GPU上,model.cuda()方法不会执行任何操作。

  • 总是将输入数据放在设备列表中的第一个设备上。

  • 在设备之间传输数据是昂贵的,应该将其作为最后的手段。

  • 优化器和梯度会被保存在GPU 0上,因此,GPU 0上使用的内存可能会比其他GPU大得多。

多节点GPU训练

在分布式训练中,每个机器上的每个GPU都有一个模型的副本,并且每个机器都会获得数据的一部分进行训练。每个模型副本在其所在的GPU上独立初始化,并在数据的一个分区上进行训练。然后,所有模型副本会彼此同步梯度更新。

这种方式可以显著加快训练速度,并且使得处理更大规模的数据集成为可能。通过将训练任务分布到多个机器和GPU上,可以同时进行多个训练任务,从而节省了训练时间。

在PyTorch中,可以使用DistributedDataParallel (DDP) 模块来实现这种分布式训练方式。它通过在每个节点上复制每个GPU上的模型并同步梯度,实现了模型的并行训练和梯度更新。代码如下:

def tng_dataloader():
     d = MNIST()
     # 4: 创建数据加载器
     # 将训练数据集分发到每个机器上
     dist_sampler = DistributedSampler(dataset)
     dataloader = DataLoader(d, shuffle=False, sampler=dist_sampler)

def main_process_entrypoint(gpu_nb):
     # 2: 设置所有机器和GPU之间的连接
     world = nb_gpus * nb_nodes
     # 初始化分布式训练环境,并指定通信后端和当前进程的排名和总共的进程数
     dist.init_process_group("nccl", rank=gpu_nb, world_size=world)

     # 3: 将模型移动到当前GPU,并使用DistributedDataParallel将模型包装起来
     # DistributedDataParallel会将模型的副本复制到每个GPU上,并确保在训练过程中同步梯度更新
     torch.cuda.set_device(gpu_nb)
     model.cuda(gpu_nb)
     model = DistributedDataParallel(model, device_ids=[gpu_nb])

     # train your model now...

if  __name__ == '__main__':
     # 1: 生成多个进程,每个进程都会调用main_process_entrypoint()函数。
     # 这样可以在每个机器上启动多个进程进行分布式训练
     mp.spawn(main_process_entrypoint, nprocs=8)

而在Lightning中,分布式训练变得更加简单。只需设置节点数量和GPU列表,Trainer类会自动处理剩下的细节。

# train on 1024 gpus across 128 nodes
trainer = Trainer(nb_gpu_nodes=128, gpus=[0, 1, 2, 3, 4, 5, 6, 7])

在单个节点上多GPU更快的训练

使用DistributedDataParallel (DDP) 在单个节点上的多个GPU上进行训练通常比使用DataParallel 更快。这是因为DDP只执行梯度同步的通信,而不是将整个模型复制到每个GPU上。

在Lightning中,可以通过将distributed_backend参数设置为’ddp’,并指定要使用的GPU数量来轻松实现在单个节点上的多GPU训练。示例如下:

# 在同一台机器上的4个GPU上进行训练,使用DDP比DataParallel更快
trainer = Trainer(distributed_backend='ddp', gpus=[0, 1, 2, 3])

在这个例子中,distributed_backend参数被设置为’ddp’来启用分布式训练,gpus参数指定要使用的GPU的索引。通过这种方式,Lightning将使用DDP来进行训练,从而更有效地利用多个GPU。

相关推荐

  1. ChatGPT

    2024-06-09 13:36:01       24 阅读
  2. PyTorch训练多任务模型技巧

    2024-06-09 13:36:01       37 阅读

最近更新

  1. Sklearn简介、安装教程、入门学习

    2024-06-09 13:36:01       0 阅读
  2. Golang - 连接MySQL执行查询脚本

    2024-06-09 13:36:01       0 阅读
  3. 【WSL开发】在WSL中部署jupyter lab并远程访问

    2024-06-09 13:36:01       0 阅读
  4. C语言FFT算法

    2024-06-09 13:36:01       0 阅读
  5. 嵌入式微处理器重点学习(二)

    2024-06-09 13:36:01       0 阅读
  6. React编写可复用图标组件

    2024-06-09 13:36:01       0 阅读
  7. LogicFlow 学习笔记——6. LogicFlow 基础 网格 Grid

    2024-06-09 13:36:01       0 阅读

热门阅读

  1. 图搜索算法A*、Dijkstra在路径规划中的应用

    2024-06-09 13:36:01       3 阅读
  2. Lambda 表达式(C++11)

    2024-06-09 13:36:01       4 阅读
  3. 我的创作纪念日

    2024-06-09 13:36:01       4 阅读
  4. Vue中 .passive 修饰符的作用

    2024-06-09 13:36:01       4 阅读
  5. 【python】flask相关包依赖关系问题

    2024-06-09 13:36:01       3 阅读
  6. HTTP常见响应状态码

    2024-06-09 13:36:01       4 阅读
  7. 深入了解Git:从数据模型到集成IDEA

    2024-06-09 13:36:01       4 阅读
  8. C语言---深入指针(4)

    2024-06-09 13:36:01       3 阅读
  9. 回溯之分割回文串

    2024-06-09 13:36:01       4 阅读
  10. 【环境搭建】2.阿里云ECS服务器 安装MySQL

    2024-06-09 13:36:01       3 阅读
  11. 2 程序的灵魂—算法-2.2 简单算法举例-【例 2.5】

    2024-06-09 13:36:01       3 阅读
  12. 【小海实习日记】金融-现货以及合约理解

    2024-06-09 13:36:01       5 阅读