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 cvt checkpoints from the original repository url https github commicrosoftcvt import argparse import json from collections import ordereddict import torch from huggingfacehub import cacheddownload hfhuburl from transformers import autoimageprocessor cvtconfig cvtforimageclassification def embeddingsidx embed embed append fcvt encoder stages idx embedding convolutionembeddings projection weight fstageidx patchembed proj weight embed append fcvt encoder stages idx embedding convolutionembeddings projection bias fstageidx patchembed proj bias embed append fcvt encoder stages idx embedding convolutionembeddings normalization weight fstageidx patchembed norm weight embed append fcvt encoder stages idx embedding convolutionembeddings normalization bias fstageidx patchembed norm bias return embed def attentionidx cnt attentionweights attentionweights append fcvt encoder stages idx layers cnt attention attention convolutionprojectionquery convolutionprojection convolution weight fstageidx blocks cnt attn convprojq conv weight attentionweights append fcvt encoder stages idx layers cnt attention attention convolutionprojectionquery convolutionprojection normalization weight fstageidx blocks cnt attn convprojq bn weight attentionweights append fcvt encoder stages idx layers cnt attention attention convolutionprojectionquery convolutionprojection normalization bias fstageidx blocks cnt attn convprojq bn bias attentionweights append fcvt encoder stages idx layers cnt attention attention convolutionprojectionquery convolutionprojection normalization runningmean fstageidx blocks cnt attn convprojq bn runningmean attentionweights append fcvt encoder stages idx layers cnt attention attention convolutionprojectionquery convolutionprojection normalization runningvar fstageidx blocks cnt attn convprojq bn runningvar attentionweights append fcvt encoder stages idx layers cnt attention attention convolutionprojectionquery convolutionprojection normalization numbatchestracked fstageidx blocks cnt attn convprojq bn numbatchestracked attentionweights append fcvt encoder stages idx layers cnt attention attention convolutionprojectionkey convolutionprojection convolution weight fstageidx blocks cnt attn convprojk conv weight attentionweights append fcvt encoder stages idx layers cnt attention attention convolutionprojectionkey convolutionprojection normalization weight fstageidx blocks cnt attn convprojk bn weight attentionweights append fcvt encoder stages idx layers cnt attention attention convolutionprojectionkey convolutionprojection normalization bias fstageidx blocks cnt attn convprojk bn bias attentionweights append fcvt encoder stages idx layers cnt attention attention convolutionprojectionkey convolutionprojection normalization runningmean fstageidx blocks cnt attn convprojk bn runningmean attentionweights append fcvt encoder stages idx layers cnt attention attention convolutionprojectionkey convolutionprojection normalization runningvar fstageidx blocks cnt attn convprojk bn runningvar attentionweights append fcvt encoder stages idx layers cnt attention attention convolutionprojectionkey convolutionprojection normalization numbatchestracked fstageidx blocks cnt attn convprojk bn numbatchestracked attentionweights append fcvt encoder stages idx layers cnt attention attention convolutionprojectionvalue convolutionprojection convolution weight fstageidx blocks cnt attn convprojv conv weight attentionweights append fcvt encoder stages idx layers cnt attention attention convolutionprojectionvalue convolutionprojection normalization weight fstageidx blocks cnt attn convprojv bn weight attentionweights append fcvt encoder stages idx layers cnt attention attention convolutionprojectionvalue convolutionprojection normalization bias fstageidx blocks cnt attn convprojv bn bias attentionweights append fcvt encoder stages idx layers cnt attention attention convolutionprojectionvalue convolutionprojection normalization runningmean fstageidx blocks cnt attn convprojv bn runningmean attentionweights append fcvt encoder stages idx layers cnt attention attention convolutionprojectionvalue convolutionprojection normalization runningvar fstageidx blocks cnt attn convprojv bn runningvar attentionweights append fcvt encoder stages idx layers cnt attention attention convolutionprojectionvalue convolutionprojection normalization numbatchestracked fstageidx blocks cnt attn convprojv bn numbatchestracked attentionweights append fcvt encoder stages idx layers cnt attention attention projectionquery weight fstageidx blocks cnt attn projq weight attentionweights append fcvt encoder stages idx layers cnt attention attention projectionquery bias fstageidx blocks cnt attn projq bias attentionweights append fcvt encoder stages idx layers cnt attention attention projectionkey weight fstageidx blocks cnt attn projk weight attentionweights append fcvt encoder stages idx layers cnt attention attention projectionkey bias fstageidx blocks cnt attn projk bias attentionweights append fcvt encoder stages idx layers cnt attention attention projectionvalue weight fstageidx blocks cnt attn projv weight attentionweights append fcvt encoder stages idx layers cnt attention attention projectionvalue bias fstageidx blocks cnt attn projv bias attentionweights append fcvt encoder stages idx layers cnt attention output dense weight fstageidx blocks cnt attn proj weight attentionweights append fcvt encoder stages idx layers cnt attention output dense bias fstageidx blocks cnt attn proj bias attentionweights append fcvt encoder stages idx layers cnt intermediate dense weight fstageidx blocks cnt mlp fc1 weight attentionweights append fcvt encoder stages idx layers cnt intermediate dense bias fstageidx blocks cnt mlp fc1 bias attentionweights append fcvt encoder stages idx layers cnt output dense weight fstageidx blocks cnt mlp fc2 weight attentionweights append fcvt encoder stages idx layers cnt output dense bias fstageidx blocks cnt mlp fc2 bias attentionweights append fcvt encoder stages idx layers cnt layernormbefore weight fstageidx blocks cnt norm1 weight attentionweights append fcvt encoder stages idx layers cnt layernormbefore bias fstageidx blocks cnt norm1 bias attentionweights append fcvt encoder stages idx layers cnt layernormafter weight fstageidx blocks cnt norm2 weight attentionweights append fcvt encoder stages idx layers cnt layernormafter bias fstageidx blocks cnt norm2 bias return attentionweights def clstokenidx token token appendfcvt encoder stages idx clstoken stage2 clstoken return token def final head head appendlayernorm weight norm weight head appendlayernorm bias norm bias head appendclassifier weight head weight head appendclassifier bias head bias return head def convertcvtcheckpointcvtmodel imagesize cvtfilename pytorchdumpfolder imglabelsfile imagenet1kid2label json numlabels 1000 repoid huggingfacelabelfiles numlabels numlabels id2label json loadopencacheddownloadhfhuburlrepoid imglabelsfile repotypedataset r id2label intk v for k v in id2label items id2label id2label label2id v k for k v in id2label items config config cvtconfignumlabelsnumlabels id2labelid2label label2idlabel2id for depth size 13 13 1210 if cvtmodel rsplit 114 6 13 config depth 1 2 10 for depth size 21 21 1416 elif cvtmodel rsplit 114 6 21 config depth 1 4 16 for wide cvt similar to wideresnet depth size 24 w24 2 2 20 else config depth 2 2 20 config numheads 3 12 16 config embeddim 192 768 1024 model cvtforimageclassificationconfig imageprocessor autoimageprocessor frompretrainedfacebookconvnextbase22422k1k imageprocessor sizeshortestedge imagesize originalweights torch loadcvtfilename maplocationtorch devicecpu huggingfaceweights ordereddict listofstatedict for idx in rangelenconfig depth if config clstokenidx listofstatedict listofstatedict clstokenidx listofstatedict listofstatedict embeddingsidx for cnt in rangeconfig depthidx listofstatedict listofstatedict attentionidx cnt listofstatedict listofstatedict final for gg in listofstatedict printgg for i in rangelenlistofstatedict huggingfaceweightslistofstatedicti0 originalweightslistofstatedicti1 model loadstatedicthuggingfaceweights model savepretrainedpytorchdumpfolder imageprocessor savepretrainedpytorchdumpfolder download the weights from zoo https 1drv msus ahixjnjblw9rzf3rmw7sslha8h eblq0al if name main parser argparse argumentparser parser addargument cvtmodel defaultcvtw24 typestr helpname of the cvt model you d like to convert parser addargument imagesize default384 typeint helpinput image size parser addargument cvtfilename defaultrcvtmodelscvtw24384x384in22k pth typestr helpinput image size parser addargument pytorchdumpfolderpath defaultnone typestr helppath to the output pytorch model directory args parser parseargs convertcvtcheckpointargs cvtmodel args imagesize args cvtfilename 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 cvt checkpoints from the original repository url https github com microsoft cvt the function helps in renaming embedding layer weights args idx stage number in original model the function helps in renaming attention block layers weights args idx stage number in original model cnt count of blocks in each stage function helps in renaming cls_token weights function helps in renaming final classification layer fucntion to convert the microsoft cvt checkpoint to huggingface checkpoint for depth size 13 13 1 2 10 for depth size 21 21 1 4 16 for wide cvt similar to wide resnet depth size 24 w24 2 2 20 download the weights from zoo https 1drv ms u s ahixjn_j blw9rzf3rmw7sslha8h e blq0al
import argparse
import json
from collections import OrderedDict

import torch
from huggingface_hub import cached_download, hf_hub_url

from transformers import AutoImageProcessor, CvtConfig, CvtForImageClassification


def embeddings(idx):
    
    embed = []
    embed.append(
        (
            f"cvt.encoder.stages.{idx}.embedding.convolution_embeddings.projection.weight",
            f"stage{idx}.patch_embed.proj.weight",
        )
    )
    embed.append(
        (
            f"cvt.encoder.stages.{idx}.embedding.convolution_embeddings.projection.bias",
            f"stage{idx}.patch_embed.proj.bias",
        )
    )
    embed.append(
        (
            f"cvt.encoder.stages.{idx}.embedding.convolution_embeddings.normalization.weight",
            f"stage{idx}.patch_embed.norm.weight",
        )
    )
    embed.append(
        (
            f"cvt.encoder.stages.{idx}.embedding.convolution_embeddings.normalization.bias",
            f"stage{idx}.patch_embed.norm.bias",
        )
    )
    return embed


def attention(idx, cnt):
    
    attention_weights = []
    attention_weights.append(
        (
            f"cvt.encoder.stages.{idx}.layers.{cnt}.attention.attention.convolution_projection_query.convolution_projection.convolution.weight",
            f"stage{idx}.blocks.{cnt}.attn.conv_proj_q.conv.weight",
        )
    )
    attention_weights.append(
        (
            f"cvt.encoder.stages.{idx}.layers.{cnt}.attention.attention.convolution_projection_query.convolution_projection.normalization.weight",
            f"stage{idx}.blocks.{cnt}.attn.conv_proj_q.bn.weight",
        )
    )
    attention_weights.append(
        (
            f"cvt.encoder.stages.{idx}.layers.{cnt}.attention.attention.convolution_projection_query.convolution_projection.normalization.bias",
            f"stage{idx}.blocks.{cnt}.attn.conv_proj_q.bn.bias",
        )
    )
    attention_weights.append(
        (
            f"cvt.encoder.stages.{idx}.layers.{cnt}.attention.attention.convolution_projection_query.convolution_projection.normalization.running_mean",
            f"stage{idx}.blocks.{cnt}.attn.conv_proj_q.bn.running_mean",
        )
    )
    attention_weights.append(
        (
            f"cvt.encoder.stages.{idx}.layers.{cnt}.attention.attention.convolution_projection_query.convolution_projection.normalization.running_var",
            f"stage{idx}.blocks.{cnt}.attn.conv_proj_q.bn.running_var",
        )
    )
    attention_weights.append(
        (
            f"cvt.encoder.stages.{idx}.layers.{cnt}.attention.attention.convolution_projection_query.convolution_projection.normalization.num_batches_tracked",
            f"stage{idx}.blocks.{cnt}.attn.conv_proj_q.bn.num_batches_tracked",
        )
    )
    attention_weights.append(
        (
            f"cvt.encoder.stages.{idx}.layers.{cnt}.attention.attention.convolution_projection_key.convolution_projection.convolution.weight",
            f"stage{idx}.blocks.{cnt}.attn.conv_proj_k.conv.weight",
        )
    )
    attention_weights.append(
        (
            f"cvt.encoder.stages.{idx}.layers.{cnt}.attention.attention.convolution_projection_key.convolution_projection.normalization.weight",
            f"stage{idx}.blocks.{cnt}.attn.conv_proj_k.bn.weight",
        )
    )
    attention_weights.append(
        (
            f"cvt.encoder.stages.{idx}.layers.{cnt}.attention.attention.convolution_projection_key.convolution_projection.normalization.bias",
            f"stage{idx}.blocks.{cnt}.attn.conv_proj_k.bn.bias",
        )
    )
    attention_weights.append(
        (
            f"cvt.encoder.stages.{idx}.layers.{cnt}.attention.attention.convolution_projection_key.convolution_projection.normalization.running_mean",
            f"stage{idx}.blocks.{cnt}.attn.conv_proj_k.bn.running_mean",
        )
    )
    attention_weights.append(
        (
            f"cvt.encoder.stages.{idx}.layers.{cnt}.attention.attention.convolution_projection_key.convolution_projection.normalization.running_var",
            f"stage{idx}.blocks.{cnt}.attn.conv_proj_k.bn.running_var",
        )
    )
    attention_weights.append(
        (
            f"cvt.encoder.stages.{idx}.layers.{cnt}.attention.attention.convolution_projection_key.convolution_projection.normalization.num_batches_tracked",
            f"stage{idx}.blocks.{cnt}.attn.conv_proj_k.bn.num_batches_tracked",
        )
    )
    attention_weights.append(
        (
            f"cvt.encoder.stages.{idx}.layers.{cnt}.attention.attention.convolution_projection_value.convolution_projection.convolution.weight",
            f"stage{idx}.blocks.{cnt}.attn.conv_proj_v.conv.weight",
        )
    )
    attention_weights.append(
        (
            f"cvt.encoder.stages.{idx}.layers.{cnt}.attention.attention.convolution_projection_value.convolution_projection.normalization.weight",
            f"stage{idx}.blocks.{cnt}.attn.conv_proj_v.bn.weight",
        )
    )
    attention_weights.append(
        (
            f"cvt.encoder.stages.{idx}.layers.{cnt}.attention.attention.convolution_projection_value.convolution_projection.normalization.bias",
            f"stage{idx}.blocks.{cnt}.attn.conv_proj_v.bn.bias",
        )
    )
    attention_weights.append(
        (
            f"cvt.encoder.stages.{idx}.layers.{cnt}.attention.attention.convolution_projection_value.convolution_projection.normalization.running_mean",
            f"stage{idx}.blocks.{cnt}.attn.conv_proj_v.bn.running_mean",
        )
    )
    attention_weights.append(
        (
            f"cvt.encoder.stages.{idx}.layers.{cnt}.attention.attention.convolution_projection_value.convolution_projection.normalization.running_var",
            f"stage{idx}.blocks.{cnt}.attn.conv_proj_v.bn.running_var",
        )
    )
    attention_weights.append(
        (
            f"cvt.encoder.stages.{idx}.layers.{cnt}.attention.attention.convolution_projection_value.convolution_projection.normalization.num_batches_tracked",
            f"stage{idx}.blocks.{cnt}.attn.conv_proj_v.bn.num_batches_tracked",
        )
    )
    attention_weights.append(
        (
            f"cvt.encoder.stages.{idx}.layers.{cnt}.attention.attention.projection_query.weight",
            f"stage{idx}.blocks.{cnt}.attn.proj_q.weight",
        )
    )
    attention_weights.append(
        (
            f"cvt.encoder.stages.{idx}.layers.{cnt}.attention.attention.projection_query.bias",
            f"stage{idx}.blocks.{cnt}.attn.proj_q.bias",
        )
    )
    attention_weights.append(
        (
            f"cvt.encoder.stages.{idx}.layers.{cnt}.attention.attention.projection_key.weight",
            f"stage{idx}.blocks.{cnt}.attn.proj_k.weight",
        )
    )
    attention_weights.append(
        (
            f"cvt.encoder.stages.{idx}.layers.{cnt}.attention.attention.projection_key.bias",
            f"stage{idx}.blocks.{cnt}.attn.proj_k.bias",
        )
    )
    attention_weights.append(
        (
            f"cvt.encoder.stages.{idx}.layers.{cnt}.attention.attention.projection_value.weight",
            f"stage{idx}.blocks.{cnt}.attn.proj_v.weight",
        )
    )
    attention_weights.append(
        (
            f"cvt.encoder.stages.{idx}.layers.{cnt}.attention.attention.projection_value.bias",
            f"stage{idx}.blocks.{cnt}.attn.proj_v.bias",
        )
    )
    attention_weights.append(
        (
            f"cvt.encoder.stages.{idx}.layers.{cnt}.attention.output.dense.weight",
            f"stage{idx}.blocks.{cnt}.attn.proj.weight",
        )
    )
    attention_weights.append(
        (
            f"cvt.encoder.stages.{idx}.layers.{cnt}.attention.output.dense.bias",
            f"stage{idx}.blocks.{cnt}.attn.proj.bias",
        )
    )
    attention_weights.append(
        (f"cvt.encoder.stages.{idx}.layers.{cnt}.intermediate.dense.weight", f"stage{idx}.blocks.{cnt}.mlp.fc1.weight")
    )
    attention_weights.append(
        (f"cvt.encoder.stages.{idx}.layers.{cnt}.intermediate.dense.bias", f"stage{idx}.blocks.{cnt}.mlp.fc1.bias")
    )
    attention_weights.append(
        (f"cvt.encoder.stages.{idx}.layers.{cnt}.output.dense.weight", f"stage{idx}.blocks.{cnt}.mlp.fc2.weight")
    )
    attention_weights.append(
        (f"cvt.encoder.stages.{idx}.layers.{cnt}.output.dense.bias", f"stage{idx}.blocks.{cnt}.mlp.fc2.bias")
    )
    attention_weights.append(
        (f"cvt.encoder.stages.{idx}.layers.{cnt}.layernorm_before.weight", f"stage{idx}.blocks.{cnt}.norm1.weight")
    )
    attention_weights.append(
        (f"cvt.encoder.stages.{idx}.layers.{cnt}.layernorm_before.bias", f"stage{idx}.blocks.{cnt}.norm1.bias")
    )
    attention_weights.append(
        (f"cvt.encoder.stages.{idx}.layers.{cnt}.layernorm_after.weight", f"stage{idx}.blocks.{cnt}.norm2.weight")
    )
    attention_weights.append(
        (f"cvt.encoder.stages.{idx}.layers.{cnt}.layernorm_after.bias", f"stage{idx}.blocks.{cnt}.norm2.bias")
    )
    return attention_weights


def cls_token(idx):
    
    token = []
    token.append((f"cvt.encoder.stages.{idx}.cls_token", "stage2.cls_token"))
    return token


def final():
    
    head = []
    head.append(("layernorm.weight", "norm.weight"))
    head.append(("layernorm.bias", "norm.bias"))
    head.append(("classifier.weight", "head.weight"))
    head.append(("classifier.bias", "head.bias"))
    return head


def convert_cvt_checkpoint(cvt_model, image_size, cvt_file_name, pytorch_dump_folder):
    
    img_labels_file = "imagenet-1k-id2label.json"
    num_labels = 1000

    repo_id = "huggingface/label-files"
    num_labels = num_labels
    id2label = json.load(open(cached_download(hf_hub_url(repo_id, img_labels_file, repo_type="dataset")), "r"))
    id2label = {int(k): v for k, v in id2label.items()}

    id2label = id2label
    label2id = {v: k for k, v in id2label.items()}

    config = config = CvtConfig(num_labels=num_labels, id2label=id2label, label2id=label2id)

    
    if cvt_model.rsplit("/", 1)[-1][4:6] == "13":
        config.depth = [1, 2, 10]

    
    elif cvt_model.rsplit("/", 1)[-1][4:6] == "21":
        config.depth = [1, 4, 16]

    
    else:
        config.depth = [2, 2, 20]
        config.num_heads = [3, 12, 16]
        config.embed_dim = [192, 768, 1024]

    model = CvtForImageClassification(config)
    image_processor = AutoImageProcessor.from_pretrained("facebook/convnext-base-224-22k-1k")
    image_processor.size["shortest_edge"] = image_size
    original_weights = torch.load(cvt_file_name, map_location=torch.device("cpu"))

    huggingface_weights = OrderedDict()
    list_of_state_dict = []

    for idx in range(len(config.depth)):
        if config.cls_token[idx]:
            list_of_state_dict = list_of_state_dict + cls_token(idx)
        list_of_state_dict = list_of_state_dict + embeddings(idx)
        for cnt in range(config.depth[idx]):
            list_of_state_dict = list_of_state_dict + attention(idx, cnt)

    list_of_state_dict = list_of_state_dict + final()
    for gg in list_of_state_dict:
        print(gg)
    for i in range(len(list_of_state_dict)):
        huggingface_weights[list_of_state_dict[i][0]] = original_weights[list_of_state_dict[i][1]]

    model.load_state_dict(huggingface_weights)
    model.save_pretrained(pytorch_dump_folder)
    image_processor.save_pretrained(pytorch_dump_folder)




if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--cvt_model",
        default="cvt-w24",
        type=str,
        help="Name of the cvt model you'd like to convert.",
    )
    parser.add_argument(
        "--image_size",
        default=384,
        type=int,
        help="Input Image Size",
    )
    parser.add_argument(
        "--cvt_file_name",
        default=r"cvtmodels\CvT-w24-384x384-IN-22k.pth",
        type=str,
        help="Input Image Size",
    )
    parser.add_argument(
        "--pytorch_dump_folder_path", default=None, type=str, help="Path to the output PyTorch model directory."
    )

    args = parser.parse_args()
    convert_cvt_checkpoint(args.cvt_model, args.image_size, args.cvt_file_name, args.pytorch_dump_folder_path)