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 hybrid 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 timm data import resolvedataconfig from timm data transformsfactory import createtransform from transformers import bitconfig vithybridconfig vithybridforimageclassification vithybridimageprocessor vithybridmodel from transformers imageutils import pilimageresampling 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 fmt off stem renamekeys appendclstoken vit embeddings clstoken renamekeys appendposembed vit embeddings positionembeddings renamekeys appendpatchembed proj weight vit embeddings patchembeddings projection weight renamekeys appendpatchembed proj bias vit embeddings patchembeddings projection bias backbone renamekeys appendpatchembed backbone stem conv weight vit embeddings patchembeddings backbone bit embedder convolution weight renamekeys appendpatchembed backbone stem norm weight vit embeddings patchembeddings backbone bit embedder norm weight renamekeys appendpatchembed backbone stem norm bias vit embeddings patchembeddings backbone bit embedder norm bias for stageidx in rangelenconfig backboneconfig depths for layeridx in rangeconfig backboneconfig depthsstageidx renamekeys appendfpatchembed backbone stages stageidx blocks layeridx conv1 weight fvit embeddings patchembeddings backbone bit encoder stages stageidx layers layeridx conv1 weight renamekeys appendfpatchembed backbone stages stageidx blocks layeridx norm1 weight fvit embeddings patchembeddings backbone bit encoder stages stageidx layers layeridx norm1 weight renamekeys appendfpatchembed backbone stages stageidx blocks layeridx norm1 bias fvit embeddings patchembeddings backbone bit encoder stages stageidx layers layeridx norm1 bias renamekeys appendfpatchembed backbone stages stageidx blocks layeridx conv2 weight fvit embeddings patchembeddings backbone bit encoder stages stageidx layers layeridx conv2 weight renamekeys appendfpatchembed backbone stages stageidx blocks layeridx norm2 weight fvit embeddings patchembeddings backbone bit encoder stages stageidx layers layeridx norm2 weight renamekeys appendfpatchembed backbone stages stageidx blocks layeridx norm2 bias fvit embeddings patchembeddings backbone bit encoder stages stageidx layers layeridx norm2 bias renamekeys appendfpatchembed backbone stages stageidx blocks layeridx conv3 weight fvit embeddings patchembeddings backbone bit encoder stages stageidx layers layeridx conv3 weight renamekeys appendfpatchembed backbone stages stageidx blocks layeridx norm3 weight fvit embeddings patchembeddings backbone bit encoder stages stageidx layers layeridx norm3 weight renamekeys appendfpatchembed backbone stages stageidx blocks layeridx norm3 bias fvit embeddings patchembeddings backbone bit encoder stages stageidx layers layeridx norm3 bias renamekeys appendfpatchembed backbone stages stageidx blocks 0 downsample conv weight fvit embeddings patchembeddings backbone bit encoder stages stageidx layers 0 downsample conv weight renamekeys appendfpatchembed backbone stages stageidx blocks 0 downsample norm weight fvit embeddings patchembeddings backbone bit encoder stages stageidx layers 0 downsample norm weight renamekeys appendfpatchembed backbone stages stageidx blocks 0 downsample norm bias fvit embeddings patchembeddings backbone bit encoder stages stageidx layers 0 downsample norm bias transformer encoder 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 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 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 fmt on 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 pushtohubfalse define default vit hybrid configuration backboneconfig bitconfig globalpaddingsame layertypebottleneck depths3 4 9 outfeaturesstage3 embeddingdynamicpaddingtrue config vithybridconfigbackboneconfigbackboneconfig imagesize384 numlabels1000 basemodel false load original model from timm timmmodel timm createmodelvitname pretrainedtrue timmmodel eval load statedict of original model remove and rename some keys statedict timmmodel statedict if basemodel removeclassificationheadstatedict renamekeys createrenamekeysconfig basemodel for src dest in renamekeys renamekeystatedict src dest readinqkvstatedict config basemodel 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 load huggingface model if vitname5 in21k model vithybridmodelconfig eval else model vithybridforimageclassificationconfig eval model loadstatedictstatedict create image processor transform createtransformresolvedataconfig modeltimmmodel timmtransforms transform transforms pillowresamplings bilinear pilimageresampling bilinear bicubic pilimageresampling bicubic nearest pilimageresampling nearest processor vithybridimageprocessor doresizetrue sizeshortestedge timmtransforms0 size resamplepillowresamplingstimmtransforms0 interpolation value docentercroptrue cropsizeheight timmtransforms1 size0 width timmtransforms1 size1 donormalizetrue imagemeantimmtransforms1 mean tolist imagestdtimmtransforms1 std tolist image prepareimg timmpixelvalues transformimage unsqueeze0 pixelvalues processorimage returntensorspt pixelvalues verify pixel values assert torch allclosetimmpixelvalues pixelvalues verify logits with torch nograd outputs modelpixelvalues logits outputs logits printpredicted class logits argmax1 item if basemodel timmpooledoutput timmmodel forwardfeaturespixelvalues assert timmpooledoutput shape outputs pooleroutput shape assert torch allclosetimmpooledoutput outputs pooleroutput atol1e3 else timmlogits timmmodelpixelvalues assert timmlogits shape outputs logits shape assert torch allclosetimmlogits outputs logits atol1e3 printlooks ok if pytorchdumpfolderpath is not none pathpytorchdumpfolderpath mkdirexistoktrue printfsaving model vitname to pytorchdumpfolderpath model savepretrainedpytorchdumpfolderpath printfsaving processor to pytorchdumpfolderpath processor savepretrainedpytorchdumpfolderpath if pushtohub printfpushing model and processor to the hub vitname model pushtohubfybelkadavitname processor pushtohubfybelkadavitname if name main parser argparse argumentparser required parameters parser addargument vitname defaultvitbaser50s16384 typestr helpname of the hybrid vit timm model you d like to convert parser addargument pytorchdumpfolderpath defaultnone typestr helppath to the output pytorch model directory parser addargument pushtohub actionstoretrue helpwhether to upload the model to the huggingface hub args parser parseargs convertvitcheckpointargs vitname args pytorchdumpfolderpath args pushtohub 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 hybrid checkpoints from the timm library here we list all keys to be renamed original name on the left our name on the right fmt off stem backbone transformer encoder encoder layers output projection 2 feedforward neural networks and 2 layernorms layernorm pooler if just the base model we should remove vit from all keys that start with vit layernorm classification head fmt on 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 hybrid configuration load original model from timm load state_dict of original model remove and rename some keys load huggingface model create image processor verify pixel values verify logits 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 timm.data import resolve_data_config
from timm.data.transforms_factory import create_transform

from transformers import (
    BitConfig,
    ViTHybridConfig,
    ViTHybridForImageClassification,
    ViTHybridImageProcessor,
    ViTHybridModel,
)
from transformers.image_utils import PILImageResampling
from transformers.utils import logging


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



def create_rename_keys(config, base_model=False):
    rename_keys = []

    
    
    rename_keys.append(("cls_token", "vit.embeddings.cls_token"))
    rename_keys.append(("pos_embed", "vit.embeddings.position_embeddings"))

    rename_keys.append(("patch_embed.proj.weight", "vit.embeddings.patch_embeddings.projection.weight"))
    rename_keys.append(("patch_embed.proj.bias", "vit.embeddings.patch_embeddings.projection.bias"))

    
    rename_keys.append(("patch_embed.backbone.stem.conv.weight", "vit.embeddings.patch_embeddings.backbone.bit.embedder.convolution.weight"))
    rename_keys.append(("patch_embed.backbone.stem.norm.weight", "vit.embeddings.patch_embeddings.backbone.bit.embedder.norm.weight"))
    rename_keys.append(("patch_embed.backbone.stem.norm.bias", "vit.embeddings.patch_embeddings.backbone.bit.embedder.norm.bias"))

    for stage_idx in range(len(config.backbone_config.depths)):
        for layer_idx in range(config.backbone_config.depths[stage_idx]):
            rename_keys.append((f"patch_embed.backbone.stages.{stage_idx}.blocks.{layer_idx}.conv1.weight", f"vit.embeddings.patch_embeddings.backbone.bit.encoder.stages.{stage_idx}.layers.{layer_idx}.conv1.weight"))
            rename_keys.append((f"patch_embed.backbone.stages.{stage_idx}.blocks.{layer_idx}.norm1.weight", f"vit.embeddings.patch_embeddings.backbone.bit.encoder.stages.{stage_idx}.layers.{layer_idx}.norm1.weight"))
            rename_keys.append((f"patch_embed.backbone.stages.{stage_idx}.blocks.{layer_idx}.norm1.bias", f"vit.embeddings.patch_embeddings.backbone.bit.encoder.stages.{stage_idx}.layers.{layer_idx}.norm1.bias"))
            rename_keys.append((f"patch_embed.backbone.stages.{stage_idx}.blocks.{layer_idx}.conv2.weight", f"vit.embeddings.patch_embeddings.backbone.bit.encoder.stages.{stage_idx}.layers.{layer_idx}.conv2.weight"))
            rename_keys.append((f"patch_embed.backbone.stages.{stage_idx}.blocks.{layer_idx}.norm2.weight", f"vit.embeddings.patch_embeddings.backbone.bit.encoder.stages.{stage_idx}.layers.{layer_idx}.norm2.weight"))
            rename_keys.append((f"patch_embed.backbone.stages.{stage_idx}.blocks.{layer_idx}.norm2.bias", f"vit.embeddings.patch_embeddings.backbone.bit.encoder.stages.{stage_idx}.layers.{layer_idx}.norm2.bias"))
            rename_keys.append((f"patch_embed.backbone.stages.{stage_idx}.blocks.{layer_idx}.conv3.weight", f"vit.embeddings.patch_embeddings.backbone.bit.encoder.stages.{stage_idx}.layers.{layer_idx}.conv3.weight"))
            rename_keys.append((f"patch_embed.backbone.stages.{stage_idx}.blocks.{layer_idx}.norm3.weight", f"vit.embeddings.patch_embeddings.backbone.bit.encoder.stages.{stage_idx}.layers.{layer_idx}.norm3.weight"))
            rename_keys.append((f"patch_embed.backbone.stages.{stage_idx}.blocks.{layer_idx}.norm3.bias", f"vit.embeddings.patch_embeddings.backbone.bit.encoder.stages.{stage_idx}.layers.{layer_idx}.norm3.bias"))

        rename_keys.append((f"patch_embed.backbone.stages.{stage_idx}.blocks.0.downsample.conv.weight", f"vit.embeddings.patch_embeddings.backbone.bit.encoder.stages.{stage_idx}.layers.0.downsample.conv.weight"))
        rename_keys.append((f"patch_embed.backbone.stages.{stage_idx}.blocks.0.downsample.norm.weight", f"vit.embeddings.patch_embeddings.backbone.bit.encoder.stages.{stage_idx}.layers.0.downsample.norm.weight"))
        rename_keys.append((f"patch_embed.backbone.stages.{stage_idx}.blocks.0.downsample.norm.bias", f"vit.embeddings.patch_embeddings.backbone.bit.encoder.stages.{stage_idx}.layers.0.downsample.norm.bias"))

    
    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"))

    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("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, push_to_hub=False):
    

    
    backbone_config = BitConfig(
        global_padding="same",
        layer_type="bottleneck",
        depths=(3, 4, 9),
        out_features=["stage3"],
        embedding_dynamic_padding=True,
    )
    config = ViTHybridConfig(backbone_config=backbone_config, image_size=384, num_labels=1000)
    base_model = False

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

    
    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)

    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 vit_name[-5:] == "in21k":
        model = ViTHybridModel(config).eval()
    else:
        model = ViTHybridForImageClassification(config).eval()
    model.load_state_dict(state_dict)

    
    transform = create_transform(**resolve_data_config({}, model=timm_model))
    timm_transforms = transform.transforms

    pillow_resamplings = {
        "bilinear": PILImageResampling.BILINEAR,
        "bicubic": PILImageResampling.BICUBIC,
        "nearest": PILImageResampling.NEAREST,
    }

    processor = ViTHybridImageProcessor(
        do_resize=True,
        size={"shortest_edge": timm_transforms[0].size},
        resample=pillow_resamplings[timm_transforms[0].interpolation.value],
        do_center_crop=True,
        crop_size={"height": timm_transforms[1].size[0], "width": timm_transforms[1].size[1]},
        do_normalize=True,
        image_mean=timm_transforms[-1].mean.tolist(),
        image_std=timm_transforms[-1].std.tolist(),
    )

    image = prepare_img()
    timm_pixel_values = transform(image).unsqueeze(0)
    pixel_values = processor(image, return_tensors="pt").pixel_values

    
    assert torch.allclose(timm_pixel_values, pixel_values)

    
    with torch.no_grad():
        outputs = model(pixel_values)
        logits = outputs.logits

    print("Predicted class:", logits.argmax(-1).item())
    if base_model:
        timm_pooled_output = timm_model.forward_features(pixel_values)
        assert timm_pooled_output.shape == outputs.pooler_output.shape
        assert torch.allclose(timm_pooled_output, outputs.pooler_output, atol=1e-3)
    else:
        timm_logits = timm_model(pixel_values)
        assert timm_logits.shape == outputs.logits.shape
        assert torch.allclose(timm_logits, outputs.logits, atol=1e-3)
    print("Looks ok!")

    if pytorch_dump_folder_path is not None:
        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 processor to {pytorch_dump_folder_path}")
        processor.save_pretrained(pytorch_dump_folder_path)

    if push_to_hub:
        print(f"Pushing model and processor to the hub {vit_name}")
        model.push_to_hub(f"ybelkada/{vit_name}")
        processor.push_to_hub(f"ybelkada/{vit_name}")


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    
    parser.add_argument(
        "--vit_name",
        default="vit_base_r50_s16_384",
        type=str,
        help="Name of the hybrid 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."
    )
    parser.add_argument(
        "--push_to_hub", action="store_true", help="Whether to upload the model to the HuggingFace hub."
    )

    args = parser.parse_args()
    convert_vit_checkpoint(args.vit_name, args.pytorch_dump_folder_path, args.push_to_hub)