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 resnet checkpoints from timm import argparse import json from dataclasses import dataclass field from functools import partial from pathlib import path from typing import list import timm import torch import torch nn as nn from huggingfacehub import hfhubdownload from torch import tensor from transformers import autoimageprocessor resnetconfig resnetforimageclassification from transformers utils import logging logging setverbosityinfo logger logging getlogger 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 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 convertweightandpushname str config resnetconfig savedirectory path pushtohub bool true printfconverting name with torch nograd frommodel timm createmodelname pretrainedtrue eval ourmodel resnetforimageclassificationconfig eval moduletransfer moduletransfersrcfrommodel destourmodel x torch randn1 3 224 224 moduletransferx assert torch allclosefrommodelx ourmodelx logits the model logits don t match the original one checkpointname fresnet joinname split resnet 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 expectedshape 1 numlabels 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 partialresnetconfig numlabelsnumlabels id2labelid2label label2idlabel2id namestoconfig resnet18 imagenetpretrainedconfig depths2 2 2 2 hiddensizes64 128 256 512 layertypebasic resnet26 imagenetpretrainedconfig depths2 2 2 2 hiddensizes256 512 1024 2048 layertypebottleneck resnet34 imagenetpretrainedconfig depths3 4 6 3 hiddensizes64 128 256 512 layertypebasic resnet50 imagenetpretrainedconfig depths3 4 6 3 hiddensizes256 512 1024 2048 layertypebottleneck resnet101 imagenetpretrainedconfig depths3 4 23 3 hiddensizes256 512 1024 2048 layertypebottleneck resnet152 imagenetpretrainedconfig depths3 8 36 3 hiddensizes256 512 1024 2048 layertypebottleneck if modelname convertweightandpushmodelname namestoconfigmodelname savedirectory pushtohub else for modelname config in namestoconfig items convertweightandpushmodelname config savedirectory pushtohub return config expectedshape 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 resnet18 26 34 50 101 152 if none all of them will the converted parser addargument pytorchdumpfolderpath defaultnone typepath requiredtrue helppath to the output pytorch model directory 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 convertweightsandpushpytorchdumpfolderpath 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 resnet checkpoints from timm 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 we can use the convnext one required parameters
import argparse
import json
from dataclasses import dataclass, field
from functools import partial
from pathlib import Path
from typing import List

import timm
import torch
import torch.nn as nn
from huggingface_hub import hf_hub_download
from torch import Tensor

from transformers import AutoImageProcessor, ResNetConfig, ResNetForImageClassification
from transformers.utils import logging


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


@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:
            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 convert_weight_and_push(name: str, config: ResNetConfig, save_directory: Path, push_to_hub: bool = True):
    print(f"Converting {name}...")
    with torch.no_grad():
        from_model = timm.create_model(name, pretrained=True).eval()
        our_model = ResNetForImageClassification(config).eval()
        module_transfer = ModuleTransfer(src=from_model, dest=our_model)
        x = torch.randn((1, 3, 224, 224))
        module_transfer(x)

    assert torch.allclose(from_model(x), our_model(x).logits), "The model logits don't match the original one."

    checkpoint_name = f"resnet{'-'.join(name.split('resnet'))}"
    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
    expected_shape = (1, num_labels)

    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(ResNetConfig, num_labels=num_labels, id2label=id2label, label2id=label2id)

    names_to_config = {
        "resnet18": ImageNetPreTrainedConfig(
            depths=[2, 2, 2, 2], hidden_sizes=[64, 128, 256, 512], layer_type="basic"
        ),
        "resnet26": ImageNetPreTrainedConfig(
            depths=[2, 2, 2, 2], hidden_sizes=[256, 512, 1024, 2048], layer_type="bottleneck"
        ),
        "resnet34": ImageNetPreTrainedConfig(
            depths=[3, 4, 6, 3], hidden_sizes=[64, 128, 256, 512], layer_type="basic"
        ),
        "resnet50": ImageNetPreTrainedConfig(
            depths=[3, 4, 6, 3], hidden_sizes=[256, 512, 1024, 2048], layer_type="bottleneck"
        ),
        "resnet101": ImageNetPreTrainedConfig(
            depths=[3, 4, 23, 3], hidden_sizes=[256, 512, 1024, 2048], layer_type="bottleneck"
        ),
        "resnet152": ImageNetPreTrainedConfig(
            depths=[3, 8, 36, 3], hidden_sizes=[256, 512, 1024, 2048], layer_type="bottleneck"
        ),
    }

    if model_name:
        convert_weight_and_push(model_name, names_to_config[model_name], save_directory, push_to_hub)
    else:
        for model_name, config in names_to_config.items():
            convert_weight_and_push(model_name, config, save_directory, push_to_hub)
    return config, expected_shape


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: resnet18,26,34,50,101,152. 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(
        "--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)
    convert_weights_and_push(pytorch_dump_folder_path, args.model_name, args.push_to_hub)