I’m new to convolutional neural network and machine learning in general, so sorry in advance if my question is really straighforward, but I have searched for a while without getting a solution to fix it. For a semantic segmentation task, I am pre-training a U-Net model of 4 convolution blocks in keras with a dataset whose patches have input_shape = (128,128,3). Then, I need to train this same model, but with a new dataset (input_shape=(128,128,5)). So far, I think I’ve managed to apply the convolution necessary for the U-Net model to accept tensors from 5 channels. I’ve saved the model with the weights in an h5 file. However, I got the error “keras load_weights() ValueError: Layer count mismatch when loading weights from file. Model expected 4 layers, found 23 saved layers” at the line indicate in bold in my second code block.

I have already searched in the keras litterature and looked in other related questions here in stack overflow, but they dosn’t fully answer my question.
After using the model.summary() method, I found out that when training the new model using the pre-trained model weights, an adapted input convolution layer and the output layer, I end up with a model of just 4 layers (the input layers, the pre-trained model as a single layer and the output layer.

Here my transfer learning model code:

pretrained_unet = load_model(".projectsunet_v2.h5")
#freezing before the layer with the same input size
pretrained_unet = freeze_up_to(pretrained_unet, "dropout_5")

## Create a new input layer for the Unet
input_unet = Input(shape=(128, 128, 5))
input_layer = Conv2D(3, (3, 3), padding='same', activation='relu', input_shape= (128,128,5))

# Convert 5-channel input to 3-channel input
filtered_unet_input = Conv2D(3, (1, 1), activation='relu')(input_unet)
filtered_unet_input = pretrained_unet(filtered_unet_input)
#filtered_unet_input = tf.convert_to_tensor(input_unet[:,:,:3])
# Combine the UNET head and Unet
combined_output = Dense(n_classes, activation="softmax")(filtered_unet_input)
combined_model = keras.Model(input_unet,combined_output)
#combined_output = pretrained_unet(filtered_unet_input)

combined_model.summary()

_________________________________________________________________
 Layer (type)                Output Shape              Param #
=================================================================
 input_1 (InputLayer)        [(None, 128, 128, 5)]     0

 conv2d_1 (Conv2D)           (None, 128, 128, 3)       18

 UNet (Functional)           (None, 128, 128, 2)       8557570

 dense (Dense)               (None, 128, 128, 2)      

 6

=================================================================
Total params: 8,557,594
Trainable params: 110,874
Non-trainable params: 8,446,720

But my pretrained model has 23 layers. So there’s a mismatch between the number of layers and the weights can’t be correctly assigned. Here my code for the pretrained model:

import tensorflow as tf
from keras.callbacks import ModelCheckpoint, EarlyStopping, TensorBoard
from keras.models import load_model, Model
from keras.layers import Input, Dense, Conv2D

# Visualization
import matplotlib.pyplot as plt
import matplotlib.gridspec as gridspec
import matplotlib.patches as patches

# Working with arrays
import numpy as np
from tabulate import tabulate

# External files with functions to load the dataset,
# create a CNN model, and a data generator.
from importlib import reload
import utils
import models
import data_generator
from data_generator import DataGenerator

reload(utils)
reload(models)
reload(data_generator)

from utils import *
from models import *
from data_generator import *


# Module to track model's metrics
import wandb
import random

def fix_gpu():
    config = tf.compat.v1.ConfigProto()
    config.gpu_options.allow_growth = True
    session = tf.compat.v1.InteractiveSession(config=config)


fix_gpu()

PROJECT_DIR = "."  # os.getcwd()
SEED = 42
color2index = {(0, 0, 0): 0,
               (255, 255, 255): 1
               }

n_classes = len(color2index)

BATCH_SIZE = 8
PATCH_SIZE = 128
STEP_SIZE = 128
EPOCHS = 7

#define the function to freeze desired layers of the UNET model
def freeze_up_to(model, freeze_layer_name):
  """Function to freeze some layers of the model

  Args:
      model (keras.Model): a keras.Model
      freeze_layer_name (str): layer name of "model". All layers up to this layer will be freezed.

  Returns:
      keras.Model: a keras.Model with some layers freezed.
  """
  # Getting layer number based on layer name
  for id_layer, layer in enumerate(model.layers):
    if layer.name == freeze_layer_name:
      layer_number = id_layer
      break

  # Froze layers
  for layer in model.layers[:layer_number]:
    layer.trainable = False

  return model

# Path to the dataset folder
DATA_PATH = join("./projects", "data/train")
print(DATA_PATH)
list_images, list_masks = read_dataset(DATA_PATH)

splits = train_val_test_dataset(list_images,
                                list_masks,
                                val_size=0.25,
                                seed=SEED)
#data generation of the train dataset (images and masks)
data_gen_train = DataGenerator(batch_size=BATCH_SIZE,
                               patch_size=PATCH_SIZE,
                               step_size=STEP_SIZE,
                               list_images=splits["images_train"],
                               list_masks=splits["masks_train"],
                               n_classes=n_classes,
                               colormap_gt=color2index
                               )
#data generation of the validation dataset (images and masks)
data_gen_val = DataGenerator(batch_size=BATCH_SIZE,
                             patch_size=PATCH_SIZE,
                             step_size=STEP_SIZE,
                             list_images=splits["images_val"],
                             list_masks=splits["masks_val"],
                             n_classes=n_classes,
                             colormap_gt=color2index
                             )

print("Number of patches for training: {}".format(len(data_gen_train) * BATCH_SIZE))
print("nNumber of patches for validation: {}".format(len(data_gen_val) * BATCH_SIZE))

a, b = data_gen_train[0]
imgs, labels = data_gen_train[0]

show_batch(imgs, labels, color2index)

unet = get_unet(img_size=PATCH_SIZE,
                n_classes=n_classes)

unet.summary()

# Checkpoint
autosave = ModelCheckpoint("./projects/unet_v2.h5",
                           mode="max",
                           save_best_only=True,
                           monitor="val_iou",
                           verbose=1)

# Early stopping
early_stopping = EarlyStopping(monitor='val_iou',
                               patience=30,
                               verbose=1,
                               mode='max')

# Train the UNET model
unet.fit(
    data_gen_train,
    validation_data=data_gen_val,
    epochs=EPOCHS,
    callbacks=[autosave, early_stopping]
)

# Save the trained UNET model
unet.save("./unet_v2.h5")

# Load the pretrained UNET model
pretrained_unet = load_model("./projects/unet_v2.h5")

#freezing before the layer with the same input size
pretrained_unet = freeze_up_to(pretrained_unet, "dropout_5")

## Create a new input layer for the Unet
input_unet = Input(shape=(128, 128, 5))
input_layer = Conv2D(3, (3,3), padding='same', activation='relu', input_shape= (128,128,5))

# Convert 5-channel input to 3-channel input
filtered_unet_input = Conv2D(3, (1, 1), activation='relu')(input_unet)
filtered_unet_input = pretrained_unet(filtered_unet_input)

# Combine the UNET head and Unet
combined_output = Dense(n_classes, activation="softmax")(filtered_unet_input)
combined_model = keras.Model(input_unet,combined_output)


# Path to the new dataset folder
NEW_DATA_PATH = join("/projects", "data/DOP")  
new_list_images, new_list_masks = read_dataset(NEW_DATA_PATH)

new_splits = train_val_test_dataset(new_list_images, new_list_masks, val_size=0.25, seed=SEED)

# Prepare data generators for the new dataset
data_gen_new_train = DataGenerator(
    batch_size=BATCH_SIZE,
    patch_size=PATCH_SIZE,
    step_size=STEP_SIZE,
    list_images=new_splits["images_train"],
    list_masks=new_splits["masks_train"],
    n_classes=n_classes,
    colormap_gt=color2index
)

data_gen_new_val = DataGenerator(
    batch_size=BATCH_SIZE,
    patch_size=PATCH_SIZE,
    step_size=STEP_SIZE,
    list_images=new_splits["images_val"],
    list_masks=new_splits["masks_val"],
    n_classes=n_classes,
    colormap_gt=color2index
)

combined_model.compile(optimizer= Adam(),
                loss="categorical_crossentropy",
                metrics=["accuracy", 
                         keras.metrics.OneHotMeanIoU(num_classes=n_classes, 
                                       name="iou")])
# Fine-tune the combined model
combined_model.fit(
    data_gen_new_train,
    validation_data=data_gen_new_val,
    epochs=8,  # You may want to use fewer epochs for fine-tuning
    callbacks=[autosave, early_stopping]
)

#Loading the weights of the trained model.
****#I get the error in this line: **
combined_model.load_weights("unet_v2.h5")**

# Evaluate the model on test data
data_gen_test = DataGenerator(
    batch_size=BATCH_SIZE,
    patch_size=PATCH_SIZE,
    step_size=STEP_SIZE,
    list_images=new_splits["images_test"],
    list_masks=new_splits["masks_test"],
    n_classes=n_classes,
    colormap_gt=color2index
)

scores_train = combined_model.evaluate(data_gen_new_train, verbose=0)
scores_val = combined_model.evaluate(data_gen_new_val, verbose=0)
scores_test = combined_model.evaluate(data_gen_test, verbose=0)

If needed, this is the method of the class Models that define the architecture of the U-Net model:

def get_unet(img_size, n_classes):
  """Function to create a U-Net architecture

  Args:
      img_size (int): size of the input image
      n_classes (int): number of classes

  Returns:
      keras.Model: a keras model created using
        the functional API
  """
  # Input
  input = Input(shape=(img_size,img_size,3))
  # Downsampling
  f1,p1 = downsampling(input, 64, times=2)
  f2,p2 = downsampling(p1, 128, times=2)
  f3,p3 = downsampling(p2, 256, times=2)
  # Bottleneck
  blottleneck = conv_block(p3, 512, times=2)
  #Upsampling
  u7 = upsampling(blottleneck, 256, layer_concat=f3)
  u8 = upsampling(u7, 128, layer_concat=f2)
  u9 = upsampling(u8, 64, layer_concat=f1)
  # Output
  output = Conv2D(filters=n_classes,
                  kernel_size=1,
                  padding="same",
                  activation="softmax")(u9)
  model = Model(inputs=input, outputs=output, name="UNet")

  model.compile(optimizer= Adam(),
                loss="categorical_crossentropy",
                metrics=["accuracy", 
                         keras.metrics.OneHotMeanIoU(num_classes=n_classes, 
                                       name="iou")])

  return model

Do you know of an alternative for loading the model not as a single layer, but as all its constituent layers? And the same for the weights? Or maybe it’s the way I am building the new model from the pre-trained one. If anybody could give me some guidance, I will really appreciate it.