(self, ignore_last_checkpoint=False, file_name=None)
| 67 | return save_data() |
| 68 | |
| 69 | def load(self, ignore_last_checkpoint=False, file_name=None): |
| 70 | save_file = os.path.join(self.cfg.OUTPUT_DIR, "last_checkpoint") |
| 71 | try: |
| 72 | with open(save_file, "r") as last_checkpoint: |
| 73 | f = last_checkpoint.read().strip() |
| 74 | except IOError: |
| 75 | self.logger.info("No checkpoint found. Initializing model from scratch") |
| 76 | if file_name is None: |
| 77 | return {} |
| 78 | |
| 79 | if ignore_last_checkpoint: |
| 80 | self.logger.info("Forced to Initialize model from scratch") |
| 81 | return {} |
| 82 | if file_name is not None: |
| 83 | f = file_name |
| 84 | |
| 85 | self.logger.info("Loading checkpoint from {}".format(f)) |
| 86 | checkpoint = torch.load(f, map_location=torch.device("cpu")) |
| 87 | for name, model in self.models.items(): |
| 88 | if name in checkpoint["models"]: |
| 89 | try: |
| 90 | model_dict = checkpoint["models"].pop(name) |
| 91 | if model_dict is not None: |
| 92 | self.models[name].load_state_dict(model_dict, strict=False) |
| 93 | else: |
| 94 | self.logger.warning("State dict for model \"%s\" is None " % name) |
| 95 | except RuntimeError as e: |
| 96 | self.logger.warning('%s\nFailed to load: %s\n%s' % ('!' * 160, name, '!' * 160)) |
| 97 | self.logger.warning('\nFailed to load: %s' % str(e)) |
| 98 | else: |
| 99 | self.logger.warning("No state dict for model: %s" % name) |
| 100 | checkpoint.pop('models') |
| 101 | if "auxiliary" in checkpoint and self.auxiliary: |
| 102 | self.logger.info("Loading auxiliary from {}".format(f)) |
| 103 | for name, item in self.auxiliary.items(): |
| 104 | try: |
| 105 | if name in checkpoint["auxiliary"]: |
| 106 | self.auxiliary[name].load_state_dict(checkpoint["auxiliary"].pop(name)) |
| 107 | if "optimizers" in checkpoint and name in checkpoint["optimizers"]: |
| 108 | self.auxiliary[name].load_state_dict(checkpoint["optimizers"].pop(name)) |
| 109 | if name in checkpoint: |
| 110 | self.auxiliary[name].load_state_dict(checkpoint.pop(name)) |
| 111 | except (IndexError, ValueError): |
| 112 | self.logger.warning('%s\nFailed to load: %s\n%s' % ('!' * 160, name, '!' * 160)) |
| 113 | checkpoint.pop('auxiliary') |
| 114 | |
| 115 | return checkpoint |
| 116 | |
| 117 | def tag_last_checkpoint(self, last_filename): |
| 118 | save_file = os.path.join(self.cfg.OUTPUT_DIR, "last_checkpoint") |
no test coverage detected