Support for fit_generator()

[rohitrawat]

Unlike the fit() method, fit_generator() does not have a 'batch_size' parameter defined. This produces a KeyError: 'batch_size' error on line 81 of tqdm_callback.py:

self.batch_count = int(ceil(self.params['nb_sample'] / self.params['batch_size']))

fit_generator() is used instead of fit() when reading image directories and augmenting images on the fly.

Here is a minimum working example that works when not using ', verbose=0, callbacks=[TQDMCallback()]', but immediately fails when used with it:

'''Trains a simple convnet on the MNIST dataset.
Gets to 99.25% test accuracy after 12 epochs
(there is still a lot of margin for parameter tuning).
16 seconds per epoch on a GRID K520 GPU.
'''

from __future__ import print_function
import numpy as np
np.random.seed(1337)  # for reproducibility

from keras.datasets import mnist
from keras.models import Sequential
from keras.layers import Dense, Dropout, Activation, Flatten
from keras.layers import Convolution2D, MaxPooling2D
from keras.utils import np_utils
from keras import backend as K

from keras.preprocessing.image import ImageDataGenerator

from keras_tqdm import TQDMCallback, TQDMNotebookCallback

batch_size = 128
nb_classes = 10
nb_epoch = 1

# input image dimensions
img_rows, img_cols = 28, 28
# number of convolutional filters to use
nb_filters = 32
# size of pooling area for max pooling
pool_size = (2, 2)
# convolution kernel size
kernel_size = (3, 3)

# the data, shuffled and split between train and test sets
(X_train, y_train), (X_test, y_test) = mnist.load_data()

if K.image_dim_ordering() == 'th':
    X_train = X_train.reshape(X_train.shape[0], 1, img_rows, img_cols)
    X_test = X_test.reshape(X_test.shape[0], 1, img_rows, img_cols)
    input_shape = (1, img_rows, img_cols)
else:
    X_train = X_train.reshape(X_train.shape[0], img_rows, img_cols, 1)
    X_test = X_test.reshape(X_test.shape[0], img_rows, img_cols, 1)
    input_shape = (img_rows, img_cols, 1)

X_train = X_train.astype('float32')
X_test = X_test.astype('float32')
X_train /= 255
X_test /= 255
print('X_train shape:', X_train.shape)
print(X_train.shape[0], 'train samples')
print(X_test.shape[0], 'test samples')

# convert class vectors to binary class matrices
Y_train = np_utils.to_categorical(y_train, nb_classes)
Y_test = np_utils.to_categorical(y_test, nb_classes)

model = Sequential()

model.add(Convolution2D(nb_filters, kernel_size[0], kernel_size[1],
                        border_mode='valid',
                        input_shape=input_shape))
model.add(Activation('relu'))
model.add(Convolution2D(nb_filters, kernel_size[0], kernel_size[1]))
model.add(Activation('relu'))
model.add(MaxPooling2D(pool_size=pool_size))
model.add(Dropout(0.25))

model.add(Flatten())
model.add(Dense(128))
model.add(Activation('relu'))
model.add(Dropout(0.5))
model.add(Dense(nb_classes))
model.add(Activation('softmax'))

model.compile(loss='categorical_crossentropy',
              optimizer='adadelta',
              metrics=['accuracy'])

# fit() works well with or without TQDMCallback():
# model.fit(X_train, Y_train, batch_size=batch_size, nb_epoch=nb_epoch,
#          validation_data=(X_test, Y_test), verbose=0, callbacks=[TQDMCallback()])

# Using an ImageDataGenerator
datagen = ImageDataGenerator(
    featurewise_center=True,
    featurewise_std_normalization=True,
    rotation_range=20,
    width_shift_range=0.2,
    height_shift_range=0.2,
    horizontal_flip=True)

# compute quantities required for featurewise normalization
# (std, mean, and principal components if ZCA whitening is applied)
datagen.fit(X_train)

# fits the model on batches with real-time data augmentation:

#####>>> This works well without using TQDMCallback():
#model.fit_generator(datagen.flow(X_train, Y_train, batch_size=32),
#                    samples_per_epoch=len(X_train), nb_epoch=nb_epoch)

#####>>> but fails when using TQDMCallback():
model.fit_generator(datagen.flow(X_train, Y_train, batch_size=32),
                    samples_per_epoch=len(X_train), nb_epoch=nb_epoch, verbose=0, callbacks=[TQDMCallback()])

score = model.evaluate(X_test, Y_test, verbose=0)
print('Test score:', score[0])
print('Test accuracy:', score[1])

[rohitrawat]

On looking at keras/engine/training.py#L1483, 'batch_size' is not included in params when it is built by fit_generator()

        callbacks.set_params({
            'nb_epoch': nb_epoch,
            'nb_sample': samples_per_epoch,
            'verbose': verbose,
            'do_validation': do_validation,
            'metrics': callback_metrics,
        })

On the other hand, _fit_loop() which is called by fit() (keras/engine/training.py#L855) does define the 'batch_size'.

        callbacks.set_params({
            'batch_size': batch_size,
            'nb_epoch': nb_epoch,
            'nb_sample': nb_train_sample,
            'verbose': verbose,
            'do_validation': do_validation,
            'metrics': callback_metrics or [],
        })

It turns out that fit_generator() does not know the batch size as it is determined by however many samples are enqueued by the generator. To fix this, we can count samples instead of batches. The Callback documentation says that number of samples in the current batch is passed as logs['size']. I have written a fix in my fork.