TensorFlow特征列使用报错：员工晋升分类模型训练问题排查

问题场景 #

使用Kaggle的HR Analytics: Employee Promotion Dataset训练TensorFlow二分类模型，判断员工是否晋升。数据集包含数值特征与分类特征，尝试通过TensorFlow Feature Columns将DataFrame转换为训练数据集时，触发ValueError: Feature (key: age) cannot have rank 0错误。

原代码 #

# Import training data from kaggle download "data"
df = pd.read_csv('data/train.csv')

# Removing NaN from dataset
df['previous_year_rating'] = df['previous_year_rating'].fillna(0)
df['education'] = df['education'].fillna('None')

#Define variable names for numerical and categorical data
categorical_vars = ['department', 'region', 'education', 'gender', 'recruitment_channel']
numeric_vars = ['no_of_trainings', 'age', 'previous_year_rating', 'length_of_service', 'awards_won?', 'avg_training_score']

# Create Feature Columns
feature_columns = []

# Categorical variables
for var in categorical_vars:
  categorical_column = tf.feature_column.categorical_column_with_vocabulary_list(
      var, df[var].unique())
  indicator_column = tf.feature_column.indicator_column(categorical_column)
  feature_columns.append(indicator_column)

# Numeric variables
for var in numeric_vars:
    feature_columns.append(tf.feature_column.numeric_column(var))


# A utility method to create a tf.data dataset from a Pandas Dataframe
def df_to_dataset(df):
    y = df.pop('is_promoted')
    ds = tf.data.Dataset.from_tensor_slices((dict(df), y))
    return ds

# Create dataset from dataframe using df_to_datas（笔误）
ds = df_to_dataset(df)

# Set random seed
tf.random.set_seed(42)

# Create Model 
model = tf.keras.Sequential([
    tf.keras.layers.DenseFeatures(feature_columns),
    tf.keras.layers.Dense(1, activation='relu'),
    tf.keras.layers.Dense(1, activation='sigmoid')
])

# Compile
model.compile(optimizer='adam',
              loss=tf.keras.losses.BinaryCrossentropy(),
              metrics=['accuracy'])

# Train 
model.fit(ds, epochs=10)

完整报错信息 #

ValueError                                Traceback (most recent call last)
Cell In [150], line 16
     11 model.compile(optimizer='adam',
     12               loss=tf.keras.losses.BinaryCrossentropy(),
     13               metrics=['accuracy'])
     15 # Train 
---> 16 model.fit(train_ds, 
     17           epochs=10)

File ~\miniconda3\envs\py310\lib\site-packages\keras\utils\traceback_utils.py:70, in filter_traceback.<locals>.error_handler(*args, **kwargs)
     67     filtered_tb = _process_traceback_frames(e.__traceback__)
     68     # To get the full stack trace, call:
     69     # `tf.debugging.disable_traceback_filtering()`
---> 70     raise e.with_traceback(filtered_tb) from None
     71 finally:
     72     del filtered_tb

File ~\AppData\Local\Temp\__autograph_generated_file3hf_ppno.py:15, in outer_factory.<locals>.inner_factory.<locals>.tf__train_function(iterator)
     13 try:
     14     do_return = True
---> 15     retval_ = ag__.converted_call(ag__.ld(step_function), (ag__.ld(self), ag__.ld(iterator)), None, fscope)
     16 except:
     17     do_return = False

ValueError: in user code:

    File "C:\Users\mjadw\miniconda3\envs\py310\lib\site-packages\keras\engine\training.py", line 1249, in train_function  *
        return step_function(self, iterator)
    File "C:\Users\mjadw\miniconda3\envs\py310\lib\site-packages\keras\engine\training.py", line 1233, in step_function  **
        outputs = model.distribute_strategy.run(run_step, args=(data,))
    File "C:\Users\mjadw\miniconda3\envs\py310\lib\site-packages\keras\engine\training.py", line 1222, in run_step  **
        outputs = model.train_step(data)
    File "C:\Users\mjadw\miniconda3\envs\py310\lib\site-packages\keras\engine\training.py", line 1023, in train_step
        y_pred = self(x, training=True)
    File "C:\Users\mjadw\miniconda3\envs\py310\lib\site-packages\keras\utils\traceback_utils.py", line 70, in error_handler
        raise e.with_traceback(filtered_tb) from None

    ValueError: Exception encountered when calling layer 'dense_features_2' (type DenseFeatures).
    
    Feature (key: age) cannot have rank 0. Given: Tensor("IteratorGetNext:0", shape=(), dtype=int64)
    
    Call arguments received by layer 'dense_features_2' (type DenseFeatures):
      • features={'employee_id': 'tf.Tensor(shape=(), dtype=int64)', 'department': 'tf.Tensor(shape=(), dtype=string)', 'region': 'tf.Tensor(shape=(), dtype=string)', 'education': 'tf.Tensor(shape=(), dtype=string)', 'gender': 'tf.Tensor(shape=(), dtype=string)', 'recruitment_channel': 'tf.Tensor(shape=(), dtype=string)', 'no_of_trainings': 'tf.Tensor(shape=(), dtype=int64)', 'age': 'tf.Tensor(shape=(), dtype=int64)', 'previous_year_rating': 'tf.Tensor(shape=(), dtype=float32)', 'length_of_service': 'tf.Tensor(shape=(), dtype=int64)', 'awards_won?': 'tf.Tensor(shape=(), dtype=int64)', 'avg_training_score': 'tf.Tensor(shape=(), dtype=int64)'}
      • cols_to_output_tensors=None
      • training=True

错误原因分析 #

1. 数据集无批量处理：tf.data.Dataset.from_tensor_slices生成的是单个样本的数据集，每个特征都是标量（rank 0），但DenseFeatures层要求输入是批量数据（至少rank 1的张量，即形状为(batch_size,)）。
2. 代码笔误：创建数据集时调用了不存在的df_to_datas函数，实际应为df_to_dataset；训练时使用了未定义的train_ds，实际应为ds。

修复后的代码 #

# Import training data from kaggle download "data"
import pandas as pd
import tensorflow as tf

df = pd.read_csv('data/train.csv')

# Removing NaN from dataset
df['previous_year_rating'] = df['previous_year_rating'].fillna(0)
df['education'] = df['education'].fillna('None')

# Define variable names for numerical and categorical data
categorical_vars = ['department', 'region', 'education', 'gender', 'recruitment_channel']
numeric_vars = ['no_of_trainings', 'age', 'previous_year_rating', 'length_of_service', 'awards_won?', 'avg_training_score']

# Create Feature Columns
feature_columns = []

# Categorical variables
for var in categorical_vars:
  categorical_column = tf.feature_column.categorical_column_with_vocabulary_list(
      var, df[var].unique())
  indicator_column = tf.feature_column.indicator_column(categorical_column)
  feature_columns.append(indicator_column)

# Numeric variables
for var in numeric_vars:
    feature_columns.append(tf.feature_column.numeric_column(var))


# A utility method to create a tf.data dataset from a Pandas Dataframe
def df_to_dataset(df, batch_size=32):
    y = df.pop('is_promoted')
    ds = tf.data.Dataset.from_tensor_slices((dict(df), y))
    # 添加shuffle和batch处理
    ds = ds.shuffle(buffer_size=len(df)).batch(batch_size)
    return ds

# 修正笔误，创建带批量的数据集
ds = df_to_dataset(df)

# Set random seed
tf.random.set_seed(42)

# Create Model 
model = tf.keras.Sequential([
    tf.keras.layers.DenseFeatures(feature_columns),
    tf.keras.layers.Dense(32, activation='relu'),  # 调整隐藏层神经元数量，避免过窄
    tf.keras.layers.Dense(1, activation='sigmoid')
])

# Compile
model.compile(optimizer='adam',
              loss=tf.keras.losses.BinaryCrossentropy(),
              metrics=['accuracy'])

# Train，修正变量名
model.fit(ds, epochs=10)

无需LabelEncoder的分类特征处理替代方案 #

除了TensorFlow Feature Columns，还可以用Keras预处理层直接嵌入模型，更灵活且无需提前编码：

方案1：Keras StringLookup + CategoryEncoding层 #

import pandas as pd
import tensorflow as tf
from tensorflow.keras.layers import Input, StringLookup, CategoryEncoding, Dense, Concatenate
from tensorflow.keras.models import Model

df = pd.read_csv('data/train.csv')
df['previous_year_rating'] = df['previous_year_rating'].fillna(0)
df['education'] = df['education'].fillna('None')

# 分离特征与标签
y = df.pop('is_promoted')

# 构建输入层
inputs = {}
for col in df.columns:
    if col in categorical_vars:
        inputs[col] = Input(shape=(1,), dtype=tf.string, name=col)
    else:
        inputs[col] = Input(shape=(1,), dtype=tf.float32, name=col)

# 处理分类特征
encoded_categorical_features = []
for col in categorical_vars:
    # 字符串转索引
    lookup = StringLookup(vocabulary=df[col].unique().tolist(), output_mode='int')
    x = lookup(inputs[col])
    # 独热编码
    encoding = CategoryEncoding(num_tokens=lookup.vocabulary_size(), output_mode='one_hot')
    encoded_categorical_features.append(encoding(x))

# 拼接所有特征
all_features = encoded_categorical_features + [inputs[col] for col in numeric_vars]
x = Concatenate()(all_features)

# 构建模型
x = Dense(32, activation='relu')(x)
output = Dense(1, activation='sigmoid')(x)
model = Model(inputs=inputs, outputs=output)

# 编译训练
model.compile(optimizer='adam', loss=tf.keras.losses.BinaryCrossentropy(), metrics=['accuracy'])

# 转换数据集并训练
ds = tf.data.Dataset.from_tensor_slices((dict(df), y)).shuffle(len(df)).batch(32)
model.fit(ds, epochs=10)

方案2：使用categorical_column_with_hash_bucket（适合高基数分类特征） #

如果分类特征的类别数量极大（比如region），可以用哈希桶替代词汇表，避免手动维护词汇：

# 替换分类特征处理部分
for var in categorical_vars:
  # 用哈希桶，指定桶的数量（根据类别数量调整）
  categorical_column = tf.feature_column.categorical_column_with_hash_bucket(var, hash_bucket_size=100)
  indicator_column = tf.feature_column.indicator_column(categorical_column)
  feature_columns.append(indicator_column)

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