(self, args)
| 989 | self.worker_init_fn = None |
| 990 | |
| 991 | def load(self, args): |
| 992 | if args.load_path == '': |
| 993 | return |
| 994 | load_path = args.load_path if args.load_single else args.load_path.replace('ckpt_task_', f'ckpt_task{self.ginfo.task_id}_') |
| 995 | |
| 996 | try: |
| 997 | checkpoint = torch.load(load_path, 'cpu') |
| 998 | except: |
| 999 | raise FileNotFoundError(f'=> no checkpoint found at {load_path}') |
| 1000 | |
| 1001 | if self.ginfo.task_rank == 0: |
| 1002 | printlog(f"Recovering from {load_path}, keys={list(checkpoint.keys())}") |
| 1003 | |
| 1004 | if 'state_dict' in checkpoint: |
| 1005 | pretrained_state_dict = checkpoint['state_dict'] |
| 1006 | else: |
| 1007 | pretrained_state_dict = checkpoint |
| 1008 | if len(args.ignore) > 0: |
| 1009 | for k in list(pretrained_state_dict.keys()): |
| 1010 | flag = False |
| 1011 | for prefix in args.ignore: |
| 1012 | if k.startswith(prefix): |
| 1013 | flag = True |
| 1014 | the_prefix = prefix |
| 1015 | break |
| 1016 | if flag: |
| 1017 | print('ignoring {} (prefix: {})'.format(k, the_prefix)) |
| 1018 | del pretrained_state_dict[k] |
| 1019 | load_state_model(self.model, pretrained_state_dict, self.ginfo) |
| 1020 | |
| 1021 | def initialize(self, args): |
| 1022 | self.create_dataset() |
no test coverage detected