Get decoder from trained autoencoder model in Keras

Question

I am training a deep autoencoder to map human faces to a 128 dimensional latent space, and then decode them back to its original 128x128x3 format.

I was hoping that after training the autoencoder, I would somehow be able to 'slice' the second half of the autoencoder, i.e. the decoder network responsible for mapping the latent space (128,) to the image space (128, 128, 3) by using the functional Keras API and autoenc_model.get_layer()

Here are the relevant layers of my model:

INPUT_SHAPE=(128,128,3)
input_img = Input(shape=INPUT_SHAPE, name='enc_input')

#1
x = Conv2D(64, (3, 3), padding='same', activation='relu')(input_img)
x = BatchNormalization()(x)

//Many Conv2D, BatchNormalization(), MaxPooling() layers
.
.
.

#Flatten
fc_input = Flatten(name='enc_output')(x)

y = Dropout(DROP_RATE)(fc_input)
y = Dense(128, activation='relu')(y)
y = Dropout(DROP_RATE)(y)
fc_output = Dense(128, activation='linear')(y)   

#Reshape
decoder_input = Reshape((8, 8, 2), name='decoder_input')(fc_output)

#Decoder part

#UnPooling-1
z = UpSampling2D()(decoder_input)
//Many Conv2D, BatchNormalization, UpSampling2D layers
.
.
.
#16
decoder_output = Conv2D(3, (3, 3), padding='same', activation='linear', name='decoder_output')(z)

autoenc_model = Model(input_img, decoder_output)

here is the notebook containing the entire model architecture.

To get the decodeer network from the trained autoencoder, I have tried using:

dec_model = Model(inputs=autoenc_model.get_layer('decoder_input').input, outputs=autoenc_model.get_layer('decoder_output').output)

and

dec_model = Model(autoenc_model.get_layer('decoder_input'), autoenc_model.get_layer('decoder_output'))

neither of which seem to work.

I need to extract the decoder layers out of the autoencoder as I want to train the entire autoencoder model first, then use the encoder and the decoder independently.

I could not find a satisfactory answer anywhere else. The Keras blog article on building autoencoders only covers how to extract the decoder for 2 layered autoencoders.

The decoder input/output shape should be: (128, ) and (128, 128, 3), which is the input shape of the 'decoder_input' and output shape of the 'decoder_output' layers respectively.

Answer 1

Couple of changes are needed:

z = UpSampling2D()(decoder_input)

to

direct_input = Input(shape=(8,8,2), name='d_input')
#UnPooling-1
z = UpSampling2D()(direct_input)

and

autoenc_model = Model(input_img, decoder_output)

to

dec_model = Model(direct_input, decoder_output)
autoenc_model = Model(input_img, dec_model(decoder_input))

Now, you can train on the auto encoder and predict using the decoder.

import numpy as np
autoenc_model.fit(np.ones((5,128,128,3)), np.ones((5,128,128,3)))
dec_model.predict(np.ones((1,8,8,2)))

You can also refer this self-contained example: https://github.com/keras-team/keras/blob/master/examples/variational_autoencoder.py

Answer 2

My solution isn't very elegant, and there are probably better solutions out there, but since no-one replied yet, I'll post it (I was actually hoping someone would so I can improve my own implementation, as you'll see below).

So what I did was built a network that can take a secondary input, directly into the latent space. Unfortunately, both inputs are obligatory, so I end up with a network that requires dummy arrays full of zeros for the 'unwanted' input (you'll see in a second).

Using Keras functional API:

image_input = Input(shape=image_shape)
conv1 = Conv2D(...,activation='relu')(image_input)
...
dense_encoder = Dense(...)(<layer>)
z_input = Input(shape=n_latent)
decoder_entry = Dense(...,activation='relu')(Add()([dense_encoder,z_input]))
...
decoder_output = Conv2DTranspose(...)


model = Model(inputs=[image_input,z_input], outputs=decoder_output)
model.compile(optimizer='adam', loss='binary_crossentropy', metrics=['accuracy'])

encoder = Model(inputs=image_input,outputs=dense_encoder)
decoder = Model(inputs=[z_input,image_input], outputs=decoder_output)

Note that you shouldn't compile the encoder and decoder.

(some code is either omitted or left with ... for you to fill in your specific needs).

Finally, to train you'll have to provide one empty array. So to train the entire auto-encoder:

images is X in this context

model.fit([images,np.zeros((len(n_latent),...))],images)

And then you can get the latent features using:

latent_features = encoder.predict(images)

Or use the decoder with latent input and dummy variables (note the order of inputs above):

decoder.predict([Z_inputs,np.zeros(shape=images.shape)])

Finally, another solution I haven't tried is build to parallel models, with the same architecture, one the autoencoder, and the second only the decoder part, and then use:

decoder_layer.set_weights(model_layer.get_weights()) 

It should work, but I haven't confirmed it. It does have the disadvantage of having to copy the weights again every time your train the autoencoder model.

So to conclude, I am aware of the many problems here, but again, I only posted this because I saw no-one else replied, and was hoping this will still be of some use to you.

Please comment if something is not clear.