如何使用TensorFlow的text_dataset_from_directory实现回归而非分类任务

调整步骤 #

你需要把原本的分类任务改成回归任务，按以下步骤修改即可：

1. 标签映射处理 #

你当前用label_mode='int'得到的标签是目录的顺序索引（比如对应1_{10分的目录，标签值是0}9），不是真实的评分值，需要加一层映射把标签转成实际评分：

def map_to_real_score(text, label_idx):
    # 将类别索引转为对应真实评分，转为float32适配回归任务要求
    real_score = tf.strings.to_number(raw_train_ds.class_names[label_idx], out_type=tf.float32)
    return text, real_score

# 对三个数据集都做映射
raw_train_ds = raw_train_ds.map(map_to_real_score, num_parallel_calls=AUTOTUNE)
val_ds = val_ds.map(map_to_real_score, num_parallel_calls=AUTOTUNE)
test_ds = test_ds.map(map_to_real_score, num_parallel_calls=AUTOTUNE)

2. 模型结构调整 #

输出层不需要分类用的softmax激活，改为输出单个神经元，用线性激活（不需要额外指定激活函数）：

# 示例，假设你用的是预训练文本模型加全连接层的结构
model = tf.keras.Sequential([
    # 这里保留你原本的文本预处理层、预训练语言模型层
    # 输出层改为单个神经元，无激活函数
    tf.keras.layers.Dense(1)
])

3. 损失与评估指标替换 #

原本分类用的交叉熵损失替换为回归任务常用的损失，编译模型时调整参数：

model.compile(
    optimizer=tf.keras.optimizers.Adam(learning_rate=2e-5),
    # 回归用均方误差损失
    loss=tf.keras.losses.MeanSquaredError(),
    # 评估指标用平均绝对误差，直观显示预测分和真实分的差值
    metrics=[tf.keras.metrics.MeanAbsoluteError()]
)

4. 预测结果后处理 #

模型预测输出的是浮点数值，你要整数评分的话，做范围裁剪和四舍五入即可：

pred_score = model.predict(["待预测的评论文本"])
# 限制评分在1~10范围内，四舍五入取整
final_score = int(tf.round(tf.clip_by_value(pred_score, 1, 10)))

如果对预测精度要求更高，也可以改用有序回归方案，不需要对标签做连续化假设，更适配评分的有序属性，调整成本也很低。

内容的提问来源于stack exchange，提问作者蘇韋文