| 69 | |
| 70 | |
| 71 | def load_pretrained_weights(model, pretrained_weights, checkpoint_key, model_name, patch_size): |
| 72 | if os.path.isfile(pretrained_weights): |
| 73 | state_dict = torch.load(pretrained_weights, map_location="cpu") |
| 74 | if checkpoint_key is not None and checkpoint_key in state_dict: |
| 75 | print(f"Take key {checkpoint_key} in provided checkpoint dict") |
| 76 | state_dict = state_dict[checkpoint_key] |
| 77 | # remove `module.` prefix |
| 78 | state_dict = {k.replace("module.", ""): v for k, v in state_dict.items()} |
| 79 | # remove `backbone.` prefix induced by multicrop wrapper |
| 80 | state_dict = {k.replace("backbone.", ""): v for k, v in state_dict.items()} |
| 81 | msg = model.load_state_dict(state_dict, strict=False) |
| 82 | print('Pretrained weights found at {} and loaded with msg: {}'.format(pretrained_weights, msg)) |
| 83 | else: |
| 84 | print("Please use the `--pretrained_weights` argument to indicate the path of the checkpoint to evaluate.") |
| 85 | url = None |
| 86 | if model_name == "vit_small" and patch_size == 16: |
| 87 | url = "dino_deitsmall16_pretrain/dino_deitsmall16_pretrain.pth" |
| 88 | elif model_name == "vit_small" and patch_size == 8: |
| 89 | url = "dino_deitsmall8_pretrain/dino_deitsmall8_pretrain.pth" |
| 90 | elif model_name == "vit_base" and patch_size == 16: |
| 91 | url = "dino_vitbase16_pretrain/dino_vitbase16_pretrain.pth" |
| 92 | elif model_name == "vit_base" and patch_size == 8: |
| 93 | url = "dino_vitbase8_pretrain/dino_vitbase8_pretrain.pth" |
| 94 | elif model_name == "xcit_small_12_p16": |
| 95 | url = "dino_xcit_small_12_p16_pretrain/dino_xcit_small_12_p16_pretrain.pth" |
| 96 | elif model_name == "xcit_small_12_p8": |
| 97 | url = "dino_xcit_small_12_p8_pretrain/dino_xcit_small_12_p8_pretrain.pth" |
| 98 | elif model_name == "xcit_medium_24_p16": |
| 99 | url = "dino_xcit_medium_24_p16_pretrain/dino_xcit_medium_24_p16_pretrain.pth" |
| 100 | elif model_name == "xcit_medium_24_p8": |
| 101 | url = "dino_xcit_medium_24_p8_pretrain/dino_xcit_medium_24_p8_pretrain.pth" |
| 102 | elif model_name == "resnet50": |
| 103 | url = "dino_resnet50_pretrain/dino_resnet50_pretrain.pth" |
| 104 | if url is not None: |
| 105 | print("Since no pretrained weights have been provided, we load the reference pretrained DINO weights.") |
| 106 | state_dict = torch.hub.load_state_dict_from_url(url="https://dl.fbaipublicfiles.com/dino/" + url) |
| 107 | model.load_state_dict(state_dict, strict=True) |
| 108 | else: |
| 109 | print("There is no reference weights available for this model => We use random weights.") |
| 110 | |
| 111 | |
| 112 | def clip_gradients(model, clip): |