Mercurial > repos > rliterman > csp2

comparison CSP2/CSP2_env/env-d9b9114564458d9d-741b3de822f2aaca6c6caa4325c4afce/lib/python3.8/site-packages/tqdm/keras.py @ 68:5028fdace37b

Find changesets by keywords (author, files, the commit message), revision number or hash, or revset expression.
planemo upload commit 2e9511a184a1ca667c7be0c6321a36dc4e3d116d
author jpayne
date Tue, 18 Mar 2025 16:23:26 -0400
parents
children
comparison
equal deleted inserted replaced
67:0e9998148a16 68:5028fdace37b
1 from copy import copy
2 from functools import partial
3
4 from .auto import tqdm as tqdm_auto
5
6 try:
7 import keras
8 except (ImportError, AttributeError) as e:
9 try:
10 from tensorflow import keras
11 except ImportError:
12 raise e
13 __author__ = {"github.com/": ["casperdcl"]}
14 __all__ = ['TqdmCallback']
15
16
17 class TqdmCallback(keras.callbacks.Callback):
18 """Keras callback for epoch and batch progress."""
19 @staticmethod
20 def bar2callback(bar, pop=None, delta=(lambda logs: 1)):
21 def callback(_, logs=None):
22 n = delta(logs)
23 if logs:
24 if pop:
25 logs = copy(logs)
26 [logs.pop(i, 0) for i in pop]
27 bar.set_postfix(logs, refresh=False)
28 bar.update(n)
29
30 return callback
31
32 def __init__(self, epochs=None, data_size=None, batch_size=None, verbose=1,
33 tqdm_class=tqdm_auto, **tqdm_kwargs):
34 """
35 Parameters
36 ----------
37 epochs : int, optional
38 data_size : int, optional
39 Number of training pairs.
40 batch_size : int, optional
41 Number of training pairs per batch.
42 verbose : int
43 0: epoch, 1: batch (transient), 2: batch. [default: 1].
44 Will be set to `0` unless both `data_size` and `batch_size`
45 are given.
46 tqdm_class : optional
47 `tqdm` class to use for bars [default: `tqdm.auto.tqdm`].
48 tqdm_kwargs : optional
49 Any other arguments used for all bars.
50 """
51 if tqdm_kwargs:
52 tqdm_class = partial(tqdm_class, **tqdm_kwargs)
53 self.tqdm_class = tqdm_class
54 self.epoch_bar = tqdm_class(total=epochs, unit='epoch')
55 self.on_epoch_end = self.bar2callback(self.epoch_bar)
56 if data_size and batch_size:
57 self.batches = batches = (data_size + batch_size - 1) // batch_size
58 else:
59 self.batches = batches = None
60 self.verbose = verbose
61 if verbose == 1:
62 self.batch_bar = tqdm_class(total=batches, unit='batch', leave=False)
63 self.on_batch_end = self.bar2callback(
64 self.batch_bar, pop=['batch', 'size'],
65 delta=lambda logs: logs.get('size', 1))
66
67 def on_train_begin(self, *_, **__):
68 params = self.params.get
69 auto_total = params('epochs', params('nb_epoch', None))
70 if auto_total is not None and auto_total != self.epoch_bar.total:
71 self.epoch_bar.reset(total=auto_total)
72
73 def on_epoch_begin(self, epoch, *_, **__):
74 if self.epoch_bar.n < epoch:
75 ebar = self.epoch_bar
76 ebar.n = ebar.last_print_n = ebar.initial = epoch
77 if self.verbose:
78 params = self.params.get
79 total = params('samples', params(
80 'nb_sample', params('steps', None))) or self.batches
81 if self.verbose == 2:
82 if hasattr(self, 'batch_bar'):
83 self.batch_bar.close()
84 self.batch_bar = self.tqdm_class(
85 total=total, unit='batch', leave=True,
86 unit_scale=1 / (params('batch_size', 1) or 1))
87 self.on_batch_end = self.bar2callback(
88 self.batch_bar, pop=['batch', 'size'],
89 delta=lambda logs: logs.get('size', 1))
90 elif self.verbose == 1:
91 self.batch_bar.unit_scale = 1 / (params('batch_size', 1) or 1)
92 self.batch_bar.reset(total=total)
93 else:
94 raise KeyError('Unknown verbosity')
95
96 def on_train_end(self, *_, **__):
97 if hasattr(self, 'batch_bar'):
98 self.batch_bar.close()
99 self.epoch_bar.close()
100
101 def display(self):
102 """Displays in the current cell in Notebooks."""
103 container = getattr(self.epoch_bar, 'container', None)
104 if container is None:
105 return
106 from .notebook import display
107 display(container)
108 batch_bar = getattr(self, 'batch_bar', None)
109 if batch_bar is not None:
110 display(batch_bar.container)
111
112 @staticmethod
113 def _implements_train_batch_hooks():
114 return True
115
116 @staticmethod
117 def _implements_test_batch_hooks():
118 return True
119
120 @staticmethod
121 def _implements_predict_batch_hooks():
122 return True