Преобразование Darknet53 дает результаты "nan" в Tensorflow 2.0

Question

Я пытаюсь конвертировать Yolo v3 в tenorflow 2.0 Я написал сетевые слои Darknet53, и я могу запустить его как на тестовых входах, так и на реальных изображениях, но результат в обоих случаях равен 'nan'

Я уже пытался масштабировать код вверх и вниз. Сначала делим изображение на 255, чтобы масштабировать от 0 до 1, как в оригинальной статье, и умножаем на 255, чтобы проверить, не слишком ли низкие значения для обработки. Я также пробовал исходное изображение без масштабирования.

Вот реальная сеть

import tensorflow as tf 

class ResnetIdentityBlock(tf.keras.Model):
    def __init__(self, kernel_size, filters):
        super(ResnetIdentityBlock, self).__init__(name='')
        filters1, filters2, filters3 = filters
        self.conv2a = tf.keras.layers.Conv2D(filters1, (1,1))
        self.bn2a = tf.keras.layers.BatchNormalization()

        self.conv2b = tf.keras.layers.Conv2D(filters2, kernel_size, padding="SAME")
        self.bn2b = tf.keras.layers.BatchNormalization()

        self.conv2c = tf.keras.layers.Conv2D(filters3, (1,1))
        self.bn2c = tf.keras.layers.BatchNormalization()

    def call(self, input_tensor, training=False):
        x = self.conv2a(input_tensor)
        x = self.bn2a(x, training=training)
        x = tf.nn.relu(x)
        print(f"after conv + bn a \n{x.shape}")
        x = self.conv2b(x)
        x = self.bn2b(x, training=training)
        x = tf.nn.relu(x)
        print(f"after conv + bn b \n{x.shape}")
        x = self.conv2c(x)
        x = self.bn2c(x, training=training)
        print(f"after conv + bn c \n{x.shape}")
        x += input_tensor
        print(f"input shape\n {input_tensor.shape}")
        print(f"final shape \n{x.shape}")
        return tf.nn.relu(x)


class Darknet53Block(tf.keras.Model):
    '''the main block of Darknet53, two conv followed by a residual'''
    def __init__(self, filters):
        super(Darknet53Block, self).__init__(name='')
        self.conv2a = tf.keras.layers.Conv2D(filters,1)
        self.bn2a = tf.keras.layers.BatchNormalization(epsilon=1e-05)

        self.conv2b = tf.keras.layers.Conv2D(filters*2, 3, padding="SAME")
        self.bn2b = tf.keras.layers.BatchNormalization(epsilon=1e-05)
    def call(self, input_tensor, training=False):
        x = self.conv2a(input_tensor)
        x = self.bn2a(x, training=training)
        x = tf.nn.leaky_relu(x, 0.1)
        x = self.conv2b(x)
        x = self.bn2b(x, training=training)
        x = tf.nn.leaky_relu(x, 0.1)
        x += input_tensor
        return x 

class Darknet53(tf.keras.Model):
    def __init__(self, *args, **kwargs):
        super(Darknet53, self).__init__(name='')

        self.bn = tf.keras.layers.BatchNormalization(epsilon=1e-05)

        self.conv2a = tf.keras.layers.Conv2D(filters=32,kernel_size=3)
        self.conv2b = tf.keras.layers.Conv2D(filters=64, kernel_size=3, strides=2)
        self.dn_a = Darknet53Block(filters=32)

        self.conv2c = tf.keras.layers.Conv2D(filters=128,kernel_size=3,  strides=2)
        self.dn_b = Darknet53Block(filters=64)

        self.conv2d = tf.keras.layers.Conv2D(filters=256, kernel_size=3, strides=2)
        self.dn_c = Darknet53Block(filters=128)

        self.conv2e = tf.keras.layers.Conv2D(filters=512, kernel_size=3, strides=2)
        self.dn_d = Darknet53Block(filters=256)

        self.conv2f = tf.keras.layers.Conv2D(filters=1024,kernel_size=3,  strides=2)
        self.dn_e = Darknet53Block(filters=512)

    def call(self, input_tensor):
        x = self.conv2a(input_tensor)
        x = self.bn(x)
        x = tf.nn.leaky_relu(x, alpha=0.1)
        x = self.conv2b(x)
        x = self.bn(x)
        x = tf.nn.leaky_relu(x, alpha=0.1)
        x = self.dn_a(x)
        x = self.conv2c(x)
        x = self.bn(x)
        x = tf.nn.leaky_relu(x, alpha=0.1)    
        for i in range(2):
            x = self.dn_b(x)
        x = self.conv2d(x)
        x = self.bn(x)
        x = tf.nn.leaky_relu(x, alpha=0.1)
        for i in range(8):
            x = self.dn_c(x)
        scale_1 = x 

        x = self.conv2e(x)
        x = self.bn(x)
        x = tf.nn.leaky_relu(x, alpha=0.1)
        for i in range(8):
            x = self.dn_d(x)
        scale_2 = x 

        x = self.conv2f(x)
        x = self.bn(x)
        x = tf.nn.leaky_relu(x, alpha=0.1)
        for i in range(4):
            x = self.dn_e(x)
        return scale_1, scale_2, x

и вот код теста

if __name__ == "__main__":

    import numpy as np
    import cv2 

    darknet = Darknet53()

    img = cv2.imread("20190413_143522_126.jpg")
    img = cv2.resize(img, (416,416))/255.0
    img = img.reshape((1,416,416,3))
    img = np.asarray(img, dtype=np.float32)
    print(img)

    result1, result2, result3 = darknet(img)
    print(result1)
    print(result2)
    print(result3)

Я ожидаю, что результат будет между 0 и 1. Размер результатов в порядке, но все они, кажется, заполнены 'nan'