Resumes training from a checkpoint at the given file path. Loads the model's state, the optimizer's state, and the scheduler's state. Args: resume_path (str): The file path to the checkpoint to resume from.
(
self,
model=None,
optim=None,
scheduler=None,
scaler=None,
)
| 263 | dist.barrier() |
| 264 | |
| 265 | def resume_checkpoint( |
| 266 | self, |
| 267 | model=None, |
| 268 | optim=None, |
| 269 | scheduler=None, |
| 270 | scaler=None, |
| 271 | ): |
| 272 | """ |
| 273 | Resumes training from a checkpoint at the given file path. |
| 274 | Loads the model's state, the optimizer's state, and the scheduler's state. |
| 275 | |
| 276 | Args: |
| 277 | resume_path (str): The file path to the checkpoint to resume from. |
| 278 | """ |
| 279 | if self.resume: |
| 280 | ckpt = os.path.join(self.output_dir, "model.pt") |
| 281 | if os.path.isfile(ckpt): |
| 282 | checkpoint = torch.load(ckpt, map_location="cpu") |
| 283 | self.start_epoch = checkpoint["epoch"] |
| 284 | # self.model.load_state_dict(checkpoint['state_dict']) |
| 285 | src_state = checkpoint["state_dict"] |
| 286 | dst_state = model.state_dict() |
| 287 | for k in dst_state.keys(): |
| 288 | if not k.startswith("module.") and "module." + k in src_state.keys(): |
| 289 | k_ddp = "module." + k |
| 290 | elif k.startswith("module.") and "module." + k not in src_state.keys(): |
| 291 | k_ddp = k.replace("module.", "", 1) |
| 292 | else: |
| 293 | k_ddp = k |
| 294 | |
| 295 | if k_ddp in src_state.keys(): |
| 296 | dst_state[k] = src_state[k_ddp] |
| 297 | else: |
| 298 | print(f"Miss key in ckpt: model: {k}, ckpt: {k_ddp}") |
| 299 | |
| 300 | model.load_state_dict(dst_state) |
| 301 | optim.load_state_dict(checkpoint["optimizer"]) |
| 302 | scheduler.load_state_dict(checkpoint["scheduler"]) |
| 303 | if scaler is not None and "scaler_state" in checkpoint: |
| 304 | scaler.load_state_dict(checkpoint["scaler_state"]) |
| 305 | |
| 306 | self.saved_ckpts = checkpoint["saved_ckpts"] |
| 307 | self.val_acc_step_or_epoch = ( |
| 308 | checkpoint["val_acc_step_or_epoch"] |
| 309 | if "val_acc_step_or_epoch" in checkpoint |
| 310 | else {} |
| 311 | ) |
| 312 | self.val_loss_step_or_epoch = ( |
| 313 | checkpoint["val_loss_step_or_epoch"] |
| 314 | if "val_loss_step_or_epoch" in checkpoint |
| 315 | else {} |
| 316 | ) |
| 317 | self.best_step_or_epoch = ( |
| 318 | checkpoint["best_step_or_epoch"] if "best_step_or_epoch" in checkpoint else "" |
| 319 | ) |
| 320 | self.start_data_split_i = ( |
| 321 | checkpoint["data_split_i"] if "data_split_i" in checkpoint else 0 |
| 322 | ) |
no test coverage detected