Простой GAN предсказывает NaN в Tensorflow после 2 шагов

Я реализую базовую GAN на основе документации Tensorflow.

После 2 шагов обучения все прогнозы от генератора будут NaN.

Я не знаю, почему это происходит, но я заметил, что градиенты 2-го слоя свертки дискриминатора все NaN с первого шага:

<tf.Tensor: shape=(5, 5, 64, 128), dtype=float32, numpy= массив([[[[нан, нан, нан, ..., нан, нан, нан], [нан, нан, нан,

Мои функции потерь:

loss = BinaryCrossentropy(from_logits=True)

def generator_loss(discriminator_output):
    loss_value = loss(tf.ones_like(discriminator_output), discriminator_output)
    print(f"Generator loss: {loss_value}")
    return loss_value

def discriminator_loss(outputs_for_reals, outputs_for_fakes):
    real_loss = loss(tf.ones_like(outputs_for_reals), outputs_for_reals)
    fake_loss = loss(tf.zeros_like(outputs_for_fakes), outputs_for_fakes)
    print(f"Disc loss (real/fake): {real_loss}/{fake_loss}")

    return real_loss + fake_loss

Тренировочный цикл:

def training_step(self, real_images):
    noise = tf.random.normal([BATCH_SIZE, NOISE_DIM])

    with tf.GradientTape() as generator_tape, tf.GradientTape() as discriminator_tape:
        generated_images = self._generator(noise, training=True)
        if True in tf.math.is_nan(generated_images)[0, 0, :, 0]:
            print("NaN result found")
        else:
            print("OK Result")

        results_for_real = self._discriminator(real_images, training=True)
        results_for_fake = self._discriminator(generated_images, training=True)

        generator_loss = self._generator_loss(results_for_fake)
        discriminator_loss = self._discriminator_loss(results_for_real, results_for_fake)

    generator_gradients = generator_tape.gradient(generator_loss,
                                                  self._generator.trainable_variables)
    discriminator_gradients = discriminator_tape.gradient(discriminator_loss,
                                                          self._discriminator.trainable_variables)


    self._generator_optimizer.apply_gradients(
        zip(generator_gradients, self._generator.trainable_variables))
    self._discriminator_optimizer.apply_gradients(
        zip(discriminator_gradients, self._discriminator.trainable_variables))

    return generator_loss, discriminator_loss, generated_images

Я строю модели точно так же, как в документации.

Что я пробовал:

• Снижение скорости обучения
• Запуск модели в разных training режимах (training=False/training=True)
• Украшаем тренировочный этап tf.function

Что бы я ни делал, вызов генератора с любыми входными данными будет производить исключительно NaN элементов. Пример вывода:

2020-12-11 18:15:42.783543: W tensorflow/stream_executor/gpu/redzone_allocator.cc:314] Internal: Invoking GPU asm compilation is supported on Cuda non-Windows platforms only
Relying on driver to perform ptx compilation. 
Modify $PATH to customize ptxas location.
This message will be only logged once.
OK Result
Generator loss: 1.5484254360198975
Disc loss (real/fake): 1.4994642734527588/0.2869953215122223
OK Result
Generator loss: 1.2899521589279175
Disc loss (real/fake): 1.3189733028411865/0.3767632842063904
Backend Qt5Agg is interactive backend. Turning interactive mode on.
NaN result found