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.
gioco9014 is a new contributor to this site. Take care in asking for clarification, commenting, and answering.
Check out our Code of Conduct.