Tensorflow - 在使用 CrossEntropy 损失训练鉴别器模型时 NaN 或接近零损失

Question

我正在尝试实现一个模型，该模型将输入作为 (q, a) 对，其中 q 是问题，a 是答案，并且 q 和 a 都是位置编码的。输出将是基于给定问题的答案的真实程度。所以这归结为一个二元分类任务，其中输出介于 0（假）和 1（真实）之间。

我的模型如下所示：

我接受两个输入，将它们连接起来，通过 RNN 传递，然后使用 sigmoid 来获得概率。 我将每个训练步骤定义为：

cross_entropy = tf.keras.losses.BinaryCrossentropy(from_logits=True)
optimizer = tf.keras.optimizers.Adam(1e-2)

@tf.function
def train_step(ip, tg, label):
    with tf.GradientTape() as tape:
        out = model([ip, tg])
        loss = cross_entropy(label, out)
        print(label, out)
        
    gradients = tape.gradient(loss, model.trainable_variables)
    optimizer.apply_gradients(zip(gradients, model.trainable_variables))
    return loss    

并使用

调用每个批次的步骤

for epoch in range(epochs):
        print("Epoch: %s"%(epoch + 1))
        batch_loss = 0.0
        
        for batch, ((ip, tg), label) in enumerate(concat_dataset.take(steps_per_epoch)):
            loss = train_step(ip, tg, label)
            batch_loss += loss
        

其中 ip, tg 是 (q, a) 对，标签为 0 或 1 表示假或真 (q, a) 样本。 当我训练模型时，我不断得到 NaN 或损失小到 1e-20 我无法弄清楚这里出了什么问题。我认为它要么是梯度爆炸，要么是梯度下降，我尝试降低和增加亚当的学习率。我也使用了 SGD，但结果相同。