CUDA：out of memory的办理方法（实测有用）

一、题目概述       

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）。

2. import torch
3. import torch.nn as nn
4. import torch.nn.functional as F
5. import torch.optim as optim
6. from torchvision import datasets, transforms
7. from torch.utils.data import DataLoader
9. # 定义一个简单的多层感知器模型
10. class SimpleMLP(nn.Module):
11.     def __init__(self):
12.         super(SimpleMLP, self).__init__()
13.         self.fc1 = nn.Linear(28 * 28, 512)  # 输入层到隐藏层
14.         self.fc2 = nn.Linear(512, 256)      # 隐藏层到隐藏层
15.         self.fc3 = nn.Linear(256, 10)       # 隐藏层到输出层
17.     def forward(self, x):
18.         x = x.view(-1, 28 * 28)  # 展平图像
19.         x = F.relu(self.fc1(x))  # 第一个隐藏层
20.         x = F.relu(self.fc2(x))  # 第二个隐藏层
21.         x = self.fc3(x)          # 输出层
22.         return x
24. # 设置数据加载器
25. transform = transforms.Compose([
26.     transforms.ToTensor(),  # 将图片转换为Tensor
27.     transforms.Normalize((0.5,), (0.5,))  # 归一化
28. ])
30. train_dataset = datasets.MNIST(root='./data', train=True, transform=transform, download=True)
31. test_dataset = datasets.MNIST(root='./data', train=False, transform=transform)
33. train_loader = DataLoader(dataset=train_dataset, batch_size=32, shuffle=True)
34. test_loader = DataLoader(

复制代码

二、修改计谋        

为了办理这些题目，我们可以接纳多种计谋。起首，调解批量巨细是一种简朴有用的方法。通过镌汰每次迭代中处置惩罚的数据量，可以减轻显存压力，但大概会影响模子训练的效率和终极性能。其次，梯度累积技术可以在不改变批量巨细时，通过累积多个小批量的梯度来模仿大批量的效果，如许可以在保持较大有用批量的同时镌汰单次迭代的显存需求。
        模子优化也是一个紧张的方向。通过剪枝、量化或知识蒸馏等技术镌汰模子巨细，大概重新筹划模子布局以镌汰参数数量，可以有用地镌汰显存的利用。肴杂精度训练，即在训练过程中利用半精度浮点数（FP16）代替全精度浮点数（FP32），不但可以镌汰显存利用，还可以加速训练过程。在内存管理方面，我们可以在PyTorch中通过删除不再必要的变量并调用torch.cuda.empty_cache()来开释未利用的显存。别的，通过设置情况变量来调解CUDA的内存分配计谋，镌汰内存碎片化，也是一个有用的本事。
        末了，分布式训练可以将模子和数据分布到多个GPU上，从而分散显存压力。这种方法固然必要更多的硬件资源，但在处置惩罚特殊大的模子或数据集时，可以显着进步训练的可行性和效率。
1. 减小批量巨细（Batch Size）

减小批量巨细是办理CUDA内存不足的简朴有用方法。以下是代码示例：

2. import torch
3. import torchvision.models as models
5. # 尝试较大的批量大小
6. model = models.resnet50().cuda()  # 将模型放到GPU上
7. input = torch.randn(32, 3, 224, 224).cuda()  # 大批量的输入数据
8. try:
9.     output = model(input)  # 尝试运行模型
10. except RuntimeError as e:
11.     if 'out of memory' in str(e):
12.         print("CUDA内存不足，尝试减少批量大小...")
13.         torch.cuda.empty_cache()  # 清理缓存
14.         input = torch.randn(16, 3, 224, 224).cuda()  # 减小批量大小后重试
15.         output = model(input)

复制代码

2. 利用梯度累积（Gradient Accumulation）

通过累积梯度，可以在不增长显存压力的情况下训练更大批量的数据。以下是代码示例：

2. optimizer.zero_grad()
3. for i in range(gradient_accumulation_steps):
4.     loss = model(input[i]).backward()
5. optimizer.step()

复制代码

3. 清理显存

在不必要变量时手动删除，并调用 torch.cuda.empty_cache() 来开释GPU内存。

2. import torch, gc
4. # 清理Python垃圾回收
5. gc.collect()
6. # 清理PyTorch缓存
7. torch.cuda.empty_cache()

复制代码

4. 利用肴杂精度训练（Mixed Precision Training）

通过利用肴杂精度训练，可以镌汰内存的利用。以下是代码示例：

2. from torch.cuda.amp import GradScaler, autocast
4. scaler = GradScaler()
6. for data, target in train_loader:
7.     optimizer.zero_grad()
8.     with autocast():
9.         output = model(data)
10.         loss = criterion(output, target)
11.     scaler.scale(loss).backward()
12.     scaler.step(optimizer)
13.     scaler.update()

复制代码

5. 设置情况变量镌汰内存碎片化

通过设置 PYTORCH_CUDA_ALLOC_CONF 情况变量，调解内存分配计谋，镌汰内存碎片化。

2. import os
3. os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'max_split_size_mb:128'

复制代码

这些代码示例提供了几种差异的方法来办理CUDA内存不足的题目，可以根据具体情况选择符合的计谋。
三、示例修改

        这个代码将包罗一些改进，比如利用更小的批量巨细来镌汰显存利用，以及利用梯度累积来模仿更大的批量巨细。别的，我还会添加一些注释来资助明确代码的每个部门。

2. import torch
3. import torch.nn as nn
4. import torch.nn.functional as F
5. import torch.optim as optim
6. from torchvision import datasets, transforms
7. from torch.utils.data import DataLoader
9. # 定义一个简单的多层感知器模型
10. class SimpleMLP(nn.Module):
11.     def __init__(self):
12.         super(SimpleMLP, self).__init__()
13.         self.fc1 = nn.Linear(28 * 28, 512)  # 输入层到隐藏层
14.         self.fc2 = nn.Linear(512, 256)      # 隐藏层到隐藏层
15.         self.fc3 = nn.Linear(256, 10)       # 隐藏层到输出层
17.     def forward(self, x):
18.         x = x.view(-1, 28 * 28)  # 展平图像
19.         x = F.relu(self.fc1(x))  # 第一个隐藏层
20.         x = F.relu(self.fc2(x))  # 第二个隐藏层
21.         x = self.fc3(x)          # 输出层
22.         return x
24. # 设置数据加载器
25. transform = transforms.Compose([
26.     transforms.ToTensor(),  # 将图片转换为Tensor
27.     transforms.Normalize((0.5,), (0.5,))  # 归一化
28. ])
30. train_dataset = datasets.MNIST(root='./data', train=True, transform=transform, download=True)
31. test_dataset = datasets.MNIST(root='./data', train=False, transform=transform)
33. # 使用更小的批量大小以减少显存使用
34. batch_size = 16
35. train_loader = DataLoader(dataset=train_dataset, batch_size=batch_size, shuffle=True)
36. test_loader = DataLoader(dataset=test_dataset, batch_size=batch_size, shuffle=False)
38. # 初始化模型、损失函数和优化器
39. model = SimpleMLP().cuda()  # 将模型移到GPU上
40. criterion = nn.CrossEntropyLoss()
41. optimizer = optim.Adam(model.parameters(), lr=0.001)
43. # 梯度累积步骤
44. gradient_accumulation_steps = 4
46. # 训练模型
47. def train_model(model, criterion, optimizer, train_loader, test_loader, epochs=5):
48.     for epoch in range(epochs):
49.         model.train()  # 设置模型为训练模式
50.         running_loss = 0.0
51.         for i, (images, labels) in enumerate(train_loader):
52.             images, labels = images.cuda(), labels.cuda()  # 将数据移到GPU上
53.             optimizer.zero_grad()  # 清空梯度
55.             # 梯度累积
56.             for _ in range(gradient_accumulation_steps):
57.                 outputs = model(images)  # 前向传播
58.                 loss = criterion(outputs, labels)  # 计算损失
59.                 loss = loss / gradient_accumulation_steps  # 将损失除以累积步数
60.                 loss.backward()  # 反向传播
62.             optimizer.step()  # 更新权重
63.             running_loss += loss.item() * gradient_accumulation_steps
65.         print(f'Epoch {epoch+1}, Loss: {running_loss/len(train_loader)}')
67.         # 测试模型
68.         model.eval()  # 设置模型为评估模式
69.         correct = 0
70.         total = 0
71.         with torch.no_grad():  # 在测试阶段不计算梯度
72.             for images, labels in test_loader:
73.                 images, labels = images.cuda(), labels.cuda()  # 将数据移到GPU上
74.                 outputs = model(images)
75.                 _, predicted = torch.max(outputs.data, 1)
76.                 total += labels.size(0)
77.                 correct += (predicted == labels).sum().item()
79.         print(f'Accuracy of the network on the test images: {100 * correct / total}%')
81. # 运行训练
82. train_model(model, criterion, optimizer, train_loader, test_loader, epochs=5)

复制代码

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

• 减小批量巨细：将批量巨细从32镌汰到16，以镌汰每次迭代中GPU必要处置惩罚的数据量。
  
• 梯度累积：通过设置gradient_accumulation_steps为4，我们可以在不增长显存压力的情况下，通过累积梯度来模仿更大的批量巨细。
  
• 数据和模子移到GPU：在训练和测试阶段，将数据和模子移到GPU上，以利用GPU的盘算能力。
  

这些改进可以资助镌汰显存的利用，同时保持模子训练的效率和性能。