Fix InstanceNorm checkpoints incompatibility (prior to 0.4)
(self, state_dict, module, keys, i=0)
| 158 | torch.save(net.cpu().state_dict(), save_path) |
| 159 | |
| 160 | def __patch_instance_norm_state_dict(self, state_dict, module, keys, i=0): |
| 161 | """Fix InstanceNorm checkpoints incompatibility (prior to 0.4)""" |
| 162 | key = keys[i] |
| 163 | if i + 1 == len(keys): # at the end, pointing to a parameter/buffer |
| 164 | if module.__class__.__name__.startswith('InstanceNorm') and \ |
| 165 | (key == 'running_mean' or key == 'running_var'): |
| 166 | if getattr(module, key) is None: |
| 167 | state_dict.pop('.'.join(keys)) |
| 168 | if module.__class__.__name__.startswith('InstanceNorm') and \ |
| 169 | (key == 'num_batches_tracked'): |
| 170 | state_dict.pop('.'.join(keys)) |
| 171 | else: |
| 172 | self.__patch_instance_norm_state_dict(state_dict, getattr(module, key), keys, i + 1) |
| 173 | |
| 174 | def load_networks(self, epoch): |
| 175 | """Load all the networks from the disk. |