Simultaneously training two CNN models - tensorflow

Question

I wish to train two CNN models in the same training loop where the input of the second model is the feature map generated by the forward pass on the first model.
def custom_training(inputd, x_train, y_train, x_val, y_val, n_epochs):
optimizer = tf.keras.optimizers.Adam()
ce_loss = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=False)
loss_train = np.zeros(shape=(n_epochs,), dtype=np.float32)
acc_train = np.zeros(shape=(n_epochs,), dtype=np.float32)
loss_val = np.zeros(shape=(n_epochs,))
acc_val = np.zeros(shape=(n_epochs,))
unet_model = tf.keras.Model(inputd, build_unet(inputd, 16))
#ds_model = DS_module()
n_batches = len(generator_train)
print(n_batches)
for epoch in range(n_epochs):
epoch_loss_avg = tf.keras.metrics.Mean() # Keeping track of the training loss
epoch_acc_avg = tf.keras.metrics.Mean() # Keeping track of the training accuracy
for batch in range(n_batches):
x, y = generator_train[batch]
with tf.GradientTape() as tape: # Forward pass
y_ = unet_model(x, training=True)
loss = ce_loss(y_true=y, y_pred=y_)
k = compute_K(y_)
print(k)
ds_model = tf.keras.Model(y_, DS_module(y_, k, num_class))
grad = tape.gradient(loss, unet_model.trainable_variables) # Backpropagation
optimizer.apply_gradients(zip(grad, unet_model.trainable_variables)) # Update network weights
epoch_loss_avg(loss)
epoch_acc_avg(accuracy_score(y_true=y, y_pred=np.argmax(y_, axis=-1)))
#generator.on_epoch_end()
loss_train[epoch] = epoch_loss_avg.result()
acc_train[epoch] = epoch_acc_avg.result()
y_ = unet_model.predict(x_val) # Validation predictions
loss_val[epoch] = ce_loss(y_true=y_val, y_pred=y_).numpy()
acc_val[epoch] = accuracy_score(y_true=y_val, y_pred=np.argmax(y_, axis=-1))
Below are the content of the two models, which consist of Unet, and another convolutional model, which will take as input the feature map generated by the Unet model:
def build_unet(input_layer, start_neurons):
# 128 -> 64
conv1 = Conv2D(start_neurons * 1, (3, 3), activation="relu", padding="same")(input_layer)
conv1 = Conv2D(start_neurons * 1, (3, 3), activation="relu", padding="same")(conv1)
pool1 = MaxPooling2D((2, 2))(conv1)
pool1 = Dropout(0.25)(pool1)
# 64 -> 32
conv2 = Conv2D(start_neurons * 2, (3, 3), activation="relu", padding="same")(pool1)
conv2 = Conv2D(start_neurons * 2, (3, 3), activation="relu", padding="same")(conv2)
pool2 = MaxPooling2D((2, 2))(conv2)
pool2 = Dropout(0.5)(pool2)
# 32 -> 16
conv3 = Conv2D(start_neurons * 4, (3, 3), activation="relu", padding="same")(pool2)
conv3 = Conv2D(start_neurons * 4, (3, 3), activation="relu", padding="same")(conv3)
pool3 = MaxPooling2D((2, 2))(conv3)
pool3 = Dropout(0.5)(pool3)
# 16 -> 8
conv4 = Conv2D(start_neurons * 8, (3, 3), activation="relu", padding="same")(pool3)
conv4 = Conv2D(start_neurons * 8, (3, 3), activation="relu", padding="same")(conv4)
pool4 = MaxPooling2D((2, 2))(conv4)
pool4 = Dropout(0.5)(pool4)
# Middle
convm = Conv2D(start_neurons * 16, (3, 3), activation="relu", padding="same")(pool4)
convm = Conv2D(start_neurons * 16, (3, 3), activation="relu", padding="same")(convm)
# 8 -> 16
deconv4 = Conv2DTranspose(start_neurons * 8, (3, 3), strides=(2, 2), padding="same")(convm)
uconv4 = concatenate([deconv4, conv4])
uconv4 = Dropout(0.5)(uconv4)
uconv4 = Conv2D(start_neurons * 8, (3, 3), activation="relu", padding="same")(uconv4)
uconv4 = Conv2D(start_neurons * 8, (3, 3), activation="relu", padding="same")(uconv4)
# 16 -> 32
deconv3 = Conv2DTranspose(start_neurons * 4, (3, 3), strides=(2, 2), padding="same")(uconv4)
uconv3 = concatenate([deconv3, conv3])
uconv3 = Dropout(0.5)(uconv3)
uconv3 = Conv2D(start_neurons * 4, (3, 3), activation="relu", padding="same")(uconv3)
uconv3 = Conv2D(start_neurons * 4, (3, 3), activation="relu", padding="same")(uconv3)
# 32 -> 64
deconv2 = Conv2DTranspose(start_neurons * 2, (3, 3), strides=(2, 2), padding="same")(uconv3)
uconv2 = concatenate([deconv2, conv2])
uconv2 = Dropout(0.5)(uconv2)
uconv2 = Conv2D(start_neurons * 2, (3, 3), activation="relu", padding="same")(uconv2)
uconv2 = Conv2D(start_neurons * 2, (3, 3), activation="relu", padding="same")(uconv2)
# 64 -> 128
deconv1 = Conv2DTranspose(start_neurons * 1, (3, 3), strides=(2, 2), padding="same")(uconv2)
uconv1 = concatenate([deconv1, conv1])
uconv1 = Dropout(0.5)(uconv1)
uconv1 = Conv2D(start_neurons * 1, (3, 3), activation="relu", padding="same")(uconv1)
uconv1 = Conv2D(start_neurons * 1, (3, 3), activation="relu", padding="same")(uconv1)
feature_maps = Conv2D(start_neurons, (1, 1), activation='relu')(uconv1)
return feature_maps
#tf.function
def DS_module(feature_map, prototypes, num_class):
ED=DS1(prototypes, 16)(feature_map)
ED_ac=DS1_activate(prototypes)(ED)
mass_prototypes=DS2(prototypes, num_class)(ED_ac)
mass_prototypes_omega=DS2_omega(prototypes, num_class)(mass_prototypes)
mass_Dempster=DS3_Dempster(prototypes, num_class)(mass_prototypes_omega)
pignistic=DM_pignistic(num_class)(mass_Dempster)
pignistic=DS3_normalize()(pignistic)
outputs = pignistic
return outputs
When I execute this code, i obtain the following error:
Error message obtained
This is the content of compute_k function, it executes PCA, Kmeans and Silhouette score on the feature map.
def compute_K(featureMap):
X = tf.reshape(featureMap, [featureMap.shape[0]*featureMap.shape[1]*featureMap.shape[2], featureMap.shape[3]])
pca = PCA(n_components=2)
principalComponents = pca.fit_transform(X)
principalDf = pd.DataFrame(data = principalComponents, columns = ['principal component 1', 'principal component 2'])
X = principalDf.to_numpy()
print(len(X))
#range_n_clusters = [3, 4, 5, 6, 7, 8, 9, 10, 12, 13, 15, 16, 17, 18, 19, 20]
range_n_clusters = [3, 4, 5, 6, 7]
silhouette_avg_n_clusters = []
max_sillhouette = 0
max_k = 0
plt.figure(figsize=(15, 7))
plt.subplot(111)
for n_clusters in range_n_clusters:
# Initialize the clusterer with n_clusters value and a random generator
# seed of 10 for reproducibility.
clusterer = KMeans(n_clusters=n_clusters, random_state=42)
cluster_labels = clusterer.fit_predict(X)
silhouette_avg = silhouette_score(X, cluster_labels)
print("For n_clusters =", n_clusters,
"The average silhouette_score is :", silhouette_avg)
silhouette_avg_n_clusters.append(silhouette_avg)
if (n_clusters == 2):
max_sillhouette = silhouette_avg
max_k = n_clusters
if (max_sillhouette< silhouette_avg):
max_sillhouette = silhouette_avg
max_k = n_clusters
return max_k