Градиентная лента Tensorflow 2 не работает должным образом

Question

В настоящее время я тренирую модель обнаружения больших объектов в Tensorflow 2 с пользовательским обучением l oop с использованием градиентной ленты . Проблема в том, что модель не улучшает потери, поскольку градиенты очень низкие. Я воспроизвел проблему в простой задаче классификации с использованием cifar10 и обнаружил, что небольшая модель отлично работает без проблем, в то время как более крупная модель (VGG16) совсем не улучшает потери. Ниже приведен код для воспроизведения проблемы.

Модель VGG16:

import tensorflow as tf
from tensorflow.keras.layers import Dense, Flatten, Conv2D, Dropout, MaxPooling2D, BatchNormalization, Input, Concatenate
import os

def create_vgg16(number_classes, include_fully=True, input_shape=(300, 300, 3), input_tensor=None):
    if input_tensor is None:
        img_input = Input(shape=input_shape)
    else:
        img_input = input_tensor
    x = Conv2D(64, (3, 3), activation='relu', padding='same', kernel_initializer='he_normal', name='conv1_1')(img_input)
    x = Conv2D(64, (3, 3), activation='relu', padding='same', kernel_initializer='he_normal', name='conv1_2')(x)
    x = MaxPooling2D(pool_size=(2, 2), strides=(2, 2), padding='same', name='pool1')(x)

    x = Conv2D(128, (3, 3), activation='relu', padding='same', kernel_initializer='he_normal', name='conv2_1')(x)
    x = Conv2D(128, (3, 3), activation='relu', padding='same', kernel_initializer='he_normal', name='conv2_2')(x)
    x = MaxPooling2D(pool_size=(2, 2), strides=(2, 2), padding='same', name='pool2')(x)

    x = Conv2D(256, (3, 3), activation='relu', padding='same', kernel_initializer='he_normal', name='conv3_1')(x)
    x = Conv2D(256, (3, 3), activation='relu', padding='same', kernel_initializer='he_normal', name='conv3_2')(x)
    x = Conv2D(256, (3, 3), activation='relu', padding='same', kernel_initializer='he_normal', name='conv3_3')(x)
    x = MaxPooling2D(pool_size=(2, 2), strides=(2, 2), padding='same', name='pool3')(x)

    x = Conv2D(512, (3, 3), activation='relu', padding='same', kernel_initializer='he_normal', name='conv4_1')(x)
    x = Conv2D(512, (3, 3), activation='relu', padding='same', kernel_initializer='he_normal', name='conv4_2')(x)
    x = Conv2D(512, (3, 3), activation='relu', padding='same', kernel_initializer='he_normal', name='conv4_3')(x)
    x = MaxPooling2D(pool_size=(2, 2), strides=(2, 2), padding='same', name='pool4')(x)

    x = Conv2D(512, (3, 3), activation='relu', padding='same', kernel_initializer='he_normal', name='conv5_1')(x)
    x = Conv2D(512, (3, 3), activation='relu', padding='same', kernel_initializer='he_normal', name='conv5_2')(x)
    x = Conv2D(512, (3, 3), activation='relu', padding='same', kernel_initializer='he_normal', name='conv5_3')(x)
    x = MaxPooling2D(pool_size=(3, 3), strides=(1, 1), padding='same', name='pool5')(x)

    if include_fully:
        x = Flatten(name='flatten')(x)
        x = Dense(4096, activation='relu', name='fc1')(x)
        x = Dense(4096, activation='relu', name='fc2')(x)
        x = Dense(number_classes, activation='softmax', name='predictions')(x)

    if input_tensor is not None:
        inputs = tf.keras.utils.get_source_inputs(input_tensor)
    else:
        inputs = img_input
    model = tf.keras.models.Model(inputs, x, name='vgg16')

    return model

Маленькая модель CNN:

def create_small_cnn(n_classes, input_shape=(32, 32, 3)):
    img_input = tf.keras.Input(shape=input_shape)
    x = tf.keras.layers.Conv2D(64, (3, 3), activation='relu', padding='same', name='conv1_1')(img_input)
    x = tf.keras.layers.Conv2D(64, (3, 3), activation='relu', padding='same', name='conv1_2')(x)
    x = tf.keras.layers.Flatten(name='flatten')(x)
    x = tf.keras.layers.Dense(16, activation='relu', name='fc1')(x)
    x = tf.keras.layers.Dense(n_classes, activation='softmax', name='softmax')(x)

    model = tf.keras.Model(img_input, x, name='small_cnn')
    return model

Обучение l oop:

def main():
    number_classes = 10
    # Load and one hot encode data
    (x_train, y_train), (x_test, y_test) = tf.keras.datasets.cifar10.load_data()
    x_train, x_test = x_train, x_test
    y_train = tf.reshape(y_train, [-1])
    y_train = tf.one_hot(y_train, number_classes).numpy()
    y_test = tf.reshape(y_test, [-1])
    y_test = tf.one_hot(y_test, number_classes).numpy()


    # Define model
    model = create_vgg16(number_classes, input_shape=(32, 32, 3))
    # model = create_small_cnn(number_classes, input_shape=(32, 32, 3))

    # Instantiate an optimizer to train the model.
    optimizer = tf.keras.optimizers.Adam(lr=0.001, beta_1=0.9, beta_2=0.999, epsilon=1e-08, decay=0.0)

    # Instantiate a loss function.s
    loss_fn = tf.keras.losses.CategoricalCrossentropy()

    # Prepare the metrics.
    train_acc_metric = tf.keras.metrics.CategoricalAccuracy()
    val_acc_metric = tf.keras.metrics.CategoricalAccuracy()

    # Prepare the training dataset.
    batch_size = 64
    train_dataset = tf.data.Dataset.from_tensor_slices(
      (tf.cast(x_train/255, tf.float32),
       tf.cast(y_train,tf.int64)))
    train_dataset = train_dataset.shuffle(buffer_size=1024).batch(batch_size)

    # Prepare the validation dataset.
    val_dataset = tf.data.Dataset.from_tensor_slices(
      (tf.cast(x_test/255, tf.float32),
       tf.cast(y_test,tf.int64)))
    val_dataset = val_dataset.shuffle(buffer_size=1024).batch(batch_size)

    model.summary()

    for epoch in range(100):
      print('Start of epoch %d' % (epoch,))
      for step, (x_batch_train, y_batch_train) in enumerate(train_dataset):
        with tf.GradientTape() as tape:
          logits = model(x_batch_train)
          loss_value = loss_fn(y_batch_train, logits)
        grads = tape.gradient(loss_value, model.trainable_variables)
        optimizer.apply_gradients(zip(grads, model.trainable_variables))

        train_acc_metric(y_batch_train[0], logits[0][:-1])

        if step % 200 == 0:
            print('Training loss (for one batch) at step %s: %s' % (step, float(loss_value)))


      # Display metrics at the end of each epoch.
      train_acc = train_acc_metric.result()
      print('Training acc over epoch: %s' % (float(train_acc),))
      # Reset training metrics at the end of each epoch
      train_acc_metric.reset_states()

      # Run a validation loop at the end of each epoch.
      for x_batch_val, y_batch_val in val_dataset:
        val_logits = model(x_batch_val)

        val_acc_metric(y_batch_val[0], val_logits[0][:-1])
      val_acc = val_acc_metric.result()
      val_acc_metric.reset_states()
      print('Validation acc: %s' % (float(val_acc),))


if __name__ == '__main__':
    main()

Если вы запустите показанный код, вы увидите, что сетевое обучение отлично работает при использовании маленькой модели CNN. Но, с другой стороны, он не работает с точно таким же набором данных с той же предварительной обработкой с использованием стандартной модели VGG16. Чтобы еще больше запутать ситуацию, модель VGG будет отлично работать при использовании model.fit вместо пользовательского обучения l oop с градиентной лентой .

Кто-нибудь имеет представление, почему это так и как решить эту проблему?