如何修改tf.data.Dataset中的数据类型以提升训练速度？

问题原因 #

你原有写法不生效是因为循环里的image_batch只是遍历过程中生成的临时副本，对它做修改不会写入原train_ds数据集对象，tf.data.Dataset的设计逻辑要求通过链式调用转换方法修改数据流向。

正确实现方式 #

使用tf.data.Dataset.map()方法对数据集的每个批次做类型转换，转换后赋值给新的数据集对象即可：

# 定义类型转换函数
def convert_dtype(image, label):
    # 按需调整转换类型，label如果不需要修改可保留原类型
    return tf.cast(image, tf.int16), label

# 对训练集做转换
train_ds = train_ds.map(convert_dtype, num_parallel_calls=tf.data.AUTOTUNE)
# 验证集可同步做转换
val_ds = val_ds.map(convert_dtype, num_parallel_calls=tf.data.AUTOTUNE)

# 转换后可追加prefetch配置进一步优化训练速度
train_ds = train_ds.prefetch(tf.data.AUTOTUNE)
val_ds = val_ds.prefetch(tf.data.AUTOTUNE)

验证与注意事项 #

• 转换完成后可执行以下代码验证类型是否修改成功：

for image_batch, labels_batch in train_ds.take(1):
    print(image_batch.dtype)

• 如果后续搭配混合精度训练，更建议转换为tf.float16类型，既能获得加速效果，也能避免模型输入类型不兼容的报错。

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