Load all the networks from the disk. Parameters: epoch (int) -- current epoch; used in the file name '%s_net_%s.pth' % (epoch, name)
(self, epoch)
| 194 | self.__patch_instance_norm_state_dict(state_dict, getattr(module, key), keys, i + 1) |
| 195 | |
| 196 | def load_networks(self, epoch): |
| 197 | """Load all the networks from the disk. |
| 198 | |
| 199 | Parameters: |
| 200 | epoch (int) -- current epoch; used in the file name '%s_net_%s.pth' % (epoch, name) |
| 201 | """ |
| 202 | for name in self.model_names: |
| 203 | if isinstance(name, str): |
| 204 | load_filename = '%s_net_%s.pth' % (epoch, name) |
| 205 | if self.opt.isTrain and self.opt.pretrained_name is not None: |
| 206 | load_dir = os.path.join(self.opt.checkpoints_dir, self.opt.pretrained_name) |
| 207 | else: |
| 208 | load_dir = self.save_dir |
| 209 | |
| 210 | load_path = os.path.join(load_dir, load_filename) |
| 211 | net = getattr(self, 'net' + name) |
| 212 | if isinstance(net, torch.nn.DataParallel): |
| 213 | net = net.module |
| 214 | print('loading the model from %s' % load_path) |
| 215 | # if you are using PyTorch newer than 0.4 (e.g., built from |
| 216 | # GitHub source), you can remove str() on self.device |
| 217 | state_dict = torch.load(load_path, map_location=str(self.device)) |
| 218 | if hasattr(state_dict, '_metadata'): |
| 219 | del state_dict._metadata |
| 220 | |
| 221 | # patch InstanceNorm checkpoints prior to 0.4 |
| 222 | # for key in list(state_dict.keys()): # need to copy keys here because we mutate in loop |
| 223 | # self.__patch_instance_norm_state_dict(state_dict, net, key.split('.')) |
| 224 | net.load_state_dict(state_dict) |
| 225 | |
| 226 | def print_networks(self, verbose): |
| 227 | """Print the total number of parameters in the network and (if verbose) network architecture |