fcn benchmark with flat outputs

fcn_bench_flat.py

from keras.models import Sequential
from keras.layers.core import Permute, Flatten, Layer
from keras.layers.convolutional import Convolution2D
from keras.layers.core import Activation
import theano
import theano.tensor as T
import datetime
import numpy as np

assert theano.config.optimizer_excluding == 'cudnn', \
    "CuDNN is slow. Add 'optimizer_excluding=cudnn' to THEANO_FLAGS"

now = datetime.datetime.now

if theano.config.floatX == 'float64':
    epsilon = 1.0e-9
else:
    epsilon = 1.0e-7


def categorical_crossentropy_2d(y_true, y_pred):
    y_pred = T.clip(y_pred, epsilon, 1.0 - epsilon)
    # scale preds so that the class probas of each sample sum to 1
    y_pred /= y_pred.sum(axis=-1, keepdims=True)
    cce = - (y_true * T.log(y_pred)).sum()
    return cce


def categorical_crossentropy_nn(y_true, y_pred):
    return - (y_true * T.log(y_pred)).sum()


class ChannelNormalization(Layer):
    def __init__(self, **kwargs):
        super(ChannelNormalization, self).__init__(**kwargs)

    def get_output(self, train=False):
        X = self.get_input(train)
        output = T.clip(X, epsilon, 1.0 - epsilon)
        return output / output.sum(axis=-1, keepdims=True)

    def get_config(self):
        config = {"name": self.__class__.__name__}
        base_config = super(ChannelNormalization, self).get_config()
        return dict(list(base_config.items()) + list(config.items()))


# input data
np.random.seed(123456)
n_images = 10
s = (850, 649)
x = np.random.rand(n_images, 1, s[0], s[1]).astype('float32')
y = np.random.rand(n_images, (s[0]-6) * (s[1]-6) * 21).astype('float32')

layers = [
    Convolution2D(64, 3, 3, activation='relu', border_mode='valid', input_shape=(1, s[0], s[1]), disable_cudnn=True),
    Convolution2D(64, 3, 3, activation='relu', border_mode='valid', disable_cudnn=True),
    Convolution2D(21, 3, 3, activation='relu', border_mode='valid', disable_cudnn=True),
    Permute((2, 3, 1)),
    Activation('softmax'),
    ChannelNormalization(),
    Flatten(),
]
model = Sequential()
for l in layers:
    model.add(l)

# compile
t = now()
model.compile(optimizer='adagrad', loss=categorical_crossentropy_nn)
print('Compilation time %s' % (now() - t))

# train
t = now()
model.fit(x, y, batch_size=1, nb_epoch=1)
print('Training time for %d images (batch_size=1): %s' % (n_images, (now() - t)))

t = now()
model.fit(x, y, batch_size=n_images, nb_epoch=1)
print('Training time for %d images (batch_size=%d): %s' % (n_images, n_images, (now() - t)))

# predict
t = now()
model.predict(x)
print('Prediction time for %d images: %s' % (n_images, (now() - t)))