CUDA：out of memory的解决方法（实测有效）

一、问题概述       

1.问题分析         

CUDA out of memory问题通常发生在深度学习训练过程中，当GPU的显存不足以容纳模型、输入数据以及中间计算结果时就会触发。这个问题可能由几个因素引起：

1. 模型和数据规模：深度学习模型尤其是大型模型，如Transformer或大型CNN，拥有大量的参数，这些参数在训练时需要被加载到GPU显存中。同时，如果批量大小（batch size）设置得过大，一次性处理的数据量也会增加，进一步加大显存的负担。

2. 内存管理：深度学习框架在处理数据和模型时，可能会因为不当的内存管理而导致显存没有被及时释放，或者由于频繁的内存分配和释放造成内存碎片化，这会减少可用的连续内存块，即使总显存使用量没有达到100%，也可能出现内存不足的情况。

3. 数据加载和预处理：如果数据加载和预处理步骤没有优化，比如加载了大量未被立即使用的数据，或者数据没有被适当地压缩或降维，也可能导致显存使用量激增。

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(),  # 将图片转换为Tensortransforms.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, autocastscaler = 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(),  # 将图片转换为Tensortransforms.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.0for 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_stepsprint(f'Epoch {epoch+1}, Loss: {running_loss/len(train_loader)}')# 测试模型model.eval()  # 设置模型为评估模式correct = 0total = 0with 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)

在这个修改后的代码中，做了以下几点改进：

1. 减小批量大小：将批量大小从32减少到16，以减少每次迭代中GPU需要处理的数据量。
2. 梯度累积：通过设置gradient_accumulation_steps为4，我们可以在不增加显存压力的情况下，通过累积梯度来模拟更大的批量大小。
3. 数据和模型移到GPU：在训练和测试阶段，将数据和模型移到GPU上，以利用GPU的计算能力。

这些改进可以帮助减少显存的使用，同时保持模型训练的效率和性能。