(self)
| 156 | self.save_checkpoint(checkpoint) |
| 157 | |
| 158 | def restore_checkpoint(self): |
| 159 | t_start = time.time() |
| 160 | if self.distributed: |
| 161 | # load the model on cpu first to avoid GPU RAM surge |
| 162 | # when loading a model checkpoint |
| 163 | # tmp = torch.load(self.continue_state_object, |
| 164 | # map_location=lambda storage, loc: storage.cuda( |
| 165 | # self.local_rank)) |
| 166 | tmp = torch.load(self.continue_state_object, map_location=torch.device("cpu")) |
| 167 | else: |
| 168 | tmp = torch.load(self.continue_state_object) |
| 169 | t_ioend = time.time() |
| 170 | self.state.model = load_model(self.state.model, tmp["model"], is_restore=True) |
| 171 | self.state.optimizer.load_state_dict(tmp["optimizer"]) |
| 172 | self.state.epoch = tmp["epoch"] + 1 |
| 173 | self.state.iteration = tmp["iteration"] |
| 174 | del tmp |
| 175 | t_end = time.time() |
| 176 | logger.info( |
| 177 | "Load checkpoint from file {}, Time usage:\n\tIO: {}, restore checkpoint: {}".format( |
| 178 | self.continue_state_object, t_ioend - t_start, t_end - t_ioend |
| 179 | ) |
| 180 | ) |
| 181 | |
| 182 | def __enter__(self): |
| 183 | return self |
no test coverage detected