在Kaggle中如何用TensorFlow借助GPU进行视频预处理与模型训练？

我有一批视频数据，用CPU提取帧耗时约5小时，打算在Kaggle平台借助GPU完成帧提取预处理和CNN模型训练，但TensorFlow无法调用GPU，PyTorch却能正常使用GPU。我尝试了以下代码检测GPU：

gpus = tf.config.experimental.list_physical_devices('GPU')
for gpu in gpus:
    print("Name:", gpu.name, "  Type:", gpu.device_type)

以及预处理设备选择代码：

def prepare_all_videos(df, root_dir):
    device = '/device:GPU:0' if len(tf.config.list_physical_devices('GPU')) > 0 else '/device:CPU:0'
    num_samples = len(df)
    ...

• 检查环境版本兼容性
  Kaggle默认的TensorFlow版本可能和CUDA版本不匹配，先运行以下命令确认版本：

  !pip list | grep tensorflow
  !nvidia-smi
  

  根据nvidia-smi显示的CUDA版本，安装对应兼容的TensorFlow版本（比如CUDA 11.8适配TensorFlow 2.12及以上）：

  !pip install tensorflow==2.15.0 --upgrade
  

• 正确配置GPU内存策略
  弃用旧的experimentalAPI，改用稳定的GPU配置方式，避免显存占用异常：

  import tensorflow as tf
  
  # 检测可用GPU
  physical_devices = tf.config.list_physical_devices('GPU')
  print(f"可用GPU数量: {len(physical_devices)}")
  
  # 启用动态显存增长（推荐）
  if physical_devices:
      try:
          for gpu in physical_devices:
              tf.config.experimental.set_memory_growth(gpu, True)
          print("GPU动态显存增长已启用")
      except RuntimeError as e:
          print(f"配置失败: {e}")
  

• 确保预处理代码实际跑在GPU上
  仅定义device变量没用，需要用tf.device()上下文管理器包裹实际的预处理逻辑，让TensorFlow把运算分配到GPU：

  def prepare_all_videos(df, root_dir):
      physical_devices = tf.config.list_physical_devices('GPU')
      target_device = '/GPU:0' if physical_devices else '/CPU:0'
      
      num_samples = len(df)
      
      # 将预处理逻辑放在GPU上下文内
      with tf.device(target_device):
          # 这里写你的视频帧提取、预处理代码
          # 例如使用tf.io读取视频、tf.image处理帧等操作
          pass
  

• 验证GPU是否正常工作
  加入简单的GPU运算测试，确认TensorFlow真的在使用GPU：

  # 查看当前默认GPU设备
  print("当前GPU设备:", tf.test.gpu_device_name())
  
  # 运行GPU运算测试
  with tf.device('/GPU:0'):
      test_tensor = tf.random.normal((1000, 1000))
      result = tf.matmul(test_tensor, test_tensor)
  print("GPU运算完成，结果形状:", result.shape)
  

内容的提问来源于stack exchange，提问作者Mr. Ghosh