TensorFlow 2.1 GPU训练报Could not find valid device for node如何解决？

问题根因 #

错误核心是tf.one_hot算子的输入、输出数据类型组合不在当前GPU支持的算子内核列表中，无法匹配到可执行的GPU设备。你的标签数组y的类型和tf.one_hot默认输出的类型组合不满足GPU算子要求，导致调用失败。

直接修复方案 #

修改自定义数据生成器中__data_generation方法的对应代码，显式指定符合GPU支持的输入输出类型：

# 原代码
y[y==4] = 3;
mask = tf.one_hot(y, 4);
Y = tf.image.resize(mask, (IMG_SIZE, IMG_SIZE));

# 修改后代码
import numpy as np
y[y==4] = 3
# 将输入标签转为GPU支持的int32类型
y = y.astype(np.int32)
# 显式指定one_hot输出为GPU支持的float32类型
mask = tf.one_hot(y, depth=4, dtype=tf.float32)
Y = tf.image.resize(mask, (IMG_SIZE, IMG_SIZE))

修改后重新执行训练代码即可，该操作不会影响训练精度和速度。

额外排查步骤 #

如果修改后仍报错，按顺序执行以下排查：

• 先验证TensorFlow能否正常识别GPU，执行以下测试代码：

import tensorflow as tf
# 输出为True则GPU识别正常
print(tf.test.is_gpu_available())
# 输出你的GPU设备名则驱动、CUDA适配正常
print(tf.test.gpu_device_name())

如果输出为False，优先检查你的TensorFlow版本、CUDA版本、cuDNN版本三者是否匹配，版本不匹配会导致GPU算子加载失败。

• 若版本无问题，可将one_hot操作调度到CPU执行（该预处理操作开销极小，不会影响整体训练速度）：

y[y==4] = 3
y = y.astype(np.int32)
with tf.device('/CPU:0'):
    mask = tf.one_hot(y, depth=4, dtype=tf.float32)
Y = tf.image.resize(mask, (IMG_SIZE, IMG_SIZE))

内容的提问来源于stack exchange，提问作者Esteban Alejandro Orozco Snche