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 yolos checkpoints from the original repository url https github comhustvlyolos import argparse import json from pathlib import path import requests import torch from huggingfacehub import hfhubdownload from pil import image from transformers import yolosconfig yolosforobjectdetection yolosimageprocessor from transformers utils import logging logging setverbosityinfo logger logging getloggername def getyolosconfigyolosname str yolosconfig config yolosconfig size of the architecture if yolosti in yolosname config hiddensize 192 config intermediatesize 768 config numhiddenlayers 12 config numattentionheads 3 config imagesize 800 1333 config usemidpositionembeddings false elif yolosname yolossdwr config hiddensize 330 config numhiddenlayers 14 config numattentionheads 6 config intermediatesize 1320 elif yoloss in yolosname config hiddensize 384 config intermediatesize 1536 config numhiddenlayers 12 config numattentionheads 6 elif yolosb in yolosname config imagesize 800 1344 config numlabels 91 repoid huggingfacelabelfiles filename cocodetectionid2label 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 return config we split up the matrix of each encoder layer into queries keys and values def readinqkvstatedict dict config yolosconfig basemodel bool false for i in rangeconfig numhiddenlayers read in weights bias of input projection layer in timm this is a single matrix bias inprojweight statedict popfblocks i attn qkv weight inprojbias statedict popfblocks i attn qkv bias next add query keys and values in that order to the state dict statedictfencoder layer i attention attention query weight inprojweight config hiddensize statedictfencoder layer i attention attention query bias inprojbias config hiddensize statedictfencoder layer i attention attention key weight inprojweight config hiddensize config hiddensize 2 statedictfencoder layer i attention attention key bias inprojbias config hiddensize config hiddensize 2 statedictfencoder layer i attention attention value weight inprojweightconfig hiddensize statedictfencoder layer i attention attention value bias inprojbiasconfig hiddensize def renamekeyname str str if backbone in name name name replacebackbone vit if clstoken in name name name replaceclstoken embeddings clstoken if dettoken in name name name replacedettoken embeddings detectiontokens if midposembed in name name name replacemidposembed encoder midpositionembeddings if posembed in name name name replaceposembed embeddings positionembeddings if patchembed proj in name name name replacepatchembed proj embeddings patchembeddings projection if blocks in name name name replaceblocks 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 classembed in name name name replaceclassembed classlabelsclassifier if bboxembed in name name name replacebboxembed bboxpredictor if vit norm in name name name replacevit norm vit layernorm return name def convertstatedictorigstatedict dict model yolosforobjectdetection dict for key in origstatedict copy keys val origstatedict popkey if qkv in key keysplit key split layernum intkeysplit2 dim model vit encoder layerlayernum attention attention allheadsize if weight in key origstatedictfvit encoder layer layernum attention attention query weight val dim origstatedictfvit encoder layer layernum attention attention key weight val dim dim 2 origstatedictfvit encoder layer layernum attention attention value weight valdim else origstatedictfvit encoder layer layernum attention attention query bias val dim origstatedictfvit encoder layer layernum attention attention key bias valdim dim 2 origstatedictfvit encoder layer layernum attention attention value bias valdim else origstatedictrenamekeykey val return origstatedict we will verify our results on an image of cute cats def prepareimg torch tensor url http images cocodataset orgval2017000000039769 jpg im image openrequests geturl streamtrue raw return im torch nograd def convertyoloscheckpoint yolosname str checkpointpath str pytorchdumpfolderpath str pushtohub bool false config getyolosconfigyolosname load original statedict statedict torch loadcheckpointpath maplocationcpumodel load model model yolosforobjectdetectionconfig model eval newstatedict convertstatedictstatedict model model loadstatedictnewstatedict check outputs on an image prepared by yolosimageprocessor size 800 if yolosname yolosti else 512 imageprocessor yolosimageprocessorformatcocodetection sizesize encoding imageprocessorimagesprepareimg returntensorspt outputs modelencoding logits predboxes outputs logits outputs predboxes expectedslicelogits expectedsliceboxes none none if yolosname yolosti expectedslicelogits torch tensor 39 5022 11 9820 17 6888 29 9574 9 9769 17 7691 42 3281 20 7200 30 6294 expectedsliceboxes torch tensor 0 4021 0 0836 0 7979 0 0184 0 2609 0 0364 0 1781 0 2004 0 2095 elif yolosname yoloss200pre expectedslicelogits torch tensor 24 0248 10 3024 14 8290 42 0392 16 8200 27 4334 27 2743 11 8154 18 7148 expectedsliceboxes torch tensor 0 2559 0 5455 0 4706 0 2989 0 7279 0 1875 0 7732 0 4017 0 4462 elif yolosname yoloss300pre expectedslicelogits torch tensor 36 2220 14 4385 23 5457 35 6970 14 7583 21 3935 31 5939 13 6042 16 8049 expectedsliceboxes torch tensor 0 7614 0 2316 0 4728 0 7168 0 4495 0 3855 0 4996 0 1466 0 9996 elif yolosname yolossdwr expectedslicelogits torch tensor 42 8668 24 1049 41 1690 34 7456 14 1274 24 9194 33 7898 12 1946 25 6495 expectedsliceboxes torch tensor 0 5587 0 2773 0 0605 0 5004 0 3014 0 9994 0 4999 0 1548 0 9994 elif yolosname yolosbase expectedslicelogits torch tensor 40 6064 24 3084 32 6447 55 1990 30 7719 35 5877 51 4311 33 3507 35 6462 expectedsliceboxes torch tensor 0 5555 0 2794 0 0655 0 9049 0 2664 0 1894 0 9183 0 1984 0 1635 else raise valueerrorfunknown yolosname yolosname assert torch allcloselogits0 3 3 expectedslicelogits atol1e4 assert torch allclosepredboxes0 3 3 expectedsliceboxes atol1e4 pathpytorchdumpfolderpath mkdirexistoktrue printfsaving model yolosname to pytorchdumpfolderpath model savepretrainedpytorchdumpfolderpath printfsaving image processor to pytorchdumpfolderpath imageprocessor savepretrainedpytorchdumpfolderpath if pushtohub modelmapping yolosti yolostiny yoloss200pre yolossmall yoloss300pre yolossmall300 yolossdwr yolossmalldwr yolosbase yolosbase printpushing to the hub modelname modelmappingyolosname imageprocessor pushtohubmodelname organizationhustvl model pushtohubmodelname organizationhustvl if name main parser argparse argumentparser required parameters parser addargument yolosname defaultyoloss200pre typestr help name of the yolos model you d like to convert should be one of yolosti yoloss200pre yoloss300pre yolossdwr yolosbase parser addargument checkpointpath defaultnone typestr helppath to the original state dict pth file parser addargument pytorchdumpfolderpath defaultnone typestr helppath to the output pytorch model directory parser addargument pushtohub actionstoretrue helpwhether or not to push the converted model to the hub args parser parseargs convertyoloscheckpointargs yolosname args checkpointpath 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 yolos checkpoints from the original repository url https github com hustvl yolos size of the architecture 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 we will verify our results on an image of cute cats copy paste tweak model s weights to our yolos structure load original state_dict load model check outputs on an image prepared by yolosimageprocessor 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 YolosConfig, YolosForObjectDetection, YolosImageProcessor
from transformers.utils import logging


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


def get_yolos_config(yolos_name: str) -> YolosConfig:
    config = YolosConfig()

    
    if "yolos_ti" in yolos_name:
        config.hidden_size = 192
        config.intermediate_size = 768
        config.num_hidden_layers = 12
        config.num_attention_heads = 3
        config.image_size = [800, 1333]
        config.use_mid_position_embeddings = False
    elif yolos_name == "yolos_s_dWr":
        config.hidden_size = 330
        config.num_hidden_layers = 14
        config.num_attention_heads = 6
        config.intermediate_size = 1320
    elif "yolos_s" in yolos_name:
        config.hidden_size = 384
        config.intermediate_size = 1536
        config.num_hidden_layers = 12
        config.num_attention_heads = 6
    elif "yolos_b" in yolos_name:
        config.image_size = [800, 1344]

    config.num_labels = 91
    repo_id = "huggingface/label-files"
    filename = "coco-detection-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()}

    return config



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


def rename_key(name: str) -> str:
    if "backbone" in name:
        name = name.replace("backbone", "vit")
    if "cls_token" in name:
        name = name.replace("cls_token", "embeddings.cls_token")
    if "det_token" in name:
        name = name.replace("det_token", "embeddings.detection_tokens")
    if "mid_pos_embed" in name:
        name = name.replace("mid_pos_embed", "encoder.mid_position_embeddings")
    if "pos_embed" in name:
        name = name.replace("pos_embed", "embeddings.position_embeddings")
    if "patch_embed.proj" in name:
        name = name.replace("patch_embed.proj", "embeddings.patch_embeddings.projection")
    if "blocks" in name:
        name = name.replace("blocks", "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 "class_embed" in name:
        name = name.replace("class_embed", "class_labels_classifier")
    if "bbox_embed" in name:
        name = name.replace("bbox_embed", "bbox_predictor")
    if "vit.norm" in name:
        name = name.replace("vit.norm", "vit.layernorm")

    return name


def convert_state_dict(orig_state_dict: dict, model: YolosForObjectDetection) -> dict:
    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[2])
            dim = model.vit.encoder.layer[layer_num].attention.attention.all_head_size
            if "weight" in key:
                orig_state_dict[f"vit.encoder.layer.{layer_num}.attention.attention.query.weight"] = val[:dim, :]
                orig_state_dict[f"vit.encoder.layer.{layer_num}.attention.attention.key.weight"] = val[
                    dim : dim * 2, :
                ]
                orig_state_dict[f"vit.encoder.layer.{layer_num}.attention.attention.value.weight"] = val[-dim:, :]
            else:
                orig_state_dict[f"vit.encoder.layer.{layer_num}.attention.attention.query.bias"] = val[:dim]
                orig_state_dict[f"vit.encoder.layer.{layer_num}.attention.attention.key.bias"] = val[dim : dim * 2]
                orig_state_dict[f"vit.encoder.layer.{layer_num}.attention.attention.value.bias"] = val[-dim:]
        else:
            orig_state_dict[rename_key(key)] = val

    return orig_state_dict



def prepare_img() -> torch.Tensor:
    url = "http://images.cocodataset.org/val2017/000000039769.jpg"
    im = Image.open(requests.get(url, stream=True).raw)
    return im


@torch.no_grad()
def convert_yolos_checkpoint(
    yolos_name: str, checkpoint_path: str, pytorch_dump_folder_path: str, push_to_hub: bool = False
):
    
    config = get_yolos_config(yolos_name)

    
    state_dict = torch.load(checkpoint_path, map_location="cpu")["model"]

    
    model = YolosForObjectDetection(config)
    model.eval()
    new_state_dict = convert_state_dict(state_dict, model)
    model.load_state_dict(new_state_dict)

    
    size = 800 if yolos_name != "yolos_ti" else 512
    image_processor = YolosImageProcessor(format="coco_detection", size=size)
    encoding = image_processor(images=prepare_img(), return_tensors="pt")
    outputs = model(**encoding)
    logits, pred_boxes = outputs.logits, outputs.pred_boxes

    expected_slice_logits, expected_slice_boxes = None, None
    if yolos_name == "yolos_ti":
        expected_slice_logits = torch.tensor(
            [[-39.5022, -11.9820, -17.6888], [-29.9574, -9.9769, -17.7691], [-42.3281, -20.7200, -30.6294]]
        )
        expected_slice_boxes = torch.tensor(
            [[0.4021, 0.0836, 0.7979], [0.0184, 0.2609, 0.0364], [0.1781, 0.2004, 0.2095]]
        )
    elif yolos_name == "yolos_s_200_pre":
        expected_slice_logits = torch.tensor(
            [[-24.0248, -10.3024, -14.8290], [-42.0392, -16.8200, -27.4334], [-27.2743, -11.8154, -18.7148]]
        )
        expected_slice_boxes = torch.tensor(
            [[0.2559, 0.5455, 0.4706], [0.2989, 0.7279, 0.1875], [0.7732, 0.4017, 0.4462]]
        )
    elif yolos_name == "yolos_s_300_pre":
        expected_slice_logits = torch.tensor(
            [[-36.2220, -14.4385, -23.5457], [-35.6970, -14.7583, -21.3935], [-31.5939, -13.6042, -16.8049]]
        )
        expected_slice_boxes = torch.tensor(
            [[0.7614, 0.2316, 0.4728], [0.7168, 0.4495, 0.3855], [0.4996, 0.1466, 0.9996]]
        )
    elif yolos_name == "yolos_s_dWr":
        expected_slice_logits = torch.tensor(
            [[-42.8668, -24.1049, -41.1690], [-34.7456, -14.1274, -24.9194], [-33.7898, -12.1946, -25.6495]]
        )
        expected_slice_boxes = torch.tensor(
            [[0.5587, 0.2773, 0.0605], [0.5004, 0.3014, 0.9994], [0.4999, 0.1548, 0.9994]]
        )
    elif yolos_name == "yolos_base":
        expected_slice_logits = torch.tensor(
            [[-40.6064, -24.3084, -32.6447], [-55.1990, -30.7719, -35.5877], [-51.4311, -33.3507, -35.6462]]
        )
        expected_slice_boxes = torch.tensor(
            [[0.5555, 0.2794, 0.0655], [0.9049, 0.2664, 0.1894], [0.9183, 0.1984, 0.1635]]
        )
    else:
        raise ValueError(f"Unknown yolos_name: {yolos_name}")

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

    Path(pytorch_dump_folder_path).mkdir(exist_ok=True)
    print(f"Saving model {yolos_name} 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:
        model_mapping = {
            "yolos_ti": "yolos-tiny",
            "yolos_s_200_pre": "yolos-small",
            "yolos_s_300_pre": "yolos-small-300",
            "yolos_s_dWr": "yolos-small-dwr",
            "yolos_base": "yolos-base",
        }

        print("Pushing to the hub...")
        model_name = model_mapping[yolos_name]
        image_processor.push_to_hub(model_name, organization="hustvl")
        model.push_to_hub(model_name, organization="hustvl")


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    
    parser.add_argument(
        "--yolos_name",
        default="yolos_s_200_pre",
        type=str,
        help=(
            "Name of the YOLOS model you'd like to convert. Should be one of 'yolos_ti', 'yolos_s_200_pre',"
            " 'yolos_s_300_pre', 'yolos_s_dWr', 'yolos_base'."
        ),
    )
    parser.add_argument(
        "--checkpoint_path", default=None, type=str, help="Path to the original state dict (.pth file)."
    )
    parser.add_argument(
        "--pytorch_dump_folder_path", default=None, type=str, help="Path to the output PyTorch model directory."
    )
    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_yolos_checkpoint(args.yolos_name, args.checkpoint_path, args.pytorch_dump_folder_path, args.push_to_hub)