python - 如何在训练期间调整 gpu 批量大小?
问题描述
令我惊讶的是,我在网上找不到任何关于如何在不停止训练的情况下动态调整 GPU 批量大小的资源。
思路如下:
1) 有一个(几乎)与正在使用的 GPU 无关的训练脚本。批量大小将动态调整,无需用户干预或调整。
2) 仍然能够指定所需的训练批量大小,即使太大而无法容纳最大的已知 GPU。
例如,假设我想使用批量大小为 4096 个图像的模型来训练模型,每个图像大小为 1024x1024。还假设我可以访问具有不同 NVidea GPU 的服务器,但我不知道会提前分配给我哪一个。(或者每个人都想使用最大的 GPU,而我要等很长时间才能成为我的任期)。
我希望我的训练脚本找到最大批量大小(假设它是每个 GPU 批次 32 个图像),并且仅在处理完所有 4096 个图像后才更新优化器(一个训练批次 = 128 个 GPU 批次)。
解决方案
有不同的方法来解决这个问题。但是,如果指定可以完成这项工作的 GPU,或者无法使用多个 GPU,那么动态调整 GPU 批量大小会很方便。
我用 pytorch 中的说明性训练示例准备了这个 repo(它应该在 TensorFlow 中类似地工作)
在下面的代码中,try/except 用于在不停止训练的情况下尝试不同的 GPU 批量大小。当批次变得太大时,它会被缩小并关闭适配。请检查 repo 以了解实现细节和可能的错误修复。
它还实现了一种称为 Batch Spoofing 的技术,该技术在进行反向传播之前执行许多前向传递。在 PyTorch 中,它只需要替换 optimizer.zero_grad()。
import torch
import torchvision
import torch.optim as optim
import torch.nn as nn
# Example of how to use it with Pytorch
if __name__ == "__main__":
# #############################################################
# 1) Initialize the dataset, model, optimizer and loss as usual.
# Initialize a fake dataset
trainset = torchvision.datasets.FakeData(size=1_000_000,
image_size=(3, 224, 224),
num_classes=1000)
# initialize the model, loss and SGD-based optimizer
resnet = torchvision.models.resnet152(pretrained=True,
progress=True)
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(resnet.parameters(), lr=0.01)
continue_training = True # criteria to stop the training
# #############################################################
# 2) Set parameters for the adaptive batch size
adapt = True # while this is true, the algorithm will perform batch adaptation
gpu_batch_size = 2 # initial gpu batch_size, it can be super small
train_batch_size = 2048 # the train batch size of desire
# Modified training loop to allow for adaptive batch size
while continue_training:
# #############################################################
# 3) Initialize dataloader and batch spoofing parameter
# Dataloader has to be reinicialized for each new batch size.
trainloader = torch.utils.data.DataLoader(trainset,
batch_size=int(gpu_batch_size),
shuffle=True)
# Number of repetitions for batch spoofing
repeat = max(1, int(train_batch_size / gpu_batch_size))
try: # This will make sure that training is not halted when the batch size is too large
# #############################################################
# 4) Epoch loop with batch spoofing
optimizer.zero_grad() # done before training because of batch spoofing.
for i, (x, y) in enumerate(trainloader):
y_pred = resnet(x)
loss = criterion(y_pred, y)
loss.backward()
# batch spoofing
if not i % repeat:
optimizer.step()
optimizer.zero_grad()
# #############################################################
# 5) Adapt batch size while no RuntimeError is rased.
# Increase batch size and get out of the loop
if adapt:
gpu_batch_size *= 2
break
# Stopping criteria for training
if i > 100:
continue_training = False
# #############################################################
# 6) After the largest batch size is found, the training progresses with the fixed batch size.
# CUDA out of memory is a RuntimeError, the moment we will get to it when our batch size is too large.
except RuntimeError as run_error:
gpu_batch_size /= 2 # resize the batch size for the biggest that works in memory
adapt = False # turn off the batch adaptation
# Number of repetitions for batch spoofing
repeat = max(1, int(train_batch_size / gpu_batch_size))
# Manual check if the RuntimeError was caused by the CUDA or something else.
print(f"---\nRuntimeError: \n{run_error}\n---\n Is it a cuda error?")
如果你有可以在 Tensorflow、Caffe 或其他方面做类似的代码,请分享!
推荐阅读
- kivy - Kivy:ScrollView 将布局挤压在一起
- python - Python数据框跳过带有关键错误的行
- elasticsearch - ElasticSearch 聚合过滤器(非嵌套)数组
- reactjs - React 功能组件通信
- azure-devops - ProductionRun 的 Azure DevOps 服务器到服务迁移时间更长?
- lua - LUA - 如何仅屏蔽字符串中的数字?
- meshlab - 在 Meshlab 中 - 当找到修剪模型的表面积时 - 我是只得到模型的外表面还是内部?
- javascript - Vue3“在'vue-router'中找不到导出'createWebHistory,createRouter'
- discord - 为什么 discord.js 不允许我在这一行中使用 .then?
- firebase - Firebase rules.MapDiff 是否知道嵌套字段的更改?