Search code examples
kerasautoencoderregularized

How to create an autoencoder where each layer of encoder should represent the same as a layer of the decoder


I want to build an autoencoder where each layer in the encoder has the same meaning as a correspondent layer in the decoder. So if the autoencoder is perfectly trained, the values of those layers should be roughly the same.

So lets say the autoencoder consists of e1 -> e2 -> e3 -> d2 -> d1, whereas e1 is the input and d1 the output. A normal autoencoder trains to have the same result in d1 as e1, but I want the additional constraint, that e2 and d2 are the same. Therefore I want an additional backpropagation path which leads from d2 to e2 and trains at the same time as the normal path from d1 to e1. (d stands for decoder, e for encoder).

I tried to use the error between e2 and d2 as a regularization term with the CustomRegularization layer from the first answer of this link https://github.com/keras-team/keras/issues/5563. I also use this for the error between e1 and d1 instead of the normal path.

The following code is written such that more than 1 intermediate layer can be handled and also uses 4 layers. In the out commented code is a normal autoencoder which only propagates from start to end.

from keras.layers import Dense
import numpy as np
from keras.datasets import mnist
from keras.models import Model
from keras.engine.topology import Layer
from keras import objectives
from keras.layers import Input
import keras
import matplotlib.pyplot as plt


#A layer which can be given as an output to force a regularization term between two layers
class CustomRegularization(Layer):
    def __init__(self, **kwargs):
        super(CustomRegularization, self).__init__(**kwargs)

    def call(self, x, mask=None):
        ld=x[0]
        rd=x[1]
        bce = objectives.binary_crossentropy(ld, rd)
        loss2 = keras.backend.sum(bce)
        self.add_loss(loss2, x)
        return bce

    def get_output_shape_for(self, input_shape):
        return (input_shape[0][0],1)


def zero_loss(y_true, y_pred):
    return keras.backend.zeros_like(y_pred)

#Create regularization layer between two corresponding layers of encoder and decoder
def buildUpDownRegularization(layerNo, input, up_layers, down_layers):
    for i in range(0, layerNo):
        input = up_layers[i](input)
    start = input
    for i in range(layerNo, len(up_layers)):
        input = up_layers[i](input)

    for j in range(0, len(down_layers) - layerNo):
        input = down_layers[j](input)
    end = input
    cr = CustomRegularization()([start, end])
    return cr


# Define shape of the network, layers, some hyperparameters and training data
sizes = [784, 400, 200, 100, 50]
up_layers = []
down_layers = []
for i in range(1, len(sizes)):
    layer = Dense(units=sizes[i], activation='sigmoid', input_dim=sizes[i-1])
    up_layers.append(layer)
for i in range(len(sizes)-2, -1, -1):
    layer = Dense(units=sizes[i], activation='sigmoid', input_dim=sizes[i+1])
    down_layers.append(layer)

batch_size = 128
num_classes = 10
epochs = 100
(x_train, y_train), (x_test, y_test) = mnist.load_data()

x_train = x_train.astype('float32')
x_test = x_test.astype('float32')
x_train /= 255
x_test /= 255
x_train = x_train.reshape([x_train.shape[0], 28*28])
x_test = x_test.reshape([x_test.shape[0], 28*28])


y_train = x_train
y_test = x_test

optimizer = keras.optimizers.Adam(lr=0.001, beta_1=0.9, beta_2=0.999, epsilon=None, decay=0.0, amsgrad=False)



"""
### Normal autoencoder like in base mnist example
model = keras.models.Sequential()
for layer in up_layers:
    model.add(layer)
for layer in down_layers:
    model.add(layer)

model.compile(optimizer=optimizer, loss=keras.backend.binary_crossentropy)
model.fit(x_train, y_train, batch_size=batch_size, epochs=epochs)

score = model.evaluate(x_test, y_test, verbose=0)
#print('Test loss:', score[0])
#print('Test accuracy:', score[1])


decoded_imgs = model.predict(x_test)


n = 10  # how many digits we will display
plt.figure(figsize=(20, 4))
for i in range(n):
    # display original
    ax = plt.subplot(2, n, i + 1)
    plt.imshow(x_test[i].reshape(28, 28))
    plt.gray()
    ax.get_xaxis().set_visible(False)
    ax.get_yaxis().set_visible(False)

    # display reconstruction
    ax = plt.subplot(2, n, i + 1 + n)
    plt.imshow(decoded_imgs[i].reshape(28, 28))
    plt.gray()
    ax.get_xaxis().set_visible(False)
    ax.get_yaxis().set_visible(False)
plt.show()

"""

### My autoencoder where each subpart is also an autoencoder

#This part is only because the model needs a path from start to end, contentwise this should do nothing
output = input = Input(shape=(sizes[0],))
for i in range(0, len(up_layers)):
    output = up_layers[i](output)
for i in range(0, len(down_layers)):
    output = down_layers[i](output)
crs = [output]
losses = [zero_loss]

#Build the regularization layer
for i in range(len(up_layers)):
    crs.append(buildUpDownRegularization(i, input, up_layers, down_layers))
    losses.append(zero_loss)


#Create and train model with adapted training data
network = Model([input], crs)
optimizer = keras.optimizers.Adam(lr=0.0001, beta_1=0.9, beta_2=0.999, epsilon=None, decay=0.0, amsgrad=False)
network.compile(loss=losses, optimizer=optimizer)

dummy_train = np.zeros([y_train.shape[0], 1])
dummy_test = np.zeros([y_test.shape[0], 1])

training_data = [y_train]
test_data = [y_test]

for i in range(len(network.outputs)-1):
    training_data.append(dummy_train)
    test_data.append(dummy_test)


network.fit(x_train, training_data, batch_size=batch_size, epochs=epochs,verbose=1, validation_data=(x_test, test_data))
score = network.evaluate(x_test, test_data, verbose=0)
print('Test loss:', score[0])
print('Test accuracy:', score[1])

decoded_imgs = network.predict(x_test)


n = 10  # how many digits we will display
plt.figure(figsize=(20, 4))
for i in range(n):
    # display original
    ax = plt.subplot(2, n, i + 1)
    plt.imshow(x_test[i].reshape(28, 28))
    plt.gray()
    ax.get_xaxis().set_visible(False)
    ax.get_yaxis().set_visible(False)

    # display reconstruction
    ax = plt.subplot(2, n, i + 1 + n)
    plt.imshow(decoded_imgs[0][i].reshape(28, 28))
    plt.gray()
    ax.get_xaxis().set_visible(False)
    ax.get_yaxis().set_visible(False)
plt.show()

If you run the code as is it will show, that the reproduction ability is no longer given in my code. I expect a similar behavior to the uncommented code, which shows a normal autoencoder.

Edit: As mentioned in the answers this works well with MSE instead of crossentropy and a lr of .01. 100 epochs with that setting produce really good results.

Edit 2: I would like that the backpropagation works as in this [image] (https://i.sstatic.net/nAGA9.jpg). So the backpropagation of the loss of a certain layer stops at the corresponding layer. I think I didn't make this clear before and I don't know if the code currently does that.

Edit 3: Although this code runs and returns a good looking solution the CustomRegularization layer is not doing what I thought it would do, therefore it does not do the same things as in the description.


Solution

  • It seems like the main issue is the use of binary cross-entropy to minimize the difference between encoder and decoder. The internal representation in the network is not going to be a single class probability like the output might be if you were classifying MNIST digits. I was able to get your network to output some reasonable-looking reconstructions with these simple changes:

    1. Using objectives.mean_squared_error instead of objectives.binary_crossentropy in the CustomRegularization class

    2. Changing number of epochs to 5

    3. Changing learning rate to .01

    Changes 2 and 3 were simply made to speed up the testing. Change 1 is the key here. Cross entropy is designed for problems where there is a binary "ground truth" variable and an estimate of that variable. However, you do not have a binary truth value in the middle of your network, only at the output layer. Thus a cross entropy loss function in the middle of the network doesn't make much sense (at least to me) -- it will be trying to measure entropy for a variable that isn't binary. Mean squared error, on the other hand, is a bit more generic and should work for this case since you are simply minimizing the difference between two real values. In essence, the middle of the network is performing regression (difference between activations in two continuous values, i.e. layers), not classification, so it needs a loss function that is appropriate for regression.

    I also want to suggest that there may be a better approach to accomplish what you want. If you really want the encoder and decoder to be exactly the same, you can share weights between them. Then they will be identical, not just highly similar, and your model will have fewer parameters to train. There is a decent explanation of shared (tied) weights autoencoders with Keras here if you're curious.

    Reading your code it does seem like it is doing what you want in your illustration, but I am not really sure how to verify that.