| 82 | torch.save(state_dict, save_path) |
| 83 | |
| 84 | def load_network(self, load_path, network, strict=True): |
| 85 | if isinstance(network, nn.DataParallel) or isinstance(network, DistributedDataParallel): |
| 86 | network = network.module |
| 87 | if os.path.exists(load_path): |
| 88 | load_net = torch.load(load_path) |
| 89 | load_net_clean = OrderedDict() # remove unnecessary 'module.' |
| 90 | for k, v in load_net.items(): |
| 91 | if k.startswith('module.'): |
| 92 | load_net_clean[k[7:]] = v |
| 93 | else: |
| 94 | load_net_clean[k] = v |
| 95 | network.load_state_dict(load_net_clean, strict=strict) |
| 96 | print("Succcefully!!!!! pretrained model has loaded!!!!!!!!!!!!!!!") |
| 97 | else: |
| 98 | print("Wrong!!!!! pretrained path not exists") |
| 99 | |
| 100 | def save_training_state(self, epoch, iter_step): |
| 101 | """Save training state during training, which will be used for resuming""" |