Loading Variable Input Size CNN model with TensorFlow

I’m working on a 3D image classification and detection project using a CNN with TensorFlow. The model is trained on 3D images the size of each image is (31,31,31,1) but I made the model input shape with variable sizes, specified as None in the input shape, because I wanted the model to accept any shape during sliding window. This works well for classification, but I encounter an issue when applying the trained model in a sliding window approach for detection on larger images.

Training Code (Summary):


model = Sequential([
    layers.Conv3D(filters=16, kernel_size=3, padding='same'),
    layers.LeakyReLU(),
    layers.BatchNormalization(),
    layers.Conv3D(filters=16, kernel_size=3, padding='same'),
    layers.LeakyReLU(),
    layers.BatchNormalization(),
    layers.MaxPool3D(pool_size=2),
    layers.BatchNormalization(),

    layers.Conv3D(filters=32, kernel_size=3, padding='same'),
    layers.LeakyReLU(),
    layers.BatchNormalization(),
    layers.Conv3D(filters=32, kernel_size=3, padding='same'),
    layers.LeakyReLU(),
    layers.BatchNormalization(),
    layers.MaxPool3D(pool_size=2),
    layers.BatchNormalization(),

    layers.Conv3D(filters=64, kernel_size=3, padding='same'),
    layers.LeakyReLU(),
    layers.BatchNormalization(),
    layers.Conv3D(filters=64, kernel_size=3, padding='same'),
    layers.LeakyReLU(),
    layers.BatchNormalization(),
    layers.MaxPool3D(pool_size=2),
    layers.BatchNormalization(),

    layers.Conv3D(filters=128, kernel_size=3, padding='same'),
    layers.LeakyReLU(),
    layers.BatchNormalization(),
    layers.Conv3D(filters=128, kernel_size=3, padding='same'),
    layers.LeakyReLU(),
    layers.BatchNormalization(),
    layers.MaxPool3D(pool_size=2),
    layers.BatchNormalization(),

    layers.GlobalAveragePooling3D(),
    layers.Dense(units=256),
    layers.LeakyReLU(),
    layers.BatchNormalization(),
    layers.Dropout(0.4),
    layers.Dense(units=1, activation="sigmoid"),
])

train_gen = DataGenerator(train_data, train_labels, base_dir, (31, 31, 31), batch_size=batch_size, shuffle=False)
test_gen = DataGenerator(val_data, val_labels, base_dir, (31, 31, 31), batch_size=batch_size, shuffle=False)

cp_callback = tf.keras.callbacks.ModelCheckpoint(filepath=checkpoint_path,
                                                 verbose=1)


log_dir="/project/checkpoints/cleancode/"+ datetime.now().strftime("%Y%m%d-%H%M%S")
tensorboard_callback = tf.keras.callbacks.TensorBoard(log_dir=log_dir, histogram_freq=1)
csv_logger = tf.keras.callbacks.CSVLogger('/project/checkpoints/cleancode/cleancode.csv')

model.compile(optimizer=opt,
              loss = tf.keras.losses.BinaryFocalCrossentropy(gamma=5, apply_class_balancing=True),
              metrics=['accuracy', BalancedAccuracy(), 'AUC',tf.keras.metrics.AUC(curve='PR',name='AUC_PR'), tfa.metrics.F1Score(num_classes=1, threshold=0.5), tf.keras.metrics.SpecificityAtSensitivity(0.5), 'Precision', 'Recall',
                        'TruePositives', 'FalsePositives', 'TrueNegatives','FalseNegatives'])
model.build(input_shape= (128,None,None,None,1))
model.summary()

history = model.fit(train_gen, validation_data = test_gen, epochs=epochs, shuffle = False , verbose = 1 ,
                    callbacks = [csv_logger, tensorboard_callback, cp_callback],
use_multiprocessing = True, workers=10)

model.save("cleancode.keras")

Sliding Window Code (Summary):

model = tf.keras.models.load_model('/project/cleancode.keras',compile=False,custom_objects={"BalancedAccuracy": BalancedAccuracy(), "F1Score": tfa.metrics.F1Score(num_classes=1, threshold=0.5)})

Error Message:

Traceback (most recent call last):
  File "/home/mustafa/.local/lib/python3.9/site-packages/IPython/core/interactiveshell.py", line 3553, in run_code
    exec(code_obj, self.user_global_ns, self.user_ns)
  File "<ipython-input-2-9c04782baac6>", line 1, in <cell line: 1>
    runfile('/home/mustafa/project/sliding_single_novoxel.py', wdir='/home/mustafa/project/')
  File "/home/mustafa/.pycharm_helpers/pydev/_pydev_bundle/pydev_umd.py", line 197, in runfile
    pydev_imports.execfile(filename, global_vars, local_vars)  # execute the script
  File "/home/mustafa/.pycharm_helpers/pydev/_pydev_imps/_pydev_execfile.py", line 18, in execfile
    exec(compile(contents+"\n", file, 'exec'), glob, loc)
  File "/home/mustafa/project/sliding_single_novoxel.py", line 103, in <module>
    model = tf.keras.models.load_model('/home/mustafa/project/cleancode.keras',compile=False,custom_objects={"BalancedAccuracy": BalancedAccuracy(), "F1Score": tfa.metrics.F1Score(num_classes=1, threshold=0.5)})
  File "/home/mustafa/.local/lib/python3.9/site-packages/keras/utils/traceback_utils.py", line 70, in error_handler
    raise e.with_traceback(filtered_tb) from None
  File "/home/mustafa/.local/lib/python3.9/site-packages/keras/layers/convolutional/base_conv.py", line 409, in _get_input_channel
    raise ValueError(
ValueError: The channel dimension of the inputs should be defined. The input_shape received is (None, None, None, None, None), where axis -1 (0-based) is the channel dimension, which found to be `None`.

Despite setting the input shape to accept variable-sized images during training, the sliding window detection script fails when loading and using the model, citing an undefined channel dimension.

How can I modify loading in the sliding window detection code to resolve this issue, ensuring the model accepts variable-sized sub-volumes during detection?

Hi @Mustafa_Mahmood, The error is due to the channel dimension mentioned in the input data does not match with training data. I have tried to train the given model with random data and it works without any error. Please refer to this gist for working code example. Thank You.

1 Like

Thank you @Kiran_Sai_Ramineni for your help.

The problem was happening after saving and loading the model to a different code for prediction.

I think the first problem was that I’m saving the model as checkpoints using

tf.keras.callbacks.ModelCheckpoint

the checkpoints were saved as (model.keras) but couldn’t be loaded normally, I managed to load it by recreating the model and building it then load the weight fo the model to the new model

model.load_weights('path to model')

even though I saved the model at the end of training using

model.save("cleancode.keras")

I was not able to load it normally, I had to build the model again and load the weight to it.