from __future__ import print_function

import time

import numpy as np
from keras.callbacks import Callback

class BatchTiming(Callback):
    """
    It measure robust stats for timing of batches and epochs.
    Useful for measuring the training process.

    For each epoch it prints median batch time and total epoch time.
    After training it prints overall median batch time and median epoch time.

    Usage: model.fit(X_train, Y_train, callbacks=[BatchTiming()])

    All times are in seconds.

    More info: https://keras.io/callbacks/
    """
    def on_train_begin(self, logs={}):
        self.all_batch_times = []
        self.all_epoch_times = []

    def on_epoch_begin(self, epoch, logs={}):
        self.epoch_batch_times = []

    def on_batch_begin(self, batch, logs={}):
        self.start_time = time.time()

    def on_batch_end(self, batch, logs={}):
        end_time = time.time()
        elapsed_time = end_time - self.start_time
        self.epoch_batch_times.append(elapsed_time)
        self.all_batch_times.append(elapsed_time)

    def on_epoch_end(self, epoch, logs={}):
        epoch_time = np.sum(self.epoch_batch_times)
        self.all_epoch_times.append(epoch_time)
        median_batch_time = np.median(self.epoch_batch_times)
        print('Epoch timing - batch (median): %0.5f, epoch: %0.5f (sec)' % \
            (median_batch_time, epoch_time))

    def on_train_end(self, logs={}):
        median_batch_time = np.median(self.all_batch_times)
        median_epoch_time = np.median(self.all_epoch_times)
        print('Overall - batch (median): %0.5f, epoch (median): %0.5f (sec)' % \
            (median_batch_time, median_epoch_time))

class SamplesPerSec(Callback):
    def __init__(self, batch_size):
        self.batch_size = batch_size

    def on_train_begin(self, logs={}):
        self.all_samples_per_sec = []

    def on_batch_begin(self, batch, logs={}):
        self.start_time = time.time()
        # self.batch_size = logs['size']

    def on_batch_end(self, batch, logs={}):
        end_time = time.time()
        elapsed_time = end_time - self.start_time
        samples_per_sec = self.batch_size / elapsed_time
        self.all_samples_per_sec.append(samples_per_sec)

    def on_epoch_end(self, epoch, logs={}):
        self.print_results()

    def print_results(self):
        print('Samples/sec: %0.2f' % np.median(self.all_samples_per_sec))

@@ -2396,8 +2396,1024 @@
 r_sec))%0A
+%0A%22%22%22%0AEnables CUDA profiling (for usage in nvprof) just for a few batches.%0A%0AThe reasons are:%0A%0A- profiling outputs are big (easily 100s MB - GBs) and repeating%0A- without a proper stop the outputs sometimes fail to save%0A%0ASince initially the TensorFlow runtime may take time to optimize the graph we%0Askip a few epochs and then enable profiling for a few batches within the next%0Aepoch.%0A%0AIt requires the %60cudaprofile%60 package.%0A%22%22%22%0Aclass CudaProfile(Callback):%0A    import cudaprofile%0A%0A    def __init__(self, warmup_epochs=0, batches_to_profile=None):%0A        self.warmup_epochs = warmup_epochs%0A        self.batches_to_profile = batches_to_profile%0A        self.enabled = False%0A%0A    def set_params(self, params):%0A        self.params = params%0A%0A    def on_epoch_begin(self, epoch, logs=%7B%7D):%0A        if epoch == self.warmup_epochs:%0A            cudaprofile.start()%0A            self.enabled = True%0A%0A    def on_batch_end(self, batch, logs=%7B%7D):%0A        if self.enabled && batch %3E= batches_to_profile:%0A            cudaprofile.stop()%0A

Add a callback for enabling CUDA profiling just for a few batches.