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 vit mae checkpoints from the original repository https github comfacebookresearchmae import argparse import requests import torch from pil import image from transformers import vitmaeconfig vitmaeforpretraining vitmaeimageprocessor def renamekeyname if clstoken in name name name replaceclstoken vit embeddings clstoken if masktoken in name name name replacemasktoken decoder masktoken if decoderposembed in name name name replacedecoderposembed decoder decoderposembed if posembed in name and decoder not in name name name replaceposembed vit embeddings positionembeddings if patchembed proj in name name name replacepatchembed proj vit embeddings patchembeddings projection if patchembed norm in name name name replacepatchembed norm vit embeddings norm if decoderblocks in name name name replacedecoderblocks decoder decoderlayers if blocks in name name name replaceblocks vit encoder layer if attn proj in name name name replaceattn proj attention output dense if attn in name name name replaceattn attention self if norm1 in name name name replacenorm1 layernormbefore if norm2 in name name name replacenorm2 layernormafter if mlp fc1 in name name name replacemlp fc1 intermediate dense if mlp fc2 in name name name replacemlp fc2 output dense if decoderembed in name name name replacedecoderembed decoder decoderembed if decodernorm in name name name replacedecodernorm decoder decodernorm if decoderpred in name name name replacedecoderpred decoder decoderpred if norm weight in name and decoder not in name name name replacenorm weight vit layernorm weight if norm bias in name and decoder not in name name name replacenorm bias vit layernorm bias return name def convertstatedictorigstatedict config for key in origstatedict copy keys val origstatedict popkey if qkv in key keysplit key split layernum intkeysplit1 if decoderblocks in key dim config decoderhiddensize prefix decoder decoderlayers if weight in key origstatedictfprefixlayernum attention attention query weight val dim origstatedictfprefixlayernum attention attention key weight valdim dim 2 origstatedictfprefixlayernum attention attention value weight valdim elif bias in key origstatedictfprefixlayernum attention attention query bias val dim origstatedictfprefixlayernum attention attention key bias valdim dim 2 origstatedictfprefixlayernum attention attention value bias valdim else dim config hiddensize prefix vit encoder layer if weight in key origstatedictfprefixlayernum attention attention query weight val dim origstatedictfprefixlayernum attention attention key weight valdim dim 2 origstatedictfprefixlayernum attention attention value weight valdim elif bias in key origstatedictfprefixlayernum attention attention query bias val dim origstatedictfprefixlayernum attention attention key bias valdim dim 2 origstatedictfprefixlayernum attention attention value bias valdim else origstatedictrenamekeykey val return origstatedict def convertvitmaecheckpointcheckpointurl pytorchdumpfolderpath config vitmaeconfig if large in checkpointurl config hiddensize 1024 config intermediatesize 4096 config numhiddenlayers 24 config numattentionheads 16 elif huge in checkpointurl config patchsize 14 config hiddensize 1280 config intermediatesize 5120 config numhiddenlayers 32 config numattentionheads 16 model vitmaeforpretrainingconfig statedict torch hub loadstatedictfromurlcheckpointurl maplocationcpumodel imageprocessor vitmaeimageprocessorsizeconfig imagesize newstatedict convertstatedictstatedict config model loadstatedictnewstatedict model eval url https userimages githubusercontent com11435359147738734196fd92f926048d5ba7ebf103d29364d jpg image image openrequests geturl streamtrue raw imageprocessor vitmaeimageprocessorsizeconfig imagesize inputs imageprocessorimagesimage returntensorspt forward pass torch manualseed2 outputs modelinputs logits outputs logits if large in checkpointurl expectedslice torch tensor 0 7309 0 7128 1 0169 1 0161 0 9058 1 1878 1 0478 0 9411 1 1911 elif huge in checkpointurl expectedslice torch tensor 1 1599 0 9199 1 2221 1 1952 0 9269 1 2307 1 2143 0 9337 1 2262 else expectedslice torch tensor 0 9192 0 8481 1 1259 1 1349 1 0034 1 2599 1 1757 1 0429 1 2726 verify logits assert torch allcloselogits0 3 3 expectedslice atol1e4 printfsaving model to pytorchdumpfolderpath model savepretrainedpytorchdumpfolderpath printfsaving image processor to pytorchdumpfolderpath imageprocessor savepretrainedpytorchdumpfolderpath if name main parser argparse argumentparser required parameters parser addargument checkpointurl defaulthttps dl fbaipublicfiles commaevisualizemaevisualizevitbase pth 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 convertvitmaecheckpointargs 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 vit mae checkpoints from the original repository https github com facebookresearch mae forward pass verify logits required parameters
import argparse

import requests
import torch
from PIL import Image

from transformers import ViTMAEConfig, ViTMAEForPreTraining, ViTMAEImageProcessor


def rename_key(name):
    if "cls_token" in name:
        name = name.replace("cls_token", "vit.embeddings.cls_token")
    if "mask_token" in name:
        name = name.replace("mask_token", "decoder.mask_token")
    if "decoder_pos_embed" in name:
        name = name.replace("decoder_pos_embed", "decoder.decoder_pos_embed")
    if "pos_embed" in name and "decoder" not in name:
        name = name.replace("pos_embed", "vit.embeddings.position_embeddings")
    if "patch_embed.proj" in name:
        name = name.replace("patch_embed.proj", "vit.embeddings.patch_embeddings.projection")
    if "patch_embed.norm" in name:
        name = name.replace("patch_embed.norm", "vit.embeddings.norm")
    if "decoder_blocks" in name:
        name = name.replace("decoder_blocks", "decoder.decoder_layers")
    if "blocks" in name:
        name = name.replace("blocks", "vit.encoder.layer")
    if "attn.proj" in name:
        name = name.replace("attn.proj", "attention.output.dense")
    if "attn" in name:
        name = name.replace("attn", "attention.self")
    if "norm1" in name:
        name = name.replace("norm1", "layernorm_before")
    if "norm2" in name:
        name = name.replace("norm2", "layernorm_after")
    if "mlp.fc1" in name:
        name = name.replace("mlp.fc1", "intermediate.dense")
    if "mlp.fc2" in name:
        name = name.replace("mlp.fc2", "output.dense")
    if "decoder_embed" in name:
        name = name.replace("decoder_embed", "decoder.decoder_embed")
    if "decoder_norm" in name:
        name = name.replace("decoder_norm", "decoder.decoder_norm")
    if "decoder_pred" in name:
        name = name.replace("decoder_pred", "decoder.decoder_pred")
    if "norm.weight" in name and "decoder" not in name:
        name = name.replace("norm.weight", "vit.layernorm.weight")
    if "norm.bias" in name and "decoder" not in name:
        name = name.replace("norm.bias", "vit.layernorm.bias")

    return name


def convert_state_dict(orig_state_dict, config):
    for key in orig_state_dict.copy().keys():
        val = orig_state_dict.pop(key)

        if "qkv" in key:
            key_split = key.split(".")
            layer_num = int(key_split[1])
            if "decoder_blocks" in key:
                dim = config.decoder_hidden_size
                prefix = "decoder.decoder_layers."
                if "weight" in key:
                    orig_state_dict[f"{prefix}{layer_num}.attention.attention.query.weight"] = val[:dim, :]
                    orig_state_dict[f"{prefix}{layer_num}.attention.attention.key.weight"] = val[dim : dim * 2, :]
                    orig_state_dict[f"{prefix}{layer_num}.attention.attention.value.weight"] = val[-dim:, :]
                elif "bias" in key:
                    orig_state_dict[f"{prefix}{layer_num}.attention.attention.query.bias"] = val[:dim]
                    orig_state_dict[f"{prefix}{layer_num}.attention.attention.key.bias"] = val[dim : dim * 2]
                    orig_state_dict[f"{prefix}{layer_num}.attention.attention.value.bias"] = val[-dim:]
            else:
                dim = config.hidden_size
                prefix = "vit.encoder.layer."
                if "weight" in key:
                    orig_state_dict[f"{prefix}{layer_num}.attention.attention.query.weight"] = val[:dim, :]
                    orig_state_dict[f"{prefix}{layer_num}.attention.attention.key.weight"] = val[dim : dim * 2, :]
                    orig_state_dict[f"{prefix}{layer_num}.attention.attention.value.weight"] = val[-dim:, :]
                elif "bias" in key:
                    orig_state_dict[f"{prefix}{layer_num}.attention.attention.query.bias"] = val[:dim]
                    orig_state_dict[f"{prefix}{layer_num}.attention.attention.key.bias"] = val[dim : dim * 2]
                    orig_state_dict[f"{prefix}{layer_num}.attention.attention.value.bias"] = val[-dim:]

        else:
            orig_state_dict[rename_key(key)] = val

    return orig_state_dict


def convert_vit_mae_checkpoint(checkpoint_url, pytorch_dump_folder_path):
    config = ViTMAEConfig()
    if "large" in checkpoint_url:
        config.hidden_size = 1024
        config.intermediate_size = 4096
        config.num_hidden_layers = 24
        config.num_attention_heads = 16
    elif "huge" in checkpoint_url:
        config.patch_size = 14
        config.hidden_size = 1280
        config.intermediate_size = 5120
        config.num_hidden_layers = 32
        config.num_attention_heads = 16

    model = ViTMAEForPreTraining(config)

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

    image_processor = ViTMAEImageProcessor(size=config.image_size)

    new_state_dict = convert_state_dict(state_dict, config)

    model.load_state_dict(new_state_dict)
    model.eval()

    url = "https://user-images.githubusercontent.com/11435359/147738734-196fd92f-9260-48d5-ba7e-bf103d29364d.jpg"

    image = Image.open(requests.get(url, stream=True).raw)
    image_processor = ViTMAEImageProcessor(size=config.image_size)
    inputs = image_processor(images=image, return_tensors="pt")

    
    torch.manual_seed(2)
    outputs = model(**inputs)
    logits = outputs.logits

    if "large" in checkpoint_url:
        expected_slice = torch.tensor(
            [[-0.7309, -0.7128, -1.0169], [-1.0161, -0.9058, -1.1878], [-1.0478, -0.9411, -1.1911]]
        )
    elif "huge" in checkpoint_url:
        expected_slice = torch.tensor(
            [[-1.1599, -0.9199, -1.2221], [-1.1952, -0.9269, -1.2307], [-1.2143, -0.9337, -1.2262]]
        )
    else:
        expected_slice = torch.tensor(
            [[-0.9192, -0.8481, -1.1259], [-1.1349, -1.0034, -1.2599], [-1.1757, -1.0429, -1.2726]]
        )

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

    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://dl.fbaipublicfiles.com/mae/visualize/mae_visualize_vit_base.pth",
        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_vit_mae_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path)