codingutf8 2022 bnrist tsinghua university tklndst nankai university 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 van checkpoints from the original repository url https github comvisualattentionnetworkvanclassification import argparse import json import sys from dataclasses import dataclass field from functools import partial from pathlib import path from typing import list import torch import torch nn as nn from huggingfacehub import cacheddownload hfhubdownload from torch import tensor from transformers import autoimageprocessor vanconfig vanforimageclassification from transformers models deprecated van modelingvan import vanlayerscaling from transformers utils import logging logging setverbosityinfo logger logging getloggername dataclass class tracker module nn module traced listnn module fielddefaultfactorylist handles list fielddefaultfactorylist def forwardhookself m inputs tensor outputs tensor hasnotsubmodules lenlistm modules 1 or isinstancem nn conv2d or isinstancem nn batchnorm2d if hasnotsubmodules if not isinstancem vanlayerscaling self traced appendm def callself x tensor for m in self module modules self handles appendm registerforwardhookself forwardhook self modulex x remove for x in self handles return self property def parametrizedself check the len of the statedict keys to see if we have learnable params return listfilterlambda x lenlistx statedict keys 0 self traced dataclass class moduletransfer src nn module dest nn module verbose int 0 srcskip list fielddefaultfactorylist destskip list fielddefaultfactorylist def callself x tensor desttraced trackerself destx parametrized srctraced trackerself srcx parametrized srctraced listfilterlambda x typex not in self srcskip srctraced desttraced listfilterlambda x typex not in self destskip desttraced if lendesttraced lensrctraced raise exception fnumbers of operations are different source module has lensrctraced operations while f destination module has lendesttraced for destm srcm in zipdesttraced srctraced destm loadstatedictsrcm statedict if self verbose 1 printftransfered fromsrcm todestm def copyparametersfrommodel nn module ourmodel nn module nn module nn parameter cannot be tracked by the tracker thus we need to manually convert them fromstatedict frommodel statedict ourstatedict ourmodel statedict config ourmodel config allkeys for stageidx in rangelenconfig hiddensizes for blockid in rangeconfig depthsstageidx fromkey fblockstageidx 1 blockid layerscale1 tokey fvan encoder stages stageidx layers blockid attentionscaling weight allkeys appendfromkey tokey fromkey fblockstageidx 1 blockid layerscale2 tokey fvan encoder stages stageidx layers blockid mlpscaling weight allkeys appendfromkey tokey for fromkey tokey in allkeys ourstatedicttokey fromstatedict popfromkey ourmodel loadstatedictourstatedict return ourmodel def convertweightandpush name str config vanconfig checkpoint str frommodel nn module savedirectory path pushtohub bool true printfdownloading weights for name checkpointpath cacheddownloadcheckpoint printfconverting name fromstatedict torch loadcheckpointpathstatedict frommodel loadstatedictfromstatedict frommodel eval with torch nograd ourmodel vanforimageclassificationconfig eval moduletransfer moduletransfersrcfrommodel destourmodel x torch randn1 3 224 224 moduletransferx ourmodel copyparametersfrommodel ourmodel if not torch allclosefrommodelx ourmodelx logits raise valueerrorthe model logits don t match the original one checkpointname name printcheckpointname if pushtohub ourmodel pushtohub repopathornamesavedirectory checkpointname commitmessageadd model usetempdirtrue we can use the convnext one imageprocessor autoimageprocessor frompretrainedfacebookconvnextbase22422k1k imageprocessor pushtohub repopathornamesavedirectory checkpointname commitmessageadd image processor usetempdirtrue printfpushed checkpointname def convertweightsandpushsavedirectory path modelname str none pushtohub bool true filename imagenet1kid2label json numlabels 1000 repoid huggingfacelabelfiles numlabels numlabels id2label json loadopenhfhubdownloadrepoid filename repotypedataset r id2label intk v for k v in id2label items id2label id2label label2id v k for k v in id2label items imagenetpretrainedconfig partialvanconfig numlabelsnumlabels id2labelid2label label2idlabel2id namestoconfig vantiny imagenetpretrainedconfig hiddensizes32 64 160 256 depths3 3 5 2 mlpratios8 8 4 4 vansmall imagenetpretrainedconfig hiddensizes64 128 320 512 depths2 2 4 2 mlpratios8 8 4 4 vanbase imagenetpretrainedconfig hiddensizes64 128 320 512 depths3 3 12 3 mlpratios8 8 4 4 vanlarge imagenetpretrainedconfig hiddensizes64 128 320 512 depths3 5 27 3 mlpratios8 8 4 4 namestooriginalmodels vantiny vantiny vansmall vansmall vanbase vanbase vanlarge vanlarge namestooriginalcheckpoints vantiny https huggingface covisualattentionnetworkvantinyoriginalresolvemainvantiny754 pth tar vansmall https huggingface covisualattentionnetworkvansmalloriginalresolvemainvansmall811 pth tar vanbase https huggingface covisualattentionnetworkvanbaseoriginalresolvemainvanbase828 pth tar vanlarge https huggingface covisualattentionnetworkvanlargeoriginalresolvemainvanlarge839 pth tar if modelname convertweightandpush modelname namestoconfigmodelname checkpointnamestooriginalcheckpointsmodelname frommodelnamestooriginalmodelsmodelname savedirectorysavedirectory pushtohubpushtohub else for modelname config in namestoconfig items convertweightandpush modelname config checkpointnamestooriginalcheckpointsmodelname frommodelnamestooriginalmodelsmodelname savedirectorysavedirectory pushtohubpushtohub if name main parser argparse argumentparser required parameters parser addargument modelname defaultnone typestr help the name of the model you wish to convert it must be one of the supported resnet architecture currently vantinysmallbaselarge if none all of them will the converted parser addargument pytorchdumpfolderpath defaultnone typepath requiredtrue helppath to the output pytorch model directory parser addargument vandir requiredtrue typepath help a path to van s original implementation directory you can download from here https github comvisualattentionnetworkvanclassification parser addargument pushtohub defaulttrue typebool requiredfalse helpif true push model and image processor to the hub args parser parseargs pytorchdumpfolderpath path args pytorchdumpfolderpath pytorchdumpfolderpath mkdirexistoktrue parentstrue vandir args vandir append the path to the parents to maskformer dir sys path appendstrvandir parent from van models van import vanbase vanlarge vansmall vantiny convertweightsandpushpytorchdumpfolderpath args modelname args pushtohub coding utf 8 2022 bnrist tsinghua university tklndst nankai university 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 van checkpoints from the original repository url https github com visual attention network van classification check the len of the state_dict keys to see if we have learnable params transfer the weights of self src to self dest by performing a forward pass using x as input under the hood we tracked all the operations in both modules nn parameter cannot be tracked by the tracker thus we need to manually convert them we can use the convnext one required parameters append the path to the parents to maskformer dir
import argparse
import json
import sys
from dataclasses import dataclass, field
from functools import partial
from pathlib import Path
from typing import List

import torch
import torch.nn as nn
from huggingface_hub import cached_download, hf_hub_download
from torch import Tensor

from transformers import AutoImageProcessor, VanConfig, VanForImageClassification
from transformers.models.deprecated.van.modeling_van import VanLayerScaling
from transformers.utils import logging


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


@dataclass
class Tracker:
    module: nn.Module
    traced: List[nn.Module] = field(default_factory=list)
    handles: list = field(default_factory=list)

    def _forward_hook(self, m, inputs: Tensor, outputs: Tensor):
        has_not_submodules = len(list(m.modules())) == 1 or isinstance(m, nn.Conv2d) or isinstance(m, nn.BatchNorm2d)
        if has_not_submodules:
            if not isinstance(m, VanLayerScaling):
                self.traced.append(m)

    def __call__(self, x: Tensor):
        for m in self.module.modules():
            self.handles.append(m.register_forward_hook(self._forward_hook))
        self.module(x)
        [x.remove() for x in self.handles]
        return self

    @property
    def parametrized(self):
        
        return list(filter(lambda x: len(list(x.state_dict().keys())) > 0, self.traced))


@dataclass
class ModuleTransfer:
    src: nn.Module
    dest: nn.Module
    verbose: int = 0
    src_skip: List = field(default_factory=list)
    dest_skip: List = field(default_factory=list)

    def __call__(self, x: Tensor):
        
        dest_traced = Tracker(self.dest)(x).parametrized
        src_traced = Tracker(self.src)(x).parametrized

        src_traced = list(filter(lambda x: type(x) not in self.src_skip, src_traced))
        dest_traced = list(filter(lambda x: type(x) not in self.dest_skip, dest_traced))

        if len(dest_traced) != len(src_traced):
            raise Exception(
                f"Numbers of operations are different. Source module has {len(src_traced)} operations while"
                f" destination module has {len(dest_traced)}."
            )

        for dest_m, src_m in zip(dest_traced, src_traced):
            dest_m.load_state_dict(src_m.state_dict())
            if self.verbose == 1:
                print(f"Transfered from={src_m} to={dest_m}")


def copy_parameters(from_model: nn.Module, our_model: nn.Module) -> nn.Module:
    
    from_state_dict = from_model.state_dict()
    our_state_dict = our_model.state_dict()
    config = our_model.config
    all_keys = []
    for stage_idx in range(len(config.hidden_sizes)):
        for block_id in range(config.depths[stage_idx]):
            from_key = f"block{stage_idx + 1}.{block_id}.layer_scale_1"
            to_key = f"van.encoder.stages.{stage_idx}.layers.{block_id}.attention_scaling.weight"

            all_keys.append((from_key, to_key))
            from_key = f"block{stage_idx + 1}.{block_id}.layer_scale_2"
            to_key = f"van.encoder.stages.{stage_idx}.layers.{block_id}.mlp_scaling.weight"

            all_keys.append((from_key, to_key))

    for from_key, to_key in all_keys:
        our_state_dict[to_key] = from_state_dict.pop(from_key)

    our_model.load_state_dict(our_state_dict)
    return our_model


def convert_weight_and_push(
    name: str,
    config: VanConfig,
    checkpoint: str,
    from_model: nn.Module,
    save_directory: Path,
    push_to_hub: bool = True,
):
    print(f"Downloading weights for {name}...")
    checkpoint_path = cached_download(checkpoint)
    print(f"Converting {name}...")
    from_state_dict = torch.load(checkpoint_path)["state_dict"]
    from_model.load_state_dict(from_state_dict)
    from_model.eval()
    with torch.no_grad():
        our_model = VanForImageClassification(config).eval()
        module_transfer = ModuleTransfer(src=from_model, dest=our_model)
        x = torch.randn((1, 3, 224, 224))
        module_transfer(x)
        our_model = copy_parameters(from_model, our_model)

    if not torch.allclose(from_model(x), our_model(x).logits):
        raise ValueError("The model logits don't match the original one.")

    checkpoint_name = name
    print(checkpoint_name)

    if push_to_hub:
        our_model.push_to_hub(
            repo_path_or_name=save_directory / checkpoint_name,
            commit_message="Add model",
            use_temp_dir=True,
        )

        
        image_processor = AutoImageProcessor.from_pretrained("facebook/convnext-base-224-22k-1k")
        image_processor.push_to_hub(
            repo_path_or_name=save_directory / checkpoint_name,
            commit_message="Add image processor",
            use_temp_dir=True,
        )

        print(f"Pushed {checkpoint_name}")


def convert_weights_and_push(save_directory: Path, model_name: str = None, push_to_hub: bool = True):
    filename = "imagenet-1k-id2label.json"
    num_labels = 1000

    repo_id = "huggingface/label-files"
    num_labels = num_labels
    id2label = json.load(open(hf_hub_download(repo_id, filename, repo_type="dataset"), "r"))
    id2label = {int(k): v for k, v in id2label.items()}

    id2label = id2label
    label2id = {v: k for k, v in id2label.items()}

    ImageNetPreTrainedConfig = partial(VanConfig, num_labels=num_labels, id2label=id2label, label2id=label2id)

    names_to_config = {
        "van-tiny": ImageNetPreTrainedConfig(
            hidden_sizes=[32, 64, 160, 256],
            depths=[3, 3, 5, 2],
            mlp_ratios=[8, 8, 4, 4],
        ),
        "van-small": ImageNetPreTrainedConfig(
            hidden_sizes=[64, 128, 320, 512],
            depths=[2, 2, 4, 2],
            mlp_ratios=[8, 8, 4, 4],
        ),
        "van-base": ImageNetPreTrainedConfig(
            hidden_sizes=[64, 128, 320, 512],
            depths=[3, 3, 12, 3],
            mlp_ratios=[8, 8, 4, 4],
        ),
        "van-large": ImageNetPreTrainedConfig(
            hidden_sizes=[64, 128, 320, 512],
            depths=[3, 5, 27, 3],
            mlp_ratios=[8, 8, 4, 4],
        ),
    }

    names_to_original_models = {
        "van-tiny": van_tiny,
        "van-small": van_small,
        "van-base": van_base,
        "van-large": van_large,
    }

    names_to_original_checkpoints = {
        "van-tiny": (
            "https://huggingface.co/Visual-Attention-Network/VAN-Tiny-original/resolve/main/van_tiny_754.pth.tar"
        ),
        "van-small": (
            "https://huggingface.co/Visual-Attention-Network/VAN-Small-original/resolve/main/van_small_811.pth.tar"
        ),
        "van-base": (
            "https://huggingface.co/Visual-Attention-Network/VAN-Base-original/resolve/main/van_base_828.pth.tar"
        ),
        "van-large": (
            "https://huggingface.co/Visual-Attention-Network/VAN-Large-original/resolve/main/van_large_839.pth.tar"
        ),
    }

    if model_name:
        convert_weight_and_push(
            model_name,
            names_to_config[model_name],
            checkpoint=names_to_original_checkpoints[model_name],
            from_model=names_to_original_models[model_name](),
            save_directory=save_directory,
            push_to_hub=push_to_hub,
        )
    else:
        for model_name, config in names_to_config.items():
            convert_weight_and_push(
                model_name,
                config,
                checkpoint=names_to_original_checkpoints[model_name],
                from_model=names_to_original_models[model_name](),
                save_directory=save_directory,
                push_to_hub=push_to_hub,
            )


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    
    parser.add_argument(
        "--model-name",
        default=None,
        type=str,
        help=(
            "The name of the model you wish to convert, it must be one of the supported resnet* architecture,"
            " currently: van-tiny/small/base/large. If `None`, all of them will the converted."
        ),
    )
    parser.add_argument(
        "--pytorch_dump_folder_path",
        default=None,
        type=Path,
        required=True,
        help="Path to the output PyTorch model directory.",
    )
    parser.add_argument(
        "--van_dir",
        required=True,
        type=Path,
        help=(
            "A path to VAN's original implementation directory. You can download from here:"
            " https://github.com/Visual-Attention-Network/VAN-Classification"
        ),
    )
    parser.add_argument(
        "--push_to_hub",
        default=True,
        type=bool,
        required=False,
        help="If True, push model and image processor to the hub.",
    )

    args = parser.parse_args()
    pytorch_dump_folder_path: Path = args.pytorch_dump_folder_path
    pytorch_dump_folder_path.mkdir(exist_ok=True, parents=True)
    van_dir = args.van_dir
    
    sys.path.append(str(van_dir.parent))
    from van.models.van import van_base, van_large, van_small, van_tiny

    convert_weights_and_push(pytorch_dump_folder_path, args.model_name, args.push_to_hub)