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 mobilenetv2 checkpoints from the tensorflowmodels library import argparse import json import re from pathlib import path import requests import torch from huggingfacehub import hfhubdownload from pil import image from transformers import mobilenetv2config mobilenetv2forimageclassification mobilenetv2forsemanticsegmentation mobilenetv2imageprocessor loadtfweightsinmobilenetv2 from transformers utils import logging logging setverbosityinfo logger logging getloggername def getmobilenetv2configmodelname config mobilenetv2configlayernormeps0 001 if quant in modelname raise valueerrorquantized models are not supported matches re matchr mobilenetv2 modelname if matches config depthmultiplier floatmatches1 config imagesize intmatches2 if modelname startswithdeeplabv3 config outputstride 8 config numlabels 21 filename pascalvocid2label json else the tensorflow version of mobilenetv2 predicts 1001 classes instead of the usual 1000 the first class index 0 is background config numlabels 1001 filename imagenet1kid2label json repoid huggingfacelabelfiles id2label json loadopenhfhubdownloadrepoid filename repotypedataset r if config numlabels 1001 id2label intk 1 v for k v in id2label items id2label0 background else id2label intk v for k v in id2label items config id2label id2label config label2id v k for k v in id2label items return config 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 convertmovilevitcheckpointmodelname checkpointpath pytorchdumpfolderpath pushtohubfalse config getmobilenetv2configmodelname load model if modelname startswithdeeplabv3 model mobilenetv2forsemanticsegmentationconfig eval else model mobilenetv2forimageclassificationconfig eval load weights from tensorflow checkpoint loadtfweightsinmobilenetv2model config checkpointpath check outputs on an image prepared by mobilenetv2imageprocessor imageprocessor mobilenetv2imageprocessor cropsizewidth config imagesize height config imagesize sizeshortestedge config imagesize 32 encoding imageprocessorimagesprepareimg returntensorspt outputs modelencoding logits outputs logits if modelname startswithdeeplabv3 assert logits shape 1 21 65 65 if modelname deeplabv3mobilenetv21 0513 expectedlogits torch tensor 17 5790 17 7581 18 3355 18 3257 18 4230 18 8973 18 6169 18 8650 19 2187 2 1595 2 0977 2 3741 2 4226 2 3028 2 6835 2 7819 2 5991 2 7706 4 2058 4 8317 4 7638 4 4136 5 0361 4 9383 4 5028 4 9644 4 8734 else raise valueerrorfunknown model name modelname assert torch allcloselogits0 3 3 3 expectedlogits atol1e4 else assert logits shape 1 1001 if modelname mobilenetv21 4224 expectedlogits torch tensor0 0181 1 0015 0 4688 elif modelname mobilenetv21 0224 expectedlogits torch tensor0 2445 1 1993 0 1905 elif modelname mobilenetv20 75160 expectedlogits torch tensor0 2482 0 4136 0 6669 elif modelname mobilenetv20 3596 expectedlogits torch tensor0 1451 0 4624 0 7192 else expectedlogits none if expectedlogits is not none assert torch allcloselogits0 3 expectedlogits atol1e4 pathpytorchdumpfolderpath mkdirexistoktrue printfsaving model modelname to pytorchdumpfolderpath model savepretrainedpytorchdumpfolderpath printfsaving image processor to pytorchdumpfolderpath imageprocessor savepretrainedpytorchdumpfolderpath if pushtohub printpushing to the hub repoid google modelname imageprocessor pushtohubrepoid model pushtohubrepoid if name main parser argparse argumentparser required parameters parser addargument modelname defaultmobilenetv21 0224 typestr helpname of the mobilenetv2 model you d like to convert should in the form mobilenetv2depthsize parser addargument checkpointpath requiredtrue typestr helppath to the original tensorflow checkpoint ckpt file parser addargument pytorchdumpfolderpath requiredtrue typestr helppath to the output pytorch model directory parser addargument pushtohub actionstoretrue helpwhether or not to push the converted model to the hub args parser parseargs convertmovilevitcheckpoint args modelname args checkpointpath args pytorchdumpfolderpath 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 mobilenetv2 checkpoints from the tensorflow models library the tensorflow version of mobilenetv2 predicts 1001 classes instead of the usual 1000 the first class index 0 is background we will verify our results on an image of cute cats copy paste tweak model s weights to our mobilenetv2 structure load model load weights from tensorflow checkpoint check outputs on an image prepared by mobilenetv2imageprocessor required parameters
import argparse
import json
import re
from pathlib import Path

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

from transformers import (
    MobileNetV2Config,
    MobileNetV2ForImageClassification,
    MobileNetV2ForSemanticSegmentation,
    MobileNetV2ImageProcessor,
    load_tf_weights_in_mobilenet_v2,
)
from transformers.utils import logging


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


def get_mobilenet_v2_config(model_name):
    config = MobileNetV2Config(layer_norm_eps=0.001)

    if "quant" in model_name:
        raise ValueError("Quantized models are not supported.")

    matches = re.match(r"^.*mobilenet_v2_([^_]*)_([^_]*)$", model_name)
    if matches:
        config.depth_multiplier = float(matches[1])
        config.image_size = int(matches[2])

    if model_name.startswith("deeplabv3_"):
        config.output_stride = 8
        config.num_labels = 21
        filename = "pascal-voc-id2label.json"
    else:
        
        
        config.num_labels = 1001
        filename = "imagenet-1k-id2label.json"

    repo_id = "huggingface/label-files"
    id2label = json.load(open(hf_hub_download(repo_id, filename, repo_type="dataset"), "r"))

    if config.num_labels == 1001:
        id2label = {int(k) + 1: v for k, v in id2label.items()}
        id2label[0] = "background"
    else:
        id2label = {int(k): v for k, v in id2label.items()}

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

    return config



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_movilevit_checkpoint(model_name, checkpoint_path, pytorch_dump_folder_path, push_to_hub=False):
    
    config = get_mobilenet_v2_config(model_name)

    
    if model_name.startswith("deeplabv3_"):
        model = MobileNetV2ForSemanticSegmentation(config).eval()
    else:
        model = MobileNetV2ForImageClassification(config).eval()

    
    load_tf_weights_in_mobilenet_v2(model, config, checkpoint_path)

    
    image_processor = MobileNetV2ImageProcessor(
        crop_size={"width": config.image_size, "height": config.image_size},
        size={"shortest_edge": config.image_size + 32},
    )
    encoding = image_processor(images=prepare_img(), return_tensors="pt")
    outputs = model(**encoding)
    logits = outputs.logits

    if model_name.startswith("deeplabv3_"):
        assert logits.shape == (1, 21, 65, 65)

        if model_name == "deeplabv3_mobilenet_v2_1.0_513":
            expected_logits = torch.tensor(
                [
                    [[17.5790, 17.7581, 18.3355], [18.3257, 18.4230, 18.8973], [18.6169, 18.8650, 19.2187]],
                    [[-2.1595, -2.0977, -2.3741], [-2.4226, -2.3028, -2.6835], [-2.7819, -2.5991, -2.7706]],
                    [[4.2058, 4.8317, 4.7638], [4.4136, 5.0361, 4.9383], [4.5028, 4.9644, 4.8734]],
                ]
            )

        else:
            raise ValueError(f"Unknown model name: {model_name}")

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

        if model_name == "mobilenet_v2_1.4_224":
            expected_logits = torch.tensor([0.0181, -1.0015, 0.4688])
        elif model_name == "mobilenet_v2_1.0_224":
            expected_logits = torch.tensor([0.2445, -1.1993, 0.1905])
        elif model_name == "mobilenet_v2_0.75_160":
            expected_logits = torch.tensor([0.2482, 0.4136, 0.6669])
        elif model_name == "mobilenet_v2_0.35_96":
            expected_logits = torch.tensor([0.1451, -0.4624, 0.7192])
        else:
            expected_logits = None

        if expected_logits is not None:
            assert torch.allclose(logits[0, :3], expected_logits, atol=1e-4)

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

    if push_to_hub:
        print("Pushing to the hub...")
        repo_id = "google/" + model_name
        image_processor.push_to_hub(repo_id)
        model.push_to_hub(repo_id)


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    
    parser.add_argument(
        "--model_name",
        default="mobilenet_v2_1.0_224",
        type=str,
        help="Name of the MobileNetV2 model you'd like to convert. Should in the form 'mobilenet_v2_<depth>_<size>'.",
    )
    parser.add_argument(
        "--checkpoint_path", required=True, type=str, help="Path to the original TensorFlow checkpoint (.ckpt file)."
    )
    parser.add_argument(
        "--pytorch_dump_folder_path", required=True, type=str, help="Path to the output PyTorch model directory."
    )
    parser.add_argument(
        "--push_to_hub", action="store_true", help="Whether or not to push the converted model to the 🤗 hub."
    )

    args = parser.parse_args()
    convert_movilevit_checkpoint(
        args.model_name, args.checkpoint_path, args.pytorch_dump_folder_path, args.push_to_hub
    )