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 poolformer checkpoints from the original repository url https github comsailsgpoolformer import argparse import json from collections import ordereddict from pathlib import path import requests import torch from huggingfacehub import hfhubdownload from pil import image from transformers import poolformerconfig poolformerforimageclassification poolformerimageprocessor from transformers utils import logging logging setverbosityinfo logger logging getloggername def replacekeywithoffsetkey offset originalname newname tofind originalname split 0 keylist key split origblocknum intkeylistkeylist indextofind 2 layernum intkeylistkeylist indextofind 1 newblocknum origblocknum offset key key replaceforigblocknum layernum originalname fblock newblocknum layernum newname return key def renamekeysstatedict newstatedict ordereddict totalembedfound patchemboffset 0 0 for key value in statedict items if key startswithnetwork key key replacenetwork poolformer encoder if proj in key works for the first embedding as well as the internal embedding layers if key endswithbias and patchembed not in key patchemboffset 1 toreplace key key findproj key key replacetoreplace fpatchembeddings totalembedfound key key replaceproj projection if key endswithbias totalembedfound 1 if patchembeddings in key key poolformer encoder key if mlp fc1 in key key replacekeywithoffsetkey patchemboffset mlp fc1 output conv1 if mlp fc2 in key key replacekeywithoffsetkey patchemboffset mlp fc2 output conv2 if norm1 in key key replacekeywithoffsetkey patchemboffset norm1 beforenorm if norm2 in key key replacekeywithoffsetkey patchemboffset norm2 afternorm if layerscale1 in key key replacekeywithoffsetkey patchemboffset layerscale1 layerscale1 if layerscale2 in key key replacekeywithoffsetkey patchemboffset layerscale2 layerscale2 if head in key key key replacehead classifier newstatedictkey value return newstatedict we will verify our results on a coco image def prepareimg url http images cocodataset orgval2017000000039769 jpg image image openrequests geturl streamtrue raw return image torch nograd def convertpoolformercheckpointmodelname checkpointpath pytorchdumpfolderpath load default poolformer configuration config poolformerconfig set attributes based on modelname repoid huggingfacelabelfiles size modelname3 config numlabels 1000 filename imagenet1kid2label json expectedshape 1 1000 set config attributes 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 if size s12 config depths 2 2 6 2 config hiddensizes 64 128 320 512 config mlpratio 4 0 croppct 0 9 elif size s24 config depths 4 4 12 4 config hiddensizes 64 128 320 512 config mlpratio 4 0 croppct 0 9 elif size s36 config depths 6 6 18 6 config hiddensizes 64 128 320 512 config mlpratio 4 0 config layerscaleinitvalue 1e6 croppct 0 9 elif size m36 config depths 6 6 18 6 config hiddensizes 96 192 384 768 config mlpratio 4 0 config layerscaleinitvalue 1e6 croppct 0 95 elif size m48 config depths 8 8 24 8 config hiddensizes 96 192 384 768 config mlpratio 4 0 config layerscaleinitvalue 1e6 croppct 0 95 else raise valueerrorfsize size not supported load image processor imageprocessor poolformerimageprocessorcroppctcroppct prepare image image prepareimg pixelvalues imageprocessorimagesimage returntensorspt pixelvalues logger infofconverting model modelname load original state dict statedict torch loadcheckpointpath maplocationtorch devicecpu rename keys statedict renamekeysstatedict create huggingface model and load state dict model poolformerforimageclassificationconfig model loadstatedictstatedict model eval define image processor imageprocessor poolformerimageprocessorcroppctcroppct pixelvalues imageprocessorimagesprepareimg returntensorspt pixelvalues forward pass outputs modelpixelvalues logits outputs logits define expected logit slices for different models if size s12 expectedslice torch tensor0 3045 0 6758 0 4869 elif size s24 expectedslice torch tensor0 4402 0 1374 0 8045 elif size s36 expectedslice torch tensor0 6080 0 5133 0 5898 elif size m36 expectedslice torch tensor0 3952 0 2263 1 2668 elif size m48 expectedslice torch tensor0 1167 0 0656 0 3423 else raise valueerrorfsize size not supported verify logits assert logits shape expectedshape assert torch allcloselogits0 3 expectedslice atol1e2 finally save model and image processor logger infofsaving pytorch model and image processor to pytorchdumpfolderpath pathpytorchdumpfolderpath mkdirexistoktrue model savepretrainedpytorchdumpfolderpath printfsaving image processor to pytorchdumpfolderpath imageprocessor savepretrainedpytorchdumpfolderpath if name main parser argparse argumentparser parser addargument modelname defaultpoolformers12 typestr helpname of the model you d like to convert parser addargument checkpointpath defaultnone typestr helppath to the original pytorch checkpoint pth file parser addargument pytorchdumpfolderpath defaultnone typestr helppath to the folder to output pytorch model args parser parseargs convertpoolformercheckpointargs modelname args checkpointpath 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 poolformer checkpoints from the original repository url https github com sail sg poolformer replaces the key by subtracting the offset from the original layer number works for the first embedding as well as the internal embedding layers we will verify our results on a coco image copy paste tweak model s weights to our poolformer structure load default poolformer configuration set attributes based on model_name set config attributes load image processor prepare image load original state dict rename keys create huggingface model and load state dict define image processor forward pass define expected logit slices for different models verify logits finally save model and image processor
import argparse
import json
from collections import OrderedDict
from pathlib import Path

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

from transformers import PoolFormerConfig, PoolFormerForImageClassification, PoolFormerImageProcessor
from transformers.utils import logging


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


def replace_key_with_offset(key, offset, original_name, new_name):
    
    to_find = original_name.split(".")[0]
    key_list = key.split(".")
    orig_block_num = int(key_list[key_list.index(to_find) - 2])
    layer_num = int(key_list[key_list.index(to_find) - 1])
    new_block_num = orig_block_num - offset

    key = key.replace(f"{orig_block_num}.{layer_num}.{original_name}", f"block.{new_block_num}.{layer_num}.{new_name}")
    return key


def rename_keys(state_dict):
    new_state_dict = OrderedDict()
    total_embed_found, patch_emb_offset = 0, 0
    for key, value in state_dict.items():
        if key.startswith("network"):
            key = key.replace("network", "poolformer.encoder")
        if "proj" in key:
            
            if key.endswith("bias") and "patch_embed" not in key:
                patch_emb_offset += 1
            to_replace = key[: key.find("proj")]
            key = key.replace(to_replace, f"patch_embeddings.{total_embed_found}.")
            key = key.replace("proj", "projection")
            if key.endswith("bias"):
                total_embed_found += 1
        if "patch_embeddings" in key:
            key = "poolformer.encoder." + key
        if "mlp.fc1" in key:
            key = replace_key_with_offset(key, patch_emb_offset, "mlp.fc1", "output.conv1")
        if "mlp.fc2" in key:
            key = replace_key_with_offset(key, patch_emb_offset, "mlp.fc2", "output.conv2")
        if "norm1" in key:
            key = replace_key_with_offset(key, patch_emb_offset, "norm1", "before_norm")
        if "norm2" in key:
            key = replace_key_with_offset(key, patch_emb_offset, "norm2", "after_norm")
        if "layer_scale_1" in key:
            key = replace_key_with_offset(key, patch_emb_offset, "layer_scale_1", "layer_scale_1")
        if "layer_scale_2" in key:
            key = replace_key_with_offset(key, patch_emb_offset, "layer_scale_2", "layer_scale_2")
        if "head" in key:
            key = key.replace("head", "classifier")
        new_state_dict[key] = value
    return new_state_dict



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

    return image


@torch.no_grad()
def convert_poolformer_checkpoint(model_name, checkpoint_path, pytorch_dump_folder_path):
    

    
    config = PoolFormerConfig()

    
    repo_id = "huggingface/label-files"
    size = model_name[-3:]
    config.num_labels = 1000
    filename = "imagenet-1k-id2label.json"
    expected_shape = (1, 1000)

    
    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 size == "s12":
        config.depths = [2, 2, 6, 2]
        config.hidden_sizes = [64, 128, 320, 512]
        config.mlp_ratio = 4.0
        crop_pct = 0.9
    elif size == "s24":
        config.depths = [4, 4, 12, 4]
        config.hidden_sizes = [64, 128, 320, 512]
        config.mlp_ratio = 4.0
        crop_pct = 0.9
    elif size == "s36":
        config.depths = [6, 6, 18, 6]
        config.hidden_sizes = [64, 128, 320, 512]
        config.mlp_ratio = 4.0
        config.layer_scale_init_value = 1e-6
        crop_pct = 0.9
    elif size == "m36":
        config.depths = [6, 6, 18, 6]
        config.hidden_sizes = [96, 192, 384, 768]
        config.mlp_ratio = 4.0
        config.layer_scale_init_value = 1e-6
        crop_pct = 0.95
    elif size == "m48":
        config.depths = [8, 8, 24, 8]
        config.hidden_sizes = [96, 192, 384, 768]
        config.mlp_ratio = 4.0
        config.layer_scale_init_value = 1e-6
        crop_pct = 0.95
    else:
        raise ValueError(f"Size {size} not supported")

    
    image_processor = PoolFormerImageProcessor(crop_pct=crop_pct)

    
    image = prepare_img()
    pixel_values = image_processor(images=image, return_tensors="pt").pixel_values

    logger.info(f"Converting model {model_name}...")

    
    state_dict = torch.load(checkpoint_path, map_location=torch.device("cpu"))

    
    state_dict = rename_keys(state_dict)

    
    model = PoolFormerForImageClassification(config)
    model.load_state_dict(state_dict)
    model.eval()

    
    image_processor = PoolFormerImageProcessor(crop_pct=crop_pct)
    pixel_values = image_processor(images=prepare_img(), return_tensors="pt").pixel_values

    
    outputs = model(pixel_values)
    logits = outputs.logits

    
    if size == "s12":
        expected_slice = torch.tensor([-0.3045, -0.6758, -0.4869])
    elif size == "s24":
        expected_slice = torch.tensor([0.4402, -0.1374, -0.8045])
    elif size == "s36":
        expected_slice = torch.tensor([-0.6080, -0.5133, -0.5898])
    elif size == "m36":
        expected_slice = torch.tensor([0.3952, 0.2263, -1.2668])
    elif size == "m48":
        expected_slice = torch.tensor([0.1167, -0.0656, -0.3423])
    else:
        raise ValueError(f"Size {size} not supported")

    
    assert logits.shape == expected_shape
    assert torch.allclose(logits[0, :3], expected_slice, atol=1e-2)

    
    logger.info(f"Saving PyTorch model and image processor to {pytorch_dump_folder_path}...")
    Path(pytorch_dump_folder_path).mkdir(exist_ok=True)
    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)


if __name__ == "__main__":
    parser = argparse.ArgumentParser()

    parser.add_argument(
        "--model_name",
        default="poolformer_s12",
        type=str,
        help="Name of the model you'd like to convert.",
    )
    parser.add_argument(
        "--checkpoint_path", default=None, type=str, help="Path to the original PyTorch checkpoint (.pth file)."
    )
    parser.add_argument(
        "--pytorch_dump_folder_path", default=None, type=str, help="Path to the folder to output PyTorch model."
    )
    args = parser.parse_args()
    convert_poolformer_checkpoint(args.model_name, args.checkpoint_path, args.pytorch_dump_folder_path)