如何从tf.data.Dataset中移除含缺失值或NaN的样本？

问题根源 #

你之前的方案有两个核心问题：

1. 自定义过滤函数逻辑错误：直接对字典类型的样本调用tf.math.is_nan，TensorFlow无法处理字典输入，导致NaN检测完全失效。
2. 操作顺序错误：先执行filter会触发字段缺失的报错，ignore_errors()还没来得及生效就中断了流程。

方案一：直接跳过所有有问题的样本（字段缺失+含NaN） #

先让ignore_errors()跳过字段缺失的样本，再过滤掉含NaN的有效样本：

# 先跳过迭代时出错的样本（比如字段4缺失）
ds = ds.ignore_errors()

# 过滤含NaN的样本
def filter_nan_samples(elem):
    # 将所有字段值堆叠成张量，统一检查是否存在NaN
    all_values = tf.stack(list(elem.values()))
    # 返回True表示保留该样本，False则过滤
    return tf.logical_not(tf.reduce_any(tf.math.is_nan(all_values)))

ds = ds.filter(filter_nan_samples)

方案二：给缺失字段设置默认值，保留更多样本 #

如果不想丢弃字段缺失的样本，可以给字段4设置默认值后再过滤NaN：

@tf.function
def fill_missing_field4(elem):
    # 检查样本字典中是否存在字段4
    keys_tensor = tf.convert_to_tensor(list(elem.keys()))
    has_field4 = tf.math.reduce_any(tf.math.equal(keys_tensor, 'field_4'))
    
    # 存在则用原字段值，不存在则设默认值（可根据需求调整为其他值）
    elem['field_4'] = tf.cond(
        has_field4,
        lambda: elem['field_4'],
        lambda: tf.constant(0.0, dtype=tf.float32)
    )
    return elem

# 先填充缺失的字段4
ds = ds.map(fill_missing_field4)

# 再过滤含NaN的样本
def filter_nan_samples(elem):
    all_values = tf.stack(list(elem.values()))
    return tf.logical_not(tf.reduce_any(tf.math.is_nan(all_values)))

ds = ds.filter(filter_nan_samples)

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