codingutf8 2021 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 beit checkpoints from the unilm repository import argparse import json from pathlib import path import requests import torch from datasets import loaddataset from huggingfacehub import hfhubdownload from pil import image from transformers import beitconfig beitforimageclassification beitformaskedimagemodeling beitforsemanticsegmentation beitimageprocessor from transformers imageutils import pilimageresampling from transformers utils import logging logging setverbosityinfo logger logging getloggername here we list all keys to be renamed original name on the left our name on the right def createrenamekeysconfig haslmheadfalse issemanticfalse prefix backbone if issemantic else renamekeys for i in rangeconfig numhiddenlayers encoder layers output projection 2 feedforward neural networks and 2 layernorms renamekeys appendfprefixblocks i norm1 weight fbeit encoder layer i layernormbefore weight renamekeys appendfprefixblocks i norm1 bias fbeit encoder layer i layernormbefore bias renamekeys append fprefixblocks i attn proj weight fbeit encoder layer i attention output dense weight renamekeys append fprefixblocks i attn proj bias fbeit encoder layer i attention output dense bias renamekeys appendfprefixblocks i norm2 weight fbeit encoder layer i layernormafter weight renamekeys appendfprefixblocks i norm2 bias fbeit encoder layer i layernormafter bias renamekeys appendfprefixblocks i mlp fc1 weight fbeit encoder layer i intermediate dense weight renamekeys appendfprefixblocks i mlp fc1 bias fbeit encoder layer i intermediate dense bias renamekeys appendfprefixblocks i mlp fc2 weight fbeit encoder layer i output dense weight renamekeys appendfprefixblocks i mlp fc2 bias fbeit encoder layer i output dense bias projection layer position embeddings renamekeys extend fprefixclstoken beit embeddings clstoken fprefixpatchembed proj weight beit embeddings patchembeddings projection weight fprefixpatchembed proj bias beit embeddings patchembeddings projection bias if haslmhead mask token shared relative position bias layernorm renamekeys extend masktoken beit embeddings masktoken relposbias relativepositionbiastable beit encoder relativepositionbias relativepositionbiastable relposbias relativepositionindex beit encoder relativepositionbias relativepositionindex norm weight layernorm weight norm bias layernorm bias elif issemantic semantic segmentation classification heads renamekeys extend decodehead convseg weight decodehead classifier weight decodehead convseg bias decodehead classifier bias auxiliaryhead convseg weight auxiliaryhead classifier weight auxiliaryhead convseg bias auxiliaryhead classifier bias else layernorm classification head renamekeys extend fcnorm weight beit pooler layernorm weight fcnorm bias beit pooler layernorm bias head weight classifier weight head bias classifier bias return renamekeys we split up the matrix of each encoder layer into queries keys and values def readinqkvstatedict config haslmheadfalse issemanticfalse for i in rangeconfig numhiddenlayers prefix backbone if issemantic else queries keys and values inprojweight statedict popfprefixblocks i attn qkv weight qbias statedict popfprefixblocks i attn qbias vbias statedict popfprefixblocks i attn vbias statedictfbeit encoder layer i attention attention query weight inprojweight config hiddensize statedictfbeit encoder layer i attention attention query bias qbias statedictfbeit encoder layer i attention attention key weight inprojweight config hiddensize config hiddensize 2 statedictfbeit encoder layer i attention attention value weight inprojweight config hiddensize statedictfbeit encoder layer i attention attention value bias vbias gamma1 and gamma2 we call them lambda because otherwise they are renamed when using frompretrained gamma1 statedict popfprefixblocks i gamma1 gamma2 statedict popfprefixblocks i gamma2 statedictfbeit encoder layer i lambda1 gamma1 statedictfbeit encoder layer i lambda2 gamma2 relativeposition bias table index if not haslmhead each layer has its own relative position bias table statedict popfprefixblocks i attn relativepositionbiastable index statedict popfprefixblocks i attn relativepositionindex statedict fbeit encoder layer i attention attention relativepositionbias relativepositionbiastable table statedict fbeit encoder layer i attention attention relativepositionbias relativepositionindex index def renamekeydct old new val dct popold dctnew val 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 convertbeitcheckpointcheckpointurl pytorchdumpfolderpath define default beit configuration config beitconfig haslmhead false issemantic false repoid huggingfacelabelfiles set config parameters based on url if checkpointurl9 4 pt22k masked image modeling config usesharedrelativepositionbias true config usemasktoken true haslmhead true elif checkpointurl9 4 ft22k intermediate finetuning on imagenet22k config userelativepositionbias true config numlabels 21841 filename imagenet22kid2label json id2label json loadopenhfhubdownloadrepoid filename repotypedataset r id2label intk v for k v in id2label items this dataset contains 21843 labels but the model only has 21841 we delete the classes as mentioned in https github comgoogleresearchbigtransferissues18 del id2label9205 del id2label15027 config id2label id2label config label2id v k for k v in id2label items elif checkpointurl8 4 to1k finetuning on imagenet1k config userelativepositionbias true config numlabels 1000 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 if 384 in checkpointurl config imagesize 384 if 512 in checkpointurl config imagesize 512 elif ade20k in checkpointurl finetuning config userelativepositionbias true config numlabels 150 filename ade20kid2label 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 640 issemantic true else raise valueerrorcheckpoint not supported url should either end with pt22k ft22k to1k or ade20k size of the architecture if base in checkpointurl pass elif large in checkpointurl config hiddensize 1024 config intermediatesize 4096 config numhiddenlayers 24 config numattentionheads 16 if ade20k in checkpointurl config imagesize 640 config outindices 7 11 15 23 else raise valueerrorshould either find base or large in checkpoint url load statedict of original model remove and rename some keys statedict torch hub loadstatedictfromurlcheckpointurl maplocationcpu checkhashtrue statedict statedictmodel if ade20k not in checkpointurl else statedictstatedict renamekeys createrenamekeysconfig haslmheadhaslmhead issemanticissemantic for src dest in renamekeys renamekeystatedict src dest readinqkvstatedict config haslmheadhaslmhead issemanticissemantic if issemantic add prefix to decoder keys for key val in statedict copy items val statedict popkey if key startswithbackbone fpn key key replacebackbone fpn fpn statedictkey val load huggingface model if checkpointurl9 4 pt22k model beitformaskedimagemodelingconfig elif ade20k in checkpointurl model beitforsemanticsegmentationconfig else model beitforimageclassificationconfig model eval model loadstatedictstatedict check outputs on an image if issemantic imageprocessor beitimageprocessorsizeconfig imagesize docentercropfalse ds loaddatasethfinternaltestingfixturesade20k splittest image image opends0file else imageprocessor beitimageprocessor sizeconfig imagesize resamplepilimageresampling bilinear docentercropfalse image prepareimg encoding imageprocessorimagesimage returntensorspt pixelvalues encodingpixelvalues outputs modelpixelvalues logits outputs logits verify logits expectedshape torch size1 1000 if checkpointurl 4 endswithbeitbasepatch16224pt22k expectedshape torch size1 196 8192 elif checkpointurl 4 endswithbeitlargepatch16224pt22k expectedshape torch size1 196 8192 elif checkpointurl 4 endswithbeitbasepatch16224pt22kft22k expectedshape torch size1 21841 expectedlogits torch tensor2 2288 2 4671 0 7395 expectedclassidx 2397 elif checkpointurl 4 endswithbeitlargepatch16224pt22kft22k expectedshape torch size1 21841 expectedlogits torch tensor1 6881 0 2787 0 5901 expectedclassidx 2396 elif checkpointurl 4 endswithbeitbasepatch16224pt22kft1k expectedlogits torch tensor0 1241 0 0798 0 6569 expectedclassidx 285 elif checkpointurl 4 endswithbeitbasepatch16224pt22kft22kto1k expectedlogits torch tensor1 2385 1 0987 1 0108 expectedclassidx 281 elif checkpointurl 4 endswithbeitbasepatch16384pt22kft22kto1k expectedlogits torch tensor1 5303 0 9484 0 3147 expectedclassidx 761 elif checkpointurl 4 endswithbeitlargepatch16224pt22kft1k expectedlogits torch tensor0 4610 0 0928 0 2086 expectedclassidx 761 elif checkpointurl 4 endswithbeitlargepatch16224pt22kft22kto1k expectedlogits torch tensor0 4804 0 6257 0 1837 expectedclassidx 761 elif checkpointurl 4 endswithbeitlargepatch16384pt22kft22kto1k expectedlogits torch tensor0 5122 0 5117 0 2113 expectedclassidx 761 elif checkpointurl 4 endswithbeitlargepatch16512pt22kft22kto1k expectedlogits torch tensor0 3062 0 7261 0 4852 expectedclassidx 761 elif checkpointurl 4 endswithbeitbasepatch16640pt22kft22ktoade20k expectedshape 1 150 160 160 expectedlogits torch tensor 4 9225 2 3954 3 0522 2 8822 1 0046 1 7561 2 9549 1 3228 2 1347 5 8168 3 4129 4 0778 3 8651 2 2214 3 0277 3 8356 2 4643 3 3535 0 0078 3 9952 4 0754 2 9856 4 6944 5 0035 3 2413 4 7813 4 9969 elif checkpointurl 4 endswithbeitlargepatch16640pt22kft22ktoade20k expectedshape 1 150 160 160 expectedlogits torch tensor 4 3305 2 3049 3 0161 2 9591 1 5305 2 2251 3 4198 1 8004 2 9062 5 8922 3 7435 4 3978 4 2063 2 7872 3 4755 4 2791 3 1874 4 1681 0 9895 4 3467 4 7663 4 2476 5 6830 6 1518 4 5550 6 2495 6 5154 else raise valueerrorcan t verify logits as model is not supported if logits shape expectedshape raise valueerrorfshape of logits not as expected logits shape expectedshape if not haslmhead if issemantic if not torch allcloselogits0 3 3 3 expectedlogits atol1e3 raise valueerrorfirst elements of logits not as expected else printpredicted class idx logits argmax1 item if not torch allcloselogits0 3 expectedlogits atol1e3 raise valueerrorfirst elements of logits not as expected if logits argmax1 item expectedclassidx raise valueerrorpredicted class index not as expected pathpytorchdumpfolderpath mkdirexistoktrue printfsaving model to pytorchdumpfolderpath model savepretrainedpytorchdumpfolderpath printfsaving image processor to pytorchdumpfolderpath imageprocessor savepretrainedpytorchdumpfolderpath if name main parser argparse argumentparser parser addargument checkpointurl defaulthttps conversationhub blob core windows netbeitsharepublicbeitbeitbasepatch16224pt22kft22kto1k pth typestr helpurl to the original pytorch checkpoint pth file parser addargument pytorchdumpfolderpath defaultnone typestr helppath to the folder to output pytorch model args parser parseargs convertbeitcheckpointargs checkpointurl args pytorchdumpfolderpath coding utf 8 2021 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 beit checkpoints from the unilm repository here we list all keys to be renamed original name on the left our name on the right encoder layers output projection 2 feedforward neural networks and 2 layernorms projection layer position embeddings mask token shared relative position bias layernorm semantic segmentation classification heads layernorm classification head we split up the matrix of each encoder layer into queries keys and values queries keys and values gamma_1 and gamma_2 we call them lambda because otherwise they are renamed when using from_pretrained relative_position bias table index each layer has its own relative position bias we will verify our results on an image of cute cats copy paste tweak model s weights to our beit structure define default beit configuration set config parameters based on url masked image modeling intermediate fine tuning on imagenet 22k this dataset contains 21843 labels but the model only has 21841 we delete the classes as mentioned in https github com google research big_transfer issues 18 fine tuning on imagenet 1k fine tuning size of the architecture load state_dict of original model remove and rename some keys add prefix to decoder keys load huggingface model check outputs on an image verify logits
import argparse
import json
from pathlib import Path

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

from transformers import (
    BeitConfig,
    BeitForImageClassification,
    BeitForMaskedImageModeling,
    BeitForSemanticSegmentation,
    BeitImageProcessor,
)
from transformers.image_utils import PILImageResampling
from transformers.utils import logging


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



def create_rename_keys(config, has_lm_head=False, is_semantic=False):
    prefix = "backbone." if is_semantic else ""

    rename_keys = []
    for i in range(config.num_hidden_layers):
        
        rename_keys.append((f"{prefix}blocks.{i}.norm1.weight", f"beit.encoder.layer.{i}.layernorm_before.weight"))
        rename_keys.append((f"{prefix}blocks.{i}.norm1.bias", f"beit.encoder.layer.{i}.layernorm_before.bias"))
        rename_keys.append(
            (f"{prefix}blocks.{i}.attn.proj.weight", f"beit.encoder.layer.{i}.attention.output.dense.weight")
        )
        rename_keys.append(
            (f"{prefix}blocks.{i}.attn.proj.bias", f"beit.encoder.layer.{i}.attention.output.dense.bias")
        )
        rename_keys.append((f"{prefix}blocks.{i}.norm2.weight", f"beit.encoder.layer.{i}.layernorm_after.weight"))
        rename_keys.append((f"{prefix}blocks.{i}.norm2.bias", f"beit.encoder.layer.{i}.layernorm_after.bias"))
        rename_keys.append((f"{prefix}blocks.{i}.mlp.fc1.weight", f"beit.encoder.layer.{i}.intermediate.dense.weight"))
        rename_keys.append((f"{prefix}blocks.{i}.mlp.fc1.bias", f"beit.encoder.layer.{i}.intermediate.dense.bias"))
        rename_keys.append((f"{prefix}blocks.{i}.mlp.fc2.weight", f"beit.encoder.layer.{i}.output.dense.weight"))
        rename_keys.append((f"{prefix}blocks.{i}.mlp.fc2.bias", f"beit.encoder.layer.{i}.output.dense.bias"))

    
    rename_keys.extend(
        [
            (f"{prefix}cls_token", "beit.embeddings.cls_token"),
            (f"{prefix}patch_embed.proj.weight", "beit.embeddings.patch_embeddings.projection.weight"),
            (f"{prefix}patch_embed.proj.bias", "beit.embeddings.patch_embeddings.projection.bias"),
        ]
    )

    if has_lm_head:
        
        rename_keys.extend(
            [
                ("mask_token", "beit.embeddings.mask_token"),
                (
                    "rel_pos_bias.relative_position_bias_table",
                    "beit.encoder.relative_position_bias.relative_position_bias_table",
                ),
                (
                    "rel_pos_bias.relative_position_index",
                    "beit.encoder.relative_position_bias.relative_position_index",
                ),
                ("norm.weight", "layernorm.weight"),
                ("norm.bias", "layernorm.bias"),
            ]
        )
    elif is_semantic:
        
        rename_keys.extend(
            [
                ("decode_head.conv_seg.weight", "decode_head.classifier.weight"),
                ("decode_head.conv_seg.bias", "decode_head.classifier.bias"),
                ("auxiliary_head.conv_seg.weight", "auxiliary_head.classifier.weight"),
                ("auxiliary_head.conv_seg.bias", "auxiliary_head.classifier.bias"),
            ]
        )
    else:
        
        rename_keys.extend(
            [
                ("fc_norm.weight", "beit.pooler.layernorm.weight"),
                ("fc_norm.bias", "beit.pooler.layernorm.bias"),
                ("head.weight", "classifier.weight"),
                ("head.bias", "classifier.bias"),
            ]
        )

    return rename_keys



def read_in_q_k_v(state_dict, config, has_lm_head=False, is_semantic=False):
    for i in range(config.num_hidden_layers):
        prefix = "backbone." if is_semantic else ""
        
        in_proj_weight = state_dict.pop(f"{prefix}blocks.{i}.attn.qkv.weight")
        q_bias = state_dict.pop(f"{prefix}blocks.{i}.attn.q_bias")
        v_bias = state_dict.pop(f"{prefix}blocks.{i}.attn.v_bias")

        state_dict[f"beit.encoder.layer.{i}.attention.attention.query.weight"] = in_proj_weight[
            : config.hidden_size, :
        ]
        state_dict[f"beit.encoder.layer.{i}.attention.attention.query.bias"] = q_bias
        state_dict[f"beit.encoder.layer.{i}.attention.attention.key.weight"] = in_proj_weight[
            config.hidden_size : config.hidden_size * 2, :
        ]
        state_dict[f"beit.encoder.layer.{i}.attention.attention.value.weight"] = in_proj_weight[
            -config.hidden_size :, :
        ]
        state_dict[f"beit.encoder.layer.{i}.attention.attention.value.bias"] = v_bias

        
        
        gamma_1 = state_dict.pop(f"{prefix}blocks.{i}.gamma_1")
        gamma_2 = state_dict.pop(f"{prefix}blocks.{i}.gamma_2")

        state_dict[f"beit.encoder.layer.{i}.lambda_1"] = gamma_1
        state_dict[f"beit.encoder.layer.{i}.lambda_2"] = gamma_2

        
        if not has_lm_head:
            
            table = state_dict.pop(f"{prefix}blocks.{i}.attn.relative_position_bias_table")
            index = state_dict.pop(f"{prefix}blocks.{i}.attn.relative_position_index")

            state_dict[
                f"beit.encoder.layer.{i}.attention.attention.relative_position_bias.relative_position_bias_table"
            ] = table
            state_dict[
                f"beit.encoder.layer.{i}.attention.attention.relative_position_bias.relative_position_index"
            ] = index


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



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_beit_checkpoint(checkpoint_url, pytorch_dump_folder_path):
    

    
    config = BeitConfig()
    has_lm_head = False
    is_semantic = False
    repo_id = "huggingface/label-files"
    
    if checkpoint_url[-9:-4] == "pt22k":
        
        config.use_shared_relative_position_bias = True
        config.use_mask_token = True
        has_lm_head = True
    elif checkpoint_url[-9:-4] == "ft22k":
        
        config.use_relative_position_bias = True
        config.num_labels = 21841
        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()}
        
        
        del id2label[9205]
        del id2label[15027]
        config.id2label = id2label
        config.label2id = {v: k for k, v in id2label.items()}
    elif checkpoint_url[-8:-4] == "to1k":
        
        config.use_relative_position_bias = True
        config.num_labels = 1000
        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 "384" in checkpoint_url:
            config.image_size = 384
        if "512" in checkpoint_url:
            config.image_size = 512
    elif "ade20k" in checkpoint_url:
        
        config.use_relative_position_bias = True
        config.num_labels = 150
        filename = "ade20k-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 = 640
        is_semantic = True
    else:
        raise ValueError("Checkpoint not supported, URL should either end with 'pt22k', 'ft22k', 'to1k' or 'ade20k'")

    
    if "base" in checkpoint_url:
        pass
    elif "large" in checkpoint_url:
        config.hidden_size = 1024
        config.intermediate_size = 4096
        config.num_hidden_layers = 24
        config.num_attention_heads = 16
        if "ade20k" in checkpoint_url:
            config.image_size = 640
            config.out_indices = [7, 11, 15, 23]
    else:
        raise ValueError("Should either find 'base' or 'large' in checkpoint URL")

    
    state_dict = torch.hub.load_state_dict_from_url(checkpoint_url, map_location="cpu", check_hash=True)
    state_dict = state_dict["model"] if "ade20k" not in checkpoint_url else state_dict["state_dict"]

    rename_keys = create_rename_keys(config, has_lm_head=has_lm_head, is_semantic=is_semantic)
    for src, dest in rename_keys:
        rename_key(state_dict, src, dest)
    read_in_q_k_v(state_dict, config, has_lm_head=has_lm_head, is_semantic=is_semantic)
    if is_semantic:
        
        for key, val in state_dict.copy().items():
            val = state_dict.pop(key)
            if key.startswith("backbone.fpn"):
                key = key.replace("backbone.fpn", "fpn")
            state_dict[key] = val

    
    if checkpoint_url[-9:-4] == "pt22k":
        model = BeitForMaskedImageModeling(config)
    elif "ade20k" in checkpoint_url:
        model = BeitForSemanticSegmentation(config)
    else:
        model = BeitForImageClassification(config)
    model.eval()
    model.load_state_dict(state_dict)

    
    if is_semantic:
        image_processor = BeitImageProcessor(size=config.image_size, do_center_crop=False)
        ds = load_dataset("hf-internal-testing/fixtures_ade20k", split="test")
        image = Image.open(ds[0]["file"])
    else:
        image_processor = BeitImageProcessor(
            size=config.image_size, resample=PILImageResampling.BILINEAR, do_center_crop=False
        )
        image = prepare_img()

    encoding = image_processor(images=image, return_tensors="pt")
    pixel_values = encoding["pixel_values"]

    outputs = model(pixel_values)
    logits = outputs.logits

    
    expected_shape = torch.Size([1, 1000])
    if checkpoint_url[:-4].endswith("beit_base_patch16_224_pt22k"):
        expected_shape = torch.Size([1, 196, 8192])
    elif checkpoint_url[:-4].endswith("beit_large_patch16_224_pt22k"):
        expected_shape = torch.Size([1, 196, 8192])
    elif checkpoint_url[:-4].endswith("beit_base_patch16_224_pt22k_ft22k"):
        expected_shape = torch.Size([1, 21841])
        expected_logits = torch.tensor([2.2288, 2.4671, 0.7395])
        expected_class_idx = 2397
    elif checkpoint_url[:-4].endswith("beit_large_patch16_224_pt22k_ft22k"):
        expected_shape = torch.Size([1, 21841])
        expected_logits = torch.tensor([1.6881, -0.2787, 0.5901])
        expected_class_idx = 2396
    elif checkpoint_url[:-4].endswith("beit_base_patch16_224_pt22k_ft1k"):
        expected_logits = torch.tensor([0.1241, 0.0798, -0.6569])
        expected_class_idx = 285
    elif checkpoint_url[:-4].endswith("beit_base_patch16_224_pt22k_ft22kto1k"):
        expected_logits = torch.tensor([-1.2385, -1.0987, -1.0108])
        expected_class_idx = 281
    elif checkpoint_url[:-4].endswith("beit_base_patch16_384_pt22k_ft22kto1k"):
        expected_logits = torch.tensor([-1.5303, -0.9484, -0.3147])
        expected_class_idx = 761
    elif checkpoint_url[:-4].endswith("beit_large_patch16_224_pt22k_ft1k"):
        expected_logits = torch.tensor([0.4610, -0.0928, 0.2086])
        expected_class_idx = 761
    elif checkpoint_url[:-4].endswith("beit_large_patch16_224_pt22k_ft22kto1k"):
        expected_logits = torch.tensor([-0.4804, 0.6257, -0.1837])
        expected_class_idx = 761
    elif checkpoint_url[:-4].endswith("beit_large_patch16_384_pt22k_ft22kto1k"):
        expected_logits = torch.tensor([[-0.5122, 0.5117, -0.2113]])
        expected_class_idx = 761
    elif checkpoint_url[:-4].endswith("beit_large_patch16_512_pt22k_ft22kto1k"):
        expected_logits = torch.tensor([-0.3062, 0.7261, 0.4852])
        expected_class_idx = 761
    elif checkpoint_url[:-4].endswith("beit_base_patch16_640_pt22k_ft22ktoade20k"):
        expected_shape = (1, 150, 160, 160)
        expected_logits = torch.tensor(
            [
                [[-4.9225, -2.3954, -3.0522], [-2.8822, -1.0046, -1.7561], [-2.9549, -1.3228, -2.1347]],
                [[-5.8168, -3.4129, -4.0778], [-3.8651, -2.2214, -3.0277], [-3.8356, -2.4643, -3.3535]],
                [[-0.0078, 3.9952, 4.0754], [2.9856, 4.6944, 5.0035], [3.2413, 4.7813, 4.9969]],
            ]
        )
    elif checkpoint_url[:-4].endswith("beit_large_patch16_640_pt22k_ft22ktoade20k"):
        expected_shape = (1, 150, 160, 160)
        expected_logits = torch.tensor(
            [
                [[-4.3305, -2.3049, -3.0161], [-2.9591, -1.5305, -2.2251], [-3.4198, -1.8004, -2.9062]],
                [[-5.8922, -3.7435, -4.3978], [-4.2063, -2.7872, -3.4755], [-4.2791, -3.1874, -4.1681]],
                [[0.9895, 4.3467, 4.7663], [4.2476, 5.6830, 6.1518], [4.5550, 6.2495, 6.5154]],
            ]
        )
    else:
        raise ValueError("Can't verify logits as model is not supported")

    if logits.shape != expected_shape:
        raise ValueError(f"Shape of logits not as expected. {logits.shape=}, {expected_shape=}")
    if not has_lm_head:
        if is_semantic:
            if not torch.allclose(logits[0, :3, :3, :3], expected_logits, atol=1e-3):
                raise ValueError("First elements of logits not as expected")
        else:
            print("Predicted class idx:", logits.argmax(-1).item())

            if not torch.allclose(logits[0, :3], expected_logits, atol=1e-3):
                raise ValueError("First elements of logits not as expected")
            if logits.argmax(-1).item() != expected_class_idx:
                raise ValueError("Predicted class index not as expected")

    Path(pytorch_dump_folder_path).mkdir(exist_ok=True)
    print(f"Saving model 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)


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

    parser.add_argument(
        "--checkpoint_url",
        default="https://conversationhub.blob.core.windows.net/beit-share-public/beit/beit_base_patch16_224_pt22k_ft22kto1k.pth",
        type=str,
        help="URL 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_beit_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path)