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 deit distilled checkpoints from the timm library import argparse import json from pathlib import path import requests import timm import torch from huggingfacehub import hfhubdownload from pil import image from transformers import deitconfig deitforimageclassificationwithteacher deitimageprocessor 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 fdeit encoder layer i layernormbefore weight renamekeys appendfblocks i norm1 bias fdeit encoder layer i layernormbefore bias renamekeys appendfblocks i attn proj weight fdeit encoder layer i attention output dense weight renamekeys appendfblocks i attn proj bias fdeit encoder layer i attention output dense bias renamekeys appendfblocks i norm2 weight fdeit encoder layer i layernormafter weight renamekeys appendfblocks i norm2 bias fdeit encoder layer i layernormafter bias renamekeys appendfblocks i mlp fc1 weight fdeit encoder layer i intermediate dense weight renamekeys appendfblocks i mlp fc1 bias fdeit encoder layer i intermediate dense bias renamekeys appendfblocks i mlp fc2 weight fdeit encoder layer i output dense weight renamekeys appendfblocks i mlp fc2 bias fdeit encoder layer i output dense bias projection layer position embeddings renamekeys extend clstoken deit embeddings clstoken disttoken deit embeddings distillationtoken patchembed proj weight deit embeddings patchembeddings projection weight patchembed proj bias deit embeddings patchembeddings projection bias posembed deit embeddings positionembeddings if basemodel layernorm pooler renamekeys extend norm weight layernorm weight norm bias layernorm bias prelogits fc weight pooler dense weight prelogits fc bias pooler dense bias if just the base model we should remove deit from all keys that start with deit renamekeys pair0 pair14 if pair1 startswithdeit else pair for pair in renamekeys else layernorm classification heads renamekeys extend norm weight deit layernorm weight norm bias deit layernorm bias head weight clsclassifier weight head bias clsclassifier bias headdist weight distillationclassifier weight headdist bias distillationclassifier 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 deit 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 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 convertdeitcheckpointdeitname pytorchdumpfolderpath define default deit configuration config deitconfig all deit models have finetuned heads basemodel false dataset finetuned on imagenet 2012 patchsize and imagesize 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 config patchsize intdeitname6 4 config imagesize intdeitname3 size of the architecture if deitname9 startswithtiny config hiddensize 192 config intermediatesize 768 config numhiddenlayers 12 config numattentionheads 3 elif deitname9 startswithsmall config hiddensize 384 config intermediatesize 1536 config numhiddenlayers 12 config numattentionheads 6 if deitname9 startswithbase pass elif deitname4 startswithlarge config hiddensize 1024 config intermediatesize 4096 config numhiddenlayers 24 config numattentionheads 16 load original model from timm timmmodel timm createmodeldeitname pretrainedtrue timmmodel eval load statedict of original model remove and rename some keys statedict timmmodel statedict renamekeys createrenamekeysconfig basemodel for src dest in renamekeys renamekeystatedict src dest readinqkvstatedict config basemodel load huggingface model model deitforimageclassificationwithteacherconfig eval model loadstatedictstatedict check outputs on an image prepared by deitimageprocessor size int 256 224 config imagesize to maintain same ratio w r t 224 images see https github comfacebookresearchdeitblobab5715372db8c6cad5740714b2216d55aeae052edatasets pyl103 imageprocessor deitimageprocessorsizesize cropsizeconfig imagesize encoding imageprocessorimagesprepareimg returntensorspt pixelvalues encodingpixelvalues outputs modelpixelvalues timmlogits timmmodelpixelvalues assert timmlogits shape outputs logits shape assert torch allclosetimmlogits outputs logits atol1e3 pathpytorchdumpfolderpath mkdirexistoktrue printfsaving model deitname to pytorchdumpfolderpath model savepretrainedpytorchdumpfolderpath printfsaving image processor to pytorchdumpfolderpath imageprocessor savepretrainedpytorchdumpfolderpath if name main parser argparse argumentparser required parameters parser addargument deitname defaultvitdeitbasedistilledpatch16224 typestr helpname of the deit timm model you d like to convert parser addargument pytorchdumpfolderpath defaultnone typestr helppath to the output pytorch model directory args parser parseargs convertdeitcheckpointargs deitname args pytorchdumpfolderpath 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 deit distilled checkpoints from the timm library 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 deit from all keys that start with deit layernorm classification heads 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 deit structure define default deit configuration all deit models have fine tuned heads dataset fine tuned on imagenet 2012 patch_size and image_size size of the architecture load original model from timm load state_dict of original model remove and rename some keys load huggingface model check outputs on an image prepared by deitimageprocessor to maintain same ratio w r t 224 images see https github com facebookresearch deit blob ab5715372db8c6cad5740714b2216d55aeae052e datasets py l103 required parameters
import argparse
import json
from pathlib import Path

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

from transformers import DeiTConfig, DeiTForImageClassificationWithTeacher, DeiTImageProcessor
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"deit.encoder.layer.{i}.layernorm_before.weight"))
        rename_keys.append((f"blocks.{i}.norm1.bias", f"deit.encoder.layer.{i}.layernorm_before.bias"))
        rename_keys.append((f"blocks.{i}.attn.proj.weight", f"deit.encoder.layer.{i}.attention.output.dense.weight"))
        rename_keys.append((f"blocks.{i}.attn.proj.bias", f"deit.encoder.layer.{i}.attention.output.dense.bias"))
        rename_keys.append((f"blocks.{i}.norm2.weight", f"deit.encoder.layer.{i}.layernorm_after.weight"))
        rename_keys.append((f"blocks.{i}.norm2.bias", f"deit.encoder.layer.{i}.layernorm_after.bias"))
        rename_keys.append((f"blocks.{i}.mlp.fc1.weight", f"deit.encoder.layer.{i}.intermediate.dense.weight"))
        rename_keys.append((f"blocks.{i}.mlp.fc1.bias", f"deit.encoder.layer.{i}.intermediate.dense.bias"))
        rename_keys.append((f"blocks.{i}.mlp.fc2.weight", f"deit.encoder.layer.{i}.output.dense.weight"))
        rename_keys.append((f"blocks.{i}.mlp.fc2.bias", f"deit.encoder.layer.{i}.output.dense.bias"))

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

    if base_model:
        
        rename_keys.extend(
            [
                ("norm.weight", "layernorm.weight"),
                ("norm.bias", "layernorm.bias"),
                ("pre_logits.fc.weight", "pooler.dense.weight"),
                ("pre_logits.fc.bias", "pooler.dense.bias"),
            ]
        )

        
        rename_keys = [(pair[0], pair[1][4:]) if pair[1].startswith("deit") else pair for pair in rename_keys]
    else:
        
        rename_keys.extend(
            [
                ("norm.weight", "deit.layernorm.weight"),
                ("norm.bias", "deit.layernorm.bias"),
                ("head.weight", "cls_classifier.weight"),
                ("head.bias", "cls_classifier.bias"),
                ("head_dist.weight", "distillation_classifier.weight"),
                ("head_dist.bias", "distillation_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 = "deit."
        
        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 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_deit_checkpoint(deit_name, pytorch_dump_folder_path):
    

    
    config = DeiTConfig()
    
    base_model = False
    
    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()}
    config.patch_size = int(deit_name[-6:-4])
    config.image_size = int(deit_name[-3:])
    
    if deit_name[9:].startswith("tiny"):
        config.hidden_size = 192
        config.intermediate_size = 768
        config.num_hidden_layers = 12
        config.num_attention_heads = 3
    elif deit_name[9:].startswith("small"):
        config.hidden_size = 384
        config.intermediate_size = 1536
        config.num_hidden_layers = 12
        config.num_attention_heads = 6
    if deit_name[9:].startswith("base"):
        pass
    elif deit_name[4:].startswith("large"):
        config.hidden_size = 1024
        config.intermediate_size = 4096
        config.num_hidden_layers = 24
        config.num_attention_heads = 16

    
    timm_model = timm.create_model(deit_name, pretrained=True)
    timm_model.eval()

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

    
    model = DeiTForImageClassificationWithTeacher(config).eval()
    model.load_state_dict(state_dict)

    
    size = int(
        (256 / 224) * config.image_size
    )  
    image_processor = DeiTImageProcessor(size=size, crop_size=config.image_size)
    encoding = image_processor(images=prepare_img(), return_tensors="pt")
    pixel_values = encoding["pixel_values"]
    outputs = model(pixel_values)

    timm_logits = timm_model(pixel_values)
    assert timm_logits.shape == outputs.logits.shape
    assert torch.allclose(timm_logits, outputs.logits, atol=1e-3)

    Path(pytorch_dump_folder_path).mkdir(exist_ok=True)
    print(f"Saving model {deit_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(
        "--deit_name",
        default="vit_deit_base_distilled_patch16_224",
        type=str,
        help="Name of the DeiT timm model 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_deit_checkpoint(args.deit_name, args.pytorch_dump_folder_path)