codingutf8 2023 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 swiftformer checkpoints from the original implementation import argparse import json from pathlib import path import requests import torch from huggingfacehub import hfhubdownload from pil import image from transformers import swiftformerconfig swiftformerforimageclassification vitimageprocessor from transformers utils import logging logging setverbosityinfo logger logging getloggername device torch devicecpu 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 def getexpectedoutputswiftformername if swiftformername swiftformerxs return torch tensor2 1703e00 2 1107e00 2 0811e00 8 8685e01 2 4360e01 elif swiftformername swiftformers return torch tensor3 9636e01 2 3478e01 1 6963e00 1 7381e00 8 6337e01 elif swiftformername swiftformerl1 return torch tensor4 2768e01 4 7429e01 1 0897e00 1 0248e00 3 5523e02 elif swiftformername swiftformerl3 return torch tensor2 5330e01 2 4211e01 6 0185e01 8 2789e01 6 0446e02 def renamekeydct old new val dct popold dctnew val def createrenamekeysstatedict renamekeys for k in statedict keys knew k if pwconv in k knew knew replace pwconv pointwiseconv if dwconv in k knew knew replace dwconv depthwiseconv if proj in k knew knew replace proj proj if patchembed in knew knew knew replacepatchembed swiftformer patchembed patchembedding if network in knew ls knew split if ls2 isdigit knew swiftformer encoder network ls1 blocks ls2 joinls3 else knew knew replacenetwork swiftformer encoder network renamekeys appendk knew return renamekeys torch nograd def convertswiftformercheckpointswiftformername pytorchdumpfolderpath originalckpt define default swiftformer configuration config swiftformerconfig dataset imagenet21k only or also 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 size of the architecture if swiftformername swiftformerxs config depths 3 3 6 4 config embeddims 48 56 112 220 elif swiftformername swiftformers config depths 3 3 9 6 config embeddims 48 64 168 224 elif swiftformername swiftformerl1 config depths 4 3 10 5 config embeddims 48 96 192 384 elif swiftformername swiftformerl3 config depths 4 4 12 6 config embeddims 64 128 320 512 load statedict of original model remove and rename some keys if originalckpt if originalckpt startswithhttps checkpoint torch hub loadstatedictfromurloriginalckpt maplocationcpu checkhashtrue else checkpoint torch loadoriginalckpt maplocationcpu statedict checkpoint renamekeys createrenamekeysstatedict for renamekeysrc renamekeydest in renamekeys renamekeystatedict renamekeysrc renamekeydest load huggingface model hfmodel swiftformerforimageclassificationconfig eval hfmodel loadstatedictstatedict prepare test inputs image prepareimg processor vitimageprocessor frompretrainedpreprocessorconfig inputs processorimagesimage returntensorspt compare outputs from both models timmlogits getexpectedoutputswiftformername hflogits hfmodelinputspixelvalues logits assert hflogits shape torch size1 1000 assert torch allclosehflogits0 0 5 timmlogits atol1e3 pathpytorchdumpfolderpath mkdirexistoktrue printfsaving model swiftformername to pytorchdumpfolderpath hfmodel savepretrainedpytorchdumpfolderpath if name main parser argparse argumentparser required parameters parser addargument swiftformername defaultswiftformerxs choicesswiftformerxs swiftformers swiftformerl1 swiftformerl3 typestr helpname of the swiftformer model you d like to convert parser addargument pytorchdumpfolderpath default convertedoutputs typestr helppath to the output pytorch model directory parser addargumentoriginalckpt defaultnone typestr helppath to the original model checkpoint args parser parseargs convertswiftformercheckpointargs swiftformername args pytorchdumpfolderpath args originalckpt coding utf 8 2023 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 swiftformer checkpoints from the original implementation we will verify our results on an image of cute cats copy paste tweak model s weights to our swiftformer structure define default swiftformer configuration dataset imagenet 21k only or also fine tuned on imagenet 2012 patch_size and image_size size of the architecture load state_dict of original model remove and rename some keys load huggingface model prepare test inputs compare outputs from both models required parameters
import argparse
import json
from pathlib import Path

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

from transformers import (
    SwiftFormerConfig,
    SwiftFormerForImageClassification,
    ViTImageProcessor,
)
from transformers.utils import logging


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

device = torch.device("cpu")



def prepare_img():
    url = "http://images.cocodataset.org/val2017/000000039769.jpg"
    im = Image.open(requests.get(url, stream=True).raw)
    return im


def get_expected_output(swiftformer_name):
    if swiftformer_name == "swiftformer_xs":
        return torch.tensor([-2.1703e00, 2.1107e00, -2.0811e00, 8.8685e-01, 2.4360e-01])

    elif swiftformer_name == "swiftformer_s":
        return torch.tensor([3.9636e-01, 2.3478e-01, -1.6963e00, -1.7381e00, -8.6337e-01])

    elif swiftformer_name == "swiftformer_l1":
        return torch.tensor([-4.2768e-01, -4.7429e-01, -1.0897e00, -1.0248e00, 3.5523e-02])

    elif swiftformer_name == "swiftformer_l3":
        return torch.tensor([-2.5330e-01, 2.4211e-01, -6.0185e-01, -8.2789e-01, -6.0446e-02])


def rename_key(dct, old, new):
    val = dct.pop(old)
    dct[new] = val


def create_rename_keys(state_dict):
    rename_keys = []
    for k in state_dict.keys():
        k_new = k
        if ".pwconv" in k:
            k_new = k_new.replace(".pwconv", ".point_wise_conv")
        if ".dwconv" in k:
            k_new = k_new.replace(".dwconv", ".depth_wise_conv")
        if ".Proj." in k:
            k_new = k_new.replace(".Proj.", ".proj.")
        if "patch_embed" in k_new:
            k_new = k_new.replace("patch_embed", "swiftformer.patch_embed.patch_embedding")
        if "network" in k_new:
            ls = k_new.split(".")
            if ls[2].isdigit():
                k_new = "swiftformer.encoder.network." + ls[1] + ".blocks." + ls[2] + "." + ".".join(ls[3:])
            else:
                k_new = k_new.replace("network", "swiftformer.encoder.network")
        rename_keys.append((k, k_new))
    return rename_keys


@torch.no_grad()
def convert_swiftformer_checkpoint(swiftformer_name, pytorch_dump_folder_path, original_ckpt):
    

    
    config = SwiftFormerConfig()

    
    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()}

    
    if swiftformer_name == "swiftformer_xs":
        config.depths = [3, 3, 6, 4]
        config.embed_dims = [48, 56, 112, 220]

    elif swiftformer_name == "swiftformer_s":
        config.depths = [3, 3, 9, 6]
        config.embed_dims = [48, 64, 168, 224]

    elif swiftformer_name == "swiftformer_l1":
        config.depths = [4, 3, 10, 5]
        config.embed_dims = [48, 96, 192, 384]

    elif swiftformer_name == "swiftformer_l3":
        config.depths = [4, 4, 12, 6]
        config.embed_dims = [64, 128, 320, 512]

    
    if original_ckpt:
        if original_ckpt.startswith("https"):
            checkpoint = torch.hub.load_state_dict_from_url(original_ckpt, map_location="cpu", check_hash=True)
        else:
            checkpoint = torch.load(original_ckpt, map_location="cpu")
    state_dict = checkpoint

    rename_keys = create_rename_keys(state_dict)
    for rename_key_src, rename_key_dest in rename_keys:
        rename_key(state_dict, rename_key_src, rename_key_dest)

    
    hf_model = SwiftFormerForImageClassification(config).eval()
    hf_model.load_state_dict(state_dict)

    
    image = prepare_img()
    processor = ViTImageProcessor.from_pretrained("preprocessor_config")
    inputs = processor(images=image, return_tensors="pt")

    
    timm_logits = get_expected_output(swiftformer_name)
    hf_logits = hf_model(inputs["pixel_values"]).logits

    assert hf_logits.shape == torch.Size([1, 1000])
    assert torch.allclose(hf_logits[0, 0:5], timm_logits, atol=1e-3)

    Path(pytorch_dump_folder_path).mkdir(exist_ok=True)
    print(f"Saving model {swiftformer_name} to {pytorch_dump_folder_path}")
    hf_model.save_pretrained(pytorch_dump_folder_path)


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    
    parser.add_argument(
        "--swiftformer_name",
        default="swiftformer_xs",
        choices=["swiftformer_xs", "swiftformer_s", "swiftformer_l1", "swiftformer_l3"],
        type=str,
        help="Name of the SwiftFormer model you'd like to convert.",
    )
    parser.add_argument(
        "--pytorch_dump_folder_path",
        default="./converted_outputs/",
        type=str,
        help="Path to the output PyTorch model directory.",
    )
    parser.add_argument("--original_ckpt", default=None, type=str, help="Path to the original model checkpoint.")

    args = parser.parse_args()
    convert_swiftformer_checkpoint(args.swiftformer_name, args.pytorch_dump_folder_path, args.original_ckpt)