codingutf8 2023 s wenhai wang enze xie xiang li dengping fan kaitao song ding liang tong lu ping luo ling shao and 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 pvt checkpoints from the original library import argparse from pathlib import path import requests import torch from pil import image from transformers import pvtconfig pvtforimageclassification pvtimageprocessor 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 def createrenamekeysconfig renamekeys for i in rangeconfig numencoderblocks remane embedings paramters renamekeys appendfposembedi 1 fpvt encoder patchembeddings i positionembeddings renamekeys appendfpatchembedi 1 proj weight fpvt encoder patchembeddings i projection weight renamekeys appendfpatchembedi 1 proj bias fpvt encoder patchembeddings i projection bias renamekeys appendfpatchembedi 1 norm weight fpvt encoder patchembeddings i layernorm weight renamekeys appendfpatchembedi 1 norm bias fpvt encoder patchembeddings i layernorm bias for j in rangeconfig depthsi rename blocks parameters renamekeys append fblocki 1 j attn q weight fpvt encoder block i j attention self query weight renamekeys append fblocki 1 j attn q bias fpvt encoder block i j attention self query bias renamekeys append fblocki 1 j attn kv weight fpvt encoder block i j attention self kv weight renamekeys appendfblocki 1 j attn kv bias fpvt encoder block i j attention self kv bias if config sequencereductionratiosi 1 renamekeys append fblocki 1 j attn norm weight fpvt encoder block i j attention self layernorm weight renamekeys append fblocki 1 j attn norm bias fpvt encoder block i j attention self layernorm bias renamekeys append fblocki 1 j attn sr weight fpvt encoder block i j attention self sequencereduction weight renamekeys append fblocki 1 j attn sr bias fpvt encoder block i j attention self sequencereduction bias renamekeys append fblocki 1 j attn proj weight fpvt encoder block i j attention output dense weight renamekeys append fblocki 1 j attn proj bias fpvt encoder block i j attention output dense bias renamekeys appendfblocki 1 j norm1 weight fpvt encoder block i j layernorm1 weight renamekeys appendfblocki 1 j norm1 bias fpvt encoder block i j layernorm1 bias renamekeys appendfblocki 1 j norm2 weight fpvt encoder block i j layernorm2 weight renamekeys appendfblocki 1 j norm2 bias fpvt encoder block i j layernorm2 bias renamekeys appendfblocki 1 j mlp fc1 weight fpvt encoder block i j mlp dense1 weight renamekeys appendfblocki 1 j mlp fc1 bias fpvt encoder block i j mlp dense1 bias renamekeys appendfblocki 1 j mlp fc2 weight fpvt encoder block i j mlp dense2 weight renamekeys appendfblocki 1 j mlp fc2 bias fpvt encoder block i j mlp dense2 bias rename cls token renamekeys extend clstoken pvt encoder patchembeddings 3 clstoken rename norm layer and classifier layer renamekeys extend norm weight pvt encoder layernorm weight norm bias pvt encoder layernorm bias head weight classifier weight head bias classifier bias return renamekeys we split up the matrix of each encoder layer into queries keys and values def readinkvstatedict config for each of the encoder blocks for i in rangeconfig numencoderblocks for j in rangeconfig depthsi read in weights bias of keys and values which is a single matrix in the original implementation kvweight statedict popfpvt encoder block i j attention self kv weight kvbias statedict popfpvt encoder block i j attention self kv bias next add keys and values in that order to the state dict statedictfpvt encoder block i j attention self key weight kvweight config hiddensizesi statedictfpvt encoder block i j attention self key bias kvbias config hiddensizesi statedictfpvt encoder block i j attention self value weight kvweight config hiddensizesi statedictfpvt encoder block i j attention self value bias kvbiasconfig hiddensizesi def renamekeydct old new val dct popold dctnew val 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 convertpvtcheckpointpvtsize pvtcheckpoint pytorchdumpfolderpath define default pvt configuration if pvtsize tiny configpath zetatechpvttiny224 elif pvtsize small configpath zetatechpvtsmall224 elif pvtsize medium configpath zetatechpvtmedium224 elif pvtsize large configpath zetatechpvtlarge224 else raise valueerrorfavailable model s size tiny small medium large but f pvtsize was given config pvtconfignameorpathconfigpath load original model from https github comwhai362pvt statedict torch loadpvtcheckpoint maplocationcpu renamekeys createrenamekeysconfig for src dest in renamekeys renamekeystatedict src dest readinkvstatedict config load huggingface model model pvtforimageclassificationconfig eval model loadstatedictstatedict check outputs on an image prepared by pvtfeatureextractor imageprocessor pvtimageprocessorsizeconfig imagesize encoding imageprocessorimagesprepareimg returntensorspt pixelvalues encodingpixelvalues outputs modelpixelvalues logits outputs logits detach cpu if pvtsize tiny expectedslicelogits torch tensor1 4192 1 9158 0 9702 elif pvtsize small expectedslicelogits torch tensor0 4353 0 1960 0 2373 elif pvtsize medium expectedslicelogits torch tensor0 2914 0 2231 0 0321 elif pvtsize large expectedslicelogits torch tensor0 3740 0 7739 0 4214 else raise valueerrorfavailable model s size tiny small medium large but f pvtsize was given assert torch allcloselogits0 3 expectedslicelogits atol1e4 pathpytorchdumpfolderpath mkdirexistoktrue printfsaving model pytorchmodel bin to pytorchdumpfolderpath model savepretrainedpytorchdumpfolderpath printfsaving image processor to pytorchdumpfolderpath imageprocessor savepretrainedpytorchdumpfolderpath if name main parser argparse argumentparser required parameters parser addargument pvtsize defaulttiny typestr helpsize of the pvt pretrained model you d like to convert parser addargument pvtcheckpoint defaultpvttiny pth typestr helpcheckpoint of the pvt pretrained model you d like to convert parser addargument pytorchdumpfolderpath defaultnone typestr helppath to the output pytorch model directory args parser parseargs convertpvtcheckpointargs pvtsize args pvtcheckpoint args pytorchdumpfolderpath coding utf 8 2023 s wenhai wang enze xie xiang li deng ping fan kaitao song ding liang tong lu ping luo ling shao and 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 pvt checkpoints from the original library here we list all keys to be renamed original name on the left our name on the right remane embedings paramters rename blocks parameters rename cls token rename norm layer and classifier layer we split up the matrix of each encoder layer into queries keys and values for each of the encoder blocks read in weights bias of keys and values which is a single matrix in the original implementation next add 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 pvt structure define default pvt configuration load original model from https github com whai362 pvt load huggingface model check outputs on an image prepared by pvtfeatureextractor required parameters
import argparse
from pathlib import Path

import requests
import torch
from PIL import Image

from transformers import PvtConfig, PvtForImageClassification, PvtImageProcessor
from transformers.utils import logging


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



def create_rename_keys(config):
    rename_keys = []
    for i in range(config.num_encoder_blocks):
        
        rename_keys.append((f"pos_embed{i + 1}", f"pvt.encoder.patch_embeddings.{i}.position_embeddings"))

        rename_keys.append((f"patch_embed{i + 1}.proj.weight", f"pvt.encoder.patch_embeddings.{i}.projection.weight"))
        rename_keys.append((f"patch_embed{i + 1}.proj.bias", f"pvt.encoder.patch_embeddings.{i}.projection.bias"))
        rename_keys.append((f"patch_embed{i + 1}.norm.weight", f"pvt.encoder.patch_embeddings.{i}.layer_norm.weight"))
        rename_keys.append((f"patch_embed{i + 1}.norm.bias", f"pvt.encoder.patch_embeddings.{i}.layer_norm.bias"))

        for j in range(config.depths[i]):
            
            rename_keys.append(
                (f"block{i + 1}.{j}.attn.q.weight", f"pvt.encoder.block.{i}.{j}.attention.self.query.weight")
            )
            rename_keys.append(
                (f"block{i + 1}.{j}.attn.q.bias", f"pvt.encoder.block.{i}.{j}.attention.self.query.bias")
            )
            rename_keys.append(
                (f"block{i + 1}.{j}.attn.kv.weight", f"pvt.encoder.block.{i}.{j}.attention.self.kv.weight")
            )
            rename_keys.append((f"block{i + 1}.{j}.attn.kv.bias", f"pvt.encoder.block.{i}.{j}.attention.self.kv.bias"))

            if config.sequence_reduction_ratios[i] > 1:
                rename_keys.append(
                    (
                        f"block{i + 1}.{j}.attn.norm.weight",
                        f"pvt.encoder.block.{i}.{j}.attention.self.layer_norm.weight",
                    )
                )
                rename_keys.append(
                    (f"block{i + 1}.{j}.attn.norm.bias", f"pvt.encoder.block.{i}.{j}.attention.self.layer_norm.bias")
                )
                rename_keys.append(
                    (
                        f"block{i + 1}.{j}.attn.sr.weight",
                        f"pvt.encoder.block.{i}.{j}.attention.self.sequence_reduction.weight",
                    )
                )
                rename_keys.append(
                    (
                        f"block{i + 1}.{j}.attn.sr.bias",
                        f"pvt.encoder.block.{i}.{j}.attention.self.sequence_reduction.bias",
                    )
                )

            rename_keys.append(
                (f"block{i + 1}.{j}.attn.proj.weight", f"pvt.encoder.block.{i}.{j}.attention.output.dense.weight")
            )
            rename_keys.append(
                (f"block{i + 1}.{j}.attn.proj.bias", f"pvt.encoder.block.{i}.{j}.attention.output.dense.bias")
            )

            rename_keys.append((f"block{i + 1}.{j}.norm1.weight", f"pvt.encoder.block.{i}.{j}.layer_norm_1.weight"))
            rename_keys.append((f"block{i + 1}.{j}.norm1.bias", f"pvt.encoder.block.{i}.{j}.layer_norm_1.bias"))

            rename_keys.append((f"block{i + 1}.{j}.norm2.weight", f"pvt.encoder.block.{i}.{j}.layer_norm_2.weight"))
            rename_keys.append((f"block{i + 1}.{j}.norm2.bias", f"pvt.encoder.block.{i}.{j}.layer_norm_2.bias"))

            rename_keys.append((f"block{i + 1}.{j}.mlp.fc1.weight", f"pvt.encoder.block.{i}.{j}.mlp.dense1.weight"))
            rename_keys.append((f"block{i + 1}.{j}.mlp.fc1.bias", f"pvt.encoder.block.{i}.{j}.mlp.dense1.bias"))
            rename_keys.append((f"block{i + 1}.{j}.mlp.fc2.weight", f"pvt.encoder.block.{i}.{j}.mlp.dense2.weight"))
            rename_keys.append((f"block{i + 1}.{j}.mlp.fc2.bias", f"pvt.encoder.block.{i}.{j}.mlp.dense2.bias"))

    
    rename_keys.extend(
        [
            ("cls_token", "pvt.encoder.patch_embeddings.3.cls_token"),
        ]
    )
    
    rename_keys.extend(
        [
            ("norm.weight", "pvt.encoder.layer_norm.weight"),
            ("norm.bias", "pvt.encoder.layer_norm.bias"),
            ("head.weight", "classifier.weight"),
            ("head.bias", "classifier.bias"),
        ]
    )

    return rename_keys



def read_in_k_v(state_dict, config):
    
    for i in range(config.num_encoder_blocks):
        for j in range(config.depths[i]):
            
            kv_weight = state_dict.pop(f"pvt.encoder.block.{i}.{j}.attention.self.kv.weight")
            kv_bias = state_dict.pop(f"pvt.encoder.block.{i}.{j}.attention.self.kv.bias")
            
            state_dict[f"pvt.encoder.block.{i}.{j}.attention.self.key.weight"] = kv_weight[: config.hidden_sizes[i], :]
            state_dict[f"pvt.encoder.block.{i}.{j}.attention.self.key.bias"] = kv_bias[: config.hidden_sizes[i]]

            state_dict[f"pvt.encoder.block.{i}.{j}.attention.self.value.weight"] = kv_weight[
                config.hidden_sizes[i] :, :
            ]
            state_dict[f"pvt.encoder.block.{i}.{j}.attention.self.value.bias"] = kv_bias[config.hidden_sizes[i] :]


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



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_pvt_checkpoint(pvt_size, pvt_checkpoint, pytorch_dump_folder_path):
    

    
    if pvt_size == "tiny":
        config_path = "Zetatech/pvt-tiny-224"
    elif pvt_size == "small":
        config_path = "Zetatech/pvt-small-224"
    elif pvt_size == "medium":
        config_path = "Zetatech/pvt-medium-224"
    elif pvt_size == "large":
        config_path = "Zetatech/pvt-large-224"
    else:
        raise ValueError(f"Available model's size: 'tiny', 'small', 'medium', 'large', but " f"'{pvt_size}' was given")
    config = PvtConfig(name_or_path=config_path)
    
    state_dict = torch.load(pvt_checkpoint, map_location="cpu")

    rename_keys = create_rename_keys(config)
    for src, dest in rename_keys:
        rename_key(state_dict, src, dest)
    read_in_k_v(state_dict, config)

    
    model = PvtForImageClassification(config).eval()
    model.load_state_dict(state_dict)

    
    image_processor = PvtImageProcessor(size=config.image_size)
    encoding = image_processor(images=prepare_img(), return_tensors="pt")
    pixel_values = encoding["pixel_values"]
    outputs = model(pixel_values)
    logits = outputs.logits.detach().cpu()

    if pvt_size == "tiny":
        expected_slice_logits = torch.tensor([-1.4192, -1.9158, -0.9702])
    elif pvt_size == "small":
        expected_slice_logits = torch.tensor([0.4353, -0.1960, -0.2373])
    elif pvt_size == "medium":
        expected_slice_logits = torch.tensor([-0.2914, -0.2231, 0.0321])
    elif pvt_size == "large":
        expected_slice_logits = torch.tensor([0.3740, -0.7739, -0.4214])
    else:
        raise ValueError(f"Available model's size: 'tiny', 'small', 'medium', 'large', but " f"'{pvt_size}' was given")

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

    Path(pytorch_dump_folder_path).mkdir(exist_ok=True)
    print(f"Saving model pytorch_model.bin 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(
        "--pvt_size",
        default="tiny",
        type=str,
        help="Size of the PVT pretrained model you'd like to convert.",
    )
    parser.add_argument(
        "--pvt_checkpoint",
        default="pvt_tiny.pth",
        type=str,
        help="Checkpoint of the PVT pretrained model 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_pvt_checkpoint(args.pvt_size, args.pvt_checkpoint, args.pytorch_dump_folder_path)