TensorFlow 量化感知训练 - 目标张量中的 NaN

如何解决TensorFlow 量化感知训练 - 目标张量中的 NaN

我正在尝试对现有的 TensorFlow+NVIDIA dali 管道实施量化感知训练 (QAT)，以执行密集预测计算机视觉任务。我的训练损失首先像往常一样减少，然后立即跳到 NaN。使用我的损失函数中的 tf.debugging API，我发现我从数据管道中的“目标”张量中获得了 NaN。奇怪的是，当 QAT 被禁用时，我没有从（相同的）数据管道中得到任何 NaN。

def custom_loss(network_output,target_tensor,...):
    tf.debugging.check_numerics(target_tensor,"target tensor numeric error",name=None)
    ...
    return loss

---

我怀疑，在某些情况下，我只是在目标张量的值超出 TensorFlow 预期范围的情况下出现溢出。如果我对目标张量进行归一化，我不会得到 NaN，这一事实加强了这种怀疑。

但话说回来，该模型似乎在 QAT 期间实际上并未进行任何量化。来自 quantization aware training comprehensive guide 中的功能示例：

# For deployment purposes,the tool adds `QuantizeLayer` after `InputLayer` so that the
# quantized model can take in float inputs instead of only uint8.

和

"请注意，生成的模型具有量化意识但未量化 （例如，权重是 float32 而不是 int8）。”

我验证了我的损失函数得到了一个带有 tf.print(target_tensor.dtype) 的 float32 目标张量。

---

有什么想法吗？