How to train multiple stacked neural networks alternatively? (Keras)

Viewed 433

Suppose I'm defining four neural network models with their own loss functions. The input of the next neural network is dependent on the output of the previous network.

Model1 -> Model2 -> Model3 -> Model4

For simplicity, let the four neural networks be shown below:

from tensorflow import keras
from tensorflow.keras import layers
import numpy as np

x = np.random.rand(300,10)
y = np.random.rand(300,1)

# Model 1
inputs1 = keras.Input(shape=(10))
x1 = layers.Dense(1, activation='relu')(inputs1)
model1 = keras.Model(inputs=inputs1, outputs=x1)
model1.compile(loss="mean_squared_error")

# Model 2
inputs2 = keras.Input(shape=(1))
x2 = layers.Dense(1, activation='relu')(inputs2)
model2 = keras.Model(inputs=inputs2, outputs=x2)
model2.compile(loss="mean_squared_error")

# Model 3
inputs3 = keras.Input(shape=(1))
x3 = layers.Dense(1, activation='relu')(inputs3)
model3 = keras.Model(inputs=inputs3, outputs=x3)
model3.compile(loss="mean_squared_error")

# Model 4
inputs4 = keras.Input(shape=(1)
x4 = layers.Dense(1, activation='relu')(inputs4)
model4 = keras.Model(inputs=inputs4, outputs=x4)
model4.compile(loss="mean_squared_error")

I would like to stack and train these models, training each model at a time while keeping the weights frozen in the remainder networks.

So in one training iteration, Model1 would be trained then the output would be passed to Model2. Model2 would be trained and then the output passed to Model3 and so forth and then another epoch would be started.

Any ideas?

1 Answers

According to your query, there are two requirements:

  1. The output of one model would be the input of the subsequent model
  2. Train each model at a time with concern loss and metrics etc.

Below is the schematic diagram of what we need to achieve. Let's consider we have 3 models (x, y, and z) where the first model's output would be the input for the next model and go on. Each of these 3 models is also self trainable. In our demonstration, we will be working with an Image dataset (mnist).

enter image description here


Let's try to achieve this.


  1. We will demonstrate one way to achieve this but hopefully, there should be a much more convenient way. We will use the MNIST data set and 3 models for example, model_1, model_2, and model_3. Though these models will individually train on MNIST's x_train and y_train (softmax), but we will get the proper feature maps from model_1 and set it as input for model_2 and so on.
from tensorflow.keras import backend as K
from tensorflow.keras import layers
from tensorflow import keras
import tensorflow as tf
import numpy as np
# Model 1
def model_1(input = (28,28,1)):
    inputs_1 = keras.Input(shape=(input), name="img1")
    x = layers.Conv2D(16, 3, activation="relu", name='conv1_model1')(inputs_1)
    x = layers.Conv2D(32, 3, activation="relu", name='conv2_model1')(x)
    x = layers.GlobalAveragePooling2D(name='gap1')(x)
    x = layers.Dense(10, activation='softmax', name='pred_1')(x)
    return keras.Model(inputs=inputs_1, outputs=x, name='model_1')

# Model 2
def model_2(input):
    inputs_2 = keras.Input(shape=(input), name='img2')
    x = layers.Conv2D(32, 3, activation="relu", name='conv1_model2')(inputs_2)
    x = layers.Conv2D(16, 3, activation="relu", name='conv2_model2')(x)
    x = layers.GlobalAveragePooling2D(name='gap2')(x)
    x = layers.Dense(10, activation='softmax', name='pred_2')(x)
    return keras.Model(inputs=inputs_2, outputs=x, name='model_2')

# Model 3
def model_3(input):
    inputs_2 = keras.Input(shape=(input), name='img3')
    x = layers.Conv2D(32, 3, activation="relu", name='conv1_model3')(inputs_2)
    x = layers.Conv2D(16, 3, activation="relu", name='conv2_model3')(x)
    x = layers.GlobalAveragePooling2D(name='gap3')(x)
    x = layers.Dense(10, activation='softmax', name='pred_3')(x)
    return keras.Model(inputs=inputs_2, outputs=x, name='model_3')

Now, we will call model_1 and take one appropriate feature map and pass it to model_2 as its input and so for the model_3.

x = model_1()
y = model_2(x.get_layer('conv2_model1').output_shape[1:]) # feat_map: (24, 24, 32)
z = model_3(y.get_layer('conv2_model2').output_shape[1:]) # feat_map: (20, 20, 16)

print(x.get_layer('conv2_model1').output_shape[1:])
print(y.get_layer('conv2_model2').output_shape[1:])
print(x.output_shape, y.output_shape, z.output_shape)
# (24, 24, 32)
# (20, 20, 16)
# (None, 10) (None, 10) (None, 10)
  1. Now we will train each model (x, y, and z) at a time. But before that here is something that needs to understand about what we've done so far. We build model x and its input would be (as we mentioned before) 28 x 28 x 1. And model y, its input is the output feature maps of layer conv2_model1 which comes from model x. In tf. keras we can do that in the following way:
# ---------- Data part ------------
# Let's first grab the data
(x_train, y_train), (_, _) = tf.keras.datasets.mnist.load_data()
x_train = tf.expand_dims(x_train, -1)
x_train = tf.divide(x_train, 255)
y_train = tf.one_hot(y_train , depth=10) 

data_x, data_y = x_train[:10], y_train[:10]
print(data_x.shape, data_y.shape)
# (10, 28, 28, 1) (10, 10)
# ---------- Modeling part ------------
x.summary()
_________________________________________________________________
Layer (type)                 Output Shape              Param #   
=================================================================
img1 (InputLayer)            [(None, 28, 28, 1)]       0         
_________________________________________________________________
conv1_model1 (Conv2D)        (None, 26, 26, 16)        160       
_________________________________________________________________
conv2_model1 (Conv2D)        (None, 24, 24, 32)        4640      
_________________________________________________________________
gap1 (GlobalAveragePooling2D (None, 32)                0         
_________________________________________________________________
pred_1 (Dense)               (None, 10)                330       
=================================================================
# model x compiling and calling fit method on it 
x.compile(
    loss = tf.keras.losses.CategoricalCrossentropy(),
    metrics = ['accuracy'],
    optimizer = tf.keras.optimizers.Adam())
x.fit(data_x, data_y, epochs=1)
# 5ms/step - loss: 2.2842 - accuracy: 0.2000

# now, we need output feature maps of layer `conv2_model1` from model `x`
# to do that, we first build a model that will give output of such model. 
# however, this was also possible to build mulit-output `x` model that 
# would give (None, 10) and output of this layer (`conv2_model1`).
pred_x_model = tf.keras.Model(x.input, x.get_layer('conv2_model1').output)

# next we pass input (data_x) to these pred_x_model 
# to get the feature maps of layer `conv2_model1`
pred_x = pred_x_model(data_x)
print(pred_x.shape)
# (10, 24, 24, 32)

As we can see how we've achieved the output feature maps of model x, especially its layer conv2_model1 (24, 24, 32). Now, we can use these feature maps as an input to the next model, which is model y. And also after training model y, we can get another feature map for another next model, in our case which is model z.

y.summary()
_________________________________________________________________
Layer (type)                 Output Shape              Param #   
=================================================================
img2 (InputLayer)            [(None, 24, 24, 32)]      0         
_________________________________________________________________
conv1_model2 (Conv2D)        (None, 22, 22, 32)        9248      
_________________________________________________________________
conv2_model2 (Conv2D)        (None, 20, 20, 16)        4624      
_________________________________________________________________
gap2 (GlobalAveragePooling2D (None, 16)                0         
_________________________________________________________________
pred_2 (Dense)               (None, 10)                170       
=================================================================

# calling model y and 
# using `pred_x` for input as it should be now for model `y`
y.compile(
    loss = tf.keras.losses.CategoricalCrossentropy(),
    metrics = ['accuracy'],
    optimizer = tf.keras.optimizers.Adam())
y.fit(pred_x, data_y, epochs=1)
# 8ms/step - loss: 2.0186 - accuracy: 0.3000

# grab the proper feature maps from model `y`'s layer (`conv2_model2`) which 
# would be the input for model `z`.
pred_y_model = tf.keras.Model(y.input, y.get_layer('conv2_model2').output)

pred_y = pred_y_model(pred_x)
print(pred_y.shape)
# (10, 20, 20, 16)

So, we have pred_y, a feature map of size (20, 20, 16) and ready to set as input for model z.

z.summary()
_________________________________________________________________
Layer (type)                 Output Shape              Param #   
=================================================================
img3 (InputLayer)            [(None, 20, 20, 16)]      0         
_________________________________________________________________
conv1_model3 (Conv2D)        (None, 18, 18, 32)        4640      
_________________________________________________________________
conv2_model3 (Conv2D)        (None, 16, 16, 16)        4624      
_________________________________________________________________
gap3 (GlobalAveragePooling2D (None, 16)                0         
_________________________________________________________________
pred_3 (Dense)               (None, 10)                170       
=================================================================

# calling model z 
# set `pred_y` as its input
z.compile(
    loss = tf.keras.losses.CategoricalCrossentropy(),
    metrics = ['accuracy'],
    optimizer = tf.keras.optimizers.Adam())
z.fit(pred_y, data_y, epochs=1)
# 5ms/step - loss: 1.9422 - accuracy: 0.3000

# as we have 3 models, our model cycling end here.
# we take the output from model z's last softmax layer
pred_z = z(pred_y)
print(pred_z.shape)
# (10, 10)

That's it.

Related