| from .plugin import Plugin |
| |
| |
| class ProgressMonitor(Plugin): |
| stat_name = 'progress' |
| |
| def __init__(self): |
| super(ProgressMonitor, self).__init__([(1, 'iteration'), (1, 'epoch')]) |
| |
| def register(self, trainer): |
| self.trainer = trainer |
| stats = self.trainer.stats.setdefault(self.stat_name, {}) |
| stats['samples_used'] = 0 |
| stats['epoch_size'] = len(trainer.dataset) |
| stats['log_iter_fields'] = [ |
| '{samples_used}/{epoch_size}', |
| '({percent:.2f}%)' |
| ] |
| |
| def iteration(self, iteration, input, *args): |
| stats = self.trainer.stats.setdefault(self.stat_name, {}) |
| stats['samples_used'] += input.size(0) |
| stats['percent'] = 100. * stats['samples_used'] / stats['epoch_size'] |
| |
| def epoch(self, *args): |
| stats = self.trainer.stats.setdefault(self.stat_name, {}) |
| stats['samples_used'] = 0 |
| stats['percent'] = 0 |
| |
