如何在Keras中用列均值填充张量NaN值并实现Lambda层？

我想用张量的列均值填充其中的NaN值，虽然知道sklearn的SimpleImputer()能轻松实现，但希望所有特征工程都在Keras或TensorFlow中完成，方便作为Lambda层加入神经网络。

我写了如下代码，但运行报错：

import tensorflow as tf
import tensorflow_transform as tft

s = tf.convert_to_tensor(df_train)

def impute_mean(tensor):
    tensor = tf.dtypes.cast(tensor, tf.float32)
    mean = tft.mean(tensor)
    tensor = tf.where(tf.math.is_nan(tensor, mean))
    return tensor

d = impute_mean(s)
d

报错信息：

---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
~\AppData\Local\Temp\ipykernel_37096\4009563778.py in <module>
     12     return tensor
     13 
---> 14 d = impute_mean(s)
     15 d

~\AppData\Local\Temp\ipykernel_37096\4009563778.py in impute_mean(tensor)
      8 def impute_mean(tensor):
      9     tensor = tf.dtypes.cast(tensor, tf.float32)
---> 10     mean = tft.mean(tensor)
     11     tensor = tf.where(tf.math.is_nan(tensor, mean))
     12     return tensor

~\Anaconda3\envs\DemandForecastEnv\lib\site-packages\tensorflow_transform\common.py in wrapped_fn(*args, **kwargs)
     71             collection.append(collections.Counter())
     72           collection[0][fn.__name__] += 1
---> 73           return fn(*args, **kwargs)
     74       else:
     75         return fn(*args, **kwargs)

~\Anaconda3\envs\DemandForecastEnv\lib\site-packages\tensorflow_transform\analyzers.py in mean(x, reduce_instance_dims, name, output_dtype)
    842   """
    843   with tf.compat.v1.name_scope(name, 'mean'):
---> 844     return _mean_and_var(x, reduce_instance_dims, output_dtype)[0]
    845 
    846 

~\Anaconda3\envs\DemandForecastEnv\lib\site-packages\tensorflow_transform\analyzers.py in _mean_and_var(x, reduce_instance_dims, output_dtype)
    909     x_mean, x_var = _apply_cacheable_combiner(
    910         WeightedMeanAndVarCombiner(output_dtype.as_numpy_dtype, output_shape),
---> 911         *combine_inputs)
    912 
    913   return x_mean, x_var

~\Anaconda3\envs\DemandForecastEnv\lib\site-packages\tensorflow_transform\analyzers.py in _apply_cacheable_combiner(combiner, *tensor_inputs)
    170   outputs_value_nodes = apply_cacheable_combine_operation(
    171       combiner, *tensor_inputs)
---> 172   return tuple(map(analyzer_nodes.wrap_as_tensor, outputs_value_nodes))  # pytype: disable=bad-return-type
    173 
    174 

~\Anaconda3\envs\DemandForecastEnv\lib\site-packages\tensorflow_transform\analyzer_nodes.py in wrap_as_tensor(output_value_node)
    320   return bind_future_as_tensor(
    321       output_value_node,
---> 322       analyzer_def.output_tensor_infos[output_value_node.value_index])
    323 
    324 

~\Anaconda3\envs\DemandForecastEnv\lib\site-packages\tensorflow_transform\analyzer_nodes.py in bind_future_as_tensor(future, tensor_info, name)
    310     return _bind_future_as_tensor_v2(future, tensor_info, name)
    311   else:
---> 312     return _bind_future_as_tensor_v1(future, tensor_info, name)
    313 
    314 

~\Anaconda3\envs\DemandForecastEnv\lib\site-packages\tensorflow_transform\analyzer_nodes.py in _bind_future_as_tensor_v1(future, tensor_info, name)
    140                               name: Optional[str] = None) -> tf.Tensor:
    141   """Bind a future value as a tensor to a TF1 graph."""
---> 142   result = tf.compat.v1.placeholder(tensor_info.dtype, tensor_info.shape, name)
    143   is_asset_filepath = tensor_info.temporary_asset_info is not None
    144   tf.compat.v1.add_to_collection(TENSOR_REPLACEMENTS,

~\Anaconda3\envs\DemandForecastEnv\lib\site-packages\tensorflow\python\ops\array_ops.py in placeholder(dtype, shape, name)
   3341   """
   3342   if context.executing_eagerly():
---> 3343     raise RuntimeError("tf.placeholder() is not compatible with "
   3344                        "eager execution.")
   3345 

RuntimeError: tf.placeholder() is not compatible with eager execution.

错误原因 #

1. 库兼容问题：tft.mean()是TensorFlow Transform（TFT）的函数，该库基于TensorFlow 1.x图模式设计，内部依赖tf.placeholder()，和当前默认的Eager Execution模式冲突，导致报错。
2. API使用错误：tf.math.is_nan()仅需传入一个张量参数，你传入了两个；tf.where()需要三个参数（条件、满足条件的替换值、不满足条件的保留值），你的写法不符合要求。

基础修正实现 #

直接用TensorFlow原生函数计算列均值，修正tf.where用法：

import tensorflow as tf

# 转换DataFrame为张量
s = tf.convert_to_tensor(df_train, dtype=tf.float32)

def impute_mean(tensor):
    # 按列计算均值，忽略NaN，keepdims保持维度匹配以便广播
    col_mean = tf.math.reduce_mean(tensor, axis=0, keepdims=True)
    # 替换NaN为对应列的均值
    imputed_tensor = tf.where(tf.math.is_nan(tensor), col_mean, tensor)
    return imputed_tensor

# 测试
d = impute_mean(s)
print(d)

封装为Keras Lambda层 #

将填充逻辑集成到神经网络，用Lambda层封装即可：

from tensorflow.keras.layers import Lambda, Input
from tensorflow.keras.models import Model

# 定义填充层
impute_layer = Lambda(impute_mean)

# 构建模型示例
input_shape = (df_train.shape[1],)
input_layer = Input(shape=input_shape)
x = impute_layer(input_layer)
# 后续可添加Dense等其他层
model = Model(inputs=input_layer, outputs=x)

进阶：训练时固定全局均值 #

如果需要在训练阶段计算训练集的全局均值，推理阶段复用该均值（避免每次推理重新计算），可以自定义Keras Layer：

class MeanImputer(tf.keras.layers.Layer):
    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self.global_mean = None

    def build(self, input_shape):
        # 初始化全局均值变量，形状为(1, 特征数)
        self.global_mean = self.add_weight(
            shape=(1, input_shape[-1]),
            initializer='zeros',
            trainable=False,
            name='global_col_mean'
        )
        super().build(input_shape)

    def call(self, inputs, training=None):
        inputs = tf.cast(inputs, tf.float32)
        if training:
            # 训练时计算当前批次均值，更新全局均值
            batch_mean = tf.math.reduce_mean(inputs, axis=0, keepdims=True)
            combined = tf.concat([self.global_mean, batch_mean], axis=0)
            self.global_mean.assign(tf.math.reduce_mean(combined, axis=0, keepdims=True))
            return tf.where(tf.math.is_nan(inputs), batch_mean, inputs)
        else:
            # 推理时用训练好的全局均值填充
            return tf.where(tf.math.is_nan(inputs), self.global_mean, inputs)

# 使用示例
imputer = MeanImputer()
input_layer = Input(shape=(df_train.shape[1],))
x = imputer(input_layer)
model = Model(inputs=input_layer, outputs=x)

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