codingutf8 2022 the huggingface inc team licensed under the apache license version 2 0 the license you may not use this file except in compliance with the license you may obtain a copy of the license at http www apache orglicenseslicense2 0 unless required by applicable law or agreed to in writing software distributed under the license is distributed on an as is basis without warranties or conditions of any kind either express or implied see the license for the specific language governing permissions and limitations under the license convert vit msn checkpoints from the original repository https github comfacebookresearchmsn import argparse import json import requests import torch from huggingfacehub import hfhubdownload from pil import image from transformers import vitimageprocessor vitmsnconfig vitmsnmodel from transformers imageutils import imagenetdefaultmean imagenetdefaultstd torch setgradenabledfalse here we list all keys to be renamed original name on the left our name on the right def createrenamekeysconfig basemodelfalse renamekeys for i in rangeconfig numhiddenlayers encoder layers output projection 2 feedforward neural networks and 2 layernorms renamekeys appendfmodule blocks i norm1 weight fvit encoder layer i layernormbefore weight renamekeys appendfmodule blocks i norm1 bias fvit encoder layer i layernormbefore bias renamekeys append fmodule blocks i attn proj weight fvit encoder layer i attention output dense weight renamekeys appendfmodule blocks i attn proj bias fvit encoder layer i attention output dense bias renamekeys appendfmodule blocks i norm2 weight fvit encoder layer i layernormafter weight renamekeys appendfmodule blocks i norm2 bias fvit encoder layer i layernormafter bias renamekeys appendfmodule blocks i mlp fc1 weight fvit encoder layer i intermediate dense weight renamekeys appendfmodule blocks i mlp fc1 bias fvit encoder layer i intermediate dense bias renamekeys appendfmodule blocks i mlp fc2 weight fvit encoder layer i output dense weight renamekeys appendfmodule blocks i mlp fc2 bias fvit encoder layer i output dense bias projection layer position embeddings renamekeys extend module clstoken vit embeddings clstoken module patchembed proj weight vit embeddings patchembeddings projection weight module patchembed proj bias vit embeddings patchembeddings projection bias module posembed vit embeddings positionembeddings if basemodel layernorm pooler renamekeys extend module norm weight layernorm weight module norm bias layernorm bias if just the base model we should remove vit from all keys that start with vit renamekeys pair0 pair14 if pair1 startswithvit else pair for pair in renamekeys else layernorm classification head renamekeys extend norm weight vit layernorm weight norm bias vit layernorm bias head weight classifier weight head bias classifier bias return renamekeys we split up the matrix of each encoder layer into queries keys and values def readinqkvstatedict config basemodelfalse for i in rangeconfig numhiddenlayers if basemodel prefix else prefix vit read in weights bias of input projection layer in timm this is a single matrix bias inprojweight statedict popfmodule blocks i attn qkv weight inprojbias statedict popfmodule blocks i attn qkv bias next add query keys and values in that order to the state dict statedictfprefixencoder layer i attention attention query weight inprojweight config hiddensize statedictfprefixencoder layer i attention attention query bias inprojbias config hiddensize statedictfprefixencoder layer i attention attention key weight inprojweight config hiddensize config hiddensize 2 statedictfprefixencoder layer i attention attention key bias inprojbias config hiddensize config hiddensize 2 statedictfprefixencoder layer i attention attention value weight inprojweight config hiddensize statedictfprefixencoder layer i attention attention value bias inprojbiasconfig hiddensize def removeclassificationheadstatedict ignorekeys head weight head bias for k in ignorekeys statedict popk none def removeprojectionheadstatedict projection head is used in the selfsupervised pretraining in msn for downstream task it s not needed ignorekeys module fc fc1 weight module fc fc1 bias module fc bn1 weight module fc bn1 bias module fc bn1 runningmean module fc bn1 runningvar module fc bn1 numbatchestracked module fc fc2 weight module fc fc2 bias module fc bn2 weight module fc bn2 bias module fc bn2 runningmean module fc bn2 runningvar module fc bn2 numbatchestracked module fc fc3 weight module fc fc3 bias for k in ignorekeys statedict popk none def renamekeydct old new val dct popold dctnew val def convertvitmsncheckpointcheckpointurl pytorchdumpfolderpath config vitmsnconfig config numlabels 1000 repoid datasetshuggingfacelabelfiles filename imagenet1kid2label json id2label json loadopenhfhubdownloadrepoid filename r id2label intk v for k v in id2label items config id2label id2label config label2id v k for k v in id2label items if s16 in checkpointurl config hiddensize 384 config intermediatesize 1536 config numattentionheads 6 elif l16 in checkpointurl config hiddensize 1024 config intermediatesize 4096 config numhiddenlayers 24 config numattentionheads 16 config hiddendropoutprob 0 1 elif b4 in checkpointurl config patchsize 4 elif l7 in checkpointurl config patchsize 7 config hiddensize 1024 config intermediatesize 4096 config numhiddenlayers 24 config numattentionheads 16 config hiddendropoutprob 0 1 model vitmsnmodelconfig statedict torch hub loadstatedictfromurlcheckpointurl maplocationcputargetencoder imageprocessor vitimageprocessorsizeconfig imagesize removeprojectionheadstatedict renamekeys createrenamekeysconfig basemodeltrue for src dest in renamekeys renamekeystatedict src dest readinqkvstatedict config basemodeltrue model loadstatedictstatedict model eval url http images cocodataset orgval2017000000039769 jpg image image openrequests geturl streamtrue raw imageprocessor vitimageprocessor sizeconfig imagesize imagemeanimagenetdefaultmean imagestdimagenetdefaultstd inputs imageprocessorimagesimage returntensorspt forward pass torch manualseed2 outputs modelinputs lasthiddenstate outputs lasthiddenstate the following colab notebook was used to generate these outputs https colab research google comgistsayakpaul3672419a04f5997827503fd84079bdd1scratchpad ipynb if s16 in checkpointurl expectedslice torch tensor1 0915 1 4876 1 1809 elif b16 in checkpointurl expectedslice torch tensor14 2889 18 9045 11 7281 elif l16 in checkpointurl expectedslice torch tensor41 5028 22 8681 45 6475 elif b4 in checkpointurl expectedslice torch tensor4 3868 5 2932 0 4137 else expectedslice torch tensor0 1792 0 6465 2 4263 verify logits assert torch allcloselasthiddenstate 0 3 expectedslice atol1e4 printfsaving model to pytorchdumpfolderpath model savepretrainedpytorchdumpfolderpath printfsaving image processor to pytorchdumpfolderpath imageprocessor savepretrainedpytorchdumpfolderpath if name main parser argparse argumentparser required parameters parser addargument checkpointurl defaulthttps dl fbaipublicfiles commsnvits16800ep pth tar typestr helpurl of the checkpoint you d like to convert parser addargument pytorchdumpfolderpath defaultnone typestr helppath to the output pytorch model directory args parser parseargs convertvitmsncheckpointargs checkpointurl args pytorchdumpfolderpath coding utf 8 2022 the huggingface inc team licensed under the apache license version 2 0 the license you may not use this file except in compliance with the license you may obtain a copy of the license at http www apache org licenses license 2 0 unless required by applicable law or agreed to in writing software distributed under the license is distributed on an as is basis without warranties or conditions of any kind either express or implied see the license for the specific language governing permissions and limitations under the license convert vit msn checkpoints from the original repository https github com facebookresearch msn here we list all keys to be renamed original name on the left our name on the right encoder layers output projection 2 feedforward neural networks and 2 layernorms projection layer position embeddings layernorm pooler if just the base model we should remove vit from all keys that start with vit layernorm classification head we split up the matrix of each encoder layer into queries keys and values read in weights bias of input projection layer in timm this is a single matrix bias next add query keys and values in that order to the state dict projection head is used in the self supervised pre training in msn for downstream task it s not needed forward pass the following colab notebook was used to generate these outputs https colab research google com gist sayakpaul 3672419a04f5997827503fd84079bdd1 scratchpad ipynb verify logits required parameters
import argparse
import json

import requests
import torch
from huggingface_hub import hf_hub_download
from PIL import Image

from transformers import ViTImageProcessor, ViTMSNConfig, ViTMSNModel
from transformers.image_utils import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD


torch.set_grad_enabled(False)



def create_rename_keys(config, base_model=False):
    rename_keys = []
    for i in range(config.num_hidden_layers):
        
        rename_keys.append((f"module.blocks.{i}.norm1.weight", f"vit.encoder.layer.{i}.layernorm_before.weight"))
        rename_keys.append((f"module.blocks.{i}.norm1.bias", f"vit.encoder.layer.{i}.layernorm_before.bias"))
        rename_keys.append(
            (f"module.blocks.{i}.attn.proj.weight", f"vit.encoder.layer.{i}.attention.output.dense.weight")
        )
        rename_keys.append((f"module.blocks.{i}.attn.proj.bias", f"vit.encoder.layer.{i}.attention.output.dense.bias"))
        rename_keys.append((f"module.blocks.{i}.norm2.weight", f"vit.encoder.layer.{i}.layernorm_after.weight"))
        rename_keys.append((f"module.blocks.{i}.norm2.bias", f"vit.encoder.layer.{i}.layernorm_after.bias"))
        rename_keys.append((f"module.blocks.{i}.mlp.fc1.weight", f"vit.encoder.layer.{i}.intermediate.dense.weight"))
        rename_keys.append((f"module.blocks.{i}.mlp.fc1.bias", f"vit.encoder.layer.{i}.intermediate.dense.bias"))
        rename_keys.append((f"module.blocks.{i}.mlp.fc2.weight", f"vit.encoder.layer.{i}.output.dense.weight"))
        rename_keys.append((f"module.blocks.{i}.mlp.fc2.bias", f"vit.encoder.layer.{i}.output.dense.bias"))

    
    rename_keys.extend(
        [
            ("module.cls_token", "vit.embeddings.cls_token"),
            ("module.patch_embed.proj.weight", "vit.embeddings.patch_embeddings.projection.weight"),
            ("module.patch_embed.proj.bias", "vit.embeddings.patch_embeddings.projection.bias"),
            ("module.pos_embed", "vit.embeddings.position_embeddings"),
        ]
    )

    if base_model:
        
        rename_keys.extend(
            [
                ("module.norm.weight", "layernorm.weight"),
                ("module.norm.bias", "layernorm.bias"),
            ]
        )

        
        rename_keys = [(pair[0], pair[1][4:]) if pair[1].startswith("vit") else pair for pair in rename_keys]
    else:
        
        rename_keys.extend(
            [
                ("norm.weight", "vit.layernorm.weight"),
                ("norm.bias", "vit.layernorm.bias"),
                ("head.weight", "classifier.weight"),
                ("head.bias", "classifier.bias"),
            ]
        )

    return rename_keys



def read_in_q_k_v(state_dict, config, base_model=False):
    for i in range(config.num_hidden_layers):
        if base_model:
            prefix = ""
        else:
            prefix = "vit."
        
        in_proj_weight = state_dict.pop(f"module.blocks.{i}.attn.qkv.weight")
        in_proj_bias = state_dict.pop(f"module.blocks.{i}.attn.qkv.bias")
        
        state_dict[f"{prefix}encoder.layer.{i}.attention.attention.query.weight"] = in_proj_weight[
            : config.hidden_size, :
        ]
        state_dict[f"{prefix}encoder.layer.{i}.attention.attention.query.bias"] = in_proj_bias[: config.hidden_size]
        state_dict[f"{prefix}encoder.layer.{i}.attention.attention.key.weight"] = in_proj_weight[
            config.hidden_size : config.hidden_size * 2, :
        ]
        state_dict[f"{prefix}encoder.layer.{i}.attention.attention.key.bias"] = in_proj_bias[
            config.hidden_size : config.hidden_size * 2
        ]
        state_dict[f"{prefix}encoder.layer.{i}.attention.attention.value.weight"] = in_proj_weight[
            -config.hidden_size :, :
        ]
        state_dict[f"{prefix}encoder.layer.{i}.attention.attention.value.bias"] = in_proj_bias[-config.hidden_size :]


def remove_classification_head_(state_dict):
    ignore_keys = ["head.weight", "head.bias"]
    for k in ignore_keys:
        state_dict.pop(k, None)


def remove_projection_head(state_dict):
    
    
    ignore_keys = [
        "module.fc.fc1.weight",
        "module.fc.fc1.bias",
        "module.fc.bn1.weight",
        "module.fc.bn1.bias",
        "module.fc.bn1.running_mean",
        "module.fc.bn1.running_var",
        "module.fc.bn1.num_batches_tracked",
        "module.fc.fc2.weight",
        "module.fc.fc2.bias",
        "module.fc.bn2.weight",
        "module.fc.bn2.bias",
        "module.fc.bn2.running_mean",
        "module.fc.bn2.running_var",
        "module.fc.bn2.num_batches_tracked",
        "module.fc.fc3.weight",
        "module.fc.fc3.bias",
    ]
    for k in ignore_keys:
        state_dict.pop(k, None)


def rename_key(dct, old, new):
    val = dct.pop(old)
    dct[new] = val


def convert_vit_msn_checkpoint(checkpoint_url, pytorch_dump_folder_path):
    config = ViTMSNConfig()
    config.num_labels = 1000

    repo_id = "datasets/huggingface/label-files"
    filename = "imagenet-1k-id2label.json"
    id2label = json.load(open(hf_hub_download(repo_id, filename), "r"))
    id2label = {int(k): v for k, v in id2label.items()}
    config.id2label = id2label
    config.label2id = {v: k for k, v in id2label.items()}

    if "s16" in checkpoint_url:
        config.hidden_size = 384
        config.intermediate_size = 1536
        config.num_attention_heads = 6
    elif "l16" in checkpoint_url:
        config.hidden_size = 1024
        config.intermediate_size = 4096
        config.num_hidden_layers = 24
        config.num_attention_heads = 16
        config.hidden_dropout_prob = 0.1
    elif "b4" in checkpoint_url:
        config.patch_size = 4
    elif "l7" in checkpoint_url:
        config.patch_size = 7
        config.hidden_size = 1024
        config.intermediate_size = 4096
        config.num_hidden_layers = 24
        config.num_attention_heads = 16
        config.hidden_dropout_prob = 0.1

    model = ViTMSNModel(config)

    state_dict = torch.hub.load_state_dict_from_url(checkpoint_url, map_location="cpu")["target_encoder"]

    image_processor = ViTImageProcessor(size=config.image_size)

    remove_projection_head(state_dict)
    rename_keys = create_rename_keys(config, base_model=True)

    for src, dest in rename_keys:
        rename_key(state_dict, src, dest)
    read_in_q_k_v(state_dict, config, base_model=True)

    model.load_state_dict(state_dict)
    model.eval()

    url = "http://images.cocodataset.org/val2017/000000039769.jpg"

    image = Image.open(requests.get(url, stream=True).raw)
    image_processor = ViTImageProcessor(
        size=config.image_size, image_mean=IMAGENET_DEFAULT_MEAN, image_std=IMAGENET_DEFAULT_STD
    )
    inputs = image_processor(images=image, return_tensors="pt")

    
    torch.manual_seed(2)
    outputs = model(**inputs)
    last_hidden_state = outputs.last_hidden_state

    
    
    if "s16" in checkpoint_url:
        expected_slice = torch.tensor([[-1.0915, -1.4876, -1.1809]])
    elif "b16" in checkpoint_url:
        expected_slice = torch.tensor([[14.2889, -18.9045, 11.7281]])
    elif "l16" in checkpoint_url:
        expected_slice = torch.tensor([[41.5028, -22.8681, 45.6475]])
    elif "b4" in checkpoint_url:
        expected_slice = torch.tensor([[-4.3868, 5.2932, -0.4137]])
    else:
        expected_slice = torch.tensor([[-0.1792, -0.6465, 2.4263]])

    
    assert torch.allclose(last_hidden_state[:, 0, :3], expected_slice, atol=1e-4)

    print(f"Saving model to {pytorch_dump_folder_path}")
    model.save_pretrained(pytorch_dump_folder_path)

    print(f"Saving image processor to {pytorch_dump_folder_path}")
    image_processor.save_pretrained(pytorch_dump_folder_path)


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    
    parser.add_argument(
        "--checkpoint_url",
        default="https://dl.fbaipublicfiles.com/msn/vits16_800ep.pth.tar",
        type=str,
        help="URL of the checkpoint you'd like to convert.",
    )
    parser.add_argument(
        "--pytorch_dump_folder_path", default=None, type=str, help="Path to the output PyTorch model directory."
    )

    args = parser.parse_args()
    convert_vit_msn_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path)