Keras中split_dataset函数处理CSV数据集时的Bug分析与解决

在Keras 3.4.0及以上版本中，开发人员发现了一个与split_dataset函数处理CSV数据集相关的Bug。这个问题主要出现在使用tf.data.experimental.make_csv_dataset加载数据集后，再调用Keras的split_dataset函数进行数据集分割时。

问题现象

当开发者尝试对通过make_csv_dataset加载的CSV数据集使用split_dataset时，会出现以下两种异常情况之一：

1. 函数调用会无限期挂起，无法完成执行
2. 输出的训练集和测试集中，列名与数据值会出现不匹配的情况（例如列名被随机重新分配）

问题重现

通过以下代码可以重现这个问题：

from keras.src.utils import split_dataset
import tensorflow as tf
import pandas as pd

# 创建测试数据
data_dict = {
    'a': [1.] * 10,
    'b': [20.] * 10,
    'c': [300.] * 10,
    'd': [4000.] * 10
}

df = pd.DataFrame(data_dict)
df.to_csv('test_data.csv', index=False)

# 加载CSV数据集
dataset = tf.data.experimental.make_csv_dataset('test_data.csv', batch_size=1)

# 尝试分割数据集 - 这里会出现问题
train, test = split_dataset(dataset, left_size=0.5, seed=1)

问题根源

经过Keras团队的分析，这个问题有两个主要原因：

1. 无限循环问题：make_csv_dataset默认会生成一个无限循环的数据集（通过设置num_epochs=None）。当split_dataset尝试处理这种无限数据集时，会导致函数挂起。

2. 列名错配问题：这个问题源于Keras的tree工具在处理OrderedDict时存在的一个Bug。在Keras 3.4.0版本中引入的一个重构（PR #19911）开始使用tree工具，而之前版本(3.3.3及以下)没有这个问题，因为它们没有使用tree工具。

解决方案

针对这两个问题，有以下解决方案：

1. 无限循环问题：在调用make_csv_dataset时显式设置num_epochs=1和shuffle=False，避免创建无限循环的数据集。

dataset = tf.data.experimental.make_csv_dataset(
    'test_data.csv', 
    batch_size=1,
    shuffle=False,
    num_epochs=1
)

2. 列名错配问题：Keras团队已经在PR #20481中修复了tree工具处理OrderedDict的Bug。用户可以通过升级到包含这个修复的Keras版本来解决这个问题。

临时解决方案

如果无法立即升级Keras版本，可以采取以下临时解决方案：

1. 将keras.src.utils.dataset_utils中的_restore_dataset_from_list函数回退到3.3.3版本的实现
2. 先将CSV数据集转换为Pandas DataFrame，再使用tf.data.Dataset.from_tensor_slices创建数据集

# 替代方案：先读取为DataFrame再转换为Dataset
df = pd.read_csv('test_data.csv')
dataset = tf.data.Dataset.from_tensor_slices(dict(df))

总结

这个Bug展示了深度学习框架中数据集处理流程的复杂性，特别是在涉及多层抽象和数据转换时。Keras团队已经识别并修复了核心问题，同时提供了多种解决方案供开发者选择。对于使用Keras处理CSV数据集的开发者，建议遵循最佳实践，明确指定数据集的epoch和shuffle参数，并保持框架版本更新。