digit.py

"""
    The digit caps model in the paper.
"""

import numpy as np
from keras.models import Sequential, Model
from keras.datasets import mnist
from keras.utils import to_categorical
from keras.layers import Input, Conv2D, Dense, Reshape
from keras_capsnet import margin_loss, conv2d_caps, DenseCapsule, Mask, CapsuleLength

img_dim = (28, 28, 1)
output_shape = 10


def load_mnist_data():

    (x_train, y_train), (x_test, y_test) = mnist.load_data()
    x_train = x_train.reshape(-1, 28, 28, 1).astype('float32')/255.0
    x_test = x_test.reshape(-1, 28, 28, 1).astype('float32')/255.0
    y_train = to_categorical(y_train.astype('float32'))
    y_test = to_categorical(y_test.astype('float32'))

    return x_train, y_train, x_test, y_test


def build_models():
    global img_dim, output_shape

    input_layer = Input(shape=(img_dim[0], img_dim[1], img_dim[2]))
    conv1 = Conv2D(filters=256, kernel_size=9, strides=1, padding='valid', activation='relu')(input_layer)
    primary_caps = conv2d_caps(
        input_layer=conv1,
        nb_filters=32,
        capsule_size=8,
        kernel_size=9,
        strides=2
    )
    digit_caps = DenseCapsule(capsule_size=16, nb_capsules=output_shape, name='densecaps')(primary_caps)
    output_layer = CapsuleLength(name='capsnet')(digit_caps)

    # Decoder part..
    true_labels_input = Input(shape=(output_shape,))  # use the true label to mask the capsule
    masked_by_true_label = Mask()([digit_caps, true_labels_input])  # for training process
    masked_by_max = Mask()(digit_caps)  # for prediction process

    decoder = Sequential(name='decoder_out')
    decoder.add(Dense(512, activation='relu', input_dim=16*output_shape))
    decoder.add(Dense(1024, activation='relu'))
    decoder.add(Dense(np.prod(img_dim), activation='sigmoid'))
    decoder.add(Reshape(target_shape=img_dim))

    train_decoder = decoder(masked_by_true_label)
    eval_decoder = decoder(masked_by_max)

    train_model = Model(
        inputs=[input_layer, true_labels_input],
        outputs=[output_layer, train_decoder]
    )

    eval_model = Model(
        inputs=input_layer,
        outputs=[eval_decoder, output_layer]
    )

    return train_model, eval_model


def main():

    x_train, y_train, x_test, y_test = load_mnist_data()
    train_model, eval_model = build_models()
    print(train_model.summary())
    train_model.compile(loss=[margin_loss, 'mse'], optimizer='adam')

    train_model.fit([x_train, y_train], [y_train, x_train], batch_size=128, epochs=5, validation_data=[[x_test, y_test], [y_test, x_test]])


if __name__ == '__main__':
    main()