cifar100-training-example.py

import tensorflow as tf
from keras.callbacks import TensorBoard
from keras.datasets import cifar100
from keras.layers import Dense, Conv2D, MaxPooling2D, GlobalAveragePooling2D, BatchNormalization, Activation, Dropout
from keras.layers import Input
from keras.models import Model
from sacred import Experiment

ex = Experiment('tb-efficiency')


@ex.config
def my_config():
    epochs = 10
    batch_size = 256
    device = '/gpu:0'


def conv2d_bn(x, filters, kernel_size=(3, 3), dropout=0.2):
    y = Conv2D(filters, kernel_size, use_bias=False)(x)
    y = BatchNormalization()(y)
    y = Activation('relu')(y)
    return y


@ex.automain
def main(epochs, batch_size, device):
    (x_train, y_train), (x_test, y_test) = cifar100.load_data()
    x_train, x_test = (x.astype(float) / 127.0 - 1
                       for x in (x_train, x_test))

    with tf.device(device):
        x = Input(shape=(32, 32, 3))
        y = conv2d_bn(x, 32)
        y = conv2d_bn(y, 32)
        y = MaxPooling2D()(y)
        y = conv2d_bn(y, 64)
        y = conv2d_bn(y, 64)
        y = MaxPooling2D()(y)
        y = conv2d_bn(y, 128)
        y = conv2d_bn(y, 128)
        y = GlobalAveragePooling2D()(y)
        y = Dropout(rate=0.5)(y)
        y = Dense(100, activation='softmax')(y)

        model = Model(x, y)
        model.compile(optimizer='adam',
                      loss='sparse_categorical_crossentropy',
                      metrics=['accuracy'])

        model.fit(x_train, y_train,
                  epochs=epochs,
                  batch_size=batch_size,
                  validation_split=1 / 3,
                  callbacks=[
                      TensorBoard(histogram_freq=1,
                                  batch_size=batch_size,
                                  write_grads=True)
                  ],
                  verbose=2)