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 levit checkpoints from timm import argparse import json from collections import ordereddict from functools import partial from pathlib import path import timm import torch from huggingfacehub import hfhubdownload from transformers import levitconfig levitforimageclassificationwithteacher levitimageprocessor from transformers utils import logging logging setverbosityinfo logger logging getlogger def convertweightandpush hiddensizes int name str config levitconfig savedirectory path pushtohub bool true printfconverting name with torch nograd if hiddensizes 128 if name1 s frommodel timm createmodellevit128s pretrainedtrue else frommodel timm createmodellevit128 pretrainedtrue if hiddensizes 192 frommodel timm createmodellevit192 pretrainedtrue if hiddensizes 256 frommodel timm createmodellevit256 pretrainedtrue if hiddensizes 384 frommodel timm createmodellevit384 pretrainedtrue frommodel eval ourmodel levitforimageclassificationwithteacherconfig eval huggingfaceweights ordereddict weights frommodel statedict ogkeys listfrommodel statedict keys newkeys listourmodel statedict keys printlenogkeys lennewkeys for i in rangelenogkeys huggingfaceweightsnewkeysi weightsogkeysi ourmodel loadstatedicthuggingfaceweights x torch randn2 3 224 224 out1 frommodelx out2 ourmodelx logits assert torch allcloseout1 out2 the model logits don t match the original one checkpointname name printcheckpointname if pushtohub ourmodel savepretrainedsavedirectory checkpointname imageprocessor levitimageprocessor imageprocessor savepretrainedsavedirectory checkpointname printfpushed checkpointname def convertweightsandpushsavedirectory path modelname str none pushtohub bool true filename imagenet1kid2label json numlabels 1000 expectedshape 1 numlabels repoid huggingfacelabelfiles numlabels numlabels id2label json loadopenhfhubdownloadrepoid filename repotypedataset r id2label intk v for k v in id2label items id2label id2label label2id v k for k v in id2label items imagenetpretrainedconfig partiallevitconfig numlabelsnumlabels id2labelid2label label2idlabel2id namestohiddensizes levit128s 128 levit128 128 levit192 192 levit256 256 levit384 384 namestoconfig levit128s imagenetpretrainedconfig hiddensizes128 256 384 numattentionheads4 6 8 depths2 3 4 keydim16 16 16 droppathrate0 levit128 imagenetpretrainedconfig hiddensizes128 256 384 numattentionheads4 8 12 depths4 4 4 keydim16 16 16 droppathrate0 levit192 imagenetpretrainedconfig hiddensizes192 288 384 numattentionheads3 5 6 depths4 4 4 keydim32 32 32 droppathrate0 levit256 imagenetpretrainedconfig hiddensizes256 384 512 numattentionheads4 6 8 depths4 4 4 keydim32 32 32 droppathrate0 levit384 imagenetpretrainedconfig hiddensizes384 512 768 numattentionheads6 9 12 depths4 4 4 keydim32 32 32 droppathrate0 1 if modelname convertweightandpush namestohiddensizesmodelname modelname namestoconfigmodelname savedirectory pushtohub else for modelname config in namestoconfig items convertweightandpushnamestohiddensizesmodelname modelname config savedirectory pushtohub return config expectedshape if name main parser argparse argumentparser required parameters parser addargument modelname defaultnone typestr helpthe name of the model you wish to convert it must be one of the supported levit architecture parser addargument pytorchdumpfolderpath defaultlevitdumpfolder typepath requiredfalse helppath to the output pytorch model directory parser addargumentpushtohub actionstoretrue helppush model and image processor to the hub parser addargument nopushtohub destpushtohub actionstorefalse helpdo not push model and image processor to the hub args parser parseargs pytorchdumpfolderpath path args pytorchdumpfolderpath pytorchdumpfolderpath mkdirexistoktrue parentstrue convertweightsandpushpytorchdumpfolderpath args modelname 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 levit checkpoints from timm required parameters
import argparse
import json
from collections import OrderedDict
from functools import partial
from pathlib import Path

import timm
import torch
from huggingface_hub import hf_hub_download

from transformers import LevitConfig, LevitForImageClassificationWithTeacher, LevitImageProcessor
from transformers.utils import logging


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


def convert_weight_and_push(
    hidden_sizes: int, name: str, config: LevitConfig, save_directory: Path, push_to_hub: bool = True
):
    print(f"Converting {name}...")

    with torch.no_grad():
        if hidden_sizes == 128:
            if name[-1] == "S":
                from_model = timm.create_model("levit_128s", pretrained=True)
            else:
                from_model = timm.create_model("levit_128", pretrained=True)
        if hidden_sizes == 192:
            from_model = timm.create_model("levit_192", pretrained=True)
        if hidden_sizes == 256:
            from_model = timm.create_model("levit_256", pretrained=True)
        if hidden_sizes == 384:
            from_model = timm.create_model("levit_384", pretrained=True)

        from_model.eval()
        our_model = LevitForImageClassificationWithTeacher(config).eval()
        huggingface_weights = OrderedDict()

        weights = from_model.state_dict()
        og_keys = list(from_model.state_dict().keys())
        new_keys = list(our_model.state_dict().keys())
        print(len(og_keys), len(new_keys))
        for i in range(len(og_keys)):
            huggingface_weights[new_keys[i]] = weights[og_keys[i]]
        our_model.load_state_dict(huggingface_weights)

        x = torch.randn((2, 3, 224, 224))
        out1 = from_model(x)
        out2 = our_model(x).logits

    assert torch.allclose(out1, out2), "The model logits don't match the original one."

    checkpoint_name = name
    print(checkpoint_name)

    if push_to_hub:
        our_model.save_pretrained(save_directory / checkpoint_name)
        image_processor = LevitImageProcessor()
        image_processor.save_pretrained(save_directory / checkpoint_name)

        print(f"Pushed {checkpoint_name}")


def convert_weights_and_push(save_directory: Path, model_name: str = None, push_to_hub: bool = True):
    filename = "imagenet-1k-id2label.json"
    num_labels = 1000
    expected_shape = (1, num_labels)

    repo_id = "huggingface/label-files"
    num_labels = num_labels
    id2label = json.load(open(hf_hub_download(repo_id, filename, 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()}

    ImageNetPreTrainedConfig = partial(LevitConfig, num_labels=num_labels, id2label=id2label, label2id=label2id)

    names_to_hidden_sizes = {
        "levit-128S": 128,
        "levit-128": 128,
        "levit-192": 192,
        "levit-256": 256,
        "levit-384": 384,
    }

    names_to_config = {
        "levit-128S": ImageNetPreTrainedConfig(
            hidden_sizes=[128, 256, 384],
            num_attention_heads=[4, 6, 8],
            depths=[2, 3, 4],
            key_dim=[16, 16, 16],
            drop_path_rate=0,
        ),
        "levit-128": ImageNetPreTrainedConfig(
            hidden_sizes=[128, 256, 384],
            num_attention_heads=[4, 8, 12],
            depths=[4, 4, 4],
            key_dim=[16, 16, 16],
            drop_path_rate=0,
        ),
        "levit-192": ImageNetPreTrainedConfig(
            hidden_sizes=[192, 288, 384],
            num_attention_heads=[3, 5, 6],
            depths=[4, 4, 4],
            key_dim=[32, 32, 32],
            drop_path_rate=0,
        ),
        "levit-256": ImageNetPreTrainedConfig(
            hidden_sizes=[256, 384, 512],
            num_attention_heads=[4, 6, 8],
            depths=[4, 4, 4],
            key_dim=[32, 32, 32],
            drop_path_rate=0,
        ),
        "levit-384": ImageNetPreTrainedConfig(
            hidden_sizes=[384, 512, 768],
            num_attention_heads=[6, 9, 12],
            depths=[4, 4, 4],
            key_dim=[32, 32, 32],
            drop_path_rate=0.1,
        ),
    }

    if model_name:
        convert_weight_and_push(
            names_to_hidden_sizes[model_name], model_name, names_to_config[model_name], save_directory, push_to_hub
        )
    else:
        for model_name, config in names_to_config.items():
            convert_weight_and_push(names_to_hidden_sizes[model_name], model_name, config, save_directory, push_to_hub)
    return config, expected_shape


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    
    parser.add_argument(
        "--model_name",
        default=None,
        type=str,
        help="The name of the model you wish to convert, it must be one of the supported Levit* architecture,",
    )
    parser.add_argument(
        "--pytorch_dump_folder_path",
        default="levit-dump-folder/",
        type=Path,
        required=False,
        help="Path to the output PyTorch model directory.",
    )
    parser.add_argument("--push_to_hub", action="store_true", help="Push model and image processor to the hub")
    parser.add_argument(
        "--no-push_to_hub",
        dest="push_to_hub",
        action="store_false",
        help="Do not push model and image processor to the hub",
    )

    args = parser.parse_args()
    pytorch_dump_folder_path: Path = args.pytorch_dump_folder_path
    pytorch_dump_folder_path.mkdir(exist_ok=True, parents=True)
    convert_weights_and_push(pytorch_dump_folder_path, args.model_name, args.push_to_hub)