Treating a SubGraph of a Neural Network as a Model In TensorFlow/Keras

Question

I am trying to train an auto encoder in tensorflow using the Keras Layer API. This API is quite nice and easy to use to setup the deep learning layers.

Just to review a quickly an autoencoder (in my mind) is a function $f(x) = z$ and its pseudo inverse \\hat{x} = f^{-1}(z) such that f(f^{-1}(x)) \\approx x. In a neural network model, you would setup a neural network with a bottleneck layer that tries to predict itself x using f^{-1}(f(x)). When the training error minimizes, you then have two components, z = f(x) is the prediction up until and including the bottleneck layer. f^{-1}(z) is the bottleneck layer to the end.

So I setup the encoder:

SZ = 6
model = tf.keras.Sequential()

model.add(layers.InputLayer(SZ))
model.add(layers.Dense(SZ))
model.add(layers.Dense(1))
model.add(layers.Dense(SZ))
model.summary()

model.compile('sgd','mse',metrics = ['accuracy'])
history= model.fit(returns.values,returns.values,epochs=100)

My difficulty here is that the weights and components (f being input+dense(SZ)+dense(1),f^{-1} being dense(1)+dense(SZ)) are trained but I do not know how to disentangle them. Is there some way to break off the two layers in the neural network and have them treated as their own separate models?

Answer 1

import tensorflow as tf
SZ=6
encoder_input = tf.keras.layers.Input(shape=(SZ,))
x = tf.keras.layers.Dense(SZ)(encoder_input)
x = tf.keras.layers.Dense(1)(x)
encoder_model = tf.keras.Model(inputs=encoder_input, outputs=x, name='encoder')

decoder_input = tf.keras.layers.Input(shape=(1,))
x2 = tf.keras.layers.Dense(SZ)(decoder_input)
decoder_model = tf.keras.Model(inputs=decoder_input, outputs=x2, name='decoder')

encoder_output = encoder_model(encoder_input)
decoder_output = decoder_model(encoder_output)

encoder_decoder_model = tf.keras.Model(inputs=encoder_input , outputs=decoder_output, name='encoder-decoder')
encoder_decoder_model.summary()

Here is the summary:

Model: "encoder-decoder"
_________________________________________________________________
Layer (type)                 Output Shape              Param #   
=================================================================
input_8 (InputLayer)         [(None, 6)]               0         
_________________________________________________________________
encoder (Model)              (None, 1)                 49        
_________________________________________________________________
decoder (Model)              (None, 6)                 12        
=================================================================
Total params: 61
Trainable params: 61
Non-trainable params: 0

you could train the encoder-decoder model and you separate encoder_model and decoder_model will be trained automatically. You could also retrieve them from your encoder_decoder model as follows:

retrieved_encoder = encoder_decoder_model.get_layer('encoder')
retrieved_encoder.summary()

it prints:

Model: "encoder"
_________________________________________________________________
Layer (type)                 Output Shape              Param #   
=================================================================
input_8 (InputLayer)         [(None, 6)]               0         
_________________________________________________________________
dense_11 (Dense)             (None, 6)                 42        
_________________________________________________________________
dense_12 (Dense)             (None, 1)                 7         
=================================================================
Total params: 49
Trainable params: 49
Non-trainable params: 0

and the decoder:

retrieved_decoder = encoder_decoder_model.get_layer('decoder')
retrieved_decoder.summary()

which prints:

Model: "decoder"
_________________________________________________________________
Layer (type)                 Output Shape              Param #   
=================================================================
input_9 (InputLayer)         [(None, 1)]               0         
_________________________________________________________________
dense_13 (Dense)             (None, 6)                 12        
=================================================================
Total params: 12
Trainable params: 12
Non-trainable params: 0