Restore trainer state. Model will get its change to update :param checkpoint: :return:
(self, checkpoint)
| 681 | return |
| 682 | |
| 683 | def restore_training_state(self, checkpoint): |
| 684 | """ |
| 685 | Restore trainer state. |
| 686 | Model will get its change to update |
| 687 | :param checkpoint: |
| 688 | :return: |
| 689 | """ |
| 690 | if self.checkpoint_callback is not None and self.checkpoint_callback is not False: |
| 691 | self.checkpoint_callback.best = checkpoint['checkpoint_callback_best'] |
| 692 | |
| 693 | self.global_step = checkpoint['global_step'] |
| 694 | self.current_epoch = checkpoint['epoch'] |
| 695 | |
| 696 | if self.testing: |
| 697 | return |
| 698 | |
| 699 | # restore the optimizers |
| 700 | optimizer_states = checkpoint['optimizer_states'] |
| 701 | for optimizer, opt_state in zip(self.optimizers, optimizer_states): |
| 702 | if optimizer is None: |
| 703 | return |
| 704 | optimizer.load_state_dict(opt_state) |
| 705 | |
| 706 | # move optimizer to GPU 1 weight at a time |
| 707 | # avoids OOM |
| 708 | if self.root_gpu is not None: |
| 709 | for state in optimizer.state.values(): |
| 710 | for k, v in state.items(): |
| 711 | if isinstance(v, torch.Tensor): |
| 712 | state[k] = v.cuda(self.root_gpu) |
| 713 | |
| 714 | # restore the lr schedulers |
| 715 | lr_schedulers = checkpoint['lr_schedulers'] |
| 716 | for scheduler, lrs_state in zip(self.lr_schedulers, lr_schedulers): |
| 717 | scheduler.load_state_dict(lrs_state) |
| 718 | |
| 719 | # -------------------- |
| 720 | # MODEL SAVE CHECKPOINT |