使用完整tf.data.Dataset训练模型出现NaN，拆分数据集后正常的问题

我构建了一个用于训练双头部神经网络的数据集，其中一个头部为LSTM，另一个为简单感知器。对数据集做了两种处理：

• 拆分为训练集和测试集，用于常规训练验证
• 保留完整数据集，用于最终全数据训练、测试与模拟

数据集处理代码 #

# 函数：将初始数据集拆分为训练集和测试集
def is_test(x, _):
    return x % int(self.val_split * 100) == 0

def is_train(x, y):
    return not is_test(x, y)

recover = lambda x, y: y
full_dataset

# 拆分测试集
test_set = full_dataset.enumerate().filter(is_test).map(recover)

# 拆分训练集
trainning_set = full_dataset.enumerate().filter(is_train).map(recover)

# 数据集批处理与优化
test_set = test_set.batch(batch_size).cache().prefetch(2)
trainning_set = trainning_set.batch(batch_size).cache().prefetch(2)

full_dataset = full_dataset.batch(batch_size).cache().prefetch(2)

数据集规格验证 #

检查三个数据集的element_spec，结果完全一致：

full_dataset:
<_PrefetchDataset element_spec=({'input1': TensorSpec(shape=(None, None, 3), dtype=tf.float32, name=None), 'input2': TensorSpec(shape=(None, 13), dtype=tf.float32, name=None)}, TensorSpec(shape=(None,), dtype=tf.float32, name=None))>

test_set: 
<_PrefetchDataset element_spec=({'input1': TensorSpec(shape=(None, None, 3), dtype=tf.float32, name=None), 'input2': TensorSpec(shape=(None, 13), dtype=tf.float32, name=None)}, TensorSpec(shape=(None,), dtype=tf.float32, name=None))>

trainning_set:
<_PrefetchDataset element_spec=({'input1': TensorSpec(shape=(None, None, 3), dtype=tf.float32, name=None), 'input2': TensorSpec(shape=(None, 13), dtype=tf.float32, name=None)}, TensorSpec(shape=(None,), dtype=tf.float32, name=None))>

训练表现差异 #

使用拆分后的训练集训练时一切正常：

model.fit(trainning_set, validation_data=data.test_set)

但使用完整数据集训练时，损失和指标全部为NaN：

model.fit(full_dataset)

训练日志输出：

Epoch 1/5
160/160 - 2s - loss: nan - nash_sutcliffe: nan - 2s/epoch - 12ms/step
Epoch 2/5
160/160 - 0s - loss: nan - nash_sutcliffe: nan - 319ms/epoch - 2ms/step
...

样本对比 #

已抽取未批处理前的test_set与full_dataset样本对比，除test_set的input1数值精度略高外，其余数据结构、类型（均为float32）完全一致：

test_set样本输出 #

for inputs, targets in test_set.take(1):
            print("Feature:", inputs)
            print("Label:", targets)

Feature: {'input1': <tf.Tensor: shape=(5, 3), dtype=float32, numpy=
array([[ 0.  , 16.12,  0.  ],
       [ 0.  , 17.42,  0.57],
       [ 0.  , 11.36, 13.97],
       [ 0.  , 10.55,  0.96],
       [ 0.  , 11.56,  0.24]], dtype=float32)>, 'input2': <tf.Tensor: shape=(13,), dtype=float32, numpy=
array([1.4391040e+02, 5.4850894e+03, 8.7901926e+00, 3.6657768e+01,
       5.4554661e+01, 9.5567673e+01, 2.0000000e+00, 5.8438915e+01,
       2.0383540e+03, 6.7381866e+01, 5.6437737e+01, 4.7759323e+00,
       0.0000000e+00], dtype=float32)>}
Label: tf.Tensor(0.91, shape=(), dtype=float32)

full_dataset样本输出 #

for inputs, targets in full_dataset.take(1):
            print("Feature:", inputs)
            print("Label:", targets)

Feature: {'input1': <tf.Tensor: shape=(5, 3), dtype=float32, numpy=
array([[0.000e+00, 9.860e+00, 0.000e+00],
       [0.000e+00, 1.308e+01, 0.000e+00],
       [0.000e+00, 1.433e+01, 1.000e-02],
       [0.000e+00, 1.630e+01, 0.000e+00],
       [0.000e+00, 1.644e+01, 0.000e+00]], dtype=float32)>, 'input2': <tf.Tensor: shape=(13,), dtype=float32, numpy=
array([1.4391040e+02, 5.4850894e+03, 8.7901926e+00, 3.6657768e+01,
       5.4554661e+01, 9.5567673e+01, 2.0000000e+00, 5.8438915e+01,
       2.0383540e+03, 6.7381866e+01, 5.6437737e+01, 4.7759323e+00,
       0.0000000e+00], dtype=float32)>}
Label: tf.Tensor(0.79, shape=(), dtype=float32)

已尝试搜索排查，但无法找到两种数据集的差异及问题根源，请求帮助分析为何完整数据集训练会出现NaN，而拆分数据集训练正常。

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