Split a tensoflow model into two parts

Question

I have a pre-trained tensorflow model for image classification. It has has some convolution and maxpooling layers followed by dense layers.

I would like to split it in two pats. The first containing the convolutions/maxpooling part, the second containing the dense layers. Then I would like to feed the first part with an image, store the result into a file/variable and then using it as input for the second part.

The idea is that I could encode on the fly images with the first part, save them to disk and then using the encoded files as input. In this way I can avoid storing the original images.

I am writing an example of a whole model (based on Coursera's/deeplearning.ai "Introduction to TensorFlow for Artificial Intelligence, Machine Learning, and Deep Learning")

import tensorflow as tf
mnist = tf.keras.datasets.fashion_mnist
(training_images, training_labels), (test_images, test_labels) = mnist.load_data()
training_images=training_images.reshape(60000, 28, 28, 1)
training_images=training_images / 255.0
test_images = test_images.reshape(10000, 28, 28, 1)
test_images=test_images/255.0
model = tf.keras.models.Sequential([
  tf.keras.layers.Conv2D(32, (3,3), activation='relu', input_shape=(28, 28, 1)),
  tf.keras.layers.MaxPooling2D(4, 4),
  tf.keras.layers.Conv2D(32, (3,3), activation='relu'),
  tf.keras.layers.MaxPooling2D(2,2),
  tf.keras.layers.Flatten(),
  tf.keras.layers.Dense(128, activation='relu'),
  tf.keras.layers.Dense(10, activation='softmax')
])
model.compile(optimizer='adam', loss='sparse_categorical_crossentropy', metrics=['accuracy'])
model.summary()
model.fit(training_images, training_labels, epochs=5)

In the example above,

• first_part would have the first 5 layers
• second_part would have the last two (and possibly an input layer)

Thank you.

Answer 1

After several attemps I have the following piece of code. I created two new networks and transferred the weights from the initial network to them.

# first part of initial model
part1_model = tf.keras.models.Sequential([
  tf.keras.layers.Conv2D(32, (3,3), activation='relu', input_shape=(28, 28, 1)),
  tf.keras.layers.MaxPooling2D(4, 4),
  tf.keras.layers.Conv2D(32, (3,3), activation='relu'),
  tf.keras.layers.MaxPooling2D(2,2),
  tf.keras.layers.Flatten()
])
part1_model.layers[0].set_weights(model.layers[0].get_weights())
part1_model.layers[1].set_weights(model.layers[1].get_weights())
part1_model.layers[2].set_weights(model.layers[2].get_weights())
part1_model.layers[3].set_weights(model.layers[3].get_weights())
part1_model.layers[4].set_weights(model.layers[4].get_weights())
part1_model.summary() 

# second part of initial model
part2_model = tf.keras.models.Sequential([
    tf.keras.layers.Dense(128, activation='relu',input_shape=(1, 128)),
  tf.keras.layers.Dense(10, activation='softmax')
])
part2_model.summary()
part2_model.layers[0].set_weights(model.layers[5].get_weights())
part2_model.layers[1].set_weights(model.layers[6].get_weights())
# predictions holds prediction for test set from initial model
predictions = model.predict(test_images)

# predictions1 holds output of first part when test set is used as input 
predictions1 = part1_model.predict(test_images)

import numpy as np
# tmp is used to transform prediction1 in a format recognizable from part2
tmp=np.zeros((10000,1,128))
for i in range(0,10000):
    tmp[i,:,:]=predictions1[i,:]

# predictions2 holds the output of second part when the result of first part (predictions1) is used as input    
predictions2 = part2_model.predict(tmp)

# check that the result of initial model is the same with the result of the two parts
ok=0
for i in range(0,10000):
    if np.argmax(predictions[i])!=np.argmax(predictions2[i]):
        print(i," False")
    else:
        ok=ok+1
print(ok)
# also sample some cases and check probabilities
print("last entry")
print(predictions2[9999])
print(predictions[9999])
print("2nd entry")
print(predictions2[1])
print(predictions[1])