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 bit checkpoints from the timm library import argparse import json from pathlib import path import requests import torch from huggingfacehub import hfhubdownload from pil import image from timm import createmodel from timm data import resolvedataconfig from timm data transformsfactory import createtransform from transformers import bitconfig bitforimageclassification bitimageprocessor from transformers imageutils import pilimageresampling from transformers utils import logging logging setverbosityinfo logger logging getloggername def getconfigmodelname repoid huggingfacelabelfiles filename imagenet1kid2label json id2label json loadopenhfhubdownloadrepoid filename repotypedataset r id2label intk v for k v in id2label items label2id v k for k v in id2label items convlayer stdconv if bit in modelname else false note that when using bit as backbone for vithybrid checkpoints one needs to additionally set config layertype bottleneck config stemtype same config convlayer stdconvsame config bitconfig convlayerconvlayer numlabels1000 id2labelid2label label2idlabel2id return config def renamekeyname if stem conv in name name name replacestem conv bit embedder convolution if blocks in name name name replaceblocks layers if head fc in name name name replacehead fc classifier 1 if name startswithnorm name bit name if bit not in name and classifier not in name name bit encoder name return name 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 convertbitcheckpointmodelname pytorchdumpfolderpath pushtohubfalse define default bit configuration config getconfigmodelname load original model from timm timmmodel createmodelmodelname pretrainedtrue timmmodel eval load statedict of original model statedict timmmodel statedict for key in statedict copy keys val statedict popkey statedictrenamekeykey val squeeze if head in key else val load huggingface model model bitforimageclassificationconfig model eval model loadstatedictstatedict create image processor transform createtransformresolvedataconfig modeltimmmodel timmtransforms transform transforms pillowresamplings bilinear pilimageresampling bilinear bicubic pilimageresampling bicubic nearest pilimageresampling nearest processor bitimageprocessor 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 printlogits logits0 3 printpredicted class model config id2labellogits argmax1 item 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 modelname and processor to pytorchdumpfolderpath model savepretrainedpytorchdumpfolderpath processor savepretrainedpytorchdumpfolderpath if pushtohub printfpushing model modelname and processor to the hub model pushtohubfybelkadamodelname processor pushtohubfybelkadamodelname if name main parser argparse argumentparser required parameters parser addargument modelname defaultresnetv250x1bitm typestr helpname of the bit 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 push the model to the hub args parser parseargs convertbitcheckpointargs modelname 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 bit checkpoints from the timm library note that when using bit as backbone for vit hybrid checkpoints one needs to additionally set config layer_type bottleneck config stem_type same config conv_layer std_conv_same we will verify our results on an image of cute cats copy paste tweak model s weights to our bit structure define default bit configuration load original model from timm load state_dict of original model load huggingface model create image processor verify pixel values verify logits 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 timm import create_model
from timm.data import resolve_data_config
from timm.data.transforms_factory import create_transform

from transformers import BitConfig, BitForImageClassification, BitImageProcessor
from transformers.image_utils import PILImageResampling
from transformers.utils import logging


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


def get_config(model_name):
    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()}
    label2id = {v: k for k, v in id2label.items()}

    conv_layer = "std_conv" if "bit" in model_name else False

    
    
    
    config = BitConfig(
        conv_layer=conv_layer,
        num_labels=1000,
        id2label=id2label,
        label2id=label2id,
    )

    return config


def rename_key(name):
    if "stem.conv" in name:
        name = name.replace("stem.conv", "bit.embedder.convolution")
    if "blocks" in name:
        name = name.replace("blocks", "layers")
    if "head.fc" in name:
        name = name.replace("head.fc", "classifier.1")
    if name.startswith("norm"):
        name = "bit." + name
    if "bit" not in name and "classifier" not in name:
        name = "bit.encoder." + name

    return name



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_bit_checkpoint(model_name, pytorch_dump_folder_path, push_to_hub=False):
    

    
    config = get_config(model_name)

    
    timm_model = create_model(model_name, pretrained=True)
    timm_model.eval()

    
    state_dict = timm_model.state_dict()
    for key in state_dict.copy().keys():
        val = state_dict.pop(key)
        state_dict[rename_key(key)] = val.squeeze() if "head" in key else val

    
    model = BitForImageClassification(config)
    model.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 = BitImageProcessor(
        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("Logits:", logits[0, :3])
    print("Predicted class:", model.config.id2label[logits.argmax(-1).item()])
    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 {model_name} and processor to {pytorch_dump_folder_path}")
        model.save_pretrained(pytorch_dump_folder_path)
        processor.save_pretrained(pytorch_dump_folder_path)

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


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    
    parser.add_argument(
        "--model_name",
        default="resnetv2_50x1_bitm",
        type=str,
        help="Name of the BiT 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 push the model to the hub.",
    )

    args = parser.parse_args()
    convert_bit_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)