Keras Functional API Example

functional.py

# Resource: https://keras.io/getting-started/functional-api-guide/

import keras.backend as K
from keras.datasets import mnist
from keras.layers import Dense, Input, Lambda
from keras.models import Model
from keras.optimizers import Adam
from keras.utils import to_categorical
import numpy as np

# Create an input layer, which allocates a tf.placeholder tensor.
input_tensor = Input(shape = (28, 28))

# I could use a Keras Flatten layer like this.
# from keras.layers import Flatten
# flattened_input = Flatten()(input_tensor)
#
# But I will use the Keras interface to low-level TF. To do this, I
# create a "Lambda" layer. Lambda layers are custom layers in which
# I can do TF stuff.
flatten_layer = Lambda(
    lambda ipt: K.reshape(ipt, (-1, 28 * 28))
)
flattened_input = flatten_layer(input_tensor)

# Create a dense layer, and feed it in the input_tensor.
# Note: it is more normal to write Dense(10, activation = 'relu')(flattened_input)
hidden_layer = Dense(10, activation = 'relu')
hidden_output = hidden_layer(flattened_input)
output_layer = Dense(10, activation = 'softmax')
classification_output = output_layer(hidden_output)

# We specify what tensors we will be fed in, and which tensors are the
# final outputs of the model.
model = Model([input_tensor], [classification_output])

# Build the optimizer
LEARNING_RATE = 0.001
optimizer = Adam(lr = LEARNING_RATE)

# "Custom" loss function
def categorical_cross_entropy(y_true, y_pred):
    # Let's see what comes into this function! They should be TF
    # Tensor objects!
    print(y_true)
    print(y_pred)

    # keras.backend is conventionally imported as K.
    # It is the Keras abstraction of TensorFlow; you're basically using
    # a nicer interface for TF here. This is lower level than the Keras
    # layers API. K functions typically manipulate TF tensors directly.
    #
    # Note I must use K.epsilon() because taking the log of a zero
    # probability gives NaN. This adds a small nudge so nothing is zero.
    # This problem wouldn't happen if I worked in logits, but I used
    # softmax at the final layer.
    return K.mean(
        K.sum(y_true * -K.log(y_pred + K.epsilon()), axis = 1)
    )

model.compile(
    # This says to use our CE function. We could have just said 'categorical_cross_entropy'
    # because Keras is nice and will figure that out for us. This is
    # more manual in case we want to have a fancy loss function.
    loss = [categorical_cross_entropy],
    # You can have multiple outputs, in which case you can specify
    # multiple loss functions. Of course, we have only one output here.
    loss_weights = [1.0],
    optimizer = optimizer,
    metrics = ['accuracy'],
)

# This just downloads the datasets.
(x_train, y_train), (x_test, y_test) = mnist.load_data()

x_mean = np.mean(x_train)
x_stddev = np.std(x_train)

# Normalize x values.
x_train = (x_train - x_mean) / x_stddev
x_test = (x_test - x_mean) / x_stddev

# One hot encode.
y_train = to_categorical(y_train, 10)
y_test = to_categorical(y_test, 10)

model.fit(
    x_train,
    y_train,
    validation_data = (x_test, y_test),
    epochs = 3
)