| 56 | torch.save(self.optimizer.state_dict(), optim_file_path) |
| 57 | |
| 58 | def load_weights(self): |
| 59 | assert os.path.isdir(self.load_weights_dir), f'\tCannot find {self.load_weights_dir}' |
| 60 | print(f'Loading a model from {self.load_weights_dir}') |
| 61 | |
| 62 | # to retrain |
| 63 | if self.pretrain and self.ddp_enable: |
| 64 | map_location = {'cuda:%d' % 0: 'cuda:%d' % (self.world_size-1)} |
| 65 | |
| 66 | for n in self.models_to_load: |
| 67 | print(f'Loading {n} weights...') |
| 68 | path = os.path.join(self.load_weights_dir, f'{n}.pth') |
| 69 | model_dict = self.models[n].state_dict() |
| 70 | |
| 71 | # distribute gpus for ddp retraining |
| 72 | if self.pretrain and self.ddp_enable: |
| 73 | pre_trained_dict = torch.load(path, map_location=map_location) |
| 74 | else: |
| 75 | pre_trained_dict = torch.load(path) |
| 76 | |
| 77 | # load parameters |
| 78 | pre_trained_dict = {k: v for k, v in pre_trained_dict.items() if k in model_dict} |
| 79 | model_dict.update(pre_trained_dict) |
| 80 | self.models[n].load_state_dict(model_dict) |
| 81 | |
| 82 | if self.mode == 'train': |
| 83 | # loading adam state |
| 84 | optim_file_path = os.path.join(self.load_weights_dir, f'{_OPTIMIZER_NAME}.pth') |
| 85 | if os.path.isfile(optim_file_path): |
| 86 | try: |
| 87 | print(f'Loading {_OPTIMIZER_NAME} weights') |
| 88 | optimizer_dict = torch.load(optim_file_path) |
| 89 | self.optimizer.load_state_dict(optimizer_dict) |
| 90 | except ValueError: |
| 91 | print(f'\tCannnot load {_OPTIMIZER_NAME} - the optimizer will be randomly initialized') |
| 92 | else: |
| 93 | print(f'\tCannot find {_OPTIMIZER_NAME} weights, so the optimizer will be randomly initialized') |