| | from copy import copy |
| | from functools import partial |
| |
|
| | from .auto import tqdm as tqdm_auto |
| |
|
| | try: |
| | import keras |
| | except (ImportError, AttributeError) as e: |
| | try: |
| | from tensorflow import keras |
| | except ImportError: |
| | raise e |
| | __author__ = {"github.com/": ["casperdcl"]} |
| | __all__ = ['TqdmCallback'] |
| |
|
| |
|
| | class TqdmCallback(keras.callbacks.Callback): |
| | """Keras callback for epoch and batch progress.""" |
| | @staticmethod |
| | def bar2callback(bar, pop=None, delta=(lambda logs: 1)): |
| | def callback(_, logs=None): |
| | n = delta(logs) |
| | if logs: |
| | if pop: |
| | logs = copy(logs) |
| | [logs.pop(i, 0) for i in pop] |
| | bar.set_postfix(logs, refresh=False) |
| | bar.update(n) |
| |
|
| | return callback |
| |
|
| | def __init__(self, epochs=None, data_size=None, batch_size=None, verbose=1, |
| | tqdm_class=tqdm_auto, **tqdm_kwargs): |
| | """ |
| | Parameters |
| | ---------- |
| | epochs : int, optional |
| | data_size : int, optional |
| | Number of training pairs. |
| | batch_size : int, optional |
| | Number of training pairs per batch. |
| | verbose : int |
| | 0: epoch, 1: batch (transient), 2: batch. [default: 1]. |
| | Will be set to `0` unless both `data_size` and `batch_size` |
| | are given. |
| | tqdm_class : optional |
| | `tqdm` class to use for bars [default: `tqdm.auto.tqdm`]. |
| | tqdm_kwargs : optional |
| | Any other arguments used for all bars. |
| | """ |
| | if tqdm_kwargs: |
| | tqdm_class = partial(tqdm_class, **tqdm_kwargs) |
| | self.tqdm_class = tqdm_class |
| | self.epoch_bar = tqdm_class(total=epochs, unit='epoch') |
| | self.on_epoch_end = self.bar2callback(self.epoch_bar) |
| | if data_size and batch_size: |
| | self.batches = batches = (data_size + batch_size - 1) // batch_size |
| | else: |
| | self.batches = batches = None |
| | self.verbose = verbose |
| | if verbose == 1: |
| | self.batch_bar = tqdm_class(total=batches, unit='batch', leave=False) |
| | self.on_batch_end = self.bar2callback( |
| | self.batch_bar, pop=['batch', 'size'], |
| | delta=lambda logs: logs.get('size', 1)) |
| |
|
| | def on_train_begin(self, *_, **__): |
| | params = self.params.get |
| | auto_total = params('epochs', params('nb_epoch', None)) |
| | if auto_total is not None and auto_total != self.epoch_bar.total: |
| | self.epoch_bar.reset(total=auto_total) |
| |
|
| | def on_epoch_begin(self, epoch, *_, **__): |
| | if self.epoch_bar.n < epoch: |
| | ebar = self.epoch_bar |
| | ebar.n = ebar.last_print_n = ebar.initial = epoch |
| | if self.verbose: |
| | params = self.params.get |
| | total = params('samples', params( |
| | 'nb_sample', params('steps', None))) or self.batches |
| | if self.verbose == 2: |
| | if hasattr(self, 'batch_bar'): |
| | self.batch_bar.close() |
| | self.batch_bar = self.tqdm_class( |
| | total=total, unit='batch', leave=True, |
| | unit_scale=1 / (params('batch_size', 1) or 1)) |
| | self.on_batch_end = self.bar2callback( |
| | self.batch_bar, pop=['batch', 'size'], |
| | delta=lambda logs: logs.get('size', 1)) |
| | elif self.verbose == 1: |
| | self.batch_bar.unit_scale = 1 / (params('batch_size', 1) or 1) |
| | self.batch_bar.reset(total=total) |
| | else: |
| | raise KeyError('Unknown verbosity') |
| |
|
| | def on_train_end(self, *_, **__): |
| | if self.verbose: |
| | self.batch_bar.close() |
| | self.epoch_bar.close() |
| |
|
| | def display(self): |
| | """Displays in the current cell in Notebooks.""" |
| | container = getattr(self.epoch_bar, 'container', None) |
| | if container is None: |
| | return |
| | from .notebook import display |
| | display(container) |
| | batch_bar = getattr(self, 'batch_bar', None) |
| | if batch_bar is not None: |
| | display(batch_bar.container) |
| |
|
| | @staticmethod |
| | def _implements_train_batch_hooks(): |
| | return True |
| |
|
| | @staticmethod |
| | def _implements_test_batch_hooks(): |
| | return True |
| |
|
| | @staticmethod |
| | def _implements_predict_batch_hooks(): |
| | return True |
| |
|
