Google Cloud训练TensorFlow模型时自定义包含TensorFlow致GPU失效求助

问题分析 #

你遇到的核心问题是：当自定义utils包的setup.py中包含TensorFlow依赖时，AI Platform预配置的GPU优化TensorFlow环境被覆盖，导致训练全程使用CPU；移除该依赖后，平台自带的GPU版本TensorFlow正常工作。这是因为AI Platform会优先安装自定义包指定的TensorFlow版本，而这个版本大概率是通用CPU版本，或者与平台GPU环境的CUDA、cuDNN配置不兼容，进而导致TensorFlow无法检测到GPU设备。

解决方案 #

以下几种方法可以确保无论自定义包是否包含TensorFlow依赖，训练都能正常使用GPU：

1. 将TensorFlow设为自定义包的可选依赖

修改utils包的setup.py，用extras_require把TensorFlow标记为开发/本地环境的可选依赖，这样AI Platform训练时只会安装基础的utils包，不会引入额外的TensorFlow，避免覆盖平台预配置的GPU版本：

from setuptools import setup, find_packages

setup(
    name="utils",
    version="0.0.1",
    packages=find_packages(),
    # 把TensorFlow设为可选依赖，本地开发时按需安装
    extras_require={
        "tf": ["tensorflow>=1.10"]
    }
)

• 本地开发需要TensorFlow时，执行pip install utils-0.0.1.tar.gz[tf]安装带依赖的版本；
• AI Platform训练时，直接使用原提交命令即可，平台会使用自带的GPU版TensorFlow。

2. 严格匹配AI Platform的TensorFlow GPU版本

如果你的utils包必须依赖TensorFlow，可以在setup.py中指定与平台runtime版本完全一致的GPU专属TensorFlow版本（TensorFlow 1.x中tensorflow-gpu是独立包）：

from setuptools import setup, find_packages

setup(
    name="utils",
    version="0.0.1",
    packages=find_packages(),
    install_requires=[
        # 与你指定的runtime-version 1.10完全匹配的GPU版本
        "tensorflow-gpu==1.10"
    ]
)

这样安装的TensorFlow版本会和AI Platform预配置的CUDA、cuDNN环境兼容，不会破坏GPU检测逻辑。

3. 添加GPU检测调试代码（辅助验证）

在训练代码开头加入以下代码，通过日志确认GPU是否被识别，方便排查问题：

# TensorFlow 1.x版本的检测代码
from tensorflow.python.client import device_lib

print("=== 本地设备列表 ===")
for device in device_lib.list_local_devices():
    print(f"设备类型: {device.device_type}, 名称: {device.name}")

如果日志中出现GPU类型的设备，说明GPU已正常识别。

4. 考虑升级runtime版本（可选）

你当前使用的runtime-version 1.10是较旧的版本，TensorFlow 1.x的GPU生态已经停止维护。如果代码兼容的话，建议升级到2.x系列的runtime版本（比如runtime-version 2.15），新版本的TensorFlow默认包含GPU支持，且环境配置更稳定，能减少这类依赖冲突问题。

总结 #

核心思路是避免自定义包的TensorFlow依赖干扰AI Platform预配置的GPU优化环境：要么把TensorFlow设为可选依赖，要么严格匹配平台的GPU版TensorFlow版本，就能保证训练正常使用指定的P100 GPU。

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