如何使用TensorFlow的flow_from_dataframe处理独热编码输出

错误原因 #

你当前的报错是因为DataFrame的age列存储的是字符串格式的独热编码，flow_from_dataframe无法直接将字符串解析为可用于训练的数值标签，即使设置class_mode="raw"也会因为数据类型不匹配失效。

解决步骤 #

步骤1：预处理age列，转换为数值数组 #

将字符串格式的独热编码转换为数值类型的numpy数组：

import numpy as np

training_data['age'] = training_data['age'].apply(
    lambda x: np.array([int(num) for num in x.strip('[]').split()], dtype=np.float32)
)

处理后可以通过print(training_data['age'].iloc[0].shape)验证，输出应为(23,)即符合要求。

步骤2：调整flow_from_dataframe参数 #

保持class_mode="raw"配置，直接使用预处理好的标签数组：

train_data_generator = idg.flow_from_dataframe(
    dataframe=training_data,
    x_col="image_name",
    y_col="age",
    class_mode="raw",
    shuffle=True,
    target_size=(227, 227),
)

可选优化方案 #

如果不需要提前存储完整独热编码，也可以直接提取类别索引存入age列，减少内存占用：

# 提取独热编码中1对应的位置下标作为类别标签
training_data['age'] = training_data['age'].apply(
    lambda x: x.strip('[]').split().index('1')
)

此时调用flow_from_dataframe时设置class_mode="categorical"即可，接口会自动完成独热编码转换，无需提前处理为数组。

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