[Tensorflow/Keras] Example how to measure average time taken per batch

main.py

# -*- coding: utf-8 -*-
"""Example how to measure average time taken per batch.
"""

import time

import numpy as np
import pandas as pd
from tensorflow.keras.models import Model
from tensorflow.keras.layers import Input, Dense
from tensorflow.keras.callbacks import Callback


EPOCHS = 1000
BATCH_SIZE = 4


class MeasureTime(Callback):
    """Measure average time taken per batch.
    """

    def __init__(self):
        self.time_start = None
        self.time_batches = []
        self.time_epochs = []

    def on_train_batch_begin(self, batch, logs=None):
        self.time_start = time.time()

    def on_train_batch_end(self, batch, logs=None):
        time_end = time.time()
        self.time_batches.append(time_end - self.time_start)

    def on_epoch_end(self, epoch, logs=None):
        # NOTE: second to millisecond multiply by 1000
        batch_avgtime = np.array(self.time_batches).mean() * 1000
        print("EPOCH:{:04d}: {:.02f}ms per batch".format(epoch, batch_avgtime))

        self.time_epochs.append(batch_avgtime)
        self.time_batches = []

    def on_train_end(self, logs=None):
        epochs_avgtime = np.array(self.time_epochs).mean()
        print("Average time taken per batch: {:.02f}ms".format(epochs_avgtime))


class FFANN_XOR:
    """Solving XOR gate with Feed-Forward Artificial Neural Network.
    """

    def __init__(self):
        self.dataset = self._generate_xor()

    def train_and_evaluate(self, callbacks):
        """Train and then evaluate

        Parameters
        ----------
        callbacks : list
            A list of ``tensorflow.keras.callbacks``.
        """

        input_data = self.dataset[['input1', 'input2']].to_numpy(dtype=np.float)
        output_data = np.reshape(self.dataset['output'].to_numpy(dtype=np.float), (-1, 1))

        model = self._create_model(input_data.shape[1], output_data.shape[1])

        model.fit(
            input_data,
            output_data,
            callbacks=callbacks,
            epochs=EPOCHS,
            shuffle=True,
            batch_size=BATCH_SIZE,
            verbose=0
        )

        test_loss = model.evaluate(input_data, output_data, verbose=0)

        print('Test test_loss: {}'.format(test_loss))
        print('Predicted output = \n{}'.format(model.predict(input_data)))
        print('Expected output = \n{}'.format(output_data))

    @staticmethod
    def _create_model(n_inputs, n_labels):
        """Create a simple FFANN model.

        Parameters
        ----------
        n_inputs : int
            Number of input neurons.
        n_labels : int
            Number of output neurons.

        Returns
        -------
        ``tensorflow.keras.models.Model``
            The created FFANN model.
        """

        input_layer = Input(shape=(n_inputs,), name='input_layer')
        hidden_layer = Dense(n_inputs * 16, activation='sigmoid', name='hidden_layer_1')(input_layer)
        hidden_layer = Dense(n_inputs * 16, activation='sigmoid', name='hidden_layer_2')(hidden_layer)
        hidden_layer = Dense(n_inputs * 16, activation='sigmoid', name='hidden_layer_3')(hidden_layer)
        output_layer = Dense(n_labels, activation='sigmoid')(hidden_layer)

        model = Model(input_layer, output_layer, name='ffann')
        model.compile(optimizer='adam', loss='mean_squared_error')
        model.summary()

        return model

    @staticmethod
    def _generate_xor():
        """Generate XOR gate dataset.

        Returns
        -------
        pd.DataFrame
            Containing XOR gate dataset.
        """

        df = pd.DataFrame(columns=['input1', 'input2', 'output'])
        df = df.append({'input1': 0, 'input2': 0, 'output': 0}, ignore_index=True)
        df = df.append({'input1': 1, 'input2': 0, 'output': 1}, ignore_index=True)
        df = df.append({'input1': 0, 'input2': 1, 'output': 1}, ignore_index=True)
        df = df.append({'input1': 1, 'input2': 1, 'output': 0}, ignore_index=True)

        return df


def main():
    ffann = FFANN_XOR()
    ffann.train_and_evaluate([MeasureTime()])


if __name__ == '__main__':
    main()