Load network. Args: load_path (str): The path of networks to be loaded. net (nn.Module): Network. strict (bool): Whether strictly loaded. param_key (str): The parameter key of loaded network. If set to None, use the root 'path'
(self, net, load_path, strict=True, param_key='params')
| 291 | load_net[k + '.ignore'] = load_net.pop(k) |
| 292 | |
| 293 | def load_network(self, net, load_path, strict=True, param_key='params'): |
| 294 | """Load network. |
| 295 | |
| 296 | Args: |
| 297 | load_path (str): The path of networks to be loaded. |
| 298 | net (nn.Module): Network. |
| 299 | strict (bool): Whether strictly loaded. |
| 300 | param_key (str): The parameter key of loaded network. If set to |
| 301 | None, use the root 'path'. |
| 302 | Default: 'params'. |
| 303 | """ |
| 304 | logger = get_root_logger() |
| 305 | net = self.get_bare_model(net) |
| 306 | load_net = torch.load(load_path, map_location=lambda storage, loc: storage) |
| 307 | if param_key is not None: |
| 308 | if param_key not in load_net and 'params' in load_net: |
| 309 | param_key = 'params' |
| 310 | logger.info('Loading: params_ema does not exist, use params.') |
| 311 | load_net = load_net[param_key] |
| 312 | logger.info(f'Loading {net.__class__.__name__} model from {load_path}, with param key: [{param_key}].') |
| 313 | # remove unnecessary 'module.' |
| 314 | for k, v in deepcopy(load_net).items(): |
| 315 | if k.startswith('module.'): |
| 316 | load_net[k[7:]] = v |
| 317 | load_net.pop(k) |
| 318 | self._print_different_keys_loading(net, load_net, strict) |
| 319 | net.load_state_dict(load_net, strict=strict) |
| 320 | |
| 321 | @master_only |
| 322 | def save_training_state(self, epoch, current_iter): |
no test coverage detected