Utility object that wrap around callbacks and progress bars.
| 659 | |
| 660 | |
| 661 | class TrainingContext(object): |
| 662 | """Utility object that wrap around callbacks and progress bars.""" |
| 663 | |
| 664 | @tf_contextlib.contextmanager |
| 665 | def on_start(self, model, callbacks=None, use_samples=False, verbose=0, |
| 666 | mode=ModeKeys.TRAIN): |
| 667 | """Provide a scope for the whole training process.""" |
| 668 | # TODO(omalleyt): Handle ProgBar as part of Callbacks once hooks are ready. |
| 669 | progbar = training_utils.get_progbar( |
| 670 | model, 'samples' if use_samples else 'steps') |
| 671 | progbar.params = callbacks.params |
| 672 | progbar.params['verbose'] = verbose |
| 673 | callbacks.model.stop_training = False |
| 674 | callbacks._call_begin_hook(mode) |
| 675 | progbar.on_train_begin() |
| 676 | |
| 677 | # Cache those two instance so that it can be used in other functions. |
| 678 | self.callbacks = callbacks |
| 679 | self.progbar = progbar |
| 680 | |
| 681 | try: |
| 682 | yield |
| 683 | finally: |
| 684 | # End of all epochs |
| 685 | self.callbacks._call_end_hook(mode) |
| 686 | |
| 687 | @tf_contextlib.contextmanager |
| 688 | def on_epoch(self, epoch=0, mode=ModeKeys.TRAIN): |
| 689 | """Provide a scope for running one epoch.""" |
| 690 | epoch_logs = {} |
| 691 | if mode == ModeKeys.TRAIN: |
| 692 | self.callbacks.on_epoch_begin(epoch, epoch_logs) |
| 693 | self.progbar.on_epoch_begin(epoch, epoch_logs) |
| 694 | try: |
| 695 | yield epoch_logs |
| 696 | finally: |
| 697 | if mode == ModeKeys.TRAIN: |
| 698 | # Epochs only apply to `fit`. |
| 699 | self.callbacks.on_epoch_end(epoch, epoch_logs) |
| 700 | self.progbar.on_epoch_end(epoch, epoch_logs) |
| 701 | |
| 702 | @tf_contextlib.contextmanager |
| 703 | def on_batch(self, step=0, mode=ModeKeys.TRAIN, size=1): |
| 704 | """Provide a scope for running one batch.""" |
| 705 | batch_logs = {'batch': step, 'size': size} |
| 706 | self.callbacks._call_batch_hook( |
| 707 | mode, 'begin', step, batch_logs) |
| 708 | self.progbar.on_batch_begin(step, batch_logs) |
| 709 | try: |
| 710 | yield batch_logs |
| 711 | finally: |
| 712 | if not batch_logs.pop('data_exhausted', False): |
| 713 | self.callbacks._call_batch_hook( |
| 714 | mode, 'end', step, batch_logs) |
| 715 | self.progbar.on_batch_end(step, batch_logs) |
no outgoing calls
no test coverage detected