Files
ai-toolkit/toolkit/progress_bar.py

29 lines
839 B
Python

from tqdm import tqdm
import time
class ToolkitProgressBar(tqdm):
def __init__(self, *args, **kwargs):
# high EMA alpha so the rate/ETA responds within a few steps
# (tqdm default of 0.3 takes ~10 steps to wash out old samples)
kwargs.setdefault('smoothing', 0.7)
super().__init__(*args, **kwargs)
self.paused = False
self.last_time = self._time()
def pause(self):
if not self.paused:
self.paused = True
self.last_time = self._time()
def unpause(self):
if self.paused:
self.paused = False
cur_t = self._time()
self.start_t += cur_t - self.last_time
self.last_print_t = cur_t
def update(self, *args, **kwargs):
if not self.paused:
super().update(*args, **kwargs)