| 9 | |
| 10 | |
| 11 | def get_state_dict(net_type: str = 'alex', version: str = '0.1'): |
| 12 | # build url |
| 13 | url = 'https://raw.githubusercontent.com/richzhang/PerceptualSimilarity/' \ |
| 14 | + f'master/lpips/weights/v{version}/{net_type}.pth' |
| 15 | |
| 16 | # download |
| 17 | old_state_dict = torch.hub.load_state_dict_from_url( |
| 18 | url, progress=True, |
| 19 | map_location=None if torch.cuda.is_available() else torch.device('cpu') |
| 20 | ) |
| 21 | |
| 22 | # rename keys |
| 23 | new_state_dict = OrderedDict() |
| 24 | for key, val in old_state_dict.items(): |
| 25 | new_key = key |
| 26 | new_key = new_key.replace('lin', '') |
| 27 | new_key = new_key.replace('model.', '') |
| 28 | new_state_dict[new_key] = val |
| 29 | |
| 30 | return new_state_dict |