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 swinv2 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 autoimageprocessor swinv2config swinv2forimageclassification def getswinv2configswinv2name config swinv2config namesplit swinv2name split modelsize namesplit1 if to in namesplit3 imgsize intnamesplit33 else imgsize intnamesplit3 if to in namesplit2 windowsize intnamesplit22 else windowsize intnamesplit26 if modelsize tiny embeddim 96 depths 2 2 6 2 numheads 3 6 12 24 elif modelsize small embeddim 96 depths 2 2 18 2 numheads 3 6 12 24 elif modelsize base embeddim 128 depths 2 2 18 2 numheads 4 8 16 32 else embeddim 192 depths 2 2 18 2 numheads 6 12 24 48 if to in swinv2name config pretrainedwindowsizes 12 12 12 6 if 22k in swinv2name and to not in swinv2name numclasses 21841 repoid huggingfacelabelfiles filename imagenet22kid2label 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 else numclasses 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 imagesize imgsize config numlabels numclasses config embeddim embeddim config depths depths config numheads numheads config windowsize windowsize return config def renamekeyname if patchembed proj in name name name replacepatchembed proj embeddings patchembeddings projection if patchembed norm in name name name replacepatchembed norm embeddings norm if layers in name name encoder name if attn proj in name name name replaceattn proj attention output dense if attn in name name name replaceattn attention self if norm1 in name name name replacenorm1 layernormbefore if norm2 in name name name replacenorm2 layernormafter if mlp fc1 in name name name replacemlp fc1 intermediate dense if mlp fc2 in name name name replacemlp fc2 output dense if qbias in name name name replaceqbias query bias if kbias in name name name replacekbias key bias if vbias in name name name replacevbias value bias if cpbmlp in name name name replacecpbmlp continuouspositionbiasmlp if name norm weight name layernorm weight if name norm bias name layernorm bias if head in name name name replacehead classifier else name swinv2 name return name def convertstatedictorigstatedict model for key in origstatedict copy keys val origstatedict popkey if mask in key continue elif qkv in key keysplit key split layernum intkeysplit1 blocknum intkeysplit3 dim model swinv2 encoder layerslayernum blocksblocknum attention self allheadsize if weight in key origstatedict fswinv2 encoder layers layernum blocks blocknum attention self query weight val dim origstatedict fswinv2 encoder layers layernum blocks blocknum attention self key weight valdim dim 2 origstatedict fswinv2 encoder layers layernum blocks blocknum attention self value weight valdim else origstatedict fswinv2 encoder layers layernum blocks blocknum attention self query bias val dim origstatedictfswinv2 encoder layers layernum blocks blocknum attention self key bias val dim dim 2 origstatedict fswinv2 encoder layers layernum blocks blocknum attention self value bias valdim else origstatedictrenamekeykey val return origstatedict def convertswinv2checkpointswinv2name pytorchdumpfolderpath timmmodel timm createmodelswinv2name pretrainedtrue timmmodel eval config getswinv2configswinv2name model swinv2forimageclassificationconfig model eval newstatedict convertstatedicttimmmodel statedict model model loadstatedictnewstatedict url http images cocodataset orgval2017000000039769 jpg imageprocessor autoimageprocessor frompretrainedmicrosoft formatswinv2name replace image image openrequests geturl streamtrue raw inputs imageprocessorimagesimage returntensorspt timmouts timmmodelinputspixelvalues hfouts modelinputs logits assert torch allclosetimmouts hfouts atol1e3 printfsaving model swinv2name to pytorchdumpfolderpath model savepretrainedpytorchdumpfolderpath printfsaving image processor to pytorchdumpfolderpath imageprocessor savepretrainedpytorchdumpfolderpath model pushtohub repopathornamepathpytorchdumpfolderpath swinv2name organizationnandwalritik commitmessageadd model if name main parser argparse argumentparser required parameters parser addargument swinv2name defaultswinv2tinypatch4window8256 typestr helpname of the swinv2 timm model you d like to convert parser addargument pytorchdumpfolderpath defaultnone typestr helppath to the output pytorch model directory args parser parseargs convertswinv2checkpointargs swinv2name 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 swinv2 checkpoints from the timm library 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 AutoImageProcessor, Swinv2Config, Swinv2ForImageClassification


def get_swinv2_config(swinv2_name):
    config = Swinv2Config()
    name_split = swinv2_name.split("_")

    model_size = name_split[1]
    if "to" in name_split[3]:
        img_size = int(name_split[3][-3:])
    else:
        img_size = int(name_split[3])
    if "to" in name_split[2]:
        window_size = int(name_split[2][-2:])
    else:
        window_size = int(name_split[2][6:])

    if model_size == "tiny":
        embed_dim = 96
        depths = (2, 2, 6, 2)
        num_heads = (3, 6, 12, 24)
    elif model_size == "small":
        embed_dim = 96
        depths = (2, 2, 18, 2)
        num_heads = (3, 6, 12, 24)
    elif model_size == "base":
        embed_dim = 128
        depths = (2, 2, 18, 2)
        num_heads = (4, 8, 16, 32)
    else:
        embed_dim = 192
        depths = (2, 2, 18, 2)
        num_heads = (6, 12, 24, 48)

    if "to" in swinv2_name:
        config.pretrained_window_sizes = (12, 12, 12, 6)

    if ("22k" in swinv2_name) and ("to" not in swinv2_name):
        num_classes = 21841
        repo_id = "huggingface/label-files"
        filename = "imagenet-22k-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()}

    else:
        num_classes = 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.image_size = img_size
    config.num_labels = num_classes
    config.embed_dim = embed_dim
    config.depths = depths
    config.num_heads = num_heads
    config.window_size = window_size

    return config


def rename_key(name):
    if "patch_embed.proj" in name:
        name = name.replace("patch_embed.proj", "embeddings.patch_embeddings.projection")
    if "patch_embed.norm" in name:
        name = name.replace("patch_embed.norm", "embeddings.norm")
    if "layers" in name:
        name = "encoder." + name
    if "attn.proj" in name:
        name = name.replace("attn.proj", "attention.output.dense")
    if "attn" in name:
        name = name.replace("attn", "attention.self")
    if "norm1" in name:
        name = name.replace("norm1", "layernorm_before")
    if "norm2" in name:
        name = name.replace("norm2", "layernorm_after")
    if "mlp.fc1" in name:
        name = name.replace("mlp.fc1", "intermediate.dense")
    if "mlp.fc2" in name:
        name = name.replace("mlp.fc2", "output.dense")
    if "q_bias" in name:
        name = name.replace("q_bias", "query.bias")
    if "k_bias" in name:
        name = name.replace("k_bias", "key.bias")
    if "v_bias" in name:
        name = name.replace("v_bias", "value.bias")
    if "cpb_mlp" in name:
        name = name.replace("cpb_mlp", "continuous_position_bias_mlp")
    if name == "norm.weight":
        name = "layernorm.weight"
    if name == "norm.bias":
        name = "layernorm.bias"

    if "head" in name:
        name = name.replace("head", "classifier")
    else:
        name = "swinv2." + name

    return name


def convert_state_dict(orig_state_dict, model):
    for key in orig_state_dict.copy().keys():
        val = orig_state_dict.pop(key)

        if "mask" in key:
            continue
        elif "qkv" in key:
            key_split = key.split(".")
            layer_num = int(key_split[1])
            block_num = int(key_split[3])
            dim = model.swinv2.encoder.layers[layer_num].blocks[block_num].attention.self.all_head_size

            if "weight" in key:
                orig_state_dict[
                    f"swinv2.encoder.layers.{layer_num}.blocks.{block_num}.attention.self.query.weight"
                ] = val[:dim, :]
                orig_state_dict[
                    f"swinv2.encoder.layers.{layer_num}.blocks.{block_num}.attention.self.key.weight"
                ] = val[dim : dim * 2, :]
                orig_state_dict[
                    f"swinv2.encoder.layers.{layer_num}.blocks.{block_num}.attention.self.value.weight"
                ] = val[-dim:, :]
            else:
                orig_state_dict[
                    f"swinv2.encoder.layers.{layer_num}.blocks.{block_num}.attention.self.query.bias"
                ] = val[:dim]
                orig_state_dict[f"swinv2.encoder.layers.{layer_num}.blocks.{block_num}.attention.self.key.bias"] = val[
                    dim : dim * 2
                ]
                orig_state_dict[
                    f"swinv2.encoder.layers.{layer_num}.blocks.{block_num}.attention.self.value.bias"
                ] = val[-dim:]
        else:
            orig_state_dict[rename_key(key)] = val

    return orig_state_dict


def convert_swinv2_checkpoint(swinv2_name, pytorch_dump_folder_path):
    timm_model = timm.create_model(swinv2_name, pretrained=True)
    timm_model.eval()

    config = get_swinv2_config(swinv2_name)
    model = Swinv2ForImageClassification(config)
    model.eval()

    new_state_dict = convert_state_dict(timm_model.state_dict(), model)
    model.load_state_dict(new_state_dict)

    url = "http://images.cocodataset.org/val2017/000000039769.jpg"

    image_processor = AutoImageProcessor.from_pretrained("microsoft/{}".format(swinv2_name.replace("_", "-")))
    image = Image.open(requests.get(url, stream=True).raw)
    inputs = image_processor(images=image, return_tensors="pt")

    timm_outs = timm_model(inputs["pixel_values"])
    hf_outs = model(**inputs).logits

    assert torch.allclose(timm_outs, hf_outs, atol=1e-3)

    print(f"Saving model {swinv2_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)

    model.push_to_hub(
        repo_path_or_name=Path(pytorch_dump_folder_path, swinv2_name),
        organization="nandwalritik",
        commit_message="Add model",
    )


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    
    parser.add_argument(
        "--swinv2_name",
        default="swinv2_tiny_patch4_window8_256",
        type=str,
        help="Name of the Swinv2 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_swinv2_checkpoint(args.swinv2_name, args.pytorch_dump_folder_path)