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 vilt checkpoints from the original github repository import argparse import json from pathlib import path import requests import torch from huggingfacehub import hfhubdownload from pil import image from transformers import berttokenizer viltconfig viltforimageandtextretrieval viltforimagesandtextclassification viltformaskedlm viltforquestionanswering viltimageprocessor viltprocessor 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 vqamodelfalse nlvrmodelfalse irtrmodelfalse renamekeys for i in rangeconfig numhiddenlayers encoder layers output projection 2 feedforward neural networks and 2 layernorms renamekeys appendftransformer blocks i norm1 weight fvilt encoder layer i layernormbefore weight renamekeys appendftransformer blocks i norm1 bias fvilt encoder layer i layernormbefore bias renamekeys append ftransformer blocks i attn proj weight fvilt encoder layer i attention output dense weight renamekeys append ftransformer blocks i attn proj bias fvilt encoder layer i attention output dense bias renamekeys appendftransformer blocks i norm2 weight fvilt encoder layer i layernormafter weight renamekeys appendftransformer blocks i norm2 bias fvilt encoder layer i layernormafter bias renamekeys append ftransformer blocks i mlp fc1 weight fvilt encoder layer i intermediate dense weight renamekeys appendftransformer blocks i mlp fc1 bias fvilt encoder layer i intermediate dense bias renamekeys appendftransformer blocks i mlp fc2 weight fvilt encoder layer i output dense weight renamekeys appendftransformer blocks i mlp fc2 bias fvilt encoder layer i output dense bias embeddings renamekeys extend text embeddings textembeddings wordembeddings weight vilt embeddings textembeddings wordembeddings weight textembeddings positionembeddings weight vilt embeddings textembeddings positionembeddings weight textembeddings positionids vilt embeddings textembeddings positionids textembeddings tokentypeembeddings weight vilt embeddings textembeddings tokentypeembeddings weight textembeddings layernorm weight vilt embeddings textembeddings layernorm weight textembeddings layernorm bias vilt embeddings textembeddings layernorm bias patch embeddings transformer clstoken vilt embeddings clstoken transformer patchembed proj weight vilt embeddings patchembeddings projection weight transformer patchembed proj bias vilt embeddings patchembeddings projection bias transformer posembed vilt embeddings positionembeddings token type embeddings tokentypeembeddings weight vilt embeddings tokentypeembeddings weight final layernorm pooler renamekeys extend transformer norm weight vilt layernorm weight transformer norm bias vilt layernorm bias pooler dense weight vilt pooler dense weight pooler dense bias vilt pooler dense bias classifier heads if vqamodel classification head renamekeys extend vqaclassifier 0 weight classifier 0 weight vqaclassifier 0 bias classifier 0 bias vqaclassifier 1 weight classifier 1 weight vqaclassifier 1 bias classifier 1 bias vqaclassifier 3 weight classifier 3 weight vqaclassifier 3 bias classifier 3 bias elif nlvrmodel classification head renamekeys extend nlvr2classifier 0 weight classifier 0 weight nlvr2classifier 0 bias classifier 0 bias nlvr2classifier 1 weight classifier 1 weight nlvr2classifier 1 bias classifier 1 bias nlvr2classifier 3 weight classifier 3 weight nlvr2classifier 3 bias classifier 3 bias else pass return renamekeys we split up the matrix of each encoder layer into queries keys and values def readinqkvstatedict config for i in rangeconfig numhiddenlayers prefix vilt read in weights bias of input projection layer in timm this is a single matrix bias inprojweight statedict popftransformer blocks i attn qkv weight inprojbias statedict popftransformer blocks i attn qkv bias next add query keys and values in that order to the state dict statedictfprefixencoder layer i attention attention query weight inprojweight config hiddensize statedictfprefixencoder layer i attention attention query bias inprojbias config hiddensize statedictfprefixencoder layer i attention attention key weight inprojweight config hiddensize config hiddensize 2 statedictfprefixencoder layer i attention attention key bias inprojbias config hiddensize config hiddensize 2 statedictfprefixencoder layer i attention attention value weight inprojweight config hiddensize statedictfprefixencoder layer i attention attention value bias inprojbiasconfig hiddensize def removeclassificationheadstatedict ignorekeys head weight head bias for k in ignorekeys statedict popk none def renamekeydct old new val dct popold dctnew val torch nograd def convertviltcheckpointcheckpointurl pytorchdumpfolderpath define configuration and initialize huggingface model config viltconfigimagesize384 patchsize32 tiewordembeddingsfalse mlmmodel false vqamodel false nlvrmodel false irtrmodel false if vqa in checkpointurl vqamodel true config numlabels 3129 repoid huggingfacelabelfiles filename vqa2id2label 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 model viltforquestionansweringconfig elif nlvr in checkpointurl nlvrmodel true config numlabels 2 config id2label 0 false 1 true config label2id v k for k v in config id2label items config modalitytypevocabsize 3 model viltforimagesandtextclassificationconfig elif irtr in checkpointurl irtrmodel true model viltforimageandtextretrievalconfig elif mlmitm in checkpointurl mlmmodel true model viltformaskedlmconfig else raise valueerrorunknown model type load statedict of original model remove and rename some keys statedict torch hub loadstatedictfromurlcheckpointurl maplocationcpustatedict renamekeys createrenamekeysconfig vqamodel nlvrmodel irtrmodel for src dest in renamekeys renamekeystatedict src dest readinqkvstatedict config if mlmmodel or irtrmodel ignorekeys itmscore fc weight itmscore fc bias for k in ignorekeys statedict popk none load state dict into huggingface model model eval if mlmmodel missingkeys unexpectedkeys model loadstatedictstatedict strictfalse assert missingkeys mlmscore decoder bias else model loadstatedictstatedict define processor imageprocessor viltimageprocessorsize384 tokenizer berttokenizer frompretrainedbertbaseuncased processor viltprocessorimageprocessor tokenizer forward pass on example inputs image text if nlvrmodel image1 image openrequests gethttps lil nlp cornell edunlvrexsex00 jpg streamtrue raw image2 image openrequests gethttps lil nlp cornell edunlvrexsex00 jpg streamtrue raw text the left image contains twice the number of dogs as the right image and at least two dogs in total are standing encoding1 processorimage1 text returntensorspt encoding2 processorimage2 text returntensorspt outputs model inputidsencoding1 inputids pixelvaluesencoding1 pixelvalues pixelvalues2encoding2 pixelvalues else image image openrequests gethttp images cocodataset orgval2017000000039769 jpg streamtrue raw if mlmmodel text a bunch of mask laying on a mask else text how many cats are there encoding processorimage text returntensorspt outputs modelencoding verify outputs if mlmmodel expectedshape torch size1 11 30522 expectedslice torch tensor12 5061 12 5123 12 5174 assert outputs logits shape expectedshape assert torch allcloseoutputs logits0 0 3 expectedslice atol1e4 verify masked token prediction equals cats predictedid outputs logits0 4 argmax1 item assert tokenizer decodepredictedid cats elif vqamodel expectedshape torch size1 3129 expectedslice torch tensor15 9495 18 1472 10 3041 assert torch allcloseoutputs logits0 3 expectedslice atol1e4 assert outputs logits shape expectedshape assert torch allcloseoutputs logits0 0 3 expectedslice atol1e4 verify vqa prediction equals 2 predictedidx outputs logits argmax1 item assert model config id2labelpredictedidx 2 elif nlvrmodel expectedshape torch size1 2 expectedslice torch tensor2 8721 2 1291 assert torch allcloseoutputs logits0 3 expectedslice atol1e4 assert outputs logits shape expectedshape pathpytorchdumpfolderpath mkdirexistoktrue printfsaving model and processor to pytorchdumpfolderpath model savepretrainedpytorchdumpfolderpath processor savepretrainedpytorchdumpfolderpath if name main parser argparse argumentparser required parameters parser addargument checkpointurl defaulthttps github comdandelinviltreleasesdownload200kvilt200kmlmitm ckpt typestr helpurl of the checkpoint you d like to convert parser addargument pytorchdumpfolderpath defaultnone typestr helppath to the output pytorch model directory args parser parseargs convertviltcheckpointargs checkpointurl 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 vilt checkpoints from the original github 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 embeddings text embeddings patch embeddings token type embeddings final layernorm pooler classifier head s classification head classification head we split up the matrix of each encoder layer into queries keys and values read in weights bias of input projection layer in timm this is a single matrix bias next add query keys and values in that order to the state dict copy paste tweak model s weights to our vilt structure define configuration and initialize huggingface model load state_dict of original model remove and rename some keys load state dict into huggingface model define processor forward pass on example inputs image text verify outputs verify masked token prediction equals cats verify vqa prediction equals 2 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 transformers import (
    BertTokenizer,
    ViltConfig,
    ViltForImageAndTextRetrieval,
    ViltForImagesAndTextClassification,
    ViltForMaskedLM,
    ViltForQuestionAnswering,
    ViltImageProcessor,
    ViltProcessor,
)
from transformers.utils import logging


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



def create_rename_keys(config, vqa_model=False, nlvr_model=False, irtr_model=False):
    rename_keys = []
    for i in range(config.num_hidden_layers):
        
        rename_keys.append((f"transformer.blocks.{i}.norm1.weight", f"vilt.encoder.layer.{i}.layernorm_before.weight"))
        rename_keys.append((f"transformer.blocks.{i}.norm1.bias", f"vilt.encoder.layer.{i}.layernorm_before.bias"))
        rename_keys.append(
            (f"transformer.blocks.{i}.attn.proj.weight", f"vilt.encoder.layer.{i}.attention.output.dense.weight")
        )
        rename_keys.append(
            (f"transformer.blocks.{i}.attn.proj.bias", f"vilt.encoder.layer.{i}.attention.output.dense.bias")
        )
        rename_keys.append((f"transformer.blocks.{i}.norm2.weight", f"vilt.encoder.layer.{i}.layernorm_after.weight"))
        rename_keys.append((f"transformer.blocks.{i}.norm2.bias", f"vilt.encoder.layer.{i}.layernorm_after.bias"))
        rename_keys.append(
            (f"transformer.blocks.{i}.mlp.fc1.weight", f"vilt.encoder.layer.{i}.intermediate.dense.weight")
        )
        rename_keys.append((f"transformer.blocks.{i}.mlp.fc1.bias", f"vilt.encoder.layer.{i}.intermediate.dense.bias"))
        rename_keys.append((f"transformer.blocks.{i}.mlp.fc2.weight", f"vilt.encoder.layer.{i}.output.dense.weight"))
        rename_keys.append((f"transformer.blocks.{i}.mlp.fc2.bias", f"vilt.encoder.layer.{i}.output.dense.bias"))

    
    rename_keys.extend(
        [
            
            ("text_embeddings.word_embeddings.weight", "vilt.embeddings.text_embeddings.word_embeddings.weight"),
            (
                "text_embeddings.position_embeddings.weight",
                "vilt.embeddings.text_embeddings.position_embeddings.weight",
            ),
            ("text_embeddings.position_ids", "vilt.embeddings.text_embeddings.position_ids"),
            (
                "text_embeddings.token_type_embeddings.weight",
                "vilt.embeddings.text_embeddings.token_type_embeddings.weight",
            ),
            ("text_embeddings.LayerNorm.weight", "vilt.embeddings.text_embeddings.LayerNorm.weight"),
            ("text_embeddings.LayerNorm.bias", "vilt.embeddings.text_embeddings.LayerNorm.bias"),
            
            ("transformer.cls_token", "vilt.embeddings.cls_token"),
            ("transformer.patch_embed.proj.weight", "vilt.embeddings.patch_embeddings.projection.weight"),
            ("transformer.patch_embed.proj.bias", "vilt.embeddings.patch_embeddings.projection.bias"),
            ("transformer.pos_embed", "vilt.embeddings.position_embeddings"),
            
            ("token_type_embeddings.weight", "vilt.embeddings.token_type_embeddings.weight"),
        ]
    )

    
    rename_keys.extend(
        [
            ("transformer.norm.weight", "vilt.layernorm.weight"),
            ("transformer.norm.bias", "vilt.layernorm.bias"),
            ("pooler.dense.weight", "vilt.pooler.dense.weight"),
            ("pooler.dense.bias", "vilt.pooler.dense.bias"),
        ]
    )

    
    if vqa_model:
        
        rename_keys.extend(
            [
                ("vqa_classifier.0.weight", "classifier.0.weight"),
                ("vqa_classifier.0.bias", "classifier.0.bias"),
                ("vqa_classifier.1.weight", "classifier.1.weight"),
                ("vqa_classifier.1.bias", "classifier.1.bias"),
                ("vqa_classifier.3.weight", "classifier.3.weight"),
                ("vqa_classifier.3.bias", "classifier.3.bias"),
            ]
        )
    elif nlvr_model:
        
        rename_keys.extend(
            [
                ("nlvr2_classifier.0.weight", "classifier.0.weight"),
                ("nlvr2_classifier.0.bias", "classifier.0.bias"),
                ("nlvr2_classifier.1.weight", "classifier.1.weight"),
                ("nlvr2_classifier.1.bias", "classifier.1.bias"),
                ("nlvr2_classifier.3.weight", "classifier.3.weight"),
                ("nlvr2_classifier.3.bias", "classifier.3.bias"),
            ]
        )
    else:
        pass

    return rename_keys



def read_in_q_k_v(state_dict, config):
    for i in range(config.num_hidden_layers):
        prefix = "vilt."
        
        in_proj_weight = state_dict.pop(f"transformer.blocks.{i}.attn.qkv.weight")
        in_proj_bias = state_dict.pop(f"transformer.blocks.{i}.attn.qkv.bias")
        
        state_dict[f"{prefix}encoder.layer.{i}.attention.attention.query.weight"] = in_proj_weight[
            : config.hidden_size, :
        ]
        state_dict[f"{prefix}encoder.layer.{i}.attention.attention.query.bias"] = in_proj_bias[: config.hidden_size]
        state_dict[f"{prefix}encoder.layer.{i}.attention.attention.key.weight"] = in_proj_weight[
            config.hidden_size : config.hidden_size * 2, :
        ]
        state_dict[f"{prefix}encoder.layer.{i}.attention.attention.key.bias"] = in_proj_bias[
            config.hidden_size : config.hidden_size * 2
        ]
        state_dict[f"{prefix}encoder.layer.{i}.attention.attention.value.weight"] = in_proj_weight[
            -config.hidden_size :, :
        ]
        state_dict[f"{prefix}encoder.layer.{i}.attention.attention.value.bias"] = in_proj_bias[-config.hidden_size :]


def remove_classification_head_(state_dict):
    ignore_keys = ["head.weight", "head.bias"]
    for k in ignore_keys:
        state_dict.pop(k, None)


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


@torch.no_grad()
def convert_vilt_checkpoint(checkpoint_url, pytorch_dump_folder_path):
    

    
    config = ViltConfig(image_size=384, patch_size=32, tie_word_embeddings=False)
    mlm_model = False
    vqa_model = False
    nlvr_model = False
    irtr_model = False
    if "vqa" in checkpoint_url:
        vqa_model = True
        config.num_labels = 3129
        repo_id = "huggingface/label-files"
        filename = "vqa2-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()}
        model = ViltForQuestionAnswering(config)
    elif "nlvr" in checkpoint_url:
        nlvr_model = True
        config.num_labels = 2
        config.id2label = {0: "False", 1: "True"}
        config.label2id = {v: k for k, v in config.id2label.items()}
        config.modality_type_vocab_size = 3
        model = ViltForImagesAndTextClassification(config)
    elif "irtr" in checkpoint_url:
        irtr_model = True
        model = ViltForImageAndTextRetrieval(config)
    elif "mlm_itm" in checkpoint_url:
        mlm_model = True
        model = ViltForMaskedLM(config)
    else:
        raise ValueError("Unknown model type")

    
    state_dict = torch.hub.load_state_dict_from_url(checkpoint_url, map_location="cpu")["state_dict"]
    rename_keys = create_rename_keys(config, vqa_model, nlvr_model, irtr_model)
    for src, dest in rename_keys:
        rename_key(state_dict, src, dest)
    read_in_q_k_v(state_dict, config)
    if mlm_model or irtr_model:
        ignore_keys = ["itm_score.fc.weight", "itm_score.fc.bias"]
        for k in ignore_keys:
            state_dict.pop(k, None)

    
    model.eval()
    if mlm_model:
        missing_keys, unexpected_keys = model.load_state_dict(state_dict, strict=False)
        assert missing_keys == ["mlm_score.decoder.bias"]
    else:
        model.load_state_dict(state_dict)

    
    image_processor = ViltImageProcessor(size=384)
    tokenizer = BertTokenizer.from_pretrained("bert-base-uncased")
    processor = ViltProcessor(image_processor, tokenizer)

    
    if nlvr_model:
        image1 = Image.open(requests.get("https://lil.nlp.cornell.edu/nlvr/exs/ex0_0.jpg", stream=True).raw)
        image2 = Image.open(requests.get("https://lil.nlp.cornell.edu/nlvr/exs/ex0_0.jpg", stream=True).raw)
        text = (
            "The left image contains twice the number of dogs as the right image, and at least two dogs in total are"
            " standing."
        )
        encoding_1 = processor(image1, text, return_tensors="pt")
        encoding_2 = processor(image2, text, return_tensors="pt")
        outputs = model(
            input_ids=encoding_1.input_ids,
            pixel_values=encoding_1.pixel_values,
            pixel_values_2=encoding_2.pixel_values,
        )
    else:
        image = Image.open(requests.get("http://images.cocodataset.org/val2017/000000039769.jpg", stream=True).raw)
        if mlm_model:
            text = "a bunch of [MASK] laying on a [MASK]."
        else:
            text = "How many cats are there?"
        encoding = processor(image, text, return_tensors="pt")
        outputs = model(**encoding)

    
    if mlm_model:
        expected_shape = torch.Size([1, 11, 30522])
        expected_slice = torch.tensor([-12.5061, -12.5123, -12.5174])
        assert outputs.logits.shape == expected_shape
        assert torch.allclose(outputs.logits[0, 0, :3], expected_slice, atol=1e-4)

        
        predicted_id = outputs.logits[0, 4, :].argmax(-1).item()
        assert tokenizer.decode([predicted_id]) == "cats"
    elif vqa_model:
        expected_shape = torch.Size([1, 3129])
        expected_slice = torch.tensor([-15.9495, -18.1472, -10.3041])
        assert torch.allclose(outputs.logits[0, :3], expected_slice, atol=1e-4)
        assert outputs.logits.shape == expected_shape
        assert torch.allclose(outputs.logits[0, 0, :3], expected_slice, atol=1e-4)

        
        predicted_idx = outputs.logits.argmax(-1).item()
        assert model.config.id2label[predicted_idx] == "2"
    elif nlvr_model:
        expected_shape = torch.Size([1, 2])
        expected_slice = torch.tensor([-2.8721, 2.1291])
        assert torch.allclose(outputs.logits[0, :3], expected_slice, atol=1e-4)
        assert outputs.logits.shape == expected_shape

    Path(pytorch_dump_folder_path).mkdir(exist_ok=True)
    print(f"Saving model and processor to {pytorch_dump_folder_path}")
    model.save_pretrained(pytorch_dump_folder_path)
    processor.save_pretrained(pytorch_dump_folder_path)


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    
    parser.add_argument(
        "--checkpoint_url",
        default="https://github.com/dandelin/ViLT/releases/download/200k/vilt_200k_mlm_itm.ckpt",
        type=str,
        help="URL of the checkpoint you'd like to convert.",
    )
    parser.add_argument(
        "--pytorch_dump_folder_path", default=None, type=str, help="Path to the output PyTorch model directory."
    )

    args = parser.parse_args()
    convert_vilt_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path)