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 checkpoints from timm and vissl import argparse import json from dataclasses import dataclass field from functools import partial from pathlib import path from typing import callable dict list tuple import timm import torch import torch nn as nn from classyvision models regnet import regnet regnetparams regnety32gf regnety64gf regnety128gf from huggingfacehub import cacheddownload hfhuburl from torch import tensor from vissl models modelhelpers import gettrunkforwardoutputs from transformers import autoimageprocessor regnetconfig regnetforimageclassification regnetmodel 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 1 srcskip list fielddefaultfactorylist destskip list fielddefaultfactorylist raiseifmismatch bool true 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 and self raiseifmismatch 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 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 nametofrommodelfuncmapdict def convertnametotimmself x str str xsplit x split return xsplit0 xsplit1 joinxsplit2 def getitemself x str callable tuplenn module dict default to timm if x not in self x self convertnametotimmx val partiallambda timm createmodelx pretrainedtrue eval none else val super getitemx return val class nametoourmodelfuncmapdict def getitemself x str callable nn module if seer in x and in1k not in x val regnetmodel else val regnetforimageclassification return val def manuallycopyvisslheadfromstatedict tostatedict keys listtuplestr str for fromkey tokey in keys tostatedicttokey fromstatedictfromkey clone printfcopied keyfromkey totokey return tostatedict def convertweightandpush name str frommodelfunc callable nn module ourmodelfunc callable nn module config regnetconfig savedirectory path pushtohub bool true printfconverting name with torch nograd frommodel fromstatedict frommodelfunc ourmodel ourmodelfuncconfig eval moduletransfer moduletransfersrcfrommodel destourmodel raiseifmismatchfalse x torch randn1 3 224 224 moduletransferx if fromstatedict is not none keys for seer in1k finetuned we have to manually copy the head if seer in name and in1k in name keys 0 clf 0 weight classifier 1 weight 0 clf 0 bias classifier 1 bias tostatedict manuallycopyvisslheadfromstatedict ourmodel statedict keys ourmodel loadstatedicttostatedict ouroutputs ourmodelx outputhiddenstatestrue ouroutput ouroutputs logits if isinstanceourmodel regnetforimageclassification else ouroutputs lasthiddenstate fromoutput frommodelx fromoutput fromoutput1 if isinstancefromoutput list else fromoutput now since i don t want to use any config files vissl seer model doesn t actually have an head so let s just check the last hidden state if seer in name and in1k in name ouroutput ouroutputs hiddenstates1 assert torch allclosefromoutput ouroutput the model logits don t match the original one if pushtohub ourmodel pushtohub repopathornamesavedirectory name commitmessageadd model usetempdirtrue size 224 if seer not in name else 384 we can use the convnext one imageprocessor autoimageprocessor frompretrainedfacebookconvnextbase22422k1k sizesize imageprocessor pushtohub repopathornamesavedirectory name commitmessageadd image processor usetempdirtrue printfpushed name def convertweightsandpushsavedirectory path modelname str none pushtohub bool true filename imagenet1kid2label json numlabels 1000 expectedshape 1 numlabels 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 regnetx002 imagenetpretrainedconfig depths1 1 4 7 hiddensizes24 56 152 368 groupswidth8 layertypex regnetx004 imagenetpretrainedconfig depths1 2 7 12 hiddensizes32 64 160 384 groupswidth16 layertypex regnetx006 imagenetpretrainedconfig depths1 3 5 7 hiddensizes48 96 240 528 groupswidth24 layertypex regnetx008 imagenetpretrainedconfig depths1 3 7 5 hiddensizes64 128 288 672 groupswidth16 layertypex regnetx016 imagenetpretrainedconfig depths2 4 10 2 hiddensizes72 168 408 912 groupswidth24 layertypex regnetx032 imagenetpretrainedconfig depths2 6 15 2 hiddensizes96 192 432 1008 groupswidth48 layertypex regnetx040 imagenetpretrainedconfig depths2 5 14 2 hiddensizes80 240 560 1360 groupswidth40 layertypex regnetx064 imagenetpretrainedconfig depths2 4 10 1 hiddensizes168 392 784 1624 groupswidth56 layertypex regnetx080 imagenetpretrainedconfig depths2 5 15 1 hiddensizes80 240 720 1920 groupswidth120 layertypex regnetx120 imagenetpretrainedconfig depths2 5 11 1 hiddensizes224 448 896 2240 groupswidth112 layertypex regnetx160 imagenetpretrainedconfig depths2 6 13 1 hiddensizes256 512 896 2048 groupswidth128 layertypex regnetx320 imagenetpretrainedconfig depths2 7 13 1 hiddensizes336 672 1344 2520 groupswidth168 layertypex y variant regnety002 imagenetpretrainedconfigdepths1 1 4 7 hiddensizes24 56 152 368 groupswidth8 regnety004 imagenetpretrainedconfig depths1 3 6 6 hiddensizes48 104 208 440 groupswidth8 regnety006 imagenetpretrainedconfig depths1 3 7 4 hiddensizes48 112 256 608 groupswidth16 regnety008 imagenetpretrainedconfig depths1 3 8 2 hiddensizes64 128 320 768 groupswidth16 regnety016 imagenetpretrainedconfig depths2 6 17 2 hiddensizes48 120 336 888 groupswidth24 regnety032 imagenetpretrainedconfig depths2 5 13 1 hiddensizes72 216 576 1512 groupswidth24 regnety040 imagenetpretrainedconfig depths2 6 12 2 hiddensizes128 192 512 1088 groupswidth64 regnety064 imagenetpretrainedconfig depths2 7 14 2 hiddensizes144 288 576 1296 groupswidth72 regnety080 imagenetpretrainedconfig depths2 4 10 1 hiddensizes168 448 896 2016 groupswidth56 regnety120 imagenetpretrainedconfig depths2 5 11 1 hiddensizes224 448 896 2240 groupswidth112 regnety160 imagenetpretrainedconfig depths2 4 11 1 hiddensizes224 448 1232 3024 groupswidth112 regnety320 imagenetpretrainedconfig depths2 5 12 1 hiddensizes232 696 1392 3712 groupswidth232 models seer https arxiv orgabs2202 08360 regnety320seer regnetconfigdepths2 5 12 1 hiddensizes232 696 1392 3712 groupswidth232 regnety640seer regnetconfigdepths2 5 12 1 hiddensizes328 984 1968 4920 groupswidth328 regnety1280seer regnetconfig depths2 7 17 1 hiddensizes528 1056 2904 7392 groupswidth264 regnety2560seer regnetconfig depths3 7 16 1 hiddensizes640 1696 2544 5088 groupswidth640 regnety10bseer imagenetpretrainedconfig depths2 7 17 1 hiddensizes2020 4040 11110 28280 groupswidth1010 finetuned on imagenet regnety320seerin1k imagenetpretrainedconfig depths2 5 12 1 hiddensizes232 696 1392 3712 groupswidth232 regnety640seerin1k imagenetpretrainedconfig depths2 5 12 1 hiddensizes328 984 1968 4920 groupswidth328 regnety1280seerin1k imagenetpretrainedconfig depths2 7 17 1 hiddensizes528 1056 2904 7392 groupswidth264 regnety2560seerin1k imagenetpretrainedconfig depths3 7 16 1 hiddensizes640 1696 2544 5088 groupswidth640 regnety10bseerin1k imagenetpretrainedconfig depths2 7 17 1 hiddensizes2020 4040 11110 28280 groupswidth1010 namestooursmodelmap nametoourmodelfuncmap namestofrommodelmap nametofrommodelfuncmap add seer weights logic def loadusingclassyvisioncheckpointurl str modelfunc callable nn module tuplenn module dict files torch hub loadstatedictfromurlcheckpointurl modeldirstrsavedirectory maplocationcpu model modelfunc check if we have a head if yes add it modelstatedict filesclassystatedictbasemodelmodel statedict modelstatedicttrunk model loadstatedictstatedict return model eval modelstatedictheads pretrained namestofrommodelmapregnety320seer partial loadusingclassyvision https dl fbaipublicfiles comvisslmodelzooseerregnet32dseerregnet32gfmodeliteration244000 torch lambda fakeregnetvisslwrapperregnety32gf namestofrommodelmapregnety640seer partial loadusingclassyvision https dl fbaipublicfiles comvisslmodelzooseerregnet64seerregnet64gfmodelfinalcheckpointphase0 torch lambda fakeregnetvisslwrapperregnety64gf namestofrommodelmapregnety1280seer partial loadusingclassyvision https dl fbaipublicfiles comvisslmodelzooswavig1bregnet128gfcnstantbs32node16sinkhorn10proto16ksyncbn64warmup8kmodelfinalcheckpointphase0 torch lambda fakeregnetvisslwrapperregnety128gf namestofrommodelmapregnety10bseer partial loadusingclassyvision https dl fbaipublicfiles comvisslmodelzooseerregnet10bmodeliteration124500conso torch lambda fakeregnetvisslwrapper regnetregnetparamsdepth27 groupwidth1010 w01744 wa620 83 wm2 52 in1k finetuned namestofrommodelmapregnety320seerin1k partial loadusingclassyvision https dl fbaipublicfiles comvisslmodelzooseerfinetunedseerregnet32finetunedin1kmodelfinalcheckpointphase78 torch lambda fakeregnetvisslwrapperregnety32gf namestofrommodelmapregnety640seerin1k partial loadusingclassyvision https dl fbaipublicfiles comvisslmodelzooseerfinetunedseerregnet64finetunedin1kmodelfinalcheckpointphase78 torch lambda fakeregnetvisslwrapperregnety64gf namestofrommodelmapregnety1280seerin1k partial loadusingclassyvision https dl fbaipublicfiles comvisslmodelzooseerfinetunedseerregnet128finetunedin1kmodelfinalcheckpointphase78 torch lambda fakeregnetvisslwrapperregnety128gf namestofrommodelmapregnety10bseerin1k partial loadusingclassyvision https dl fbaipublicfiles comvisslmodelzooseerfinetunedseer10bfinetunedin1kmodelphase28conso torch lambda fakeregnetvisslwrapper regnetregnetparamsdepth27 groupwidth1010 w01744 wa620 83 wm2 52 if modelname convertweightandpush modelname namestofrommodelmapmodelname namestooursmodelmapmodelname namestoconfigmodelname savedirectory pushtohub else for modelname config in namestoconfig items convertweightandpush modelname namestofrommodelmapmodelname namestooursmodelmapmodelname 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 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 checkpoints from timm and vissl 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 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 a dictionary with some additional logic to return a function that creates the correct original model default to timm a dictionary with some additional logic to return the correct hugging face regnet class reference for seer in1k finetuned we have to manually copy the head now since i don t want to use any config files vissl seer model doesn t actually have an head so let s just check the last hidden state we can use the convnext one y variant models seer https arxiv org abs 2202 08360 finetuned on imagenet add seer weights logic check if we have a head if yes add it pretrained in1k finetuned required parameters
import argparse
import json
from dataclasses import dataclass, field
from functools import partial
from pathlib import Path
from typing import Callable, Dict, List, Tuple

import timm
import torch
import torch.nn as nn
from classy_vision.models.regnet import RegNet, RegNetParams, RegNetY32gf, RegNetY64gf, RegNetY128gf
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.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 = 1
    src_skip: List = field(default_factory=list)
    dest_skip: List = field(default_factory=list)
    raise_if_mismatch: bool = True

    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) and self.raise_if_mismatch:
            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}")


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 NameToFromModelFuncMap(dict):
    

    def convert_name_to_timm(self, x: str) -> str:
        x_split = x.split("-")
        return x_split[0] + x_split[1] + "_" + "".join(x_split[2:])

    def __getitem__(self, x: str) -> Callable[[], Tuple[nn.Module, Dict]]:
        
        if x not in self:
            x = self.convert_name_to_timm(x)
            val = partial(lambda: (timm.create_model(x, pretrained=True).eval(), None))

        else:
            val = super().__getitem__(x)

        return val


class NameToOurModelFuncMap(dict):
    

    def __getitem__(self, x: str) -> Callable[[], nn.Module]:
        if "seer" in x and "in1k" not in x:
            val = RegNetModel
        else:
            val = RegNetForImageClassification
        return val


def manually_copy_vissl_head(from_state_dict, to_state_dict, keys: List[Tuple[str, str]]):
    for from_key, to_key in keys:
        to_state_dict[to_key] = from_state_dict[from_key].clone()
        print(f"Copied key={from_key} to={to_key}")
    return to_state_dict


def convert_weight_and_push(
    name: str,
    from_model_func: Callable[[], nn.Module],
    our_model_func: Callable[[], nn.Module],
    config: RegNetConfig,
    save_directory: Path,
    push_to_hub: bool = True,
):
    print(f"Converting {name}...")
    with torch.no_grad():
        from_model, from_state_dict = from_model_func()
        our_model = our_model_func(config).eval()
        module_transfer = ModuleTransfer(src=from_model, dest=our_model, raise_if_mismatch=False)
        x = torch.randn((1, 3, 224, 224))
        module_transfer(x)

    if from_state_dict is not None:
        keys = []
        
        if "seer" in name and "in1k" in name:
            keys = [("0.clf.0.weight", "classifier.1.weight"), ("0.clf.0.bias", "classifier.1.bias")]
        to_state_dict = manually_copy_vissl_head(from_state_dict, our_model.state_dict(), keys)
        our_model.load_state_dict(to_state_dict)

    our_outputs = our_model(x, output_hidden_states=True)
    our_output = (
        our_outputs.logits if isinstance(our_model, RegNetForImageClassification) else our_outputs.last_hidden_state
    )

    from_output = from_model(x)
    from_output = from_output[-1] if isinstance(from_output, list) else from_output

    
    if "seer" in name and "in1k" in name:
        our_output = our_outputs.hidden_states[-1]

    assert torch.allclose(from_output, our_output), "The model logits don't match the original one."

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

        size = 224 if "seer" not in name else 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 / name,
            commit_message="Add image processor",
            use_temp_dir=True,
        )

        print(f"Pushed {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(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-x-002": ImageNetPreTrainedConfig(
            depths=[1, 1, 4, 7], hidden_sizes=[24, 56, 152, 368], groups_width=8, layer_type="x"
        ),
        "regnet-x-004": ImageNetPreTrainedConfig(
            depths=[1, 2, 7, 12], hidden_sizes=[32, 64, 160, 384], groups_width=16, layer_type="x"
        ),
        "regnet-x-006": ImageNetPreTrainedConfig(
            depths=[1, 3, 5, 7], hidden_sizes=[48, 96, 240, 528], groups_width=24, layer_type="x"
        ),
        "regnet-x-008": ImageNetPreTrainedConfig(
            depths=[1, 3, 7, 5], hidden_sizes=[64, 128, 288, 672], groups_width=16, layer_type="x"
        ),
        "regnet-x-016": ImageNetPreTrainedConfig(
            depths=[2, 4, 10, 2], hidden_sizes=[72, 168, 408, 912], groups_width=24, layer_type="x"
        ),
        "regnet-x-032": ImageNetPreTrainedConfig(
            depths=[2, 6, 15, 2], hidden_sizes=[96, 192, 432, 1008], groups_width=48, layer_type="x"
        ),
        "regnet-x-040": ImageNetPreTrainedConfig(
            depths=[2, 5, 14, 2], hidden_sizes=[80, 240, 560, 1360], groups_width=40, layer_type="x"
        ),
        "regnet-x-064": ImageNetPreTrainedConfig(
            depths=[2, 4, 10, 1], hidden_sizes=[168, 392, 784, 1624], groups_width=56, layer_type="x"
        ),
        "regnet-x-080": ImageNetPreTrainedConfig(
            depths=[2, 5, 15, 1], hidden_sizes=[80, 240, 720, 1920], groups_width=120, layer_type="x"
        ),
        "regnet-x-120": ImageNetPreTrainedConfig(
            depths=[2, 5, 11, 1], hidden_sizes=[224, 448, 896, 2240], groups_width=112, layer_type="x"
        ),
        "regnet-x-160": ImageNetPreTrainedConfig(
            depths=[2, 6, 13, 1], hidden_sizes=[256, 512, 896, 2048], groups_width=128, layer_type="x"
        ),
        "regnet-x-320": ImageNetPreTrainedConfig(
            depths=[2, 7, 13, 1], hidden_sizes=[336, 672, 1344, 2520], groups_width=168, layer_type="x"
        ),
        
        "regnet-y-002": ImageNetPreTrainedConfig(depths=[1, 1, 4, 7], hidden_sizes=[24, 56, 152, 368], groups_width=8),
        "regnet-y-004": ImageNetPreTrainedConfig(
            depths=[1, 3, 6, 6], hidden_sizes=[48, 104, 208, 440], groups_width=8
        ),
        "regnet-y-006": ImageNetPreTrainedConfig(
            depths=[1, 3, 7, 4], hidden_sizes=[48, 112, 256, 608], groups_width=16
        ),
        "regnet-y-008": ImageNetPreTrainedConfig(
            depths=[1, 3, 8, 2], hidden_sizes=[64, 128, 320, 768], groups_width=16
        ),
        "regnet-y-016": ImageNetPreTrainedConfig(
            depths=[2, 6, 17, 2], hidden_sizes=[48, 120, 336, 888], groups_width=24
        ),
        "regnet-y-032": ImageNetPreTrainedConfig(
            depths=[2, 5, 13, 1], hidden_sizes=[72, 216, 576, 1512], groups_width=24
        ),
        "regnet-y-040": ImageNetPreTrainedConfig(
            depths=[2, 6, 12, 2], hidden_sizes=[128, 192, 512, 1088], groups_width=64
        ),
        "regnet-y-064": ImageNetPreTrainedConfig(
            depths=[2, 7, 14, 2], hidden_sizes=[144, 288, 576, 1296], groups_width=72
        ),
        "regnet-y-080": ImageNetPreTrainedConfig(
            depths=[2, 4, 10, 1], hidden_sizes=[168, 448, 896, 2016], groups_width=56
        ),
        "regnet-y-120": ImageNetPreTrainedConfig(
            depths=[2, 5, 11, 1], hidden_sizes=[224, 448, 896, 2240], groups_width=112
        ),
        "regnet-y-160": ImageNetPreTrainedConfig(
            depths=[2, 4, 11, 1], hidden_sizes=[224, 448, 1232, 3024], groups_width=112
        ),
        "regnet-y-320": ImageNetPreTrainedConfig(
            depths=[2, 5, 12, 1], hidden_sizes=[232, 696, 1392, 3712], groups_width=232
        ),
        
        "regnet-y-320-seer": RegNetConfig(depths=[2, 5, 12, 1], hidden_sizes=[232, 696, 1392, 3712], groups_width=232),
        "regnet-y-640-seer": RegNetConfig(depths=[2, 5, 12, 1], hidden_sizes=[328, 984, 1968, 4920], groups_width=328),
        "regnet-y-1280-seer": RegNetConfig(
            depths=[2, 7, 17, 1], hidden_sizes=[528, 1056, 2904, 7392], groups_width=264
        ),
        "regnet-y-2560-seer": RegNetConfig(
            depths=[3, 7, 16, 1], hidden_sizes=[640, 1696, 2544, 5088], groups_width=640
        ),
        "regnet-y-10b-seer": ImageNetPreTrainedConfig(
            depths=[2, 7, 17, 1], hidden_sizes=[2020, 4040, 11110, 28280], groups_width=1010
        ),
        
        "regnet-y-320-seer-in1k": ImageNetPreTrainedConfig(
            depths=[2, 5, 12, 1], hidden_sizes=[232, 696, 1392, 3712], groups_width=232
        ),
        "regnet-y-640-seer-in1k": ImageNetPreTrainedConfig(
            depths=[2, 5, 12, 1], hidden_sizes=[328, 984, 1968, 4920], groups_width=328
        ),
        "regnet-y-1280-seer-in1k": ImageNetPreTrainedConfig(
            depths=[2, 7, 17, 1], hidden_sizes=[528, 1056, 2904, 7392], groups_width=264
        ),
        "regnet-y-2560-seer-in1k": ImageNetPreTrainedConfig(
            depths=[3, 7, 16, 1], hidden_sizes=[640, 1696, 2544, 5088], groups_width=640
        ),
        "regnet-y-10b-seer-in1k": ImageNetPreTrainedConfig(
            depths=[2, 7, 17, 1], hidden_sizes=[2020, 4040, 11110, 28280], groups_width=1010
        ),
    }

    names_to_ours_model_map = NameToOurModelFuncMap()
    names_to_from_model_map = NameToFromModelFuncMap()
    

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

    
    names_to_from_model_map["regnet-y-320-seer"] = partial(
        load_using_classy_vision,
        "https://dl.fbaipublicfiles.com/vissl/model_zoo/seer_regnet32d/seer_regnet32gf_model_iteration244000.torch",
        lambda: FakeRegNetVisslWrapper(RegNetY32gf()),
    )

    names_to_from_model_map["regnet-y-640-seer"] = partial(
        load_using_classy_vision,
        "https://dl.fbaipublicfiles.com/vissl/model_zoo/seer_regnet64/seer_regnet64gf_model_final_checkpoint_phase0.torch",
        lambda: FakeRegNetVisslWrapper(RegNetY64gf()),
    )

    names_to_from_model_map["regnet-y-1280-seer"] = partial(
        load_using_classy_vision,
        "https://dl.fbaipublicfiles.com/vissl/model_zoo/swav_ig1b_regnet128Gf_cnstant_bs32_node16_sinkhorn10_proto16k_syncBN64_warmup8k/model_final_checkpoint_phase0.torch",
        lambda: FakeRegNetVisslWrapper(RegNetY128gf()),
    )

    names_to_from_model_map["regnet-y-10b-seer"] = partial(
        load_using_classy_vision,
        "https://dl.fbaipublicfiles.com/vissl/model_zoo/seer_regnet10B/model_iteration124500_conso.torch",
        lambda: FakeRegNetVisslWrapper(
            RegNet(RegNetParams(depth=27, group_width=1010, w_0=1744, w_a=620.83, w_m=2.52))
        ),
    )

    
    names_to_from_model_map["regnet-y-320-seer-in1k"] = partial(
        load_using_classy_vision,
        "https://dl.fbaipublicfiles.com/vissl/model_zoo/seer_finetuned/seer_regnet32_finetuned_in1k_model_final_checkpoint_phase78.torch",
        lambda: FakeRegNetVisslWrapper(RegNetY32gf()),
    )

    names_to_from_model_map["regnet-y-640-seer-in1k"] = partial(
        load_using_classy_vision,
        "https://dl.fbaipublicfiles.com/vissl/model_zoo/seer_finetuned/seer_regnet64_finetuned_in1k_model_final_checkpoint_phase78.torch",
        lambda: FakeRegNetVisslWrapper(RegNetY64gf()),
    )

    names_to_from_model_map["regnet-y-1280-seer-in1k"] = partial(
        load_using_classy_vision,
        "https://dl.fbaipublicfiles.com/vissl/model_zoo/seer_finetuned/seer_regnet128_finetuned_in1k_model_final_checkpoint_phase78.torch",
        lambda: FakeRegNetVisslWrapper(RegNetY128gf()),
    )

    names_to_from_model_map["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",
        lambda: FakeRegNetVisslWrapper(
            RegNet(RegNetParams(depth=27, group_width=1010, w_0=1744, w_a=620.83, w_m=2.52))
        ),
    )

    if model_name:
        convert_weight_and_push(
            model_name,
            names_to_from_model_map[model_name],
            names_to_ours_model_map[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,
                names_to_from_model_map[model_name],
                names_to_ours_model_map[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 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)