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 regnet 10b checkpoints vissl you need to install a specific version of classy vision pip install githttps github comfrancescosaveriozuppichiniclassyvision gitconvertweights import argparse import json import os import re from collections import ordereddict from dataclasses import dataclass field from functools import partial from pathlib import path from pprint import pprint from typing import dict list tuple import torch import torch nn as nn from classyvision models regnet import regnet regnetparams from huggingfacehub import cacheddownload hfhuburl from torch import tensor from vissl models modelhelpers import gettrunkforwardoutputs from transformers import autoimageprocessor regnetconfig regnetforimageclassification regnetmodel from transformers modelingutils import pretrainedmodel from transformers utils import logging logging setverbosityinfo logger logging getlogger dataclass class tracker module nn module traced listnn module fielddefaultfactorylist handles list fielddefaultfactorylist name2module dictstr nn module fielddefaultfactoryordereddict def forwardhookself m inputs tensor outputs tensor name str hasnotsubmodules lenlistm modules 1 or isinstancem nn conv2d or isinstancem nn batchnorm2d if hasnotsubmodules self traced appendm self name2modulename m def callself x tensor for name m in self module namedmodules self handles appendm registerforwardhookpartialself forwardhook namename 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 k v for k v in self name2module items if lenlistv statedict keys 0 class fakeregnetvisslwrappernn module def initself model nn module super init featureblocks listtuplestr nn module get the stem featureblocks appendconv1 model stem get all the feature blocks for k v in model trunkoutput namedchildren assert k startswithblock funexpected layer name k blockindex lenfeatureblocks 1 featureblocks appendfresblockindex v self featureblocks nn moduledictfeatureblocks def forwardself x tensor return gettrunkforwardoutputs x outfeatkeysnone featureblocksself featureblocks class fakeregnetparamsregnetparams def getexpandedparamsself return 8 2 2 8 1 0 8 2 7 8 1 0 8 2 17 8 1 0 8 2 1 8 1 0 def getfromtoourkeysmodelname str dictstr str create our model with small weights ourconfig regnetconfigdepths2 7 17 1 hiddensizes8 8 8 8 groupswidth8 if in1k in modelname ourmodel regnetforimageclassificationourconfig else ourmodel regnetmodelourconfig create from model with small weights frommodel fakeregnetvisslwrapper regnetfakeregnetparamsdepth27 groupwidth1010 w01744 wa620 83 wm2 52 with torch nograd frommodel frommodel eval ourmodel ourmodel eval x torch randn1 3 32 32 trace both desttracker trackerourmodel desttraced desttrackerx parametrized pprintdesttracker name2module srctracker trackerfrommodel srctraced srctrackerx parametrized convert the keys module dict to keys params def toparamsdictdictwithmodules paramsdict ordereddict for name module in dictwithmodules items for paramname param in module statedict items paramsdictfname paramname param return paramsdict fromtoourskeys srcstatedict toparamsdictsrctraced dststatedict toparamsdictdesttraced for srckey srcparam destkey destparam in zipsrcstatedict items dststatedict items fromtoourskeyssrckey destkey logger infofsrckey destkey if in1k was in the modelname it means it must have a classification head was finetuned if in1k in modelname fromtoourskeys0 clf 0 weight classifier 1 weight fromtoourskeys0 clf 0 bias classifier 1 bias return fromtoourskeys def convertweightsandpushsavedirectory path modelname str none pushtohub bool true filename imagenet1kid2label json numlabels 1000 repoid huggingfacelabelfiles numlabels numlabels id2label json loadopencacheddownloadhfhuburlrepoid filename repotypedataset r id2label intk v for k v in id2label items id2label id2label label2id v k for k v in id2label items imagenetpretrainedconfig partialregnetconfig numlabelsnumlabels id2labelid2label label2idlabel2id namestoconfig regnety10bseer imagenetpretrainedconfig depths2 7 17 1 hiddensizes2020 4040 11110 28280 groupswidth1010 finetuned on imagenet regnety10bseerin1k imagenetpretrainedconfig depths2 7 17 1 hiddensizes2020 4040 11110 28280 groupswidth1010 add seer weights logic def loadusingclassyvisioncheckpointurl str tupledict dict files torch hub loadstatedictfromurlcheckpointurl modeldirstrsavedirectory maplocationcpu check if we have a head if yes add it modelstatedict filesclassystatedictbasemodelmodel return modelstatedicttrunk modelstatedictheads namestofrommodel regnety10bseer partial loadusingclassyvision https dl fbaipublicfiles comvisslmodelzooseerregnet10bmodeliteration124500conso torch regnety10bseerin1k partial loadusingclassyvision https dl fbaipublicfiles comvisslmodelzooseerfinetunedseer10bfinetunedin1kmodelphase28conso torch fromtoourskeys getfromtoourkeysmodelname if not savedirectory fmodelname pth exists logger infoloading original statedict fromstatedicttrunk fromstatedicthead namestofrommodelmodelname fromstatedict fromstatedicttrunk if in1k in modelname add the head fromstatedict fromstatedicttrunk fromstatedicthead logger infodone convertedstatedict notusedkeys listfromstatedict keys regex r block part this is interesting so the original checkpoints have block0 1part in each key name we remove it for key in fromstatedict keys remove the weird block0 1part from the key srckey re subregex key now srckey from the model checkpoints is the one we got from the original model after tracing so use it to get the correct destination key destkey fromtoourskeyssrckey store the parameter with our key convertedstatedictdestkey fromstatedictkey notusedkeys removekey check that all keys have been updated assert lennotusedkeys 0 fsome keys where not used joinnotusedkeys logger infofthe following keys were not used joinnotusedkeys save our state dict to disk torch saveconvertedstatedict savedirectory fmodelname pth del convertedstatedict else logger infothe statedict was already stored on disk if pushtohub logger infoftoken is os environ hftoken logger infoloading our model create our model ourconfig namestoconfigmodelname ourmodelfunc regnetmodel if in1k in modelname ourmodelfunc regnetforimageclassification ourmodel ourmodelfuncourconfig place our model to the meta device so remove all the weights ourmodel totorch devicemeta logger infoloading statedict in our model load state dict statedictkeys ourmodel statedict keys pretrainedmodel loadpretrainedmodellowmem ourmodel statedictkeys savedirectory fmodelname pth logger infofinally pushing push it to hub ourmodel pushtohub repopathornamesavedirectory modelname commitmessageadd model outputdirsavedirectory modelname size 384 we can use the convnext one imageprocessor autoimageprocessor frompretrainedfacebookconvnextbase22422k1k sizesize imageprocessor pushtohub repopathornamesavedirectory modelname commitmessageadd image processor outputdirsavedirectory modelname 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 regnet architecture currently regnetx regnety 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 regnet 10b checkpoints vissl you need to install a specific version of classy vision pip install git https github com francescosaveriozuppichini classyvision git convert_weights check the len of the state_dict keys to see if we have learnable params fake wrapper for regnet that mimics what vissl does without the need to pass a config file get the stem get all the feature blocks used to instantiace a regnet model from classy vision with the same depth as the 10b one but with super small parameters so we can trace it in memory returns a dictionary that maps from original model s key our implementation s keys create our model with small weights create from model with small weights trace both convert the keys module dict to keys params if in1k was in the model_name it means it must have a classification head was finetuned finetuned on imagenet add seer weights logic check if we have a head if yes add it add the head this is interesting so the original checkpoints have block 0 1 part in each key name we remove it remove the weird block 0 1 part from the key now src_key from the model checkpoints is the one we got from the original model after tracing so use it to get the correct destination key store the parameter with our key check that all keys have been updated save our state dict to disk create our model place our model to the meta device so remove all the weights load state dict push it to hub we can use the convnext one required parameters
import argparse
import json
import os
import re
from collections import OrderedDict
from dataclasses import dataclass, field
from functools import partial
from pathlib import Path
from pprint import pprint
from typing import Dict, List, Tuple

import torch
import torch.nn as nn
from classy_vision.models.regnet import RegNet, RegNetParams
from huggingface_hub import cached_download, hf_hub_url
from torch import Tensor
from vissl.models.model_helpers import get_trunk_forward_outputs

from transformers import AutoImageProcessor, RegNetConfig, RegNetForImageClassification, RegNetModel
from transformers.modeling_utils import PreTrainedModel
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)
    name2module: Dict[str, nn.Module] = field(default_factory=OrderedDict)

    def _forward_hook(self, m, inputs: Tensor, outputs: Tensor, name: str):
        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)
            self.name2module[name] = m

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

    @property
    def parametrized(self):
        
        return {k: v for k, v in self.name2module.items() if len(list(v.state_dict().keys())) > 0}


class FakeRegNetVisslWrapper(nn.Module):
    

    def __init__(self, model: nn.Module):
        super().__init__()

        feature_blocks: List[Tuple[str, nn.Module]] = []
        
        feature_blocks.append(("conv1", model.stem))
        
        for k, v in model.trunk_output.named_children():
            assert k.startswith("block"), f"Unexpected layer name {k}"
            block_index = len(feature_blocks) + 1
            feature_blocks.append((f"res{block_index}", v))

        self._feature_blocks = nn.ModuleDict(feature_blocks)

    def forward(self, x: Tensor):
        return get_trunk_forward_outputs(
            x,
            out_feat_keys=None,
            feature_blocks=self._feature_blocks,
        )


class FakeRegNetParams(RegNetParams):
    

    def get_expanded_params(self):
        return [(8, 2, 2, 8, 1.0), (8, 2, 7, 8, 1.0), (8, 2, 17, 8, 1.0), (8, 2, 1, 8, 1.0)]


def get_from_to_our_keys(model_name: str) -> Dict[str, str]:
    

    
    our_config = RegNetConfig(depths=[2, 7, 17, 1], hidden_sizes=[8, 8, 8, 8], groups_width=8)
    if "in1k" in model_name:
        our_model = RegNetForImageClassification(our_config)
    else:
        our_model = RegNetModel(our_config)
    
    from_model = FakeRegNetVisslWrapper(
        RegNet(FakeRegNetParams(depth=27, group_width=1010, w_0=1744, w_a=620.83, w_m=2.52))
    )

    with torch.no_grad():
        from_model = from_model.eval()
        our_model = our_model.eval()

        x = torch.randn((1, 3, 32, 32))
        
        dest_tracker = Tracker(our_model)
        dest_traced = dest_tracker(x).parametrized

        pprint(dest_tracker.name2module)
        src_tracker = Tracker(from_model)
        src_traced = src_tracker(x).parametrized

    
    def to_params_dict(dict_with_modules):
        params_dict = OrderedDict()
        for name, module in dict_with_modules.items():
            for param_name, param in module.state_dict().items():
                params_dict[f"{name}.{param_name}"] = param
        return params_dict

    from_to_ours_keys = {}

    src_state_dict = to_params_dict(src_traced)
    dst_state_dict = to_params_dict(dest_traced)

    for (src_key, src_param), (dest_key, dest_param) in zip(src_state_dict.items(), dst_state_dict.items()):
        from_to_ours_keys[src_key] = dest_key
        logger.info(f"{src_key} -> {dest_key}")
    
    if "in1k" in model_name:
        from_to_ours_keys["0.clf.0.weight"] = "classifier.1.weight"
        from_to_ours_keys["0.clf.0.bias"] = "classifier.1.bias"

    return from_to_ours_keys


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(cached_download(hf_hub_url(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(RegNetConfig, num_labels=num_labels, id2label=id2label, label2id=label2id)

    names_to_config = {
        "regnet-y-10b-seer": ImageNetPreTrainedConfig(
            depths=[2, 7, 17, 1], hidden_sizes=[2020, 4040, 11110, 28280], groups_width=1010
        ),
        
        "regnet-y-10b-seer-in1k": ImageNetPreTrainedConfig(
            depths=[2, 7, 17, 1], hidden_sizes=[2020, 4040, 11110, 28280], groups_width=1010
        ),
    }

    
    def load_using_classy_vision(checkpoint_url: str) -> Tuple[Dict, Dict]:
        files = torch.hub.load_state_dict_from_url(checkpoint_url, model_dir=str(save_directory), map_location="cpu")
        
        model_state_dict = files["classy_state_dict"]["base_model"]["model"]
        return model_state_dict["trunk"], model_state_dict["heads"]

    names_to_from_model = {
        "regnet-y-10b-seer": partial(
            load_using_classy_vision,
            "https://dl.fbaipublicfiles.com/vissl/model_zoo/seer_regnet10B/model_iteration124500_conso.torch",
        ),
        "regnet-y-10b-seer-in1k": partial(
            load_using_classy_vision,
            "https://dl.fbaipublicfiles.com/vissl/model_zoo/seer_finetuned/seer_10b_finetuned_in1k_model_phase28_conso.torch",
        ),
    }

    from_to_ours_keys = get_from_to_our_keys(model_name)

    if not (save_directory / f"{model_name}.pth").exists():
        logger.info("Loading original state_dict.")
        from_state_dict_trunk, from_state_dict_head = names_to_from_model[model_name]()
        from_state_dict = from_state_dict_trunk
        if "in1k" in model_name:
            
            from_state_dict = {**from_state_dict_trunk, **from_state_dict_head}
        logger.info("Done!")

        converted_state_dict = {}

        not_used_keys = list(from_state_dict.keys())
        regex = r"\.block.-part."
        
        for key in from_state_dict.keys():
            
            src_key = re.sub(regex, "", key)
            
            dest_key = from_to_ours_keys[src_key]
            
            converted_state_dict[dest_key] = from_state_dict[key]
            not_used_keys.remove(key)
        
        assert len(not_used_keys) == 0, f"Some keys where not used {','.join(not_used_keys)}"

        logger.info(f"The following keys were not used: {','.join(not_used_keys)}")

        
        torch.save(converted_state_dict, save_directory / f"{model_name}.pth")

        del converted_state_dict
    else:
        logger.info("The state_dict was already stored on disk.")
    if push_to_hub:
        logger.info(f"Token is {os.environ['HF_TOKEN']}")
        logger.info("Loading our model.")
        
        our_config = names_to_config[model_name]
        our_model_func = RegNetModel
        if "in1k" in model_name:
            our_model_func = RegNetForImageClassification
        our_model = our_model_func(our_config)
        
        our_model.to(torch.device("meta"))
        logger.info("Loading state_dict in our model.")
        
        state_dict_keys = our_model.state_dict().keys()
        PreTrainedModel._load_pretrained_model_low_mem(
            our_model, state_dict_keys, [save_directory / f"{model_name}.pth"]
        )
        logger.info("Finally, pushing!")
        
        our_model.push_to_hub(
            repo_path_or_name=save_directory / model_name,
            commit_message="Add model",
            output_dir=save_directory / model_name,
        )
        size = 384
        
        image_processor = AutoImageProcessor.from_pretrained("facebook/convnext-base-224-22k-1k", size=size)
        image_processor.push_to_hub(
            repo_path_or_name=save_directory / model_name,
            commit_message="Add image processor",
            output_dir=save_directory / model_name,
        )


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 regnet* architecture,"
            " currently: regnetx-*, regnety-*. 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)