(self, network, network_label, epoch_label, save_dir='')
| 61 | |
| 62 | # helper loading function that can be used by subclasses |
| 63 | def load_network(self, network, network_label, epoch_label, save_dir=''): |
| 64 | self.resolve_version() |
| 65 | save_filename = '%s_net_%s.pth' % (epoch_label, network_label) |
| 66 | if not save_dir: |
| 67 | save_dir = self.save_dir |
| 68 | save_path = os.path.join(save_dir, save_filename) |
| 69 | if not os.path.isfile(save_path): |
| 70 | print('%s not exists yet!' % save_path) |
| 71 | if 'G0' in network_label: |
| 72 | raise('Generator must exist!') |
| 73 | else: |
| 74 | #network.load_state_dict(torch.load(save_path)) |
| 75 | try: |
| 76 | network.load_state_dict(torch.load(save_path)) |
| 77 | except: |
| 78 | pretrained_dict = torch.load(save_path) |
| 79 | model_dict = network.state_dict() |
| 80 | |
| 81 | ### printout layers in pretrained model |
| 82 | initialized = set() |
| 83 | for k, v in pretrained_dict.items(): |
| 84 | initialized.add(k.split('.')[0]) |
| 85 | #print('pretrained model has following layers: ') |
| 86 | #print(sorted(initialized)) |
| 87 | |
| 88 | try: |
| 89 | pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict} |
| 90 | network.load_state_dict(pretrained_dict) |
| 91 | print('Pretrained network %s has excessive layers; Only loading layers that are used' % network_label) |
| 92 | except: |
| 93 | print('Pretrained network %s has fewer layers; The following are not initialized:' % network_label) |
| 94 | if sys.version_info >= (3,0): |
| 95 | not_initialized = set() |
| 96 | else: |
| 97 | from sets import Set |
| 98 | not_initialized = Set() |
| 99 | for k, v in pretrained_dict.items(): |
| 100 | if v.size() == model_dict[k].size(): |
| 101 | model_dict[k] = v |
| 102 | |
| 103 | for k, v in model_dict.items(): |
| 104 | if k not in pretrained_dict or v.size() != pretrained_dict[k].size(): |
| 105 | not_initialized.add(k.split('.')[0]) |
| 106 | print(sorted(not_initialized)) |
| 107 | network.load_state_dict(model_dict) |
| 108 | |
| 109 | def concat(self, tensors, dim=0): |
| 110 | if tensors[0] is not None and tensors[1] is not None: |
no test coverage detected