Title: Adding convolutional layers freezes training
Tags: python,tensorflow,keras,deep-learning,conv-neural-network

Question: I'm trying to train an image classification model, there are six classes and the training dataset is roughly seven thousand images. My issue is that the simple model behaves fine, and can get up to 85% accuracy, but the deeper model freezes and can't exceed 5% accuracy. I'm able to increase the size of the dense layers, and increase the number of kernels in the convolutional layer, both of which increase accuracy, but as soon as I add a second convolutional layer the model gets stuck with a loss value of exactly 1.7918. If it was varying between 1.9 and 1.6 I would suspect a local minimum, but it stays at exactly 1.7918 across multiple epochs.
Below is the code for the simple model.
import tensorflow as tf
import pathlib
import os
import matplotlib.pyplot as plt
import keras
from tensorflow.keras import utils
from tensorflow.keras import layers
import pickle
from tensorflow.keras import callbacks
from keras.optimizers import gradient_descent_v2

checkpoint_dir = "./training_checkpoints"
checkpoint_prefix = os.path.join(checkpoint_dir, "ckpt_{epoch}")

callbacks = [
    tf.keras.callbacks.ModelCheckpoint(filepath=checkpoint_prefix, save_weights_only=True),
    tf.keras.callbacks.TensorBoard(log_dir='./logs')
]

gpus = tf.config.experimental.list_physical_devices('GPU')
for gpu in gpus:
    tf.config.experimental.set_memory_growth(gpu, True)

strategy = tf.distribute.MirroredStrategy()

my_dir = "/workspace/"

train_dir = os.path.join(my_dir, 'data/train')
trainT = utils.image_dataset_from_directory(train_dir, labels="inferred", label_mode="categorical", class_names=None, color_mode="grayscale", image_size=(802,393), shuffle=True, seed=2023, batch_size=32)
trainT = trainT.apply(tf.data.experimental.ignore_errors())

test_dir = os.path.join(my_dir, 'data/test')
testT = utils.image_dataset_from_directory(test_dir, labels="inferred", label_mode="categorical", class_names=None, color_mode="grayscale", image_size=(802,393), shuffle=True, seed=2023, batch_size=32)

evalu_dir = os.path.join(my_dir, 'data/evaluation')
evaluT = utils.image_dataset_from_directory(evalu_dir, labels="inferred", label_mode="categorical", class_names=None, color_mode="grayscale", image_size=(802,393), shuffle=True, seed=2023, batch_size=32)

with strategy.scope():
    model = keras.Sequential(
        [
            layers.Rescaling(1./255),
            layers.Conv2D(32, 32, activation='relu', padding="valid", data_format='channels_last'),
            layers.MaxPooling2D((13,13), strides=1),
            layers.Flatten(),
            layers.Dense(25, activation='relu', use_bias=True, kernel_initializer='glorot_uniform'),
            layers.Dense(6, activation='relu', use_bias=True, kernel_initializer='glorot_uniform'),
            layers.Activation('softmax')
        ]
    )

    opt = gradient_descent_v2.SGD(lr=0.01)

    model.compile(optimizer=opt,
                 loss=tf.keras.losses.CategoricalCrossentropy(),
                 metrics=[tf.keras.metrics.CategoricalAccuracy(),
                         tf.keras.metrics.MeanIoU(num_classes=6)])

    history = model.fit(x=trainT, epochs=20, validation_data=evaluT, callbacks=callbacks)

model.save("fitModel")

with open('/trainHistory', 'wb') as file_pi:
    pickle.dump(history.history, file_pi)

The code that I've changed for the slightly deeper model is below, anything I've not included is the same from the simple model.
    model = keras.Sequential(
        [
            layers.Rescaling(1./255),
            layers.Conv2D(32, 32, activation='relu', padding="valid", data_format='channels_last'),
            layers.MaxPooling2D((13,13), strides=1),
            layers.Conv2D(32, 7, activation='relu', padding="valid", data_format='channels_last'),
            layers.MaxPooling2D((4,4), strides=1),
            layers.Flatten(),
            layers.Dense(25, activation='relu', use_bias=True, kernel_initializer='glorot_uniform'),
            layers.Dense(6, activation='relu', use_bias=True, kernel_initializer='glorot_uniform'),
            layers.Activation('softmax')
        ]
    )

    opt = tf.keras.optimizers.SGD(learning_rate=0.1)

The increased learning rate is an attempt to get out of a local minimum, though obviously this didn't work.
I've tried altering the learning rate, I've tried using the adam optimizer, I've tried removing the pooling layer between the two convolutional layers, and I've tried removing the pooling layer after the second convolutional layer.

Answer: Changing the kernel size from 32 to 5 allowed me to add extra layers. I ended up giving the first layer a size of 5x5, and the second a size of 3x3. This out performed the original model. Changing the optimizer did nothing to solve this problem, although it did help with convergence.