一、问题概述
1.问题分析
CUDA out of memory问题通常发生在深度学习训练过程中，当GPU的显存不足以容纳模型、输入数据以及中间计算结果时就会触发。这个问题可能由几个因素引起：
模型和数据规模：深度学习模型尤其是大型模型，如Transformer或大型CNN，拥有大量的参数，这些参数在训练时需要被加载到GPU显存中。同时，如果批量大小（batch size）设置得过大，一次性处理的数据量也会增加，进一步加大显存的负担。
内存管理：深度学习框架在处理数据和模型时，可能会因为不当的内存管理而导致显存没有被及时释放，或者由于频繁的内存分配和释放造成内存碎片化，这会减少可用的连续内存块，即使总显存使用量没有达到100%，也可能出现内存不足的情况。
数据加载和预处理：如果数据加载和预处理步骤没有优化，比如加载了大量未被立即使用的数据，或者数据没有被适当地压缩或降维，也可能导致显存使用量激增。
2.报错代码
3.示例代码
这里提供一个简单的深度学习示例代码，使用PyTorch框架来实现一个简单的多层感知器（MLP）模型，用于分类MNIST数据集中的手写数字。MNIST是一个包含60,000个训练样本和10,000个测试样本的数据库，每个样本都是一个28x28的灰度图像，代表一个手写数字（0到9）。
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torchvision import datasets, transforms
from torch.utils.data import DataLoader
# 定义一个简单的多层感知器模型
class SimpleMLP(nn.Module):
def __init__(self):
super(SimpleMLP, self).__init__()
self.fc1 = nn.Linear(28 * 28, 512) # 输入层到隐藏层
self.fc2 = nn.Linear(512, 256) # 隐藏层到隐藏层
self.fc3 = nn.Linear(256, 10) # 隐藏层到输出层
def forward(self, x):
x = x.view(-1, 28 * 28) # 展平图像
x = F.relu(self.fc1(x)) # 第一个隐藏层
x = F.relu(self.fc2(x)) # 第二个隐藏层
x = self.fc3(x) # 输出层
return x
# 设置数据加载器
transform = transforms.Compose([
transforms.ToTensor(), # 将图片转换为Tensor
transforms.Normalize((0.5,), (0.5,)) # 归一化
])
train_dataset = datasets.MNIST(root='./data', train=True, transform=transform, download=True)
test_dataset = datasets.MNIST(root='./data', train=False, transform=transform)
train_loader = DataLoader(dataset=train_dataset, batch_size=32, shuffle=True)
test_loader = DataLoader(二、修改策略
为了解决这些问题，我们可以采取多种策略。首先，调整批量大小是一种简单有效的方法。通过减少每次迭代中处理的数据量，可以减轻显存压力，但可能会影响模型训练的效率和最终性能。其次，梯度累积技术可以在不改变批量大小时，通过累积多个小批量的梯度来模拟大批量的效果，这样可以在保持较大有效批量的同时减少单次迭代的显存需求。
模型优化也是一个重要的方向。通过剪枝、量化或知识蒸馏等技术减少模型大小，或者重新设计模型结构以减少参数数量，可以有效地减少显存的使用。混合精度训练，即在训练过程中使用半精度浮点数（FP16）代替全精度浮点数（FP32），不仅可以减少显存使用，还可以加速训练过程。在内存管理方面，我们可以在PyTorch中通过删除不再需要的变量并调用torch.cuda.empty_cache()来释放未使用的显存。此外，通过设置环境变量来调整CUDA的内存分配策略，减少内存碎片化，也是一个有效的手段。
最后，分布式训练可以将模型和数据分布到多个GPU上，从而分散显存压力。这种方法虽然需要更多的硬件资源，但在处理特别大的模型或数据集时，可以显著提高训练的可行性和效率。
1. 减小批量大小（Batch Size）
减小批量大小是解决CUDA内存不足的简单有效方法。以下是代码示例：
import torch
import torchvision.models as models
# 尝试较大的批量大小
model = models.resnet50().cuda() # 将模型放到GPU上
input = torch.randn(32, 3, 224, 224).cuda() # 大批量的输入数据
try:
output = model(input) # 尝试运行模型
except RuntimeError as e:
if 'out of memory' in str(e):
print("CUDA内存不足，尝试减少批量大小...")
torch.cuda.empty_cache() # 清理缓存
input = torch.randn(16, 3, 224, 224).cuda() # 减小批量大小后重试
output = model(input)2. 使用梯度累积（Gradient Accumulation）
通过累积梯度，可以在不增加显存压力的情况下训练更大批量的数据。以下是代码示例：
optimizer.zero_grad()
for i in range(gradient_accumulation_steps):
loss = model(input[i]).backward()
optimizer.step()3. 清理显存
在不需要变量时手动删除，并调用 torch.cuda.empty_cache() 来释放GPU内存。
import torch, gc
# 清理Python垃圾回收
gc.collect()
# 清理PyTorch缓存
torch.cuda.empty_cache()4. 使用混合精度训练（Mixed Precision Training）
通过使用混合精度训练，可以减少内存的使用。以下是代码示例：
from torch.cuda.amp import GradScaler, autocast
scaler = GradScaler()
for data, target in train_loader:
optimizer.zero_grad()
with autocast():
output = model(data)
loss = criterion(output, target)
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()5. 设置环境变量减少内存碎片化
通过设置 PYTORCH_CUDA_ALLOC_CONF 环境变量，调整内存分配策略，减少内存碎片化。
import os
os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'max_split_size_mb:128'这些代码示例提供了几种不同的方法来解决CUDA内存不足的问题，可以根据具体情况选择合适的策略。
三、示例修改
这个代码将包括一些改进，比如使用更小的批量大小来减少显存使用，以及使用梯度累积来模拟更大的批量大小。此外，我还会添加一些注释来帮助理解代码的每个部分。
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torchvision import datasets, transforms
from torch.utils.data import DataLoader
# 定义一个简单的多层感知器模型
class SimpleMLP(nn.Module):
def __init__(self):
super(SimpleMLP, self).__init__()
self.fc1 = nn.Linear(28 * 28, 512) # 输入层到隐藏层
self.fc2 = nn.Linear(512, 256) # 隐藏层到隐藏层
self.fc3 = nn.Linear(256, 10) # 隐藏层到输出层
def forward(self, x):
x = x.view(-1, 28 * 28) # 展平图像
x = F.relu(self.fc1(x)) # 第一个隐藏层
x = F.relu(self.fc2(x)) # 第二个隐藏层
x = self.fc3(x) # 输出层
return x
# 设置数据加载器
transform = transforms.Compose([
transforms.ToTensor(), # 将图片转换为Tensor
transforms.Normalize((0.5,), (0.5,)) # 归一化
])
train_dataset = datasets.MNIST(root='./data', train=True, transform=transform, download=True)
test_dataset = datasets.MNIST(root='./data', train=False, transform=transform)
# 使用更小的批量大小以减少显存使用
batch_size = 16
train_loader = DataLoader(dataset=train_dataset, batch_size=batch_size, shuffle=True)
test_loader = DataLoader(dataset=test_dataset, batch_size=batch_size, shuffle=False)
# 初始化模型、损失函数和优化器
model = SimpleMLP().cuda() # 将模型移到GPU上
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)
# 梯度累积步骤
gradient_accumulation_steps = 4
# 训练模型
def train_model(model, criterion, optimizer, train_loader, test_loader, epochs=5):
for epoch in range(epochs):
model.train() # 设置模型为训练模式
running_loss = 0.0
for i, (images, labels) in enumerate(train_loader):
images, labels = images.cuda(), labels.cuda() # 将数据移到GPU上
optimizer.zero_grad() # 清空梯度
# 梯度累积
for _ in range(gradient_accumulation_steps):
outputs = model(images) # 前向传播
loss = criterion(outputs, labels) # 计算损失
loss = loss / gradient_accumulation_steps # 将损失除以累积步数
loss.backward() # 反向传播
optimizer.step() # 更新权重
running_loss += loss.item() * gradient_accumulation_steps
print(f'Epoch {epoch+1}, Loss: {running_loss/len(train_loader)}')
# 测试模型
model.eval() # 设置模型为评估模式
correct = 0
total = 0
with torch.no_grad(): # 在测试阶段不计算梯度
for images, labels in test_loader:
images, labels = images.cuda(), labels.cuda() # 将数据移到GPU上
outputs = model(images)
_, predicted = torch.max(outputs.data, 1)
total += labels.size(0)
correct += (predicted == labels).sum().item()
print(f'Accuracy of the network on the test images: {100 * correct / total}%')
# 运行训练
train_model(model, criterion, optimizer, train_loader, test_loader, epochs=5)在这个修改后的代码中，做了以下几点改进：
- 减小批量大小：将批量大小从32减少到16，以减少每次迭代中GPU需要处理的数据量。
- 梯度累积：通过设置
gradient_accumulation_steps为4，我们可以在不增加显存压力的情况下，通过累积梯度来模拟更大的批量大小。 - 数据和模型移到GPU：在训练和测试阶段，将数据和模型移到GPU上，以利用GPU的计算能力。
这些改进可以帮助减少显存的使用，同时保持模型训练的效率和性能。