TensorFlow Dataset删除含NaN行及维度越界问题解决

错误原因分析 #

你遇到的「slice index 3 of dimension 0 out of bounds」错误，核心问题在于：

• 你使用map + tf.boolean_mask处理批量数据时，会导致每个batch的行数动态变化（比如原batch是8行，过滤后可能变成5行甚至0行）。
• TensorFlow Dataset的map操作默认要求输出张量的形状固定，当后续split_feature_label尝试对形状不固定的张量做切片时（比如空张量），就会触发索引越界。

解决方案 #

根据你的数据处理流程，推荐两种可行方案：

方案1：提前在单个样本级别过滤NaN（优先推荐） #

从源头过滤掉含NaN的样本，后续所有操作的样本都是有效的，避免批量处理时的形状不一致问题。

import tensorflow as tf
import numpy as np

# 生成带NaN的测试数据
data = np.random.rand(100, 4)
data[10:15, :] = np.nan  # 整行NaN
data[20, 2] = np.nan      # 部分元素NaN

# 1. 转为Dataset后先过滤无效样本
ds = tf.data.Dataset.from_tensor_slices(data)
# 过滤规则：样本中所有元素都不是NaN才保留（对应pandas的~np.isnan(ds).any(axis=1)）
ds = ds.filter(lambda x: tf.reduce_all(~tf.math.is_nan(x)))

# 2. 执行你的window、flat_map、batch、shuffle流程
ds = ds.window(2, shift=1, drop_remainder=True)
ds = ds.flat_map(lambda window: window.batch(2))
ds = ds.batch(8)
ds = ds.shuffle(10)

# 3. 特征标签拆分
def split_feature_label(x):
    features = x[:, :-1]  # 前3列作为特征
    label = x[:, -1]      # 最后1列作为标签
    return features, label

ds = ds.map(split_feature_label)

# 验证迭代
for feat, lbl in ds:
    print(f"特征形状：{feat.shape}，标签形状：{lbl.shape}")

方案2：批量过滤后重新整理数据 #

如果必须在batch之后处理，可通过flat_map将过滤后的行拆分为单个样本，再重新batch，确保输出形状固定。

import tensorflow as tf
import numpy as np

# 生成带NaN的测试数据
data = np.random.rand(100, 4)
data[10:15, :] = np.nan  # 整行NaN
data[20, 2] = np.nan      # 部分元素NaN

# 1. 执行你的window、flat_map、batch、shuffle流程
ds = tf.data.Dataset.from_tensor_slices(data)
ds = ds.window(2, shift=1, drop_remainder=True)
ds = ds.flat_map(lambda window: window.batch(2))
ds = ds.batch(8)
ds = ds.shuffle(10)

# 2. 批量过滤NaN行并重新整理
ds = ds.flat_map(lambda batch: tf.data.Dataset.from_tensor_slices(
    tf.boolean_mask(batch, tf.reduce_all(~tf.math.is_nan(batch), axis=-1))
))
# 重新batch为固定大小
ds = ds.batch(8)

# 3. 特征标签拆分
def split_feature_label(x):
    features = x[:, :-1]
    label = x[:, -1]
    return features, label

ds = ds.map(split_feature_label)

# 验证迭代
for feat, lbl in ds:
    print(f"特征形状：{feat.shape}，标签形状：{lbl.shape}")

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