这种跨环境的显存差异其实很常见，我之前帮朋友排查过类似问题，核心原因主要来自这几个方面：

• 框架版本默认配置变化：PyTorch 1.8.1对比1.7、TensorFlow 2.6对比2.4，新版本会默认启用一些性能优化特性，但这些特性往往会占用更多显存。比如PyTorch 1.8对CUDA后端的内存管理做了调整，默认预分配的显存比例更高；TensorFlow 2.6默认开启了更多内存缓存策略，XLA编译的默认行为也有变化。
• Kaggle环境的隐性优化：Kaggle的Notebook是经过优化的容器环境，框架会自动启用显存友好的配置——比如PyTorch默认关闭torch.backends.cudnn.benchmark（避免预加载大量卷积算法占用显存），TensorFlow默认开启内存增长模式，而且环境完全隔离，没有其他后台进程抢占显存。
• 本地GPU的显存预分配逻辑：RTX3090有24GB大显存，PyTorch和TensorFlow会默认尝试预占更多显存（接近满容量），而Kaggle的16GB显存会触发框架的自动限制，预分配量被严格控制在安全阈值内。
• 本地后台进程占用：你的本地环境可能有其他占用显存的程序（比如NVIDIA控制面板后台服务、未关闭的AI训练进程、甚至视频渲染工具），这些都会额外占用显存，导致整体占用居高不下。

要把本地环境的显存效率对齐Kaggle，你可以按以下步骤操作：

1. 对齐框架版本（最直接的办法） #

把本地的PyTorch降级到1.7，TensorFlow降级到2.4，完全匹配Kaggle的环境版本，直接消除版本差异带来的默认配置变化：

# 降级PyTorch到1.7（适配CUDA 11.0为例，可根据你的CUDA版本调整）
pip install torch==1.7.0 torchvision==0.8.0 torchaudio==0.7.0

# 降级TensorFlow到2.4
pip install tensorflow==2.4.0

2. 调整PyTorch显存配置 #

如果不想降级版本，手动修改PyTorch的内存管理设置：

• 限制显存预分配比例：在代码开头添加

  import torch
  # 限制当前进程最多使用80%的GPU显存（可根据需要调整）
  torch.cuda.set_per_process_memory_fraction(0.8, device=0)
  # 关闭cudnn基准测试，避免预加载大量卷积算法占用显存
  torch.backends.cudnn.benchmark = False
  

• 启用梯度检查点：针对大模型，用梯度检查点减少显存占用

  from torch.utils.checkpoint import checkpoint
  # 将模型前向传播用checkpoint包裹
  output = checkpoint(model, input)
  

• 训练间隙清理显存：在每个epoch结束后调用

  torch.cuda.empty_cache()
  

3. 调整TensorFlow显存配置 #

手动修改TensorFlow的显存策略，模拟Kaggle的环境：

• 开启内存增长模式：让TensorFlow根据实际需要逐步分配显存，而非一次性预占全部

  import tensorflow as tf
  gpus = tf.config.list_physical_devices('GPU')
  if gpus:
      try:
          for gpu in gpus:
              tf.config.experimental.set_memory_growth(gpu, True)
      except RuntimeError as e:
          print(e)
  

• 强制限制显存容量：直接模拟Kaggle的16GB显存限制

  gpus = tf.config.list_physical_devices('GPU')
  if gpus:
      tf.config.experimental.set_virtual_device_configuration(
          gpus[0],
          [tf.config.experimental.VirtualDeviceConfiguration(memory_limit=16384)]  # 16GB
      )
  

4. 清理本地后台进程 #

用nvidia-smi命令查看当前GPU的显存占用情况，找出占用显存的后台进程并关闭：

nvidia-smi

找到PID列对应的进程，用kill -9 <PID>命令关闭不必要的进程。

5. 关闭不必要的框架特性 #

比如TensorFlow的XLA编译如果默认开启，可能会增加显存占用，可以临时关闭：

tf.config.optimizer.set_jit(False)

内容的提问来源于stack exchange，提问作者Mobassir Hossen