codingutf8 2021 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 checkpoints trained with the dino method import argparse import json from pathlib import path import requests import torch from huggingfacehub import hfhubdownload from pil import image from transformers import vitconfig vitforimageclassification vitimageprocessor vitmodel from transformers utils import logging logging setverbosityinfo logger logging getloggername 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 appendfblocks i norm1 weight fvit encoder layer i layernormbefore weight renamekeys appendfblocks i norm1 bias fvit encoder layer i layernormbefore bias renamekeys appendfblocks i attn proj weight fvit encoder layer i attention output dense weight renamekeys appendfblocks i attn proj bias fvit encoder layer i attention output dense bias renamekeys appendfblocks i norm2 weight fvit encoder layer i layernormafter weight renamekeys appendfblocks i norm2 bias fvit encoder layer i layernormafter bias renamekeys appendfblocks i mlp fc1 weight fvit encoder layer i intermediate dense weight renamekeys appendfblocks i mlp fc1 bias fvit encoder layer i intermediate dense bias renamekeys appendfblocks i mlp fc2 weight fvit encoder layer i output dense weight renamekeys appendfblocks i mlp fc2 bias fvit encoder layer i output dense bias projection layer position embeddings renamekeys extend clstoken vit embeddings clstoken patchembed proj weight vit embeddings patchembeddings projection weight patchembed proj bias vit embeddings patchembeddings projection bias posembed vit embeddings positionembeddings if basemodel layernorm pooler renamekeys extend norm weight layernorm weight 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 popfblocks i attn qkv weight inprojbias statedict popfblocks 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 renamekeydct old new val dct popold dctnew val we will verify our results on an image of cute cats def prepareimg url http images cocodataset orgval2017000000039769 jpg im image openrequests geturl streamtrue raw return im torch nograd def convertvitcheckpointmodelname pytorchdumpfolderpath basemodeltrue define default vit configuration config vitconfig patchsize if modelname1 8 config patchsize 8 set labels if required if not basemodel config numlabels 1000 repoid huggingfacelabelfiles filename imagenet1kid2label json id2label json loadopenhfhubdownloadrepoid filename repotypedataset r id2label intk v for k v in id2label items config id2label id2label config label2id v k for k v in id2label items size of the architecture if modelname in dinovits8 dinovits16 config hiddensize 384 config intermediatesize 1536 config numhiddenlayers 12 config numattentionheads 6 load original model from torch hub originalmodel torch hub loadfacebookresearchdino main modelname originalmodel eval load statedict of original model remove and rename some keys statedict originalmodel statedict if basemodel removeclassificationheadstatedict renamekeys createrenamekeysconfig basemodelbasemodel for src dest in renamekeys renamekeystatedict src dest readinqkvstatedict config basemodel load huggingface model if basemodel model vitmodelconfig addpoolinglayerfalse eval else model vitforimageclassificationconfig eval model loadstatedictstatedict check outputs on an image prepared by vitimageprocessor imageprocessor vitimageprocessor encoding imageprocessorimagesprepareimg returntensorspt pixelvalues encodingpixelvalues outputs modelpixelvalues if basemodel finalhiddenstateclstoken originalmodelpixelvalues assert torch allclosefinalhiddenstateclstoken outputs lasthiddenstate 0 atol1e1 else logits originalmodelpixelvalues assert logits shape outputs logits shape assert torch allcloselogits outputs logits atol1e3 pathpytorchdumpfolderpath mkdirexistoktrue printfsaving model modelname to pytorchdumpfolderpath model savepretrainedpytorchdumpfolderpath printfsaving image processor to pytorchdumpfolderpath imageprocessor savepretrainedpytorchdumpfolderpath if name main parser argparse argumentparser required parameters parser addargument modelname defaultdinovitb16 typestr helpname of the model trained with dino you d like to convert parser addargument pytorchdumpfolderpath defaultnone typestr helppath to the output pytorch model directory parser addargument basemodel actionstoretrue helpwhether to only convert the base model no projection head weights parser setdefaultsbasemodeltrue args parser parseargs convertvitcheckpointargs modelname args pytorchdumpfolderpath args basemodel coding utf 8 2021 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 checkpoints trained with the dino method 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 we will verify our results on an image of cute cats copy paste tweak model s weights to our vit structure define default vit configuration patch_size set labels if required size of the architecture load original model from torch hub load state_dict of original model remove and rename some keys load huggingface model check outputs on an image prepared by vitimageprocessor required parameters
import argparse
import json
from pathlib import Path

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

from transformers import ViTConfig, ViTForImageClassification, ViTImageProcessor, ViTModel
from transformers.utils import logging


logging.set_verbosity_info()
logger = logging.get_logger(__name__)



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

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

    if base_model:
        
        rename_keys.extend(
            [
                ("norm.weight", "layernorm.weight"),
                ("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"blocks.{i}.attn.qkv.weight")
        in_proj_bias = state_dict.pop(f"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 rename_key(dct, old, new):
    val = dct.pop(old)
    dct[new] = val



def prepare_img():
    url = "http://images.cocodataset.org/val2017/000000039769.jpg"
    im = Image.open(requests.get(url, stream=True).raw)
    return im


@torch.no_grad()
def convert_vit_checkpoint(model_name, pytorch_dump_folder_path, base_model=True):
    

    
    config = ViTConfig()
    
    if model_name[-1] == "8":
        config.patch_size = 8
    
    if not base_model:
        config.num_labels = 1000
        repo_id = "huggingface/label-files"
        filename = "imagenet-1k-id2label.json"
        id2label = json.load(open(hf_hub_download(repo_id, filename, repo_type="dataset"), "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 model_name in ["dino_vits8", "dino_vits16"]:
        config.hidden_size = 384
        config.intermediate_size = 1536
        config.num_hidden_layers = 12
        config.num_attention_heads = 6

    
    original_model = torch.hub.load("facebookresearch/dino:main", model_name)
    original_model.eval()

    
    state_dict = original_model.state_dict()
    if base_model:
        remove_classification_head_(state_dict)
    rename_keys = create_rename_keys(config, base_model=base_model)
    for src, dest in rename_keys:
        rename_key(state_dict, src, dest)
    read_in_q_k_v(state_dict, config, base_model)

    
    if base_model:
        model = ViTModel(config, add_pooling_layer=False).eval()
    else:
        model = ViTForImageClassification(config).eval()
    model.load_state_dict(state_dict)

    
    image_processor = ViTImageProcessor()
    encoding = image_processor(images=prepare_img(), return_tensors="pt")
    pixel_values = encoding["pixel_values"]
    outputs = model(pixel_values)

    if base_model:
        final_hidden_state_cls_token = original_model(pixel_values)
        assert torch.allclose(final_hidden_state_cls_token, outputs.last_hidden_state[:, 0, :], atol=1e-1)
    else:
        logits = original_model(pixel_values)
        assert logits.shape == outputs.logits.shape
        assert torch.allclose(logits, outputs.logits, atol=1e-3)

    Path(pytorch_dump_folder_path).mkdir(exist_ok=True)
    print(f"Saving model {model_name} 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(
        "--model_name",
        default="dino_vitb16",
        type=str,
        help="Name of the model trained with DINO you'd like to convert.",
    )
    parser.add_argument(
        "--pytorch_dump_folder_path", default=None, type=str, help="Path to the output PyTorch model directory."
    )
    parser.add_argument(
        "--base_model",
        action="store_true",
        help="Whether to only convert the base model (no projection head weights).",
    )

    parser.set_defaults(base_model=True)
    args = parser.parse_args()
    convert_vit_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.base_model)