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 timesformer checkpoints from the original repository https github commcgnjutimesformer import argparse import json import gdown import numpy as np import torch from huggingfacehub import hfhubdownload from transformers import timesformerconfig timesformerforvideoclassification videomaeimageprocessor def gettimesformerconfigmodelname config timesformerconfig if large in modelname config numframes 96 if hr in modelname config numframes 16 config imagesize 448 repoid huggingfacelabelfiles if k400 in modelname config numlabels 400 filename kinetics400id2label json elif k600 in modelname config numlabels 600 filename kinetics600id2label json elif ssv2 in modelname config numlabels 174 filename somethingsomethingv2id2label json else raise valueerrormodel name should either contain k400 k600 or ssv2 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 renamekeyname if encoder in name name name replaceencoder if clstoken in name name name replaceclstoken timesformer embeddings clstoken if posembed in name name name replaceposembed timesformer embeddings positionembeddings if timeembed in name name name replacetimeembed timesformer embeddings timeembeddings if patchembed proj in name name name replacepatchembed proj timesformer embeddings patchembeddings projection if patchembed norm in name name name replacepatchembed norm timesformer embeddings norm if blocks in name name name replaceblocks timesformer encoder layer if attn proj in name name name replaceattn proj attention output dense if attn in name and bias not in name and temporal not in name name name replaceattn attention self if attn in name and temporal not in name name name replaceattn attention attention if temporalnorm1 in name name name replacetemporalnorm1 temporallayernorm if temporalattn proj in name name name replacetemporalattn temporalattention output dense if temporalfc in name name name replacetemporalfc temporaldense if norm1 in name and temporal not 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 norm weight in name and fc not in name and temporal not in name name name replacenorm weight timesformer layernorm weight if norm bias in name and fc not in name and temporal not in name name name replacenorm bias timesformer layernorm bias if head in name name name replacehead classifier return name def convertstatedictorigstatedict config for key in origstatedict copy keys val origstatedict popkey if key startswithmodel key key replacemodel if qkv in key keysplit key split layernum intkeysplit1 prefix timesformer encoder layer if temporal in key postfix temporalattention attention qkv else postfix attention attention qkv if weight in key origstatedictfprefixlayernumpostfixweight val else origstatedictfprefixlayernumpostfixbias val 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 converttimesformercheckpointcheckpointurl pytorchdumpfolderpath modelname pushtohub config gettimesformerconfigmodelname model timesformerforvideoclassificationconfig download original checkpoint hosted on google drive output pytorchmodel bin gdown cacheddownloadcheckpointurl output quietfalse files torch loadoutput maplocationcpu if model in files statedict filesmodel elif module in files statedict filesmodule else statedict filesmodelstate 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 8 returntensorspt outputs modelinputs logits outputs logits modelnames kinetics400 checkpoints hr high resolution input of 448px instead of 224px timesformerbasefinetunedk400 timesformerlargefinetunedk400 timesformerhrfinetunedk400 kinetics600 checkpoints hr high resolution input of 448px instead of 224px timesformerbasefinetunedk600 timesformerlargefinetunedk600 timesformerhrfinetunedk600 somethingsomethingv2 checkpoints hr high resolution input of 448px instead of 224px timesformerbasefinetunedssv2 timesformerlargefinetunedssv2 timesformerhrfinetunedssv2 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 timesformerbasefinetunedk400 expectedshape torch size1 400 expectedslice torch tensor0 3016 0 7713 0 4205 elif modelname timesformerbasefinetunedk600 expectedshape torch size1 600 expectedslice torch tensor0 7267 0 7466 3 2404 elif modelname timesformerbasefinetunedssv2 expectedshape torch size1 174 expectedslice torch tensor0 9059 0 6433 3 1457 elif modelname timesformerlargefinetunedk400 expectedshape torch size1 400 expectedslice torch tensor0 0 0 elif modelname timesformerlargefinetunedk600 expectedshape torch size1 600 expectedslice torch tensor0 0 0 elif modelname timesformerlargefinetunedssv2 expectedshape torch size1 174 expectedslice torch tensor0 0 0 elif modelname timesformerhrfinetunedk400 expectedshape torch size1 400 expectedslice torch tensor0 9617 3 7311 3 7708 elif modelname timesformerhrfinetunedk600 expectedshape torch size1 600 expectedslice torch tensor2 5273 0 7127 1 8848 elif modelname timesformerhrfinetunedssv2 expectedshape torch size1 174 expectedslice torch tensor3 6756 0 7513 0 7180 else raise valueerrorfmodel name not supported should be one of modelnames verify logits assert logits shape expectedshape assert torch allcloselogits0 3 expectedslice atol1e4 printlogits ok if pytorchdumpfolderpath is not none printfsaving model and image processor to pytorchdumpfolderpath imageprocessor savepretrainedpytorchdumpfolderpath model savepretrainedpytorchdumpfolderpath if pushtohub printpushing to the hub model pushtohubffcakyonmodelname if name main parser argparse argumentparser required parameters parser addargument checkpointurl defaulthttps drive google comu1uc id17yvuyp9l4mnhpick5zo6k3uooy1ka5lexportdownload 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 default typestr helppath to the output pytorch model directory parser addargumentmodelname defaulttimesformerbasefinetunedk400 typestr helpname of the model parser addargument pushtohub actionstoretrue helpwhether or not to push the converted model to the hub args parser parseargs converttimesformercheckpoint args 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 timesformer checkpoints from the original repository https github com mcg nju timesformer 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 hr high resolution input of 448px instead of 224px kinetics 600 checkpoints hr high resolution input of 448px instead of 224px something something v2 checkpoints hr high resolution input of 448px instead of 224px 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 verify logits required parameters
import argparse
import json

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

from transformers import TimesformerConfig, TimesformerForVideoClassification, VideoMAEImageProcessor


def get_timesformer_config(model_name):
    config = TimesformerConfig()

    if "large" in model_name:
        config.num_frames = 96

    if "hr" in model_name:
        config.num_frames = 16
        config.image_size = 448

    repo_id = "huggingface/label-files"
    if "k400" in model_name:
        config.num_labels = 400
        filename = "kinetics400-id2label.json"
    elif "k600" in model_name:
        config.num_labels = 600
        filename = "kinetics600-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 'k400', 'k600' or 'ssv2'.")
    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 rename_key(name):
    if "encoder." in name:
        name = name.replace("encoder.", "")
    if "cls_token" in name:
        name = name.replace("cls_token", "timesformer.embeddings.cls_token")
    if "pos_embed" in name:
        name = name.replace("pos_embed", "timesformer.embeddings.position_embeddings")
    if "time_embed" in name:
        name = name.replace("time_embed", "timesformer.embeddings.time_embeddings")
    if "patch_embed.proj" in name:
        name = name.replace("patch_embed.proj", "timesformer.embeddings.patch_embeddings.projection")
    if "patch_embed.norm" in name:
        name = name.replace("patch_embed.norm", "timesformer.embeddings.norm")
    if "blocks" in name:
        name = name.replace("blocks", "timesformer.encoder.layer")
    if "attn.proj" in name:
        name = name.replace("attn.proj", "attention.output.dense")
    if "attn" in name and "bias" not in name and "temporal" not in name:
        name = name.replace("attn", "attention.self")
    if "attn" in name and "temporal" not in name:
        name = name.replace("attn", "attention.attention")
    if "temporal_norm1" in name:
        name = name.replace("temporal_norm1", "temporal_layernorm")
    if "temporal_attn.proj" in name:
        name = name.replace("temporal_attn", "temporal_attention.output.dense")
    if "temporal_fc" in name:
        name = name.replace("temporal_fc", "temporal_dense")
    if "norm1" in name and "temporal" not 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 "norm.weight" in name and "fc" not in name and "temporal" not in name:
        name = name.replace("norm.weight", "timesformer.layernorm.weight")
    if "norm.bias" in name and "fc" not in name and "temporal" not in name:
        name = name.replace("norm.bias", "timesformer.layernorm.bias")
    if "head" 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("model."):
            key = key.replace("model.", "")

        if "qkv" in key:
            key_split = key.split(".")
            layer_num = int(key_split[1])
            prefix = "timesformer.encoder.layer."
            if "temporal" in key:
                postfix = ".temporal_attention.attention.qkv."
            else:
                postfix = ".attention.attention.qkv."
            if "weight" in key:
                orig_state_dict[f"{prefix}{layer_num}{postfix}weight"] = val
            else:
                orig_state_dict[f"{prefix}{layer_num}{postfix}bias"] = val
        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_timesformer_checkpoint(checkpoint_url, pytorch_dump_folder_path, model_name, push_to_hub):
    config = get_timesformer_config(model_name)

    model = TimesformerForVideoClassification(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"]
    elif "module" in files:
        state_dict = files["module"]
    else:
        state_dict = files["model_state"]
    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[:8], return_tensors="pt")

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

    model_names = [
        
        "timesformer-base-finetuned-k400",
        "timesformer-large-finetuned-k400",
        "timesformer-hr-finetuned-k400",
        
        "timesformer-base-finetuned-k600",
        "timesformer-large-finetuned-k600",
        "timesformer-hr-finetuned-k600",
        
        "timesformer-base-finetuned-ssv2",
        "timesformer-large-finetuned-ssv2",
        "timesformer-hr-finetuned-ssv2",
    ]

    
    if model_name == "timesformer-base-finetuned-k400":
        expected_shape = torch.Size([1, 400])
        expected_slice = torch.tensor([-0.3016, -0.7713, -0.4205])
    elif model_name == "timesformer-base-finetuned-k600":
        expected_shape = torch.Size([1, 600])
        expected_slice = torch.tensor([-0.7267, -0.7466, 3.2404])
    elif model_name == "timesformer-base-finetuned-ssv2":
        expected_shape = torch.Size([1, 174])
        expected_slice = torch.tensor([-0.9059, 0.6433, -3.1457])
    elif model_name == "timesformer-large-finetuned-k400":
        expected_shape = torch.Size([1, 400])
        expected_slice = torch.tensor([0, 0, 0])
    elif model_name == "timesformer-large-finetuned-k600":
        expected_shape = torch.Size([1, 600])
        expected_slice = torch.tensor([0, 0, 0])
    elif model_name == "timesformer-large-finetuned-ssv2":
        expected_shape = torch.Size([1, 174])
        expected_slice = torch.tensor([0, 0, 0])
    elif model_name == "timesformer-hr-finetuned-k400":
        expected_shape = torch.Size([1, 400])
        expected_slice = torch.tensor([-0.9617, -3.7311, -3.7708])
    elif model_name == "timesformer-hr-finetuned-k600":
        expected_shape = torch.Size([1, 600])
        expected_slice = torch.tensor([2.5273, 0.7127, 1.8848])
    elif model_name == "timesformer-hr-finetuned-ssv2":
        expected_shape = torch.Size([1, 174])
        expected_slice = torch.tensor([-3.6756, -0.7513, 0.7180])
    else:
        raise ValueError(f"Model name not supported. Should be one of {model_names}")

    
    assert logits.shape == expected_shape
    assert torch.allclose(logits[0, :3], expected_slice, atol=1e-4)
    print("Logits 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(f"fcakyon/{model_name}")


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    
    parser.add_argument(
        "--checkpoint_url",
        default="https://drive.google.com/u/1/uc?id=17yvuYp9L4mn-HpIcK5Zo6K3UoOy1kA5l&export=download",
        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="",
        type=str,
        help="Path to the output PyTorch model directory.",
    )
    parser.add_argument("--model_name", default="timesformer-base-finetuned-k400", 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_timesformer_checkpoint(
        args.checkpoint_url, args.pytorch_dump_folder_path, args.model_name, args.push_to_hub
    )