服务器训练LSTM模型时GPU内存占用过高及GPU切换问题求助

一、切换至空闲的GPU1运行模型 #

这是最直接的解决方式，GPU1当前完全空闲，操作简单：

• 命令行启动时指定GPU
  运行训练脚本时，通过环境变量限定TensorFlow仅可见GPU1：

  CUDA_VISIBLE_DEVICES=1 python your_lstm_training_script.py
  

• 在代码中提前指定GPU
  在Python脚本的最开头（导入TensorFlow之前）添加以下代码，强制TensorFlow使用GPU1：

  import os
  os.environ["CUDA_VISIBLE_DEVICES"] = "1"
  # 之后再导入TensorFlow和其他库
  import tensorflow as tf
  

• 手动指定设备上下文
  如果需要更精细控制代码运行设备，用tf.device上下文管理器包裹模型训练逻辑：

  import tensorflow as tf
  
  with tf.device('/device:GPU:1'):
      # 模型定义、数据加载、训练步骤都放在这里
      dataset = tf.keras.utils.timeseries_dataset_from_array(...)
      model = tf.keras.Sequential([...]) # 你的LSTM模型
      model.fit(dataset, epochs=...)
  

二、重置GPU0的内存占用 #

如果需要清空GPU0的残留占用，尝试以下方法（优先推荐前两种，避免内核重启问题）：

• 使用nvidia-smi重置GPU
  需root权限，执行后会重置指定GPU状态，清空所有内存占用：

  # 仅重置GPU0
  sudo nvidia-smi --gpu-reset -i 0
  # 重置所有GPU（可选）
  # sudo nvidia-smi --gpu-reset
  

• 用fuser彻底清理GPU进程
  比单独杀PID更彻底，可避免内核重启问题：

  # 列出所有占用GPU0的进程
  sudo fuser -v /dev/nvidia0
  # 直接杀死这些进程
  sudo fuser -k /dev/nvidia0
  

• 优化TensorFlow内存分配策略
  即使使用GPU0，也可以让TensorFlow动态分配内存，避免一次性占满GPU导致训练缓慢：

  import tensorflow as tf
  
  gpus = tf.config.list_physical_devices('GPU')
  if gpus:
      try:
          # 开启动态内存增长，TensorFlow会按需申请内存
          for gpu in gpus:
              tf.config.experimental.set_memory_growth(gpu, True)
      except RuntimeError as e:
          print(e)
  

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