Loads pretrained weights, and downloads if loading for the first time.
(model, model_name, load_fc=True, advprop=False)
| 321 | |
| 322 | |
| 323 | def load_pretrained_weights(model, model_name, load_fc=True, advprop=False): |
| 324 | """ Loads pretrained weights, and downloads if loading for the first time. """ |
| 325 | # AutoAugment or Advprop (different preprocessing) |
| 326 | url_map_ = url_map_advprop if advprop else url_map |
| 327 | state_dict = model_zoo.load_url(url_map_[model_name]) |
| 328 | if load_fc: |
| 329 | model.load_state_dict(state_dict) |
| 330 | else: |
| 331 | state_dict.pop('_fc.weight') |
| 332 | state_dict.pop('_fc.bias') |
| 333 | res = model.load_state_dict(state_dict, strict=False) |
| 334 | assert set(res.missing_keys) == set(['_fc.weight', '_fc.bias']), 'issue loading pretrained weights' |
| 335 | print('Loaded pretrained weights for {}'.format(model_name)) |