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 convnext checkpoints from the original repository url https github comfacebookresearchconvnext import argparse import json from pathlib import path import requests import torch from huggingfacehub import hfhubdownload from pil import image from transformers import convnextconfig convnextforimageclassification convnextimageprocessor from transformers utils import logging logging setverbosityinfo logger logging getloggername def getconvnextconfigcheckpointurl config convnextconfig if tiny in checkpointurl depths 3 3 9 3 hiddensizes 96 192 384 768 if small in checkpointurl depths 3 3 27 3 hiddensizes 96 192 384 768 if base in checkpointurl depths 3 3 27 3 hiddensizes 128 256 512 1024 if large in checkpointurl depths 3 3 27 3 hiddensizes 192 384 768 1536 if xlarge in checkpointurl depths 3 3 27 3 hiddensizes 256 512 1024 2048 if 1k in checkpointurl numlabels 1000 filename imagenet1kid2label json expectedshape 1 1000 else numlabels 21841 filename imagenet22kid2label json expectedshape 1 21841 repoid huggingfacelabelfiles config numlabels numlabels id2label json loadopenhfhubdownloadrepoid filename repotypedataset r id2label intk v for k v in id2label items if 1k not in checkpointurl this dataset contains 21843 labels but the model only has 21841 we delete the classes as mentioned in https github comgoogleresearchbigtransferissues18 del id2label9205 del id2label15027 config id2label id2label config label2id v k for k v in id2label items config hiddensizes hiddensizes config depths depths return config expectedshape def renamekeyname if downsamplelayers 0 0 in name name name replacedownsamplelayers 0 0 embeddings patchembeddings if downsamplelayers 0 1 in name name name replacedownsamplelayers 0 1 embeddings norm we rename to layernorm later on if downsamplelayers 1 0 in name name name replacedownsamplelayers 1 0 stages 1 downsamplinglayer 0 if downsamplelayers 1 1 in name name name replacedownsamplelayers 1 1 stages 1 downsamplinglayer 1 if downsamplelayers 2 0 in name name name replacedownsamplelayers 2 0 stages 2 downsamplinglayer 0 if downsamplelayers 2 1 in name name name replacedownsamplelayers 2 1 stages 2 downsamplinglayer 1 if downsamplelayers 3 0 in name name name replacedownsamplelayers 3 0 stages 3 downsamplinglayer 0 if downsamplelayers 3 1 in name name name replacedownsamplelayers 3 1 stages 3 downsamplinglayer 1 if stages in name and downsamplinglayer not in name stages 0 0 for instance should be renamed to stages 0 layers 0 name name lenstages 0 layers namelenstages 0 if stages in name name name replacestages encoder stages if norm in name name name replacenorm layernorm if gamma in name name name replacegamma layerscaleparameter if head in name name name replacehead classifier return name 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 convertconvnextcheckpointcheckpointurl pytorchdumpfolderpath define convnext configuration based on url config expectedshape getconvnextconfigcheckpointurl load original statedict from url statedict torch hub loadstatedictfromurlcheckpointurlmodel rename keys for key in statedict copy keys val statedict popkey statedictrenamekeykey val add prefix to all keys expect classifier head for key in statedict copy keys val statedict popkey if not key startswithclassifier key convnext key statedictkey val load huggingface model model convnextforimageclassificationconfig model loadstatedictstatedict model eval check outputs on an image prepared by convnextimageprocessor size 224 if 224 in checkpointurl else 384 imageprocessor convnextimageprocessorsizesize pixelvalues imageprocessorimagesprepareimg returntensorspt pixelvalues logits modelpixelvalues logits note the logits below were obtained without center cropping if checkpointurl https dl fbaipublicfiles comconvnextconvnexttiny1k224ema pth expectedlogits torch tensor0 1210 0 6605 0 1918 elif checkpointurl https dl fbaipublicfiles comconvnextconvnextsmall1k224ema pth expectedlogits torch tensor0 4473 0 1847 0 6365 elif checkpointurl https dl fbaipublicfiles comconvnextconvnextbase1k224ema pth expectedlogits torch tensor0 4525 0 7539 0 0308 elif checkpointurl https dl fbaipublicfiles comconvnextconvnextbase1k384 pth expectedlogits torch tensor0 3561 0 6350 0 0384 elif checkpointurl https dl fbaipublicfiles comconvnextconvnextlarge1k224ema pth expectedlogits torch tensor0 4174 0 0989 0 1489 elif checkpointurl https dl fbaipublicfiles comconvnextconvnextlarge1k384 pth expectedlogits torch tensor0 2513 0 1349 0 1613 elif checkpointurl https dl fbaipublicfiles comconvnextconvnextbase22k224 pth expectedlogits torch tensor1 2980 0 3631 0 1198 elif checkpointurl https dl fbaipublicfiles comconvnextconvnextlarge22k224 pth expectedlogits torch tensor1 2963 0 1227 0 1723 elif checkpointurl https dl fbaipublicfiles comconvnextconvnextxlarge22k224 pth expectedlogits torch tensor1 7956 0 8390 0 2820 elif checkpointurl https dl fbaipublicfiles comconvnextconvnextbase22k1k224 pth expectedlogits torch tensor0 2822 0 0502 0 0878 elif checkpointurl https dl fbaipublicfiles comconvnextconvnextbase22k1k384 pth expectedlogits torch tensor0 5672 0 0730 0 4348 elif checkpointurl https dl fbaipublicfiles comconvnextconvnextlarge22k1k224 pth expectedlogits torch tensor0 2681 0 2365 0 6246 elif checkpointurl https dl fbaipublicfiles comconvnextconvnextlarge22k1k384 pth expectedlogits torch tensor0 2642 0 3931 0 5116 elif checkpointurl https dl fbaipublicfiles comconvnextconvnextxlarge22k1k224ema pth expectedlogits torch tensor0 6677 0 1873 0 8379 elif checkpointurl https dl fbaipublicfiles comconvnextconvnextxlarge22k1k384ema pth expectedlogits torch tensor0 7749 0 2967 0 6444 else raise valueerrorfunknown url checkpointurl assert torch allcloselogits0 3 expectedlogits atol1e3 assert logits shape expectedshape pathpytorchdumpfolderpath mkdirexistoktrue printfsaving model to pytorchdumpfolderpath model savepretrainedpytorchdumpfolderpath printfsaving image processor to pytorchdumpfolderpath imageprocessor savepretrainedpytorchdumpfolderpath printpushing model to the hub modelname convnext if tiny in checkpointurl modelname tiny elif small in checkpointurl modelname small elif base in checkpointurl modelname base elif xlarge in checkpointurl modelname xlarge elif large in checkpointurl modelname large if 224 in checkpointurl modelname 224 elif 384 in checkpointurl modelname 384 if 22k in checkpointurl and 1k not in checkpointurl modelname 22k if 22k in checkpointurl and 1k in checkpointurl modelname 22k1k model pushtohub repopathornamepathpytorchdumpfolderpath modelname organizationnielsr commitmessageadd model if name main parser argparse argumentparser required parameters parser addargument checkpointurl defaulthttps dl fbaipublicfiles comconvnextconvnexttiny1k224ema pth typestr helpurl of the original convnext checkpoint you d like to convert parser addargument pytorchdumpfolderpath defaultnone typestr requiredtrue helppath to the output pytorch model directory args parser parseargs convertconvnextcheckpointargs checkpointurl args pytorchdumpfolderpath 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 convnext checkpoints from the original repository url https github com facebookresearch convnext this dataset contains 21843 labels but the model only has 21841 we delete the classes as mentioned in https github com google research big_transfer issues 18 we rename to layernorm later on stages 0 0 for instance should be renamed to stages 0 layers 0 we will verify our results on an image of cute cats copy paste tweak model s weights to our convnext structure define convnext configuration based on url load original state_dict from url rename keys add prefix to all keys expect classifier head load huggingface model check outputs on an image prepared by convnextimageprocessor note the logits below were obtained without center cropping required parameters
import argparse
import json
from pathlib import Path

import requests
import torch
from huggingface_hub import hf_hub_download
from PIL import Image

from transformers import ConvNextConfig, ConvNextForImageClassification, ConvNextImageProcessor
from transformers.utils import logging


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


def get_convnext_config(checkpoint_url):
    config = ConvNextConfig()

    if "tiny" in checkpoint_url:
        depths = [3, 3, 9, 3]
        hidden_sizes = [96, 192, 384, 768]
    if "small" in checkpoint_url:
        depths = [3, 3, 27, 3]
        hidden_sizes = [96, 192, 384, 768]
    if "base" in checkpoint_url:
        depths = [3, 3, 27, 3]
        hidden_sizes = [128, 256, 512, 1024]
    if "large" in checkpoint_url:
        depths = [3, 3, 27, 3]
        hidden_sizes = [192, 384, 768, 1536]
    if "xlarge" in checkpoint_url:
        depths = [3, 3, 27, 3]
        hidden_sizes = [256, 512, 1024, 2048]

    if "1k" in checkpoint_url:
        num_labels = 1000
        filename = "imagenet-1k-id2label.json"
        expected_shape = (1, 1000)
    else:
        num_labels = 21841
        filename = "imagenet-22k-id2label.json"
        expected_shape = (1, 21841)

    repo_id = "huggingface/label-files"
    config.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()}
    if "1k" not in checkpoint_url:
        
        
        del id2label[9205]
        del id2label[15027]
    config.id2label = id2label
    config.label2id = {v: k for k, v in id2label.items()}
    config.hidden_sizes = hidden_sizes
    config.depths = depths

    return config, expected_shape


def rename_key(name):
    if "downsample_layers.0.0" in name:
        name = name.replace("downsample_layers.0.0", "embeddings.patch_embeddings")
    if "downsample_layers.0.1" in name:
        name = name.replace("downsample_layers.0.1", "embeddings.norm")  
    if "downsample_layers.1.0" in name:
        name = name.replace("downsample_layers.1.0", "stages.1.downsampling_layer.0")
    if "downsample_layers.1.1" in name:
        name = name.replace("downsample_layers.1.1", "stages.1.downsampling_layer.1")
    if "downsample_layers.2.0" in name:
        name = name.replace("downsample_layers.2.0", "stages.2.downsampling_layer.0")
    if "downsample_layers.2.1" in name:
        name = name.replace("downsample_layers.2.1", "stages.2.downsampling_layer.1")
    if "downsample_layers.3.0" in name:
        name = name.replace("downsample_layers.3.0", "stages.3.downsampling_layer.0")
    if "downsample_layers.3.1" in name:
        name = name.replace("downsample_layers.3.1", "stages.3.downsampling_layer.1")
    if "stages" in name and "downsampling_layer" not in name:
        
        name = name[: len("stages.0")] + ".layers" + name[len("stages.0") :]
    if "stages" in name:
        name = name.replace("stages", "encoder.stages")
    if "norm" in name:
        name = name.replace("norm", "layernorm")
    if "gamma" in name:
        name = name.replace("gamma", "layer_scale_parameter")
    if "head" in name:
        name = name.replace("head", "classifier")

    return name



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_convnext_checkpoint(checkpoint_url, pytorch_dump_folder_path):
    

    
    config, expected_shape = get_convnext_config(checkpoint_url)
    
    state_dict = torch.hub.load_state_dict_from_url(checkpoint_url)["model"]
    
    for key in state_dict.copy().keys():
        val = state_dict.pop(key)
        state_dict[rename_key(key)] = val
    
    for key in state_dict.copy().keys():
        val = state_dict.pop(key)
        if not key.startswith("classifier"):
            key = "convnext." + key
        state_dict[key] = val

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

    
    size = 224 if "224" in checkpoint_url else 384
    image_processor = ConvNextImageProcessor(size=size)
    pixel_values = image_processor(images=prepare_img(), return_tensors="pt").pixel_values

    logits = model(pixel_values).logits

    
    if checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_tiny_1k_224_ema.pth":
        expected_logits = torch.tensor([-0.1210, -0.6605, 0.1918])
    elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_small_1k_224_ema.pth":
        expected_logits = torch.tensor([-0.4473, -0.1847, -0.6365])
    elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_base_1k_224_ema.pth":
        expected_logits = torch.tensor([0.4525, 0.7539, 0.0308])
    elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_base_1k_384.pth":
        expected_logits = torch.tensor([0.3561, 0.6350, -0.0384])
    elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_large_1k_224_ema.pth":
        expected_logits = torch.tensor([0.4174, -0.0989, 0.1489])
    elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_large_1k_384.pth":
        expected_logits = torch.tensor([0.2513, -0.1349, -0.1613])
    elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_base_22k_224.pth":
        expected_logits = torch.tensor([1.2980, 0.3631, -0.1198])
    elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_large_22k_224.pth":
        expected_logits = torch.tensor([1.2963, 0.1227, 0.1723])
    elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_xlarge_22k_224.pth":
        expected_logits = torch.tensor([1.7956, 0.8390, 0.2820])
    elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_base_22k_1k_224.pth":
        expected_logits = torch.tensor([-0.2822, -0.0502, -0.0878])
    elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_base_22k_1k_384.pth":
        expected_logits = torch.tensor([-0.5672, -0.0730, -0.4348])
    elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_large_22k_1k_224.pth":
        expected_logits = torch.tensor([0.2681, 0.2365, 0.6246])
    elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_large_22k_1k_384.pth":
        expected_logits = torch.tensor([-0.2642, 0.3931, 0.5116])
    elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_xlarge_22k_1k_224_ema.pth":
        expected_logits = torch.tensor([-0.6677, -0.1873, -0.8379])
    elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_xlarge_22k_1k_384_ema.pth":
        expected_logits = torch.tensor([-0.7749, -0.2967, -0.6444])
    else:
        raise ValueError(f"Unknown URL: {checkpoint_url}")

    assert torch.allclose(logits[0, :3], expected_logits, atol=1e-3)
    assert logits.shape == expected_shape

    Path(pytorch_dump_folder_path).mkdir(exist_ok=True)
    print(f"Saving model 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)

    print("Pushing model to the hub...")
    model_name = "convnext"
    if "tiny" in checkpoint_url:
        model_name += "-tiny"
    elif "small" in checkpoint_url:
        model_name += "-small"
    elif "base" in checkpoint_url:
        model_name += "-base"
    elif "xlarge" in checkpoint_url:
        model_name += "-xlarge"
    elif "large" in checkpoint_url:
        model_name += "-large"
    if "224" in checkpoint_url:
        model_name += "-224"
    elif "384" in checkpoint_url:
        model_name += "-384"
    if "22k" in checkpoint_url and "1k" not in checkpoint_url:
        model_name += "-22k"
    if "22k" in checkpoint_url and "1k" in checkpoint_url:
        model_name += "-22k-1k"

    model.push_to_hub(
        repo_path_or_name=Path(pytorch_dump_folder_path, model_name),
        organization="nielsr",
        commit_message="Add model",
    )


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    
    parser.add_argument(
        "--checkpoint_url",
        default="https://dl.fbaipublicfiles.com/convnext/convnext_tiny_1k_224_ema.pth",
        type=str,
        help="URL of the original ConvNeXT checkpoint you'd like to convert.",
    )
    parser.add_argument(
        "--pytorch_dump_folder_path",
        default=None,
        type=str,
        required=True,
        help="Path to the output PyTorch model directory.",
    )

    args = parser.parse_args()
    convert_convnext_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path)