使用tfdf转换数据集报错：label未在轴中找到

问题背景 #

在Google Colab中训练TensorFlow决策森林时，执行以下操作后转换TF数据集报错：

• 加载数据集：

  dataset_df = pd.read_csv("/content/drive/MyDrive/Medical_Error_Test_Data-_1_.csv")
  

• 填充NaN值：

  labels = dataset_df.keys()
  for label in labels:
      dataset_df[label] = dataset_df[label].fillna(0)
  

• 整数编码所有列：

  classes = []
  for label in labels:
      classes.append(dataset_df[label].unique().tolist())
  print(f"Label classes: {classes}")
  for classs, label in zip(classes, labels):
      dataset_df[label] = dataset_df[label].map(classs.index)
  

• 拆分训练测试集：

  def split_dataset(dataset, test_ratio=0.30):
      """Splits a panda dataframe in two."""
      test_indices = np.random.rand(len(dataset)) < test_ratio
      return dataset[~test_indices], dataset[test_indices]
  
  train_ds_pd, test_ds_pd = split_dataset(dataset_df)
  

• 转换为TF数据集时执行：

  train_ds = tfdf.keras.pd_dataframe_to_tf_dataset(train_ds_pd, label='Sub System')
  

  报错：KeyError: "['Sub System'] not found in axis"，且尝试所有列均出现该错误。

问题原因 #

最可能的原因是DataFrame的列名存在不可见的空白字符（如前后空格、制表符），或者CSV文件加载时列名被错误解析，导致代码中指定的列名与实际DataFrame的列名不匹配。

解决步骤 #

1. 验证实际列名
   执行以下代码查看训练集的真实列名，确认是否存在空白字符：

   # 打印所有列名
   print("训练集列名列表：", train_ds_pd.columns.tolist())
   # 查看每个列名的详细信息（含长度）
   for col in train_ds_pd.columns:
       print(f"列名：'{col}'，字符长度：{len(col)}")
   

   若输出中列名带有前后空格（如' Sub System'或'Sub System '），则说明列名与代码中指定的'Sub System'不一致。

2. 修正列名
   去除所有列名前后的空白字符：

   train_ds_pd.columns = train_ds_pd.columns.str.strip()
   test_ds_pd.columns = test_ds_pd.columns.str.strip()
   

   之后重新执行pd_dataframe_to_tf_dataset即可。

3. 额外排查项

   • 若上述方法无效，检查CSV文件的表头是否正确，可通过print(dataset_df.head())查看数据集前几行，确认列名是否与预期一致。
   • 确认拆分数据集的函数未丢失列（当前split_dataset函数逻辑正常，不会导致列丢失）。

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