| 123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122 |
- 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
|
