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 videomae checkpoints from the original repository https github commcgnjuvideomae import argparse import json import gdown import numpy as np import torch from huggingfacehub import hfhubdownload from transformers import videomaeconfig videomaeforpretraining videomaeforvideoclassification videomaeimageprocessor def getvideomaeconfigmodelname config videomaeconfig setarchitectureconfigsmodelname config if finetuned not in modelname config usemeanpooling false if finetuned in modelname repoid huggingfacelabelfiles if kinetics in modelname config numlabels 400 filename kinetics400id2label json elif ssv2 in modelname config numlabels 174 filename somethingsomethingv2id2label json else raise valueerrormodel name should either contain kinetics or ssv2 in case it s finetuned 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 return config def setarchitectureconfigsmodelname config if small in modelname config hiddensize 384 config intermediatesize 1536 config numhiddenlayers 12 config numattentionheads 16 config decodernumhiddenlayers 12 config decodernumattentionheads 3 config decoderhiddensize 192 config decoderintermediatesize 768 elif large in modelname config hiddensize 1024 config intermediatesize 4096 config numhiddenlayers 24 config numattentionheads 16 config decodernumhiddenlayers 12 config decodernumattentionheads 8 config decoderhiddensize 512 config decoderintermediatesize 2048 elif huge in modelname config hiddensize 1280 config intermediatesize 5120 config numhiddenlayers 32 config numattentionheads 16 config decodernumhiddenlayers 12 config decodernumattentionheads 8 config decoderhiddensize 640 config decoderintermediatesize 2560 elif base not in modelname raise valueerror model name should include either small base large or huge def renamekeyname if encoder in name name name replaceencoder if clstoken in name name name replaceclstoken videomae embeddings clstoken if decoderposembed in name name name replacedecoderposembed decoder decoderposembed if posembed in name and decoder not in name name name replaceposembed videomae embeddings positionembeddings if patchembed proj in name name name replacepatchembed proj videomae embeddings patchembeddings projection if patchembed norm in name name name replacepatchembed norm videomae embeddings norm if decoder blocks in name name name replacedecoder blocks decoder decoderlayers if blocks in name name name replaceblocks videomae encoder layer if attn proj in name name name replaceattn proj attention output dense if attn in name and bias not in name name name replaceattn attention self if attn in name name name replaceattn attention attention 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 and fc not in name name name replacenorm weight videomae layernorm weight if norm bias in name and decoder not in name and fc not in name name name replacenorm bias videomae layernorm bias if head in name and decoder not in name name name replacehead classifier return name def convertstatedictorigstatedict config for key in origstatedict copy keys val origstatedict popkey if key startswithencoder key key replaceencoder if qkv in key keysplit key split if key startswithdecoder blocks dim config decoderhiddensize layernum intkeysplit2 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 else dim config hiddensize layernum intkeysplit1 prefix videomae 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 else origstatedictrenamekeykey val return origstatedict we will verify our results on a video of eating spaghetti frame indices used 164 168 172 176 181 185 189 193 198 202 206 210 215 219 223 227 def preparevideo file hfhubdownload repoidhfinternaltestingspaghettivideo filenameeatingspaghetti npy repotypedataset video np loadfile return listvideo def convertvideomaecheckpointcheckpointurl pytorchdumpfolderpath modelname pushtohub config getvideomaeconfigmodelname if finetuned in modelname model videomaeforvideoclassificationconfig else model videomaeforpretrainingconfig download original checkpoint hosted on google drive output pytorchmodel bin gdown cacheddownloadcheckpointurl output quietfalse files torch loadoutput maplocationcpu if model in files statedict filesmodel else statedict filesmodule newstatedict convertstatedictstatedict config model loadstatedictnewstatedict model eval verify model on basic input imageprocessor videomaeimageprocessorimagemean0 5 0 5 0 5 imagestd0 5 0 5 0 5 video preparevideo inputs imageprocessorvideo returntensorspt if finetuned not in modelname localpath hfhubdownloadrepoidhfinternaltestingboolmaskedpos filenameboolmaskedpos pt inputsboolmaskedpos torch loadlocalpath outputs modelinputs logits outputs logits modelnames videomaesmallfinetunedkinetics videomaesmallfinetunedssv2 kinetics400 checkpoints short pretrained only for 800 epochs instead of 1600 videomaebaseshort videomaebaseshortfinetunedkinetics videomaebase videomaebasefinetunedkinetics videomaelarge videomaelargefinetunedkinetics videomaehugefinetunedkinetics somethingsomethingv2 checkpoints short pretrained only for 800 epochs instead of 2400 videomaebaseshortssv2 videomaebaseshortfinetunedssv2 videomaebasessv2 videomaebasefinetunedssv2 note logits were tested with imagemean and imagestd equal to 0 5 0 5 0 5 and 0 5 0 5 0 5 if modelname videomaesmallfinetunedkinetics expectedshape torch size1 400 expectedslice torch tensor0 9291 0 4061 0 9307 elif modelname videomaesmallfinetunedssv2 expectedshape torch size1 174 expectedslice torch tensor0 2671 0 4689 0 8235 elif modelname videomaebase expectedshape torch size1 1408 1536 expectedslice torch tensor0 7739 0 7968 0 7089 0 6701 0 7487 0 6209 0 4287 0 5158 0 4773 elif modelname videomaebaseshort expectedshape torch size1 1408 1536 expectedslice torch tensor0 7994 0 9612 0 8508 0 7401 0 8958 0 8302 0 5862 0 7468 0 7325 we verified the loss both for normalized and unnormalized targets for this one expectedloss torch tensor0 5142 if config normpixloss else torch tensor0 6469 elif modelname videomaelarge expectedshape torch size1 1408 1536 expectedslice torch tensor0 7149 0 7997 0 6966 0 6768 0 7869 0 6948 0 5139 0 6221 0 5605 elif modelname videomaelargefinetunedkinetics expectedshape torch size1 400 expectedslice torch tensor0 0771 0 0011 0 3625 elif modelname videomaehugefinetunedkinetics expectedshape torch size1 400 expectedslice torch tensor0 2433 0 1632 0 4894 elif modelname videomaebaseshortfinetunedkinetics expectedshape torch size1 400 expectedslice torch tensor0 6588 0 0990 0 2493 elif modelname videomaebasefinetunedkinetics expectedshape torch size1 400 expectedslice torch tensor0 3669 0 0688 0 2421 elif modelname videomaebaseshortssv2 expectedshape torch size1 1408 1536 expectedslice torch tensor0 4712 0 5296 0 5786 0 2278 0 2729 0 4026 0 0352 0 0730 0 2506 elif modelname videomaebaseshortfinetunedssv2 expectedshape torch size1 174 expectedslice torch tensor0 0537 0 1539 0 3266 elif modelname videomaebasessv2 expectedshape torch size1 1408 1536 expectedslice torch tensor0 8131 0 8727 0 8546 0 7366 0 9377 0 8870 0 5935 0 8874 0 8564 elif modelname videomaebasefinetunedssv2 expectedshape torch size1 174 expectedslice torch tensor0 1961 0 8337 0 6389 else raise valueerrorfmodel name not supported should be one of modelnames verify logits assert logits shape expectedshape if finetuned in modelname assert torch allcloselogits0 3 expectedslice atol1e4 else printlogits logits0 3 3 assert torch allcloselogits0 3 3 expectedslice atol1e4 printlogits ok verify loss if applicable if modelname videomaebaseshort loss outputs loss assert torch allcloseloss expectedloss atol1e4 printloss ok if pytorchdumpfolderpath is not none printfsaving model and image processor to pytorchdumpfolderpath imageprocessor savepretrainedpytorchdumpfolderpath model savepretrainedpytorchdumpfolderpath if pushtohub printpushing to the hub model pushtohubmodelname organizationnielsr if name main parser argparse argumentparser required parameters parser addargument checkpointurl defaulthttps drive google comu1uc id1tehlyskjb755tj65ptsrafug2llswqe1amp exportdownloadamp confirmtamp uuidaa3276ebfb7e482aadecdc7171df14c4 typestr help url of the original pytorch checkpoint on google drive you d like to convert should be a direct download link parser addargument pytorchdumpfolderpath defaultusersnielsroggedocumentsvideomaetest typestr helppath to the output pytorch model directory parser addargumentmodelname defaultvideomaebase typestr helpname of the model parser addargument pushtohub actionstoretrue helpwhether or not to push the converted model to the hub args parser parseargs convertvideomaecheckpointargs checkpointurl args pytorchdumpfolderpath args modelname 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 videomae checkpoints from the original repository https github com mcg nju videomae we will verify our results on a video of eating spaghetti frame indices used 164 168 172 176 181 185 189 193 198 202 206 210 215 219 223 227 download original checkpoint hosted on google drive verify model on basic input kinetics 400 checkpoints short pretrained only for 800 epochs instead of 1600 something something v2 checkpoints short pretrained only for 800 epochs instead of 2400 note logits were tested with image_mean and image_std equal to 0 5 0 5 0 5 and 0 5 0 5 0 5 we verified the loss both for normalized and unnormalized targets for this one verify logits verify loss if applicable required parameters
import argparse
import json

import gdown
import numpy as np
import torch
from huggingface_hub import hf_hub_download

from transformers import (
    VideoMAEConfig,
    VideoMAEForPreTraining,
    VideoMAEForVideoClassification,
    VideoMAEImageProcessor,
)


def get_videomae_config(model_name):
    config = VideoMAEConfig()

    set_architecture_configs(model_name, config)

    if "finetuned" not in model_name:
        config.use_mean_pooling = False

    if "finetuned" in model_name:
        repo_id = "huggingface/label-files"
        if "kinetics" in model_name:
            config.num_labels = 400
            filename = "kinetics400-id2label.json"
        elif "ssv2" in model_name:
            config.num_labels = 174
            filename = "something-something-v2-id2label.json"
        else:
            raise ValueError("Model name should either contain 'kinetics' or 'ssv2' in case it's fine-tuned.")
        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()}

    return config


def set_architecture_configs(model_name, config):
    if "small" in model_name:
        config.hidden_size = 384
        config.intermediate_size = 1536
        config.num_hidden_layers = 12
        config.num_attention_heads = 16
        config.decoder_num_hidden_layers = 12
        config.decoder_num_attention_heads = 3
        config.decoder_hidden_size = 192
        config.decoder_intermediate_size = 768
    elif "large" in model_name:
        config.hidden_size = 1024
        config.intermediate_size = 4096
        config.num_hidden_layers = 24
        config.num_attention_heads = 16
        config.decoder_num_hidden_layers = 12
        config.decoder_num_attention_heads = 8
        config.decoder_hidden_size = 512
        config.decoder_intermediate_size = 2048
    elif "huge" in model_name:
        config.hidden_size = 1280
        config.intermediate_size = 5120
        config.num_hidden_layers = 32
        config.num_attention_heads = 16
        config.decoder_num_hidden_layers = 12
        config.decoder_num_attention_heads = 8
        config.decoder_hidden_size = 640
        config.decoder_intermediate_size = 2560
    elif "base" not in model_name:
        raise ValueError('Model name should include either "small", "base", "large", or "huge"')


def rename_key(name):
    if "encoder." in name:
        name = name.replace("encoder.", "")
    if "cls_token" in name:
        name = name.replace("cls_token", "videomae.embeddings.cls_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", "videomae.embeddings.position_embeddings")
    if "patch_embed.proj" in name:
        name = name.replace("patch_embed.proj", "videomae.embeddings.patch_embeddings.projection")
    if "patch_embed.norm" in name:
        name = name.replace("patch_embed.norm", "videomae.embeddings.norm")
    if "decoder.blocks" in name:
        name = name.replace("decoder.blocks", "decoder.decoder_layers")
    if "blocks" in name:
        name = name.replace("blocks", "videomae.encoder.layer")
    if "attn.proj" in name:
        name = name.replace("attn.proj", "attention.output.dense")
    if "attn" in name and "bias" not in name:
        name = name.replace("attn", "attention.self")
    if "attn" in name:
        name = name.replace("attn", "attention.attention")
    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 and "fc" not in name:
        name = name.replace("norm.weight", "videomae.layernorm.weight")
    if "norm.bias" in name and "decoder" not in name and "fc" not in name:
        name = name.replace("norm.bias", "videomae.layernorm.bias")
    if "head" in name and "decoder" not in name:
        name = name.replace("head", "classifier")

    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 key.startswith("encoder."):
            key = key.replace("encoder.", "")

        if "qkv" in key:
            key_split = key.split(".")
            if key.startswith("decoder.blocks"):
                dim = config.decoder_hidden_size
                layer_num = int(key_split[2])
                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:, :]
            else:
                dim = config.hidden_size
                layer_num = int(key_split[1])
                prefix = "videomae.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:, :]
        else:
            orig_state_dict[rename_key(key)] = val

    return orig_state_dict




def prepare_video():
    file = hf_hub_download(
        repo_id="hf-internal-testing/spaghetti-video", filename="eating_spaghetti.npy", repo_type="dataset"
    )
    video = np.load(file)
    return list(video)


def convert_videomae_checkpoint(checkpoint_url, pytorch_dump_folder_path, model_name, push_to_hub):
    config = get_videomae_config(model_name)

    if "finetuned" in model_name:
        model = VideoMAEForVideoClassification(config)
    else:
        model = VideoMAEForPreTraining(config)

    
    output = "pytorch_model.bin"
    gdown.cached_download(checkpoint_url, output, quiet=False)
    files = torch.load(output, map_location="cpu")
    if "model" in files:
        state_dict = files["model"]
    else:
        state_dict = files["module"]
    new_state_dict = convert_state_dict(state_dict, config)

    model.load_state_dict(new_state_dict)
    model.eval()

    
    image_processor = VideoMAEImageProcessor(image_mean=[0.5, 0.5, 0.5], image_std=[0.5, 0.5, 0.5])
    video = prepare_video()
    inputs = image_processor(video, return_tensors="pt")

    if "finetuned" not in model_name:
        local_path = hf_hub_download(repo_id="hf-internal-testing/bool-masked-pos", filename="bool_masked_pos.pt")
        inputs["bool_masked_pos"] = torch.load(local_path)

    outputs = model(**inputs)
    logits = outputs.logits

    model_names = [
        "videomae-small-finetuned-kinetics",
        "videomae-small-finetuned-ssv2",
        
        "videomae-base-short",
        "videomae-base-short-finetuned-kinetics",
        "videomae-base",
        "videomae-base-finetuned-kinetics",
        "videomae-large",
        "videomae-large-finetuned-kinetics",
        "videomae-huge-finetuned-kinetics",
        
        "videomae-base-short-ssv2",
        "videomae-base-short-finetuned-ssv2",
        "videomae-base-ssv2",
        "videomae-base-finetuned-ssv2",
    ]

    
    if model_name == "videomae-small-finetuned-kinetics":
        expected_shape = torch.Size([1, 400])
        expected_slice = torch.tensor([-0.9291, -0.4061, -0.9307])
    elif model_name == "videomae-small-finetuned-ssv2":
        expected_shape = torch.Size([1, 174])
        expected_slice = torch.tensor([0.2671, -0.4689, -0.8235])
    elif model_name == "videomae-base":
        expected_shape = torch.Size([1, 1408, 1536])
        expected_slice = torch.tensor([[0.7739, 0.7968, 0.7089], [0.6701, 0.7487, 0.6209], [0.4287, 0.5158, 0.4773]])
    elif model_name == "videomae-base-short":
        expected_shape = torch.Size([1, 1408, 1536])
        expected_slice = torch.tensor([[0.7994, 0.9612, 0.8508], [0.7401, 0.8958, 0.8302], [0.5862, 0.7468, 0.7325]])
        
        expected_loss = torch.tensor([0.5142]) if config.norm_pix_loss else torch.tensor([0.6469])
    elif model_name == "videomae-large":
        expected_shape = torch.Size([1, 1408, 1536])
        expected_slice = torch.tensor([[0.7149, 0.7997, 0.6966], [0.6768, 0.7869, 0.6948], [0.5139, 0.6221, 0.5605]])
    elif model_name == "videomae-large-finetuned-kinetics":
        expected_shape = torch.Size([1, 400])
        expected_slice = torch.tensor([0.0771, 0.0011, -0.3625])
    elif model_name == "videomae-huge-finetuned-kinetics":
        expected_shape = torch.Size([1, 400])
        expected_slice = torch.tensor([0.2433, 0.1632, -0.4894])
    elif model_name == "videomae-base-short-finetuned-kinetics":
        expected_shape = torch.Size([1, 400])
        expected_slice = torch.tensor([0.6588, 0.0990, -0.2493])
    elif model_name == "videomae-base-finetuned-kinetics":
        expected_shape = torch.Size([1, 400])
        expected_slice = torch.tensor([0.3669, -0.0688, -0.2421])
    elif model_name == "videomae-base-short-ssv2":
        expected_shape = torch.Size([1, 1408, 1536])
        expected_slice = torch.tensor([[0.4712, 0.5296, 0.5786], [0.2278, 0.2729, 0.4026], [0.0352, 0.0730, 0.2506]])
    elif model_name == "videomae-base-short-finetuned-ssv2":
        expected_shape = torch.Size([1, 174])
        expected_slice = torch.tensor([-0.0537, -0.1539, -0.3266])
    elif model_name == "videomae-base-ssv2":
        expected_shape = torch.Size([1, 1408, 1536])
        expected_slice = torch.tensor([[0.8131, 0.8727, 0.8546], [0.7366, 0.9377, 0.8870], [0.5935, 0.8874, 0.8564]])
    elif model_name == "videomae-base-finetuned-ssv2":
        expected_shape = torch.Size([1, 174])
        expected_slice = torch.tensor([0.1961, -0.8337, -0.6389])
    else:
        raise ValueError(f"Model name not supported. Should be one of {model_names}")

    
    assert logits.shape == expected_shape
    if "finetuned" in model_name:
        assert torch.allclose(logits[0, :3], expected_slice, atol=1e-4)
    else:
        print("Logits:", logits[0, :3, :3])
        assert torch.allclose(logits[0, :3, :3], expected_slice, atol=1e-4)
    print("Logits ok!")

    
    if model_name == "videomae-base-short":
        loss = outputs.loss
        assert torch.allclose(loss, expected_loss, atol=1e-4)
        print("Loss ok!")

    if pytorch_dump_folder_path is not None:
        print(f"Saving model and image processor to {pytorch_dump_folder_path}")
        image_processor.save_pretrained(pytorch_dump_folder_path)
        model.save_pretrained(pytorch_dump_folder_path)

    if push_to_hub:
        print("Pushing to the hub...")
        model.push_to_hub(model_name, organization="nielsr")


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    
    parser.add_argument(
        "--checkpoint_url",
        default="https://drive.google.com/u/1/uc?id=1tEhLyskjb755TJ65ptsrafUG2llSwQE1&amp;export=download&amp;confirm=t&amp;uuid=aa3276eb-fb7e-482a-adec-dc7171df14c4",
        type=str,
        help=(
            "URL of the original PyTorch checkpoint (on Google Drive) you'd like to convert. Should be a direct"
            " download link."
        ),
    )
    parser.add_argument(
        "--pytorch_dump_folder_path",
        default="/Users/nielsrogge/Documents/VideoMAE/Test",
        type=str,
        help="Path to the output PyTorch model directory.",
    )
    parser.add_argument("--model_name", default="videomae-base", type=str, help="Name of the model.")
    parser.add_argument(
        "--push_to_hub", action="store_true", help="Whether or not to push the converted model to the 🤗 hub."
    )

    args = parser.parse_args()
    convert_videomae_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path, args.model_name, args.push_to_hub)