codingutf8 2020 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 detr checkpoints with timm backbone import argparse import json from collections import ordereddict from pathlib import path import requests import torch from huggingfacehub import hfhubdownload from pil import image from transformers import detrconfig detrforobjectdetection detrforsegmentation detrimageprocessor 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 renamekeys for i in range6 encoder layers output projection 2 feedforward neural networks and 2 layernorms renamekeys append ftransformer encoder layers i selfattn outproj weight fencoder layers i selfattn outproj weight renamekeys append ftransformer encoder layers i selfattn outproj bias fencoder layers i selfattn outproj bias renamekeys appendftransformer encoder layers i linear1 weight fencoder layers i fc1 weight renamekeys appendftransformer encoder layers i linear1 bias fencoder layers i fc1 bias renamekeys appendftransformer encoder layers i linear2 weight fencoder layers i fc2 weight renamekeys appendftransformer encoder layers i linear2 bias fencoder layers i fc2 bias renamekeys append ftransformer encoder layers i norm1 weight fencoder layers i selfattnlayernorm weight renamekeys appendftransformer encoder layers i norm1 bias fencoder layers i selfattnlayernorm bias renamekeys appendftransformer encoder layers i norm2 weight fencoder layers i finallayernorm weight renamekeys appendftransformer encoder layers i norm2 bias fencoder layers i finallayernorm bias decoder layers 2 times output projection 2 feedforward neural networks and 3 layernorms renamekeys append ftransformer decoder layers i selfattn outproj weight fdecoder layers i selfattn outproj weight renamekeys append ftransformer decoder layers i selfattn outproj bias fdecoder layers i selfattn outproj bias renamekeys append ftransformer decoder layers i multiheadattn outproj weight fdecoder layers i encoderattn outproj weight renamekeys append ftransformer decoder layers i multiheadattn outproj bias fdecoder layers i encoderattn outproj bias renamekeys appendftransformer decoder layers i linear1 weight fdecoder layers i fc1 weight renamekeys appendftransformer decoder layers i linear1 bias fdecoder layers i fc1 bias renamekeys appendftransformer decoder layers i linear2 weight fdecoder layers i fc2 weight renamekeys appendftransformer decoder layers i linear2 bias fdecoder layers i fc2 bias renamekeys append ftransformer decoder layers i norm1 weight fdecoder layers i selfattnlayernorm weight renamekeys appendftransformer decoder layers i norm1 bias fdecoder layers i selfattnlayernorm bias renamekeys append ftransformer decoder layers i norm2 weight fdecoder layers i encoderattnlayernorm weight renamekeys append ftransformer decoder layers i norm2 bias fdecoder layers i encoderattnlayernorm bias renamekeys appendftransformer decoder layers i norm3 weight fdecoder layers i finallayernorm weight renamekeys appendftransformer decoder layers i norm3 bias fdecoder layers i finallayernorm bias convolutional projection query embeddings layernorm of decoder class and bounding box heads renamekeys extend inputproj weight inputprojection weight inputproj bias inputprojection bias queryembed weight querypositionembeddings weight transformer decoder norm weight decoder layernorm weight transformer decoder norm bias decoder layernorm bias classembed weight classlabelsclassifier weight classembed bias classlabelsclassifier bias bboxembed layers 0 weight bboxpredictor layers 0 weight bboxembed layers 0 bias bboxpredictor layers 0 bias bboxembed layers 1 weight bboxpredictor layers 1 weight bboxembed layers 1 bias bboxpredictor layers 1 bias bboxembed layers 2 weight bboxpredictor layers 2 weight bboxembed layers 2 bias bboxpredictor layers 2 bias def renamekeystatedict old new val statedict popold statedictnew val def renamebackbonekeysstatedict newstatedict ordereddict for key value in statedict items if backbone 0 body in key newkey key replacebackbone 0 body backbone convencoder model newstatedictnewkey value else newstatedictkey value return newstatedict def readinqkvstatedict ispanopticfalse prefix if ispanoptic prefix detr first transformer encoder for i in range6 read in weights bias of input projection layer in pytorch s multiheadattention this is a single matrix bias inprojweight statedict popfprefixtransformer encoder layers i selfattn inprojweight inprojbias statedict popfprefixtransformer encoder layers i selfattn inprojbias next add query keys and values in that order to the state dict statedictfencoder layers i selfattn qproj weight inprojweight 256 statedictfencoder layers i selfattn qproj bias inprojbias 256 statedictfencoder layers i selfattn kproj weight inprojweight256 512 statedictfencoder layers i selfattn kproj bias inprojbias256 512 statedictfencoder layers i selfattn vproj weight inprojweight256 statedictfencoder layers i selfattn vproj bias inprojbias256 next transformer decoder which is a bit more complex because it also includes crossattention for i in range6 read in weights bias of input projection layer of selfattention inprojweight statedict popfprefixtransformer decoder layers i selfattn inprojweight inprojbias statedict popfprefixtransformer decoder layers i selfattn inprojbias next add query keys and values in that order to the state dict statedictfdecoder layers i selfattn qproj weight inprojweight 256 statedictfdecoder layers i selfattn qproj bias inprojbias 256 statedictfdecoder layers i selfattn kproj weight inprojweight256 512 statedictfdecoder layers i selfattn kproj bias inprojbias256 512 statedictfdecoder layers i selfattn vproj weight inprojweight256 statedictfdecoder layers i selfattn vproj bias inprojbias256 read in weights bias of input projection layer of crossattention inprojweightcrossattn statedict pop fprefixtransformer decoder layers i multiheadattn inprojweight inprojbiascrossattn statedict popfprefixtransformer decoder layers i multiheadattn inprojbias next add query keys and values in that order of crossattention to the state dict statedictfdecoder layers i encoderattn qproj weight inprojweightcrossattn 256 statedictfdecoder layers i encoderattn qproj bias inprojbiascrossattn 256 statedictfdecoder layers i encoderattn kproj weight inprojweightcrossattn256 512 statedictfdecoder layers i encoderattn kproj bias inprojbiascrossattn256 512 statedictfdecoder layers i encoderattn vproj weight inprojweightcrossattn256 statedictfdecoder layers i encoderattn vproj bias inprojbiascrossattn256 we will verify our results on an image of cute cats def prepareimg url http images cocodataset orgval2017000000039769 jpg im image openrequests geturl streamtrue raw return im torch nograd def convertdetrcheckpointmodelname pytorchdumpfolderpath load default config config detrconfig set backbone and dilation attributes if resnet101 in modelname config backbone resnet101 if dc5 in modelname config dilation true ispanoptic panoptic in modelname if ispanoptic config numlabels 250 else 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 load image processor format cocopanoptic if ispanoptic else cocodetection imageprocessor detrimageprocessorformatformat prepare image img prepareimg encoding imageprocessorimagesimg returntensorspt pixelvalues encodingpixelvalues logger infofconverting model modelname load original model from torch hub detr torch hub loadfacebookresearchdetr modelname pretrainedtrue eval statedict detr statedict rename keys for src dest in renamekeys if ispanoptic src detr src renamekeystatedict src dest statedict renamebackbonekeysstatedict query key and value matrices need special treatment readinqkvstatedict ispanopticispanoptic important we need to prepend a prefix to each of the base model keys as the head models use different attributes for them prefix detr model if ispanoptic else model for key in statedict copy keys if ispanoptic if key startswithdetr and not key startswithclasslabelsclassifier and not key startswithbboxpredictor val statedict popkey statedictdetr model key4 val elif classlabelsclassifier in key or bboxpredictor in key val statedict popkey statedictdetr key val elif key startswithbboxattention or key startswithmaskhead continue else val statedict popkey statedictprefix key val else if not key startswithclasslabelsclassifier and not key startswithbboxpredictor val statedict popkey statedictprefix key val finally create huggingface model and load state dict model detrforsegmentationconfig if ispanoptic else detrforobjectdetectionconfig model loadstatedictstatedict model eval verify our conversion originaloutputs detrpixelvalues outputs modelpixelvalues assert torch allcloseoutputs logits originaloutputspredlogits atol1e4 assert torch allcloseoutputs predboxes originaloutputspredboxes atol1e4 if ispanoptic assert torch allcloseoutputs predmasks originaloutputspredmasks atol1e4 save model and image processor logger infofsaving pytorch model and image processor to pytorchdumpfolderpath pathpytorchdumpfolderpath mkdirexistoktrue model savepretrainedpytorchdumpfolderpath imageprocessor savepretrainedpytorchdumpfolderpath if name main parser argparse argumentparser parser addargument modelname defaultdetrresnet50 typestr helpname of the detr model you d like to convert parser addargument pytorchdumpfolderpath defaultnone typestr helppath to the folder to output pytorch model args parser parseargs convertdetrcheckpointargs modelname args pytorchdumpfolderpath coding utf 8 2020 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 detr checkpoints with timm backbone 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 decoder layers 2 times output projection 2 feedforward neural networks and 3 layernorms convolutional projection query embeddings layernorm of decoder class and bounding box heads first transformer encoder read in weights bias of input projection layer in pytorch s multiheadattention this is a single matrix bias next add query keys and values in that order to the state dict next transformer decoder which is a bit more complex because it also includes cross attention read in weights bias of input projection layer of self attention next add query keys and values in that order to the state dict read in weights bias of input projection layer of cross attention next add query keys and values in that order of cross attention to the state dict we will verify our results on an image of cute cats copy paste tweak model s weights to our detr structure load default config set backbone and dilation attributes load image processor prepare image load original model from torch hub rename keys query key and value matrices need special treatment important we need to prepend a prefix to each of the base model keys as the head models use different attributes for them finally create huggingface model and load state dict verify our conversion save model and image processor
import argparse
import json
from collections import OrderedDict
from pathlib import Path

import requests
import torch
from huggingface_hub import hf_hub_download
from PIL import Image

from transformers import DetrConfig, DetrForObjectDetection, DetrForSegmentation, DetrImageProcessor
from transformers.utils import logging


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


rename_keys = []
for i in range(6):
    
    rename_keys.append(
        (f"transformer.encoder.layers.{i}.self_attn.out_proj.weight", f"encoder.layers.{i}.self_attn.out_proj.weight")
    )
    rename_keys.append(
        (f"transformer.encoder.layers.{i}.self_attn.out_proj.bias", f"encoder.layers.{i}.self_attn.out_proj.bias")
    )
    rename_keys.append((f"transformer.encoder.layers.{i}.linear1.weight", f"encoder.layers.{i}.fc1.weight"))
    rename_keys.append((f"transformer.encoder.layers.{i}.linear1.bias", f"encoder.layers.{i}.fc1.bias"))
    rename_keys.append((f"transformer.encoder.layers.{i}.linear2.weight", f"encoder.layers.{i}.fc2.weight"))
    rename_keys.append((f"transformer.encoder.layers.{i}.linear2.bias", f"encoder.layers.{i}.fc2.bias"))
    rename_keys.append(
        (f"transformer.encoder.layers.{i}.norm1.weight", f"encoder.layers.{i}.self_attn_layer_norm.weight")
    )
    rename_keys.append((f"transformer.encoder.layers.{i}.norm1.bias", f"encoder.layers.{i}.self_attn_layer_norm.bias"))
    rename_keys.append((f"transformer.encoder.layers.{i}.norm2.weight", f"encoder.layers.{i}.final_layer_norm.weight"))
    rename_keys.append((f"transformer.encoder.layers.{i}.norm2.bias", f"encoder.layers.{i}.final_layer_norm.bias"))
    
    rename_keys.append(
        (f"transformer.decoder.layers.{i}.self_attn.out_proj.weight", f"decoder.layers.{i}.self_attn.out_proj.weight")
    )
    rename_keys.append(
        (f"transformer.decoder.layers.{i}.self_attn.out_proj.bias", f"decoder.layers.{i}.self_attn.out_proj.bias")
    )
    rename_keys.append(
        (
            f"transformer.decoder.layers.{i}.multihead_attn.out_proj.weight",
            f"decoder.layers.{i}.encoder_attn.out_proj.weight",
        )
    )
    rename_keys.append(
        (
            f"transformer.decoder.layers.{i}.multihead_attn.out_proj.bias",
            f"decoder.layers.{i}.encoder_attn.out_proj.bias",
        )
    )
    rename_keys.append((f"transformer.decoder.layers.{i}.linear1.weight", f"decoder.layers.{i}.fc1.weight"))
    rename_keys.append((f"transformer.decoder.layers.{i}.linear1.bias", f"decoder.layers.{i}.fc1.bias"))
    rename_keys.append((f"transformer.decoder.layers.{i}.linear2.weight", f"decoder.layers.{i}.fc2.weight"))
    rename_keys.append((f"transformer.decoder.layers.{i}.linear2.bias", f"decoder.layers.{i}.fc2.bias"))
    rename_keys.append(
        (f"transformer.decoder.layers.{i}.norm1.weight", f"decoder.layers.{i}.self_attn_layer_norm.weight")
    )
    rename_keys.append((f"transformer.decoder.layers.{i}.norm1.bias", f"decoder.layers.{i}.self_attn_layer_norm.bias"))
    rename_keys.append(
        (f"transformer.decoder.layers.{i}.norm2.weight", f"decoder.layers.{i}.encoder_attn_layer_norm.weight")
    )
    rename_keys.append(
        (f"transformer.decoder.layers.{i}.norm2.bias", f"decoder.layers.{i}.encoder_attn_layer_norm.bias")
    )
    rename_keys.append((f"transformer.decoder.layers.{i}.norm3.weight", f"decoder.layers.{i}.final_layer_norm.weight"))
    rename_keys.append((f"transformer.decoder.layers.{i}.norm3.bias", f"decoder.layers.{i}.final_layer_norm.bias"))


rename_keys.extend(
    [
        ("input_proj.weight", "input_projection.weight"),
        ("input_proj.bias", "input_projection.bias"),
        ("query_embed.weight", "query_position_embeddings.weight"),
        ("transformer.decoder.norm.weight", "decoder.layernorm.weight"),
        ("transformer.decoder.norm.bias", "decoder.layernorm.bias"),
        ("class_embed.weight", "class_labels_classifier.weight"),
        ("class_embed.bias", "class_labels_classifier.bias"),
        ("bbox_embed.layers.0.weight", "bbox_predictor.layers.0.weight"),
        ("bbox_embed.layers.0.bias", "bbox_predictor.layers.0.bias"),
        ("bbox_embed.layers.1.weight", "bbox_predictor.layers.1.weight"),
        ("bbox_embed.layers.1.bias", "bbox_predictor.layers.1.bias"),
        ("bbox_embed.layers.2.weight", "bbox_predictor.layers.2.weight"),
        ("bbox_embed.layers.2.bias", "bbox_predictor.layers.2.bias"),
    ]
)


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


def rename_backbone_keys(state_dict):
    new_state_dict = OrderedDict()
    for key, value in state_dict.items():
        if "backbone.0.body" in key:
            new_key = key.replace("backbone.0.body", "backbone.conv_encoder.model")
            new_state_dict[new_key] = value
        else:
            new_state_dict[key] = value

    return new_state_dict


def read_in_q_k_v(state_dict, is_panoptic=False):
    prefix = ""
    if is_panoptic:
        prefix = "detr."

    
    for i in range(6):
        
        in_proj_weight = state_dict.pop(f"{prefix}transformer.encoder.layers.{i}.self_attn.in_proj_weight")
        in_proj_bias = state_dict.pop(f"{prefix}transformer.encoder.layers.{i}.self_attn.in_proj_bias")
        
        state_dict[f"encoder.layers.{i}.self_attn.q_proj.weight"] = in_proj_weight[:256, :]
        state_dict[f"encoder.layers.{i}.self_attn.q_proj.bias"] = in_proj_bias[:256]
        state_dict[f"encoder.layers.{i}.self_attn.k_proj.weight"] = in_proj_weight[256:512, :]
        state_dict[f"encoder.layers.{i}.self_attn.k_proj.bias"] = in_proj_bias[256:512]
        state_dict[f"encoder.layers.{i}.self_attn.v_proj.weight"] = in_proj_weight[-256:, :]
        state_dict[f"encoder.layers.{i}.self_attn.v_proj.bias"] = in_proj_bias[-256:]
    
    for i in range(6):
        
        in_proj_weight = state_dict.pop(f"{prefix}transformer.decoder.layers.{i}.self_attn.in_proj_weight")
        in_proj_bias = state_dict.pop(f"{prefix}transformer.decoder.layers.{i}.self_attn.in_proj_bias")
        
        state_dict[f"decoder.layers.{i}.self_attn.q_proj.weight"] = in_proj_weight[:256, :]
        state_dict[f"decoder.layers.{i}.self_attn.q_proj.bias"] = in_proj_bias[:256]
        state_dict[f"decoder.layers.{i}.self_attn.k_proj.weight"] = in_proj_weight[256:512, :]
        state_dict[f"decoder.layers.{i}.self_attn.k_proj.bias"] = in_proj_bias[256:512]
        state_dict[f"decoder.layers.{i}.self_attn.v_proj.weight"] = in_proj_weight[-256:, :]
        state_dict[f"decoder.layers.{i}.self_attn.v_proj.bias"] = in_proj_bias[-256:]
        
        in_proj_weight_cross_attn = state_dict.pop(
            f"{prefix}transformer.decoder.layers.{i}.multihead_attn.in_proj_weight"
        )
        in_proj_bias_cross_attn = state_dict.pop(f"{prefix}transformer.decoder.layers.{i}.multihead_attn.in_proj_bias")
        
        state_dict[f"decoder.layers.{i}.encoder_attn.q_proj.weight"] = in_proj_weight_cross_attn[:256, :]
        state_dict[f"decoder.layers.{i}.encoder_attn.q_proj.bias"] = in_proj_bias_cross_attn[:256]
        state_dict[f"decoder.layers.{i}.encoder_attn.k_proj.weight"] = in_proj_weight_cross_attn[256:512, :]
        state_dict[f"decoder.layers.{i}.encoder_attn.k_proj.bias"] = in_proj_bias_cross_attn[256:512]
        state_dict[f"decoder.layers.{i}.encoder_attn.v_proj.weight"] = in_proj_weight_cross_attn[-256:, :]
        state_dict[f"decoder.layers.{i}.encoder_attn.v_proj.bias"] = in_proj_bias_cross_attn[-256:]



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

    return im


@torch.no_grad()
def convert_detr_checkpoint(model_name, pytorch_dump_folder_path):
    

    
    config = DetrConfig()
    
    if "resnet101" in model_name:
        config.backbone = "resnet101"
    if "dc5" in model_name:
        config.dilation = True
    is_panoptic = "panoptic" in model_name
    if is_panoptic:
        config.num_labels = 250
    else:
        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()}

    
    format = "coco_panoptic" if is_panoptic else "coco_detection"
    image_processor = DetrImageProcessor(format=format)

    
    img = prepare_img()
    encoding = image_processor(images=img, return_tensors="pt")
    pixel_values = encoding["pixel_values"]

    logger.info(f"Converting model {model_name}...")

    
    detr = torch.hub.load("facebookresearch/detr", model_name, pretrained=True).eval()
    state_dict = detr.state_dict()
    
    for src, dest in rename_keys:
        if is_panoptic:
            src = "detr." + src
        rename_key(state_dict, src, dest)
    state_dict = rename_backbone_keys(state_dict)
    
    read_in_q_k_v(state_dict, is_panoptic=is_panoptic)
    
    prefix = "detr.model." if is_panoptic else "model."
    for key in state_dict.copy().keys():
        if is_panoptic:
            if (
                key.startswith("detr")
                and not key.startswith("class_labels_classifier")
                and not key.startswith("bbox_predictor")
            ):
                val = state_dict.pop(key)
                state_dict["detr.model" + key[4:]] = val
            elif "class_labels_classifier" in key or "bbox_predictor" in key:
                val = state_dict.pop(key)
                state_dict["detr." + key] = val
            elif key.startswith("bbox_attention") or key.startswith("mask_head"):
                continue
            else:
                val = state_dict.pop(key)
                state_dict[prefix + key] = val
        else:
            if not key.startswith("class_labels_classifier") and not key.startswith("bbox_predictor"):
                val = state_dict.pop(key)
                state_dict[prefix + key] = val
    
    model = DetrForSegmentation(config) if is_panoptic else DetrForObjectDetection(config)
    model.load_state_dict(state_dict)
    model.eval()
    
    original_outputs = detr(pixel_values)
    outputs = model(pixel_values)
    assert torch.allclose(outputs.logits, original_outputs["pred_logits"], atol=1e-4)
    assert torch.allclose(outputs.pred_boxes, original_outputs["pred_boxes"], atol=1e-4)
    if is_panoptic:
        assert torch.allclose(outputs.pred_masks, original_outputs["pred_masks"], atol=1e-4)

    
    logger.info(f"Saving PyTorch model and image processor to {pytorch_dump_folder_path}...")
    Path(pytorch_dump_folder_path).mkdir(exist_ok=True)
    model.save_pretrained(pytorch_dump_folder_path)
    image_processor.save_pretrained(pytorch_dump_folder_path)


if __name__ == "__main__":
    parser = argparse.ArgumentParser()

    parser.add_argument(
        "--model_name", default="detr_resnet50", type=str, help="Name of the DETR model you'd like to convert."
    )
    parser.add_argument(
        "--pytorch_dump_folder_path", default=None, type=str, help="Path to the folder to output PyTorch model."
    )
    args = parser.parse_args()
    convert_detr_checkpoint(args.model_name, args.pytorch_dump_folder_path)