活动公告

系统通知
通知：端午节限时活动来袭！
06-18 23:43
系统通知
通知：微软邮箱更换提醒
06-14 00:00
系统通知
通知：本站资源由网友上传分享，如有违规等问题请到版务模块进行投诉，资源失效请在帖子内回复要求补档，会尽快处理！
10-23 09:31
>社区 > 科技 > 经验交流 >帖子
本版发帖返回

Python释放显存实用技术在深度学习项目开发中如何有效管理GPU内存资源防止显存泄漏和程序崩溃的完整指南包括手动释放方法自动内存管理策略和常见问题解决方案

这是个蓝孩子 当前离线 威震华夏关云长
SunJu_FaceMall

3万

主题

3077

科技点

3万

积分

执行版主

碾压王

积分
32876

塔罗立华奏
执行版主 发表于 2025-9-29 15:30:00 | 显示全部楼层 |阅读模式

马上注册，结交更多好友，享用更多功能，让你轻松玩转社区。

您需要 登录 才可以下载或查看，没有账号？立即注册

x
引言

在深度学习项目开发中，GPU内存（显存）管理是一个至关重要的环节。随着模型规模的不断扩大和数据集的日益增长，GPU显存往往成为限制模型训练和推理的瓶颈。不当的内存管理不仅会导致程序运行效率低下，还可能引发显存泄漏，最终导致程序崩溃。本文将全面介绍Python中释放显存的实用技术，帮助开发者有效管理GPU内存资源，防止显存泄漏和程序崩溃。

GPU内存基础

GPU显存的工作原理

GPU显存（Video RAM，简称VRAM）是专门用于存储GPU计算所需数据的内存。与系统内存（RAM）不同，GPU显存具有更高的带宽和更低的延迟，专门为并行计算优化。在深度学习中，模型参数、中间激活值、梯度以及批处理数据都需要存储在GPU显存中。

Python与GPU内存的交互

Python通过深度学习框架（如PyTorch、TensorFlow）与GPU进行交互。这些框架提供了高级API，使开发者能够轻松地将计算任务和数据分配到GPU上。然而，这种抽象也意味着开发者对底层内存管理的控制有限，容易导致内存问题。
  2. # 检查GPU可用性
  3. import torch
  4. if torch.cuda.is_available():
  5.     print(f"可用GPU数量: {torch.cuda.device_count()}")
  6.     print(f"当前GPU: {torch.cuda.current_device()}")
  7.     print(f"GPU名称: {torch.cuda.get_device_name(torch.cuda.current_device())}")
  8.     print(f"GPU显存总量: {torch.cuda.get_device_properties(torch.cuda.current_device()).total_memory / 1024**3:.2f} GB")
复制代码

深度学习框架中的内存管理

PyTorch内存管理机制

PyTorch使用缓存分配器（Caching Allocator）来管理GPU内存。这种机制通过维护一个内存池来减少内存分配和释放的开销。当PyTorch需要分配内存时，它会首先检查内存池中是否有足够大的空闲块；如果没有，它会向操作系统申请新的内存。当释放内存时，PyTorch不会立即将内存返回给操作系统，而是将其保留在内存池中以备将来使用。
  2. import torch
  4. # 查看PyTorch内存分配统计
  5. def print_memory_stats():
  6.     allocated = torch.cuda.memory_allocated() / 1024**3  # GB
  7.     cached = torch.cuda.memory_reserved() / 1024**3  # GB
  8.     print(f"已分配显存: {allocated:.2f} GB")
  9.     print(f"缓存显存: {cached:.2f} GB")
  11. # 创建张量并查看内存使用
  12. a = torch.randn(10000, 10000, device='cuda')
  13. print_memory_stats()
  15. # 删除张量
  16. del a
  17. torch.cuda.empty_cache()  # 清空缓存
  18. print_memory_stats()
复制代码

TensorFlow内存管理机制

TensorFlow也使用类似的内存管理策略。在TensorFlow 2.x中，默认启用即时执行（Eager Execution），内存管理更加动态。TensorFlow会根据需要自动分配和释放内存，但也提供了手动控制内存分配的选项。
  2. import tensorflow as tf
  4. # 配置GPU内存增长选项，而不是一次性分配所有内存
  5. gpus = tf.config.experimental.list_physical_devices('GPU')
  6. if gpus:
  7.     try:
  8.         for gpu in gpus:
  9.             tf.config.experimental.set_memory_growth(gpu, True)
  10.     except RuntimeError as e:
  11.         print(e)
  13. # 查看内存使用情况
  14. def print_tf_memory_stats():
  15.     for gpu in gpus:
  16.         memory_info = tf.config.experimental.get_memory_info(gpu.name)
  17.         print(f"GPU {gpu.name}:")
  18.         print(f"  当前使用: {memory_info['current'] / 1024**3:.2f} GB")
  19.         print(f"  峰值使用: {memory_info['peak'] / 1024**3:.2f} GB")
  21. # 创建张量并查看内存使用
  22. a = tf.random.normal((10000, 10000))
  23. print_tf_memory_stats()
  25. # 删除张量
  26. del a
  27. print_tf_memory_stats()
复制代码

手动释放显存的方法

使用垃圾回收机制

Python的垃圾回收机制可以自动释放不再使用的对象，但在GPU内存管理中，仅依靠垃圾回收可能不够及时。我们可以手动触发垃圾回收来释放内存。
  2. import gc
  3. import torch
  5. # 创建一些大的张量
  6. a = torch.randn(10000, 10000, device='cuda')
  7. b = torch.randn(10000, 10000, device='cuda')
  8. print(f"分配前: {torch.cuda.memory_allocated() / 1024**3:.2f} GB")
  10. # 删除引用
  11. del a
  12. del b
  14. # 手动触发垃圾回收
  15. gc.collect()
  17. # 清空PyTorch的缓存分配器
  18. torch.cuda.empty_cache()
  20. print(f"释放后: {torch.cuda.memory_allocated() / 1024**3:.2f} GB")
复制代码

使用上下文管理器

上下文管理器可以确保资源在使用后被正确释放，即使在发生异常的情况下也是如此。我们可以创建自定义的上下文管理器来管理GPU内存。
  2. import torch
  3. from contextlib import contextmanager
  5. @contextmanager
  6. def gpu_memory_context(device='cuda'):
  7.     try:
  8.         # 记录初始内存状态
  9.         initial_memory = torch.cuda.memory_allocated(device=device)
  10.         yield
  11.     finally:
  12.         # 清理并释放内存
  13.         gc.collect()
  14.         torch.cuda.empty_cache()
  15.         final_memory = torch.cuda.memory_allocated(device=device)
  16.         print(f"内存变化: {(initial_memory - final_memory) / 1024**3:.2f} GB")
  18. # 使用上下文管理器
  19. with gpu_memory_context():
  20.     a = torch.randn(5000, 5000, device='cuda')
  21.     b = torch.randn(5000, 5000, device='cuda')
  22.     c = torch.matmul(a, b)
  23.     # 在这里进行计算
  24.     # 退出上下文时，内存会被自动释放
复制代码

显式释放模型和张量

在深度学习中，模型和张量是占用显存的主要对象。显式地删除不再需要的模型和张量，并释放其占用的显存，可以有效管理内存。
  2. import torch
  3. import torch.nn as nn
  5. # 定义一个简单的模型
  6. class SimpleModel(nn.Module):
  7.     def __init__(self):
  8.         super(SimpleModel, self).__init__()
  9.         self.linear = nn.Linear(1000, 1000)
  10.    
  11.     def forward(self, x):
  12.         return self.linear(x)
  14. # 创建模型并移至GPU
  15. model = SimpleModel().cuda()
  16. input = torch.randn(100, 1000).cuda()
  18. # 查看内存使用
  19. print(f"模型创建后: {torch.cuda.memory_allocated() / 1024**3:.2f} GB")
  21. # 使用模型进行前向传播
  22. output = model(input)
  23. loss = output.sum()
  24. loss.backward()
  26. # 查看内存使用
  27. print(f"反向传播后: {torch.cuda.memory_allocated() / 1024**3:.2f} GB")
  29. # 释放模型和梯度
  30. del model, input, output, loss
  31. torch.cuda.empty_cache()
  33. # 查看内存使用
  34. print(f"释放后: {torch.cuda.memory_allocated() / 1024**3:.2f} GB")
复制代码

使用内存映射文件处理大数据

对于非常大的数据集，可以考虑使用内存映射文件（Memory-mapped files）技术，这样可以避免一次性将所有数据加载到GPU显存中。
  2. import torch
  3. import numpy as np
  5. # 创建一个大的numpy数组并保存到磁盘
  6. large_array = np.random.rand(10000, 10000)
  7. np.save('large_array.npy', large_array)
  9. # 使用内存映射方式加载数组
  10. memmap_array = np.load('large_array.npy', mmap_mode='r')
  12. # 创建PyTorch张量，但不立即复制到GPU
  13. tensor = torch.from_numpy(memmap_array)
  15. # 分批处理数据
  16. batch_size = 1000
  17. for i in range(0, len(tensor), batch_size):
  18.     batch = tensor[i:i+batch_size].cuda()  # 只将当前批次加载到GPU
  19.     # 在这里处理批次数据
  20.     result = torch.mean(batch)
  21.     print(f"批次 {i//batch_size}: {result.item()}")
  22.    
  23.     # 释放当前批次的GPU内存
  24.     del batch
  25.     torch.cuda.empty_cache()
复制代码

自动内存管理策略

PyTorch的自动内存管理

PyTorch提供了几种自动内存管理的策略，包括内存池、梯度检查点和自动混合精度等。

PyTorch的缓存分配器会自动管理内存池，但我们可以通过调整环境变量来优化其行为。
  2. import os
  3. import torch
  5. # 设置PyTorch内存分配器的环境变量
  6. # PYTORCH_CUDA_ALLOC_CONF可以控制内存分配器的行为
  7. # 例如，可以设置最大内存池大小
  8. os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'max_split_size_mb:128'
  10. # 现在PyTorch会使用更小的内存块，减少内存碎片
  11. a = torch.randn(5000, 5000, device='cuda')
  12. b = torch.randn(5000, 5000, device='cuda')
  13. c = torch.matmul(a, b)
  15. # 查看内存使用情况
  16. print(f"内存使用: {torch.cuda.memory_allocated() / 1024**3:.2f} GB")
复制代码

梯度检查点是一种用计算换内存的技术，它在前向传播时不保存所有中间激活值，而是在反向传播时重新计算它们。
  2. import torch
  3. import torch.nn as nn
  4. from torch.utils.checkpoint import checkpoint
  6. # 定义一个深度模型
  7. class DeepModel(nn.Module):
  8.     def __init__(self):
  9.         super(DeepModel, self).__init__()
  10.         self.layers = nn.ModuleList([nn.Linear(1000, 1000) for _ in range(10)])
  11.    
  12.     def forward(self, x):
  13.         for layer in self.layers:
  14.             # 使用checkpoint包装每一层
  15.             x = checkpoint(layer, x)
  16.         return x
  18. # 创建模型
  19. model = DeepModel().cuda()
  20. input = torch.randn(100, 1000).cuda()
  22. # 查看内存使用
  23. print(f"使用梯度检查点前: {torch.cuda.memory_allocated() / 1024**3:.2f} GB")
  25. # 前向传播
  26. output = model(input)
  27. loss = output.sum()
  28. loss.backward()
  30. # 查看内存使用
  31. print(f"使用梯度检查点后: {torch.cuda.memory_allocated() / 1024**3:.2f} GB")
复制代码

自动混合精度可以减少内存使用，同时保持模型性能。它使用半精度浮点数（FP16）进行计算，同时在必要时使用全精度浮点数（FP32）保持数值稳定性。
  2. import torch
  3. import torch.nn as nn
  4. from torch.cuda.amp import autocast, GradScaler
  6. # 定义模型
  7. model = nn.Sequential(
  8.     nn.Linear(1000, 1000),
  9.     nn.ReLU(),
  10.     nn.Linear(1000, 1000),
  11.     nn.ReLU(),
  12.     nn.Linear(1000, 10)
  13. ).cuda()
  15. # 创建GradScaler用于梯度缩放
  16. scaler = GradScaler()
  18. # 生成输入数据
  19. input = torch.randn(100, 1000).cuda()
  20. target = torch.randint(0, 10, (100,)).cuda()
  22. # 查看内存使用
  23. print(f"AMP前: {torch.cuda.memory_allocated() / 1024**3:.2f} GB")
  25. # 使用自动混合精度进行训练
  26. optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
  28. with autocast():
  29.     output = model(input)
  30.     loss = nn.CrossEntropyLoss()(output, target)
  32. # 缩放损失并反向传播
  33. scaler.scale(loss).backward()
  35. # 缩放梯度并更新参数
  36. scaler.step(optimizer)
  37. scaler.update()
  39. # 查看内存使用
  40. print(f"AMP后: {torch.cuda.memory_allocated() / 1024**3:.2f} GB")
复制代码

TensorFlow的自动内存管理

TensorFlow也提供了多种自动内存管理策略，包括内存增长选项、虚拟设备和内存优化器等。

默认情况下，TensorFlow会一次性分配GPU上几乎所有可用的内存。我们可以启用内存增长选项，使TensorFlow只在需要时分配内存。
  2. import tensorflow as tf
  4. # 获取所有GPU设备
  5. gpus = tf.config.experimental.list_physical_devices('GPU')
  7. if gpus:
  8.     try:
  9.         # 启用内存增长
  10.         for gpu in gpus:
  11.             tf.config.experimental.set_memory_growth(gpu, True)
  12.         print("内存增长已启用")
  13.     except RuntimeError as e:
  14.         print(f"内存增长设置失败: {e}")
  16. # 查看内存使用情况
  17. for gpu in gpus:
  18.     memory_info = tf.config.experimental.get_memory_info(gpu.name)
  19.     print(f"GPU {gpu.name}:")
  20.     print(f"  当前使用: {memory_info['current'] / 1024**3:.2f} GB")
  21.     print(f"  峰值使用: {memory_info['peak'] / 1024**3:.2f} GB")
复制代码

我们可以配置虚拟设备，限制TensorFlow可以使用的GPU内存量。
  2. import tensorflow as tf
  4. # 获取所有GPU设备
  5. gpus = tf.config.experimental.list_physical_devices('GPU')
  7. if gpus:
  8.     try:
  9.         # 限制第一个GPU只能使用4GB内存
  10.         tf.config.experimental.set_virtual_device_configuration(
  11.             gpus[0],
  12.             [tf.config.experimental.VirtualDeviceConfiguration(memory_limit=4096)]
  13.         )
  14.         print("虚拟设备配置已设置")
  15.     except RuntimeError as e:
  16.         print(f"虚拟设备配置失败: {e}")
  18. # 查看内存使用情况
  19. for gpu in gpus:
  20.     memory_info = tf.config.experimental.get_memory_info(gpu.name)
  21.     print(f"GPU {gpu.name}:")
  22.     print(f"  当前使用: {memory_info['current'] / 1024**3:.2f} GB")
  23.     print(f"  峰值使用: {memory_info['peak'] / 1024**3:.2f} GB")
复制代码

TensorFlow提供了多种内存优化技术，包括模型优化、梯度累积和混合精度训练等。
  2. import tensorflow as tf
  4. # 启用混合精度训练
  5. policy = tf.keras.mixed_precision.Policy('mixed_float16')
  6. tf.keras.mixed_precision.set_global_policy(policy)
  8. # 定义模型
  9. model = tf.keras.Sequential([
  10.     tf.keras.layers.Dense(1000, activation='relu'),
  11.     tf.keras.layers.Dense(1000, activation='relu'),
  12.     tf.keras.layers.Dense(10)
  13. ])
  15. # 编译模型
  16. model.compile(
  17.     optimizer='adam',
  18.     loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
  19.     metrics=['accuracy']
  20. )
  22. # 生成数据
  23. x_train = tf.random.normal((1000, 100))
  24. y_train = tf.random.uniform((1000,), maxval=10, dtype=tf.int32)
  26. # 查看内存使用
  27. gpus = tf.config.experimental.list_physical_devices('GPU')
  28. if gpus:
  29.     memory_info = tf.config.experimental.get_memory_info(gpus[0].name)
  30.     print(f"训练前内存使用: {memory_info['current'] / 1024**3:.2f} GB")
  32. # 训练模型
  33. model.fit(x_train, y_train, epochs=2)
  35. # 查看内存使用
  36. if gpus:
  37.     memory_info = tf.config.experimental.get_memory_info(gpus[0].name)
  38.     print(f"训练后内存使用: {memory_info['current'] / 1024**3:.2f} GB")
复制代码

常见问题及解决方案

显存泄漏的识别与解决

显存泄漏是指程序在运行过程中持续占用越来越多的显存，而不释放不再使用的内存。这最终会导致程序因显存不足而崩溃。
  2. import torch
  3. import time
  4. import matplotlib.pyplot as plt
  6. def check_memory_leak():
  7.     memory_usage = []
  8.     timestamps = []
  9.    
  10.     # 记录初始内存
  11.     initial_memory = torch.cuda.memory_allocated()
  12.     print(f"初始内存使用: {initial_memory / 1024**3:.2f} GB")
  13.    
  14.     # 模拟可能泄漏的操作
  15.     tensors = []
  16.     for i in range(10):
  17.         # 创建一些张量并添加到列表中
  18.         a = torch.randn(1000, 1000, device='cuda')
  19.         b = torch.randn(1000, 1000, device='cuda')
  20.         c = torch.matmul(a, b)
  21.         tensors.append(c)
  22.         
  23.         # 记录内存使用
  24.         current_memory = torch.cuda.memory_allocated()
  25.         memory_usage.append(current_memory / 1024**3)
  26.         timestamps.append(i)
  27.         
  28.         print(f"步骤 {i}: {current_memory / 1024**3:.2f} GB")
  29.         time.sleep(0.5)
  30.    
  31.     # 绘制内存使用图
  32.     plt.figure(figsize=(10, 6))
  33.     plt.plot(timestamps, memory_usage, 'b-')
  34.     plt.xlabel('步骤')
  35.     plt.ylabel('内存使用 (GB)')
  36.     plt.title('内存使用趋势')
  37.     plt.grid(True)
  38.     plt.show()
  39.    
  40.     # 释放张量
  41.     del tensors
  42.     torch.cuda.empty_cache()
  43.    
  44.     final_memory = torch.cuda.memory_allocated()
  45.     print(f"最终内存使用: {final_memory / 1024**3:.2f} GB")
  47. # 运行检查
  48. check_memory_leak()
复制代码
  2. import torch
  3. import gc
  4. import weakref
  6. # 解决方案1: 确保删除不再需要的对象
  7. def memory_leak_solution_1():
  8.     tensors = []
  9.     for i in range(10):
  10.         a = torch.randn(1000, 1000, device='cuda')
  11.         b = torch.randn(1000, 1000, device='cuda')
  12.         c = torch.matmul(a, b)
  13.         tensors.append(c)
  14.         
  15.         # 定期清理不再需要的张量
  16.         if i % 3 == 0:
  17.             # 删除前几个张量
  18.             for j in range(min(3, len(tensors))):
  19.                 del tensors[j]
  20.             tensors = tensors[3:]
  21.             
  22.             # 手动触发垃圾回收和清空缓存
  23.             gc.collect()
  24.             torch.cuda.empty_cache()
  25.             
  26.             print(f"步骤 {i}: 清理后内存使用 {torch.cuda.memory_allocated() / 1024**3:.2f} GB")
  28. # 解决方案2: 使用弱引用避免循环引用
  29. def memory_leak_solution_2():
  30.     class TensorHolder:
  31.         def __init__(self, tensor):
  32.             self.tensor = tensor
  33.    
  34.     holders = []
  35.     for i in range(10):
  36.         a = torch.randn(1000, 1000, device='cuda')
  37.         b = torch.randn(1000, 1000, device='cuda')
  38.         c = torch.matmul(a, b)
  39.         
  40.         # 使用弱引用
  41.         holder = TensorHolder(c)
  42.         weak_holder = weakref.ref(holder)
  43.         holders.append(weak_holder)
  44.         
  45.         print(f"步骤 {i}: 内存使用 {torch.cuda.memory_allocated() / 1024**3:.2f} GB")
  46.    
  47.     # 清理
  48.     del holders
  49.     gc.collect()
  50.     torch.cuda.empty_cache()
  51.     print(f"清理后内存使用 {torch.cuda.memory_allocated() / 1024**3:.2f} GB")
  53. # 运行解决方案
  54. print("解决方案1:")
  55. memory_leak_solution_1()
  57. print("\n解决方案2:")
  58. memory_leak_solution_2()
复制代码

CUDA out of memory错误处理

“CUDA out of memory”是深度学习开发中常见的错误，通常发生在尝试分配的内存超过可用显存时。
  2. import torch
  3. import torch.nn as nn
  5. class DynamicBatchTrainer:
  6.     def __init__(self, model, initial_batch_size=32, min_batch_size=1):
  7.         self.model = model
  8.         self.current_batch_size = initial_batch_size
  9.         self.min_batch_size = min_batch_size
  10.         
  11.     def train_step(self, data, target, optimizer, criterion):
  12.         # 尝试使用当前批处理大小
  13.         while self.current_batch_size >= self.min_batch_size:
  14.             try:
  15.                 # 分割数据以适应当前批处理大小
  16.                 for i in range(0, len(data), self.current_batch_size):
  17.                     batch_data = data[i:i+self.current_batch_size]
  18.                     batch_target = target[i:i+self.current_batch_size]
  19.                     
  20.                     # 前向传播
  21.                     output = self.model(batch_data)
  22.                     loss = criterion(output, batch_target)
  23.                     
  24.                     # 反向传播
  25.                     optimizer.zero_grad()
  26.                     loss.backward()
  27.                     optimizer.step()
  28.                
  29.                 # 如果成功，增加批处理大小以获得更好的性能
  30.                 self.current_batch_size = min(self.current_batch_size * 2, len(data))
  31.                 return loss.item()
  32.                
  33.             except RuntimeError as e:
  34.                 if 'out of memory' in str(e):
  35.                     # 如果内存不足，减少批处理大小并重试
  36.                     torch.cuda.empty_cache()
  37.                     self.current_batch_size = max(self.current_batch_size // 2, self.min_batch_size)
  38.                     print(f"内存不足，减少批处理大小至 {self.current_batch_size}")
  39.                 else:
  40.                     # 如果是其他错误，重新抛出
  41.                     raise e
  42.         
  43.         # 如果达到最小批处理大小仍然失败，抛出错误
  44.         raise RuntimeError(f"即使使用最小批处理大小 {self.min_batch_size} 也无法训练")
  46. # 使用示例
  47. model = nn.Sequential(
  48.     nn.Linear(1000, 1000),
  49.     nn.ReLU(),
  50.     nn.Linear(1000, 10)
  51. ).cuda()
  53. optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
  54. criterion = nn.CrossEntropyLoss()
  56. # 生成大数据集
  57. data = torch.randn(1000, 1000).cuda()
  58. target = torch.randint(0, 10, (1000,)).cuda()
  60. # 创建动态批处理训练器
  61. trainer = DynamicBatchTrainer(model, initial_batch_size=256, min_batch_size=1)
  63. # 训练
  64. for epoch in range(3):
  65.     loss = trainer.train_step(data, target, optimizer, criterion)
  66.     print(f"Epoch {epoch}, Loss: {loss:.4f}, Batch Size: {trainer.current_batch_size}")
复制代码

梯度累积是一种技术，允许我们使用较小的批处理大小模拟较大的批处理大小，从而减少内存使用。
  2. import torch
  3. import torch.nn as nn
  5. class GradientAccumulator:
  6.     def __init__(self, model, accumulation_steps=4):
  7.         self.model = model
  8.         self.accumulation_steps = accumulation_steps
  9.         self.current_step = 0
  10.         
  11.     def train_step(self, data, target, optimizer, criterion):
  12.         # 前向传播
  13.         output = self.model(data)
  14.         loss = criterion(output, target)
  15.         
  16.         # 归一化损失以考虑累积
  17.         loss = loss / self.accumulation_steps
  18.         
  19.         # 反向传播
  20.         loss.backward()
  21.         
  22.         # 增加步骤计数
  23.         self.current_step += 1
  24.         
  25.         # 如果达到累积步骤，更新参数
  26.         if self.current_step % self.accumulation_steps == 0:
  27.             optimizer.step()
  28.             optimizer.zero_grad()
  29.             return True, loss.item() * self.accumulation_steps
  30.         
  31.         return False, loss.item() * self.accumulation_steps
  33. # 使用示例
  34. model = nn.Sequential(
  35.     nn.Linear(1000, 1000),
  36.     nn.ReLU(),
  37.     nn.Linear(1000, 10)
  38. ).cuda()
  40. optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
  41. criterion = nn.CrossEntropyLoss()
  43. # 生成大数据集
  44. data = torch.randn(1000, 1000).cuda()
  45. target = torch.randint(0, 10, (1000,)).cuda()
  47. # 创建梯度累积器
  48. accumulator = GradientAccumulator(model, accumulation_steps=4)
  50. # 训练
  51. batch_size = 32
  52. for epoch in range(3):
  53.     total_loss = 0
  54.     update_count = 0
  55.    
  56.     for i in range(0, len(data), batch_size):
  57.         batch_data = data[i:i+batch_size]
  58.         batch_target = target[i:i+batch_size]
  59.         
  60.         updated, loss = accumulator.train_step(batch_data, batch_target, optimizer, criterion)
  61.         total_loss += loss
  62.         
  63.         if updated:
  64.             update_count += 1
  65.    
  66.     avg_loss = total_loss / (len(data) // batch_size)
  67.     print(f"Epoch {epoch}, Loss: {avg_loss:.4f}, Updates: {update_count}")
复制代码

多GPU环境下的内存管理

在多GPU环境中，内存管理变得更加复杂，需要考虑数据并行、模型并行和负载均衡等问题。
  2. import torch
  3. import torch.nn as nn
  4. from torch.nn.parallel import DistributedDataParallel as DDP
  5. import torch.distributed as dist
  6. import torch.multiprocessing as mp
  7. import os
  9. def setup(rank, world_size):
  10.     os.environ['MASTER_ADDR'] = 'localhost'
  11.     os.environ['MASTER_PORT'] = '12355'
  12.    
  13.     # 初始化进程组
  14.     dist.init_process_group(
  15.         backend='nccl',
  16.         rank=rank,
  17.         world_size=world_size
  18.     )
  20. def cleanup():
  21.     dist.destroy_process_group()
  23. class ToyModel(nn.Module):
  24.     def __init__(self):
  25.         super(ToyModel, self).__init__()
  26.         self.net1 = nn.Linear(1000, 1000)
  27.         self.relu = nn.ReLU()
  28.         self.net2 = nn.Linear(1000, 10)
  29.         
  30.     def forward(self, x):
  31.         return self.net2(self.relu(self.net1(x)))
  33. def train(rank, world_size):
  34.     print(f"在GPU {rank} 上运行 DDP 示例")
  35.     setup(rank, world_size)
  36.    
  37.     # 创建模型并移至当前GPU
  38.     model = ToyModel().to(rank)
  39.    
  40.     # 使用DDP包装模型
  41.     ddp_model = DDP(model, device_ids=[rank])
  42.    
  43.     # 创建优化器
  44.     optimizer = torch.optim.SGD(ddp_model.parameters(), lr=0.001)
  45.    
  46.     # 生成数据
  47.     inputs = torch.randn(2000, 1000).to(rank)
  48.     labels = torch.randn(2000, 10).to(rank)
  49.    
  50.     # 训练循环
  51.     for epoch in range(2):
  52.         # 在每个GPU上处理数据的一部分
  53.         inputs_shard = inputs[rank::world_size]
  54.         labels_shard = labels[rank::world_size]
  55.         
  56.         # 前向传播
  57.         outputs = ddp_model(inputs_shard)
  58.         loss = torch.nn.functional.mse_loss(outputs, labels_shard)
  59.         
  60.         # 反向传播
  61.         optimizer.zero_grad()
  62.         loss.backward()
  63.         optimizer.step()
  64.         
  65.         # 打印内存使用情况
  66.         allocated = torch.cuda.memory_allocated(rank) / 1024**3
  67.         print(f"GPU {rank}, Epoch {epoch}, Loss: {loss.item()}, Memory: {allocated:.2f} GB")
  68.    
  69.     # 清理
  70.     cleanup()
  72. def main():
  73.     world_size = torch.cuda.device_count()
  74.     print(f"发现 {world_size} 个GPU")
  75.     if world_size > 1:
  76.         mp.spawn(train, args=(world_size,), nprocs=world_size, join=True)
  77.     else:
  78.         print("需要多个GPU才能运行此示例")
  80. if __name__ == '__main__':
  81.     main()
复制代码
  2. import torch
  3. import torch.nn as nn
  4. from torch.nn.parallel import DistributedDataParallel as DDP
  5. import torch.distributed as dist
  6. import torch.multiprocessing as mp
  7. import os
  9. def setup(rank, world_size):
  10.     os.environ['MASTER_ADDR'] = 'localhost'
  11.     os.environ['MASTER_PORT'] = '12355'
  12.    
  13.     # 初始化进程组
  14.     dist.init_process_group(
  15.         backend='nccl',
  16.         rank=rank,
  17.         world_size=world_size
  18.     )
  20. def cleanup():
  21.     dist.destroy_process_group()
  23. class ParallelModel(nn.Module):
  24.     def __init__(self, rank, world_size):
  25.         super(ParallelModel, self).__init__()
  26.         self.rank = rank
  27.         self.world_size = world_size
  28.         
  29.         # 将模型的不同层分配到不同的GPU
  30.         if rank == 0:
  31.             # 第一个GPU处理前半部分
  32.             self.part1 = nn.Sequential(
  33.                 nn.Linear(1000, 2000),
  34.                 nn.ReLU(),
  35.                 nn.Linear(2000, 2000),
  36.                 nn.ReLU()
  37.             )
  38.         else:
  39.             # 第二个GPU处理后半部分
  40.             self.part2 = nn.Sequential(
  41.                 nn.Linear(2000, 2000),
  42.                 nn.ReLU(),
  43.                 nn.Linear(2000, 10)
  44.             )
  45.    
  46.     def forward(self, x):
  47.         # 第一个GPU处理前向传播的第一部分
  48.         if self.rank == 0:
  49.             x = self.part1(x)
  50.             # 发送到下一个GPU
  51.             dist.send(x, dst=1)
  52.             # 接收最终结果
  53.             dist.recv(x, src=1)
  54.             return x
  55.         else:
  56.             # 第二个GPU接收中间结果
  57.             dist.recv(x, src=0)
  58.             # 处理后半部分
  59.             x = self.part2(x)
  60.             # 发送回第一个GPU
  61.             dist.send(x, dst=0)
  62.             return x
  64. def train(rank, world_size):
  65.     print(f"在GPU {rank} 上运行模型并行示例")
  66.     setup(rank, world_size)
  67.    
  68.     # 创建模型
  69.     model = ParallelModel(rank, world_size).to(rank)
  70.    
  71.     # 只在主进程创建优化器
  72.     if rank == 0:
  73.         optimizer = torch.optim.SGD(model.parameters(), lr=0.001)
  74.    
  75.     # 生成数据
  76.     if rank == 0:
  77.         inputs = torch.randn(64, 1000).to(rank)
  78.         labels = torch.randn(64, 10).to(rank)
  79.    
  80.     # 训练循环
  81.     for epoch in range(2):
  82.         if rank == 0:
  83.             # 前向传播
  84.             outputs = model(inputs)
  85.             loss = torch.nn.functional.mse_loss(outputs, labels)
  86.             
  87.             # 反向传播
  88.             optimizer.zero_grad()
  89.             loss.backward()
  90.             optimizer.step()
  91.             
  92.             # 打印内存使用情况
  93.             allocated = torch.cuda.memory_allocated(rank) / 1024**3
  94.             print(f"GPU {rank}, Epoch {epoch}, Loss: {loss.item()}, Memory: {allocated:.2f} GB")
  95.         else:
  96.             # 其他GPU只参与前向和反向传播
  97.             outputs = model(inputs)
  98.             # 反向传播会自动处理
  99.             outputs.sum().backward()
  100.             
  101.             # 打印内存使用情况
  102.             allocated = torch.cuda.memory_allocated(rank) / 1024**3
  103.             print(f"GPU {rank}, Epoch {epoch}, Memory: {allocated:.2f} GB")
  104.    
  105.     # 清理
  106.     cleanup()
  108. def main():
  109.     world_size = 2  # 使用2个GPU
  110.     if torch.cuda.device_count() >= world_size:
  111.         mp.spawn(train, args=(world_size,), nprocs=world_size, join=True)
  112.     else:
  113.         print(f"需要至少 {world_size} 个GPU才能运行此示例")
  115. if __name__ == '__main__':
  116.     main()
复制代码

最佳实践

监控GPU内存使用

有效管理GPU内存的第一步是了解内存的使用情况。以下是一些监控GPU内存使用的方法：
  2. import torch
  3. import time
  4. import psutil
  5. import GPUtil
  6. from datetime import datetime
  8. class GPUMemoryMonitor:
  9.     def __init__(self, interval=1.0):
  10.         self.interval = interval
  11.         self.running = False
  12.         self.log_data = []
  13.         
  14.     def start(self):
  15.         self.running = True
  16.         self.log_data = []
  17.         print("开始监控GPU内存使用...")
  18.         
  19.         while self.running:
  20.             timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
  21.             
  22.             # 获取GPU信息
  23.             gpus = GPUtil.getGPUs()
  24.             for i, gpu in enumerate(gpus):
  25.                 memory_used = gpu.memoryUsed
  26.                 memory_total = gpu.memoryTotal
  27.                 memory_percent = gpu.memoryUtil * 100
  28.                 gpu_load = gpu.load * 100
  29.                
  30.                 log_entry = {
  31.                     'timestamp': timestamp,
  32.                     'gpu_id': i,
  33.                     'memory_used': memory_used,
  34.                     'memory_total': memory_total,
  35.                     'memory_percent': memory_percent,
  36.                     'gpu_load': gpu_load
  37.                 }
  38.                
  39.                 self.log_data.append(log_entry)
  40.                
  41.                 print(f"[{timestamp}] GPU {i}: 内存 {memory_used}/{memory_total} MB ({memory_percent:.1f}%), 负载 {gpu_load:.1f}%")
  42.             
  43.             time.sleep(self.interval)
  44.    
  45.     def stop(self):
  46.         self.running = False
  47.         print("停止监控GPU内存使用")
  48.    
  49.     def get_memory_usage_plot(self):
  50.         import matplotlib.pyplot as plt
  51.         import pandas as pd
  52.         
  53.         if not self.log_data:
  54.             print("没有可用的监控数据")
  55.             return
  56.         
  57.         # 转换为DataFrame
  58.         df = pd.DataFrame(self.log_data)
  59.         
  60.         # 绘制内存使用图
  61.         plt.figure(figsize=(12, 6))
  62.         
  63.         for gpu_id in df['gpu_id'].unique():
  64.             gpu_data = df[df['gpu_id'] == gpu_id]
  65.             plt.plot(gpu_data['timestamp'], gpu_data['memory_percent'], label=f'GPU {gpu_id}')
  66.         
  67.         plt.xlabel('时间')
  68.         plt.ylabel('内存使用率 (%)')
  69.         plt.title('GPU内存使用率随时间变化')
  70.         plt.legend()
  71.         plt.grid(True)
  72.         plt.xticks(rotation=45)
  73.         plt.tight_layout()
  74.         plt.show()
  76. # 使用示例
  77. monitor = GPUMemoryMonitor(interval=0.5)
  79. # 在单独的线程中启动监控
  80. import threading
  81. monitor_thread = threading.Thread(target=monitor.start)
  82. monitor_thread.start()
  84. # 模拟一些GPU操作
  85. try:
  86.     a = torch.randn(5000, 5000, device='cuda')
  87.     time.sleep(2)
  88.     b = torch.randn(5000, 5000, device='cuda')
  89.     time.sleep(2)
  90.     c = torch.matmul(a, b)
  91.     time.sleep(2)
  92.     del a, b, c
  93.     torch.cuda.empty_cache()
  94.     time.sleep(2)
  95. finally:
  96.     # 停止监控
  97.     monitor.stop()
  98.     monitor_thread.join()
  99.    
  100.     # 显示内存使用图
  101.     monitor.get_memory_usage_plot()
复制代码

内存优化策略

以下是一些有效的内存优化策略：
  2. import torch
  3. import torch.nn as nn
  4. from torch.utils.checkpoint import checkpoint
  5. from torch.cuda.amp import autocast, GradScaler
  7. class MemoryOptimizedTrainer:
  8.     def __init__(self, model, device='cuda'):
  9.         self.model = model.to(device)
  10.         self.device = device
  11.         self.scaler = GradScaler()
  12.         
  13.     def train_with_gradient_checkpointing(self, data_loader, optimizer, criterion, epochs=1):
  14.         """使用梯度检查点进行训练，减少内存使用"""
  15.         self.model.train()
  16.         
  17.         for epoch in range(epochs):
  18.             total_loss = 0
  19.             for batch_idx, (data, target) in enumerate(data_loader):
  20.                 data, target = data.to(self.device), target.to(self.device)
  21.                
  22.                 def custom_forward(*inputs):
  23.                     """自定义前向函数，用于梯度检查点"""
  24.                     return self.model(inputs[0])
  25.                
  26.                 # 使用梯度检查点
  27.                 output = checkpoint(custom_forward, data)
  28.                 loss = criterion(output, target)
  29.                
  30.                 # 反向传播
  31.                 optimizer.zero_grad()
  32.                 loss.backward()
  33.                 optimizer.step()
  34.                
  35.                 total_loss += loss.item()
  36.                
  37.                 # 打印内存使用情况
  38.                 if batch_idx % 10 == 0:
  39.                     memory_used = torch.cuda.memory_allocated(self.device) / 1024**3
  40.                     print(f'Epoch: {epoch}, Batch: {batch_idx}, Loss: {loss.item():.4f}, Memory: {memory_used:.2f} GB')
  41.             
  42.             avg_loss = total_loss / len(data_loader)
  43.             print(f'Epoch {epoch}, Average Loss: {avg_loss:.4f}')
  44.    
  45.     def train_with_mixed_precision(self, data_loader, optimizer, criterion, epochs=1):
  46.         """使用混合精度进行训练，减少内存使用并提高速度"""
  47.         self.model.train()
  48.         
  49.         for epoch in range(epochs):
  50.             total_loss = 0
  51.             for batch_idx, (data, target) in enumerate(data_loader):
  52.                 data, target = data.to(self.device), target.to(self.device)
  53.                
  54.                 # 使用自动混合精度
  55.                 with autocast():
  56.                     output = self.model(data)
  57.                     loss = criterion(output, target)
  58.                
  59.                 # 缩放损失并反向传播
  60.                 optimizer.zero_grad()
  61.                 self.scaler.scale(loss).backward()
  62.                
  63.                 # 缩放梯度并更新参数
  64.                 self.scaler.step(optimizer)
  65.                 self.scaler.update()
  66.                
  67.                 total_loss += loss.item()
  68.                
  69.                 # 打印内存使用情况
  70.                 if batch_idx % 10 == 0:
  71.                     memory_used = torch.cuda.memory_allocated(self.device) / 1024**3
  72.                     print(f'Epoch: {epoch}, Batch: {batch_idx}, Loss: {loss.item():.4f}, Memory: {memory_used:.2f} GB')
  73.             
  74.             avg_loss = total_loss / len(data_loader)
  75.             print(f'Epoch {epoch}, Average Loss: {avg_loss:.4f}')
  76.    
  77.     def train_with_gradient_accumulation(self, data_loader, optimizer, criterion, accumulation_steps=4, epochs=1):
  78.         """使用梯度累积进行训练，模拟更大的批处理大小"""
  79.         self.model.train()
  80.         
  81.         for epoch in range(epochs):
  82.             total_loss = 0
  83.             optimizer.zero_grad()
  84.             
  85.             for batch_idx, (data, target) in enumerate(data_loader):
  86.                 data, target = data.to(self.device), target.to(self.device)
  87.                
  88.                 # 前向传播
  89.                 with autocast():
  90.                     output = self.model(data)
  91.                     loss = criterion(output, target)
  92.                
  93.                 # 归一化损失
  94.                 loss = loss / accumulation_steps
  95.                
  96.                 # 反向传播
  97.                 self.scaler.scale(loss).backward()
  98.                
  99.                 total_loss += loss.item() * accumulation_steps
  100.                
  101.                 # 如果达到累积步骤，更新参数
  102.                 if (batch_idx + 1) % accumulation_steps == 0:
  103.                     # 缩放梯度并更新参数
  104.                     self.scaler.step(optimizer)
  105.                     self.scaler.update()
  106.                     optimizer.zero_grad()
  107.                
  108.                 # 打印内存使用情况
  109.                 if batch_idx % 10 == 0:
  110.                     memory_used = torch.cuda.memory_allocated(self.device) / 1024**3
  111.                     print(f'Epoch: {epoch}, Batch: {batch_idx}, Loss: {loss.item() * accumulation_steps:.4f}, Memory: {memory_used:.2f} GB')
  112.             
  113.             # 处理剩余的批次
  114.             if len(data_loader) % accumulation_steps != 0:
  115.                 self.scaler.step(optimizer)
  116.                 self.scaler.update()
  117.                 optimizer.zero_grad()
  118.             
  119.             avg_loss = total_loss / len(data_loader)
  120.             print(f'Epoch {epoch}, Average Loss: {avg_loss:.4f}')
  122. # 使用示例
  123. # 创建一个简单的模型
  124. model = nn.Sequential(
  125.     nn.Linear(1000, 2000),
  126.     nn.ReLU(),
  127.     nn.Linear(2000, 2000),
  128.     nn.ReLU(),
  129.     nn.Linear(2000, 10)
  130. )
  132. # 创建优化器和损失函数
  133. optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
  134. criterion = nn.CrossEntropyLoss()
  136. # 创建数据加载器
  137. from torch.utils.data import DataLoader, TensorDataset
  138. data = torch.randn(1000, 1000)
  139. target = torch.randint(0, 10, (1000,))
  140. dataset = TensorDataset(data, target)
  141. data_loader = DataLoader(dataset, batch_size=32, shuffle=True)
  143. # 创建内存优化训练器
  144. trainer = MemoryOptimizedTrainer(model)
  146. # 使用不同的优化策略进行训练
  147. print("使用梯度检查点训练:")
  148. trainer.train_with_gradient_checkpointing(data_loader, optimizer, criterion, epochs=1)
  150. print("\n使用混合精度训练:")
  151. trainer.train_with_mixed_precision(data_loader, optimizer, criterion, epochs=1)
  153. print("\n使用梯度累积训练:")
  154. trainer.train_with_gradient_accumulation(data_loader, optimizer, criterion, accumulation_steps=4, epochs=1)
复制代码

内存泄漏预防

预防内存泄漏比解决内存泄漏更容易。以下是一些预防内存泄漏的最佳实践：
  2. import torch
  3. import gc
  4. import weakref
  5. from contextlib import contextmanager
  7. class MemoryLeakPrevention:
  8.     @staticmethod
  9.     @contextmanager
  10.     def memory_context(device='cuda'):
  11.         """上下文管理器，确保在退出时清理内存"""
  12.         try:
  13.             # 记录初始内存状态
  14.             initial_memory = torch.cuda.memory_allocated(device=device)
  15.             yield
  16.         finally:
  17.             # 清理内存
  18.             gc.collect()
  19.             torch.cuda.empty_cache(device=device)
  20.             final_memory = torch.cuda.memory_allocated(device=device)
  21.             print(f"内存变化: {(initial_memory - final_memory) / 1024**3:.2f} GB")
  22.    
  23.     @staticmethod
  24.     def safe_model_loading(model_path, device='cuda'):
  25.         """安全加载模型，避免内存泄漏"""
  26.         with MemoryLeakPrevention.memory_context(device):
  27.             # 首先在CPU上加载模型
  28.             model = torch.load(model_path, map_location='cpu')
  29.             
  30.             # 然后移动到目标设备
  31.             model = model.to(device)
  32.             
  33.             return model
  34.    
  35.     @staticmethod
  36.     def safe_tensor_creation(size, device='cuda', dtype=torch.float32):
  37.         """安全创建张量，避免内存泄漏"""
  38.         with MemoryLeakPrevention.memory_context(device):
  39.             return torch.randn(size, device=device, dtype=dtype)
  40.    
  41.     @staticmethod
  42.     def safe_training_loop(model, data_loader, optimizer, criterion, epochs=1, device='cuda'):
  43.         """安全的训练循环，避免内存泄漏"""
  44.         model = model.to(device)
  45.         
  46.         for epoch in range(epochs):
  47.             model.train()
  48.             total_loss = 0
  49.             
  50.             for batch_idx, (data, target) in enumerate(data_loader):
  51.                 data, target = data.to(device), target.to(device)
  52.                
  53.                 # 前向传播
  54.                 output = model(data)
  55.                 loss = criterion(output, target)
  56.                
  57.                 # 反向传播
  58.                 optimizer.zero_grad()
  59.                 loss.backward()
  60.                 optimizer.step()
  61.                
  62.                 total_loss += loss.item()
  63.                
  64.                 # 定期清理内存
  65.                 if batch_idx % 100 == 0:
  66.                     with MemoryLeakPrevention.memory_context(device):
  67.                         pass  # 上下文管理器会自动清理内存
  68.             
  69.             avg_loss = total_loss / len(data_loader)
  70.             print(f'Epoch {epoch}, Average Loss: {avg_loss:.4f}')
  71.    
  72.     @staticmethod
  73.     def safe_inference(model, data_loader, device='cuda'):
  74.         """安全的推理过程，避免内存泄漏"""
  75.         model = model.to(device)
  76.         model.eval()
  77.         
  78.         results = []
  79.         
  80.         with torch.no_grad():
  81.             for batch_idx, (data, _) in enumerate(data_loader):
  82.                 data = data.to(device)
  83.                
  84.                 # 前向传播
  85.                 output = model(data)
  86.                
  87.                 # 保存结果
  88.                 results.append(output.cpu())
  89.                
  90.                 # 定期清理内存
  91.                 if batch_idx % 100 == 0:
  92.                     with MemoryLeakPrevention.memory_context(device):
  93.                         pass  # 上下文管理器会自动清理内存
  94.         
  95.         return torch.cat(results)
  97. # 使用示例
  98. # 创建一个简单的模型
  99. model = torch.nn.Sequential(
  100.     torch.nn.Linear(1000, 100),
  101.     torch.nn.ReLU(),
  102.     torch.nn.Linear(100, 10)
  103. )
  105. # 创建数据加载器
  106. from torch.utils.data import DataLoader, TensorDataset
  107. data = torch.randn(1000, 1000)
  108. target = torch.randint(0, 10, (1000,))
  109. dataset = TensorDataset(data, target)
  110. data_loader = DataLoader(dataset, batch_size=32, shuffle=True)
  112. # 创建优化器和损失函数
  113. optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
  114. criterion = torch.nn.CrossEntropyLoss()
  116. # 使用安全的训练循环
  117. print("使用安全的训练循环:")
  118. MemoryLeakPrevention.safe_training_loop(model, data_loader, optimizer, criterion, epochs=1)
  120. # 使用安全的推理过程
  121. print("\n使用安全的推理过程:")
  122. results = MemoryLeakPrevention.safe_inference(model, data_loader)
  123. print(f"推理结果形状: {results.shape}")
  125. # 安全创建张量
  126. print("\n安全创建张量:")
  127. tensor = MemoryLeakPrevention.safe_tensor_creation((1000, 1000))
  128. print(f"张量形状: {tensor.shape}")
复制代码

结论

在深度学习项目开发中，有效管理GPU内存资源是确保程序稳定运行的关键。本文全面介绍了Python中释放显存的实用技术，包括手动释放方法、自动内存管理策略和常见问题解决方案。

通过掌握这些技术，开发者可以：

1. 更有效地管理GPU内存资源，避免内存不足的问题
2. 防止显存泄漏，提高程序的稳定性
3. 在有限的硬件资源上训练更大的模型
4. 优化内存使用，提高程序运行效率

在实际应用中，建议开发者结合多种技术，根据具体项目需求选择最适合的内存管理策略。同时，定期监控内存使用情况，及时发现和解决内存问题，也是确保程序稳定运行的重要措施。

随着深度学习技术的不断发展，GPU内存管理技术也在不断进步。开发者应保持学习，关注最新的内存管理技术和最佳实践，以应对日益复杂的深度学习项目需求。
「七転び八起き（ななころびやおき）」
回复

使用道具 举报

返回列表 发新帖
您需要登录后才可以回帖 登录 | 立即注册
高级模式
B Color Image Link Quote Code Smilies

本版积分规则