(self, args)
| 264 | pin_memory=False, sampler=self.sampler, worker_init_fn=self.worker_init_fn) |
| 265 | |
| 266 | def load(self, args): |
| 267 | if args.load_path == '': |
| 268 | return |
| 269 | load_path = args.load_path if args.load_single else args.load_path.replace('ckpt_task_', f'ckpt_task{self.ginfo.task_id}_') |
| 270 | |
| 271 | try: |
| 272 | checkpoint = torch.load(load_path, 'cpu') |
| 273 | except: |
| 274 | raise FileNotFoundError(f'=> no checkpoint found at {load_path}') |
| 275 | |
| 276 | if self.ginfo.task_rank == 0: |
| 277 | printlog(f"Recovering from {load_path}, keys={list(checkpoint.keys())}") |
| 278 | |
| 279 | if 'state_dict' in checkpoint: |
| 280 | pretrained_state_dict = checkpoint['state_dict'] |
| 281 | else: |
| 282 | pretrained_state_dict = checkpoint |
| 283 | |
| 284 | ignores = args.ignore + self.config.get('load_ignore', []) |
| 285 | if len(ignores) > 0: |
| 286 | for k in list(pretrained_state_dict.keys()): |
| 287 | flag = False |
| 288 | for prefix in ignores: |
| 289 | if k.startswith(prefix): |
| 290 | flag = True |
| 291 | the_prefix = prefix |
| 292 | break |
| 293 | if flag: |
| 294 | print('ignoring {} (prefix: {})'.format(k, the_prefix)) |
| 295 | del pretrained_state_dict[k] |
| 296 | load_state_model(self.model, pretrained_state_dict, self.ginfo) |
| 297 | if args.finetune and not args.recover: |
| 298 | return |
| 299 | if 'optimizer' in checkpoint: |
| 300 | load_state_optimizer(self.optimizer, checkpoint['optimizer'], self.ginfo) |
| 301 | self.last_iter = checkpoint['step'] - 1 |
| 302 | elif args.recover: |
| 303 | self.last_iter = checkpoint['step'] - 1 |
| 304 | |
| 305 | def pre_run(self): |
| 306 | tmp = self.tmp |
no test coverage detected