TensorFlow图执行中将tf.dense张量转为tf.one_hot张量的问题

使用TensorFlow 2.11版本，基于TFRecords构建tf.data pipeline训练双输入分类器时，遇到无法将仅包含标量的稠密Tensor转换为tf.one_hot向量的问题。

相关代码 #

数据集加载代码 #

# 获取所有训练记录的绝对路径
training_names= [record_path+'/'+rec for rec in os.listdir(record_path) if rec.startswith('train')]
# 加载TF数据集
train_dataset = tf.data.TFRecordDataset(training_names[1])
train_dataset = train_dataset.map(return_xy)

映射函数return_xy #

def return_xy(example_proto):
    # 解析样本
    sample= parse_function(example_proto)
    # 解码图像1
    encoded_image1 = sample['image/encoded_1']
    decoded_image1 = decode_image(encoded_image1)
    # 解码图像2
    encoded_image2 = sample['image/encoded_2']
    decoded_image2 = decode_image(encoded_image2)
    # 解码标签
    print(f'image/object/class/{level}: {sample["image/object/class/"+level]}')
    class_label = tf.sparse.to_dense(sample["image/object/class/"+level])
    print(f'type of class label :{type(class_label)}')
    print(class_label)
    # 尝试转换为深度26的onehot -> 如何提取值或直接转换？
    label_onehot=tf.one_hot(class_label,26)
    # 调整图像大小
    input_left=tf.image.resize(decoded_image1,[416, 416])
    input_right=tf.image.resize(decoded_image2,[416, 416])
    return {'input_3res1':input_left, 'input_5res2':input_right} ,  label_onehot

当前输出 #

运行上述代码后，标签相关输出如下：

image/object/class/'+level: SparseTensor(indices=Tensor("ParseSingleExample/ParseExample/ParseExampleV2:14", shape=(None, 1), dtype=int64), values=Tensor("ParseSingleExample/ParseExample/ParseExampleV2:31", shape=(None,), dtype=int64), dense_shape=Tensor("ParseSingleExample/ParseExample/ParseExampleV2:48", shape=(1,), dtype=int64))
type of class label :<class 'tensorflow.python.framework.ops.Tensor'>
Tensor("SparseToDense:0", shape=(None,), dtype=int64)

急切执行下的验证 #

在急切执行模式下，能正确提取到标签值：

raw_dataset = tf.data.TFRecordDataset([rec_file])
parsed_dataset = raw_dataset.map(parse_function) # 仅解析
for sample in parsed_dataset:
    class_label=tf.sparse.to_dense(sample['image/object/class/label_level3'])[0]
    print(f'type of class label :{type(class_label)}')
    print(f'labels  from labelmap :{class_label}')

输出：

type of class label :<class 'tensorflow.python.framework.ops.EagerTensor'>
labels  from labelmap :7

测试验证 #

若将随机数传入tf.one_hot（如tf.one_hot(randint, 26)），模型可正常启动训练（无意义但验证了pipeline的其他流程正常）。

核心问题 #

如何将形状为(None,)、类型为int64的Tensor("SparseToDense:0")转换为形状为(26,)、类型为float32的Tensor("one_hot:0")？

已尝试的方法 #

• 在data.map(parse_xy)中调用.numpy()：仅对Eager Tensor有效，执行失败；
• 启用急切执行：tf.data函数始终以编译图模式运行，该设置无效；
• 使用tf.pyfunc：返回形状未知的张量，无法通过.set_shape()修改。

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