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 and nondistilled deit checkpoints from the timm library import argparse from pathlib import path import requests import timm import torch from pil import image from timm data import imagenetinfo inferimagenetsubset from transformers import deitimageprocessor 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 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 convertvitcheckpointvitname pytorchdumpfolderpath define default vit configuration config vitconfig basemodel false load original model from timm timmmodel timm createmodelvitname pretrainedtrue timmmodel eval detect unsupported vit models in transformers fcnorm is present if not isinstancegetattrtimmmodel fcnorm none torch nn identity raise valueerrorfvitname is not supported in transformers because of the presence of fcnorm use of global average pooling in combination or without class token if getattrtimmmodel globalpool none avg raise valueerrorfvitname is not supported in transformers because of use of global average pooling clip style vit with normpre layer present if clip in vitname and not isinstancegetattrtimmmodel normpre none torch nn identity raise valueerror fvitname is not supported in transformers because it s a clip style vit with normpre layer siglip style vit with attnpool layer present if siglip in vitname and getattrtimmmodel globalpool none map raise valueerror fvitname is not supported in transformers because it s a siglip style vit with attnpool use of layer scale in vit model blocks if not isinstancegetattrtimmmodel blocks0 ls1 none torch nn identity or not isinstance getattrtimmmodel blocks0 ls2 none torch nn identity raise valueerrorfvitname is not supported in transformers because it uses a layer scale in its blocks hybrid resnetvits if not isinstancetimmmodel patchembed timm layers patchembed raise valueerrorfvitname is not supported in transformers because it is a hybrid resnetvit get patch size and image size from the patch embedding submodule config patchsize timmmodel patchembed patchsize0 config imagesize timmmodel patchembed imgsize0 retrieve architecturespecific parameters from the timm model config hiddensize timmmodel embeddim config intermediatesize timmmodel blocks0 mlp fc1 outfeatures config numhiddenlayers lentimmmodel blocks config numattentionheads timmmodel blocks0 attn numheads check whether the model has a classification head or not if timmmodel numclasses 0 config numlabels timmmodel numclasses infer imagenet subset from timm model imagenetsubset inferimagenetsubsettimmmodel datasetinfo imagenetinfoimagenetsubset config id2label i datasetinfo indextolabelnamei for i in rangedatasetinfo numclasses config label2id v k for k v in config id2label items else printfvitname is going to be converted as a feature extractor only basemodel true load statedict of original model statedict timmmodel statedict remove and rename some keys in the state dict if basemodel removeclassificationheadstatedict renamekeys createrenamekeysconfig basemodel 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 vitimageprocessordeitimageprocessor if deit in vitname imageprocessor deitimageprocessorsizeconfig imagesize else imageprocessor vitimageprocessorsizeconfig imagesize encoding imageprocessorimagesprepareimg returntensorspt pixelvalues encodingpixelvalues outputs modelpixelvalues if basemodel timmpooledoutput timmmodel forwardfeaturespixelvalues assert timmpooledoutput shape outputs lasthiddenstate shape assert torch allclosetimmpooledoutput outputs lasthiddenstate atol1e1 else timmlogits timmmodelpixelvalues assert timmlogits shape outputs logits shape assert torch allclosetimmlogits outputs logits atol1e3 pathpytorchdumpfolderpath mkdirexistoktrue printfsaving model vitname to pytorchdumpfolderpath model savepretrainedpytorchdumpfolderpath printfsaving image processor to pytorchdumpfolderpath imageprocessor savepretrainedpytorchdumpfolderpath if name main parser argparse argumentparser required parameters parser addargument vitname defaultvitbasepatch16224 typestr helpname of the vit timm model you d like to convert parser addargument pytorchdumpfolderpath defaultnone typestr helppath to the output pytorch model directory args parser parseargs convertvitcheckpointargs vitname 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 vit and non distilled deit 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 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 load original model from timm detect unsupported vit models in transformers fc_norm is present use of global average pooling in combination or without class token clip style vit with norm_pre layer present siglip style vit with attn_pool layer present use of layer scale in vit model blocks hybrid resnet vits get patch size and image size from the patch embedding submodule retrieve architecture specific parameters from the timm model check whether the model has a classification head or not infer imagenet subset from timm model load state_dict of original model remove and rename some keys in the state dict load huggingface model check outputs on an image prepared by vitimageprocessor deitimageprocessor required parameters
import argparse
from pathlib import Path

import requests
import timm
import torch
from PIL import Image
from timm.data import ImageNetInfo, infer_imagenet_subset

from transformers import DeiTImageProcessor, 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(vit_name, pytorch_dump_folder_path):
    

    
    config = ViTConfig()
    base_model = False

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

    
    
    if not isinstance(getattr(timm_model, "fc_norm", None), torch.nn.Identity):
        raise ValueError(f"{vit_name} is not supported in transformers because of the presence of fc_norm.")

    
    if getattr(timm_model, "global_pool", None) == "avg":
        raise ValueError(f"{vit_name} is not supported in transformers because of use of global average pooling.")

    
    if "clip" in vit_name and not isinstance(getattr(timm_model, "norm_pre", None), torch.nn.Identity):
        raise ValueError(
            f"{vit_name} is not supported in transformers because it's a CLIP style ViT with norm_pre layer."
        )

    
    if "siglip" in vit_name and getattr(timm_model, "global_pool", None) == "map":
        raise ValueError(
            f"{vit_name} is not supported in transformers because it's a SigLIP style ViT with attn_pool."
        )

    
    if not isinstance(getattr(timm_model.blocks[0], "ls1", None), torch.nn.Identity) or not isinstance(
        getattr(timm_model.blocks[0], "ls2", None), torch.nn.Identity
    ):
        raise ValueError(f"{vit_name} is not supported in transformers because it uses a layer scale in its blocks.")

    
    if not isinstance(timm_model.patch_embed, timm.layers.PatchEmbed):
        raise ValueError(f"{vit_name} is not supported in transformers because it is a hybrid ResNet-ViT.")

    
    config.patch_size = timm_model.patch_embed.patch_size[0]
    config.image_size = timm_model.patch_embed.img_size[0]

    
    config.hidden_size = timm_model.embed_dim
    config.intermediate_size = timm_model.blocks[0].mlp.fc1.out_features
    config.num_hidden_layers = len(timm_model.blocks)
    config.num_attention_heads = timm_model.blocks[0].attn.num_heads

    
    if timm_model.num_classes != 0:
        config.num_labels = timm_model.num_classes
        
        imagenet_subset = infer_imagenet_subset(timm_model)
        dataset_info = ImageNetInfo(imagenet_subset)
        config.id2label = {i: dataset_info.index_to_label_name(i) for i in range(dataset_info.num_classes())}
        config.label2id = {v: k for k, v in config.id2label.items()}
    else:
        print(f"{vit_name} is going to be converted as a feature extractor only.")
        base_model = True

    
    state_dict = timm_model.state_dict()

    
    if base_model:
        remove_classification_head_(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)

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

    
    if "deit" in vit_name:
        image_processor = DeiTImageProcessor(size=config.image_size)
    else:
        image_processor = ViTImageProcessor(size=config.image_size)
    encoding = image_processor(images=prepare_img(), return_tensors="pt")
    pixel_values = encoding["pixel_values"]
    outputs = model(pixel_values)

    if base_model:
        timm_pooled_output = timm_model.forward_features(pixel_values)
        assert timm_pooled_output.shape == outputs.last_hidden_state.shape
        assert torch.allclose(timm_pooled_output, outputs.last_hidden_state, atol=1e-1)
    else:
        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 {vit_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(
        "--vit_name",
        default="vit_base_patch16_224",
        type=str,
        help="Name of the ViT 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_vit_checkpoint(args.vit_name, args.pytorch_dump_folder_path)