如何在仅CPU实例上避免TensorFlow GPU镜像的cuDNN等注册错误

以下是几种可行的解决方案，可根据你的场景选择：

• 在推理启动脚本中强制TensorFlow使用CPU
  在SageMaker推理服务的启动脚本（如serve）开头添加代码，彻底禁用TensorFlow对GPU的检测：

  import os
  os.environ['CUDA_VISIBLE_DEVICES'] = '-1'
  import tensorflow as tf
  tf.config.set_visible_devices([], 'GPU')
  

  这段代码会让TensorFlow完全忽略GPU设备，避免初始化CUDA组件时触发错误。

• 在SageMaker端点配置中动态注入环境变量
  部署端点时通过env参数注入强制CPU模式的环境变量，无需修改镜像：

  from sagemaker.model import Model
  
  model = Model(
      image_uri="你的GPU镜像URI",
      role="你的SageMaker角色ARN",
      env={
          "CUDA_VISIBLE_DEVICES": "-1",
          "TF_CPP_MIN_LOG_LEVEL": "3",
          "TF_FORCE_GPU_ALLOW_GROWTH": "false"
      }
  )
  model.deploy(instance_type="ml.c5.xlarge", initial_instance_count=1)
  

  训练时使用GPU实例启动训练任务，无需设置这些变量，同一个镜像即可兼容两种场景。

• 在启动脚本中添加环境自动检测逻辑
  让脚本自动识别当前实例是否有GPU，无GPU时自动切换到CPU模式：

  import os
  import tensorflow as tf
  
  # 检查是否存在物理GPU设备
  gpus = tf.config.list_physical_devices('GPU')
  if not gpus:
      os.environ['CUDA_VISIBLE_DEVICES'] = '-1'
      tf.config.set_visible_devices([], 'GPU')
  

  这种方式无需手动配置环境变量，镜像会自适应运行环境。

• 在Dockerfile中添加CPU模式的默认配置（可选）
  如果需要镜像默认以CPU模式启动（训练时再手动覆盖），可以在Dockerfile中添加：

  ENV CUDA_VISIBLE_DEVICES="-1"
  ENV TF_CPP_MIN_LOG_LEVEL="3"
  

  训练GPU任务时，通过环境变量覆盖CUDA_VISIBLE_DEVICES为默认值即可启用GPU支持。

内容的提问来源于stack exchange，提问作者Lukas Baliūnas