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 dit checkpoints from the unilm repository import argparse import json from pathlib import path import requests import torch from huggingfacehub import hfhubdownload from pil import image from transformers import beitconfig beitforimageclassification beitformaskedimagemodeling 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 fprefixposembed beit embeddings positionembeddings if haslmhead mask token layernorm renamekeys extend masktoken beit embeddings masktoken norm weight layernorm weight norm bias layernorm 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 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 convertditcheckpointcheckpointurl pytorchdumpfolderpath pushtohubfalse define default beit configuration haslmhead false if rvlcdip in checkpointurl else true config beitconfiguseabsolutepositionembeddingstrue usemasktokenhaslmhead size of the architecture if large in checkpointurl or ditl in checkpointurl config hiddensize 1024 config intermediatesize 4096 config numhiddenlayers 24 config numattentionheads 16 labels if rvlcdip in checkpointurl config numlabels 16 repoid huggingfacelabelfiles filename rvlcdipid2label 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 load statedict of original model remove and rename some keys statedict torch hub loadstatedictfromurlcheckpointurl maplocationcpumodel renamekeys createrenamekeysconfig haslmheadhaslmhead for src dest in renamekeys renamekeystatedict src dest readinqkvstatedict config haslmheadhaslmhead load huggingface model model beitformaskedimagemodelingconfig if haslmhead else beitforimageclassificationconfig model eval model loadstatedictstatedict check outputs on an image imageprocessor beitimageprocessor sizeconfig imagesize resamplepilimageresampling bilinear docentercropfalse image prepareimg encoding imageprocessorimagesimage returntensorspt pixelvalues encodingpixelvalues outputs modelpixelvalues logits outputs logits verify logits expectedshape 1 16 if rvlcdip in checkpointurl else 1 196 8192 assert logits shape torch sizeexpectedshape shape of logits not as expected pathpytorchdumpfolderpath mkdirexistoktrue printfsaving model to pytorchdumpfolderpath model savepretrainedpytorchdumpfolderpath printfsaving image processor to pytorchdumpfolderpath imageprocessor savepretrainedpytorchdumpfolderpath if pushtohub if haslmhead modelname ditbase if base in checkpointurl else ditlarge else modelname ditbasefinetunedrvlcdip if ditb in checkpointurl else ditlargefinetunedrvlcdip imageprocessor pushtohub repopathornamepathpytorchdumpfolderpath modelname organizationnielsr commitmessageadd image processor usetempdirtrue model pushtohub repopathornamepathpytorchdumpfolderpath modelname organizationnielsr commitmessageadd model usetempdirtrue if name main parser argparse argumentparser parser addargument checkpointurl defaulthttps layoutlm blob core windows netditditptsditbase224p16500k62d53a pth typestr helpurl to the original pytorch checkpoint pth file parser addargument pytorchdumpfolderpath defaultnone typestr helppath to the folder to output pytorch model parser addargument pushtohub actionstoretrue args parser parseargs convertditcheckpointargs checkpointurl 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 dit 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 layernorm 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 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 size of the architecture labels load state_dict of original model remove and rename some keys load huggingface model check outputs on an image verify logits
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 BeitConfig, BeitForImageClassification, BeitForMaskedImageModeling, 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"),
            (f"{prefix}pos_embed", "beit.embeddings.position_embeddings"),
        ]
    )

    if has_lm_head:
        
        rename_keys.extend(
            [
                ("mask_token", "beit.embeddings.mask_token"),
                ("norm.weight", "layernorm.weight"),
                ("norm.bias", "layernorm.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


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

    
    has_lm_head = False if "rvlcdip" in checkpoint_url else True
    config = BeitConfig(use_absolute_position_embeddings=True, use_mask_token=has_lm_head)

    
    if "large" in checkpoint_url or "dit-l" in checkpoint_url:
        config.hidden_size = 1024
        config.intermediate_size = 4096
        config.num_hidden_layers = 24
        config.num_attention_heads = 16

    
    if "rvlcdip" in checkpoint_url:
        config.num_labels = 16
        repo_id = "huggingface/label-files"
        filename = "rvlcdip-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()}

    
    state_dict = torch.hub.load_state_dict_from_url(checkpoint_url, map_location="cpu")["model"]

    rename_keys = create_rename_keys(config, has_lm_head=has_lm_head)
    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)

    
    model = BeitForMaskedImageModeling(config) if has_lm_head else BeitForImageClassification(config)
    model.eval()
    model.load_state_dict(state_dict)

    
    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 = [1, 16] if "rvlcdip" in checkpoint_url else [1, 196, 8192]
    assert logits.shape == torch.Size(expected_shape), "Shape of logits 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 push_to_hub:
        if has_lm_head:
            model_name = "dit-base" if "base" in checkpoint_url else "dit-large"
        else:
            model_name = "dit-base-finetuned-rvlcdip" if "dit-b" in checkpoint_url else "dit-large-finetuned-rvlcdip"
        image_processor.push_to_hub(
            repo_path_or_name=Path(pytorch_dump_folder_path, model_name),
            organization="nielsr",
            commit_message="Add image processor",
            use_temp_dir=True,
        )
        model.push_to_hub(
            repo_path_or_name=Path(pytorch_dump_folder_path, model_name),
            organization="nielsr",
            commit_message="Add model",
            use_temp_dir=True,
        )


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

    parser.add_argument(
        "--checkpoint_url",
        default="https://layoutlm.blob.core.windows.net/dit/dit-pts/dit-base-224-p16-500k-62d53a.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."
    )
    parser.add_argument(
        "--push_to_hub",
        action="store_true",
    )
    args = parser.parse_args()
    convert_dit_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path, args.push_to_hub)