UnifiedConversionAPI_MobileNet.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "alternative-gospel",
   "metadata": {},
   "outputs": [],
   "source": [
    "import coremltools as ct"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "removed-semiconductor",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Download class labels (from a separate file)\n",
    "import urllib\n",
    "label_url = 'https://storage.googleapis.com/download.tensorflow.org/data/ImageNetLabels.txt'\n",
    "class_labels = urllib.request.urlopen(label_url).read().splitlines()\n",
    "class_labels = class_labels[1:] # remove the first class which is background\n",
    "assert len(class_labels) == 1000\n",
    "\n",
    "# make sure entries of class_labels are strings\n",
    "for i, label in enumerate(class_labels):\n",
    "  if isinstance(label, bytes):\n",
    "    class_labels[i] = label.decode(\"utf8\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "dramatic-raise",
   "metadata": {},
   "outputs": [],
   "source": [
    "classifier_config = ct.ClassifierConfig(class_labels)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "separated-athletics",
   "metadata": {},
   "source": [
    "# TensorFlow (tf.keras)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "established-favorite",
   "metadata": {},
   "outputs": [],
   "source": [
    "from tensorflow.keras.applications import mobilenet_v2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "liquid-gossip",
   "metadata": {},
   "outputs": [],
   "source": [
    "tf_model = mobilenet_v2.MobileNetV2(include_top=True, weights='imagenet')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "final-sarah",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Running TensorFlow Graph Passes: 100%|██████████| 5/5 [00:00<00:00,  7.74 passes/s]\n",
      "Converting Frontend ==> MIL Ops: 100%|██████████| 428/428 [00:01<00:00, 291.40 ops/s] \n",
      "Running MIL optimization passes: 100%|██████████| 18/18 [00:01<00:00, 14.93 passes/s]\n",
      "Translating MIL ==> MLModel Ops: 100%|██████████| 751/751 [00:00<00:00, 1124.17 ops/s]\n"
     ]
    }
   ],
   "source": [
    "image_input = ct.ImageType(scale=2/255, bias=[-1, -1, -1])\n",
    "\n",
    "mlmodel_tf = ct.convert(tf_model,\n",
    "                        inputs=[image_input],\n",
    "                        classifier_config=classifier_config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "inside-lover",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set feature description (these show up as comments in Xcode)\n",
    "mlmodel_tf.input_description[\"input_1\"] = \"Input image to be classified\"\n",
    "mlmodel_tf.output_description[\"classLabel\"] = \"Most likely image category\"\n",
    "mlmodel_tf.output_description[\"Identity\"] = \"Probability of each image category\"\n",
    "mlmodel_tf.short_description = \"MobileNet v2 converted from TensorFlow\"\n",
    "mlmodel_tf.version = \"1.0\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "early-currency",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Save model\n",
    "mlmodel_tf.save(\"MobileNetV2_TF.mlmodel\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "piano-colonial",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "input {\n",
      "  name: \"input_1\"\n",
      "  shortDescription: \"Input image to be classified\"\n",
      "  type {\n",
      "    imageType {\n",
      "      width: 224\n",
      "      height: 224\n",
      "      colorSpace: RGB\n",
      "      imageSizeRange {\n",
      "        widthRange {\n",
      "          lowerBound: 224\n",
      "          upperBound: 224\n",
      "        }\n",
      "        heightRange {\n",
      "          lowerBound: 224\n",
      "          upperBound: 224\n",
      "        }\n",
      "      }\n",
      "    }\n",
      "  }\n",
      "}\n",
      "output {\n",
      "  name: \"Identity\"\n",
      "  shortDescription: \"Probability of each image category\"\n",
      "  type {\n",
      "    dictionaryType {\n",
      "      stringKeyType {\n",
      "      }\n",
      "    }\n",
      "  }\n",
      "}\n",
      "output {\n",
      "  name: \"classLabel\"\n",
      "  shortDescription: \"Most likely image category\"\n",
      "  type {\n",
      "    stringType {\n",
      "    }\n",
      "  }\n",
      "}\n",
      "predictedFeatureName: \"classLabel\"\n",
      "predictedProbabilitiesName: \"Identity\"\n",
      "metadata {\n",
      "  shortDescription: \"MobileNet v2 converted from TensorFlow\"\n",
      "  versionString: \"1.0\"\n",
      "  userDefined {\n",
      "    key: \"com.github.apple.coremltools.source\"\n",
      "    value: \"tensorflow==2.1.0\"\n",
      "  }\n",
      "  userDefined {\n",
      "    key: \"com.github.apple.coremltools.version\"\n",
      "    value: \"4.1\"\n",
      "  }\n",
      "}\n",
      "\n"
     ]
    }
   ],
   "source": [
    "print(mlmodel_tf)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "significant-smoke",
   "metadata": {},
   "source": [
    "# PyTorch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "cubic-dimension",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torchvision"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "aware-scheme",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch.nn as nn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "effective-works",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Wrapper class for passing the softmax function through the output of the pre-trained model\n",
    "class WrappedMobileNetV2(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(WrappedMobileNetV2, self).__init__()\n",
    "        self.model = torchvision.models.mobilenet_v2(pretrained=True).eval()\n",
    "        \n",
    "    def forward(self, x):\n",
    "        m = nn.Softmax(dim=1)\n",
    "        res = m(self.model(x))\n",
    "        return res"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "greater-prince",
   "metadata": {},
   "outputs": [],
   "source": [
    "traceable_model = WrappedMobileNetV2().eval()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "eligible-disaster",
   "metadata": {},
   "outputs": [],
   "source": [
    "example_input = torch.rand(1, 3, 224, 224) \n",
    "traced_model = torch.jit.trace(traceable_model, example_input)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "working-application",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Converting Frontend ==> MIL Ops: 100%|█████████▉| 386/387 [00:01<00:00, 377.46 ops/s]\n",
      "Running MIL optimization passes: 100%|██████████| 18/18 [00:00<00:00, 51.25 passes/s]\n",
      "Translating MIL ==> MLModel Ops: 100%|██████████| 706/706 [00:00<00:00, 765.32 ops/s] \n"
     ]
    }
   ],
   "source": [
    "image_input = ct.ImageType(name=\"input_1\", shape=example_input.shape, \n",
    "                           scale=2/255, bias=[-1, -1, -1])\n",
    "\n",
    "mlmodel_torch = ct.convert(traced_model,\n",
    "                           inputs=[image_input],\n",
    "                           classifier_config=classifier_config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "front-calculation",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "input {\n",
      "  name: \"input_1\"\n",
      "  type {\n",
      "    imageType {\n",
      "      width: 224\n",
      "      height: 224\n",
      "      colorSpace: RGB\n",
      "    }\n",
      "  }\n",
      "}\n",
      "output {\n",
      "  name: \"649\"\n",
      "  type {\n",
      "    dictionaryType {\n",
      "      stringKeyType {\n",
      "      }\n",
      "    }\n",
      "  }\n",
      "}\n",
      "output {\n",
      "  name: \"classLabel\"\n",
      "  type {\n",
      "    stringType {\n",
      "    }\n",
      "  }\n",
      "}\n",
      "predictedFeatureName: \"classLabel\"\n",
      "predictedProbabilitiesName: \"649\"\n",
      "metadata {\n",
      "  userDefined {\n",
      "    key: \"com.github.apple.coremltools.source\"\n",
      "    value: \"torch==1.6.0\"\n",
      "  }\n",
      "  userDefined {\n",
      "    key: \"com.github.apple.coremltools.version\"\n",
      "    value: \"4.1\"\n",
      "  }\n",
      "}\n",
      "\n"
     ]
    }
   ],
   "source": [
    "print(mlmodel_torch)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "id": "secure-questionnaire",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'649'"
      ]
     },
     "execution_count": 39,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "m = str(mlmodel_torch)\n",
    "target = 'predictedProbabilitiesName'\n",
    "m = m[m.find(target)+len(target):]\n",
    "node_name = m.split(\"\\\"\", 2)[1]\n",
    "node_name"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "id": "disabled-mechanism",
   "metadata": {},
   "outputs": [],
   "source": [
    "spec = mlmodel_torch.get_spec()\n",
    "\n",
    "ct.utils.rename_feature(spec, node_name, \"Identity\")\n",
    "mlmodel_torch = ct.models.MLModel(spec)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "id": "running-hearing",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set feature description (these show up as comments in Xcode)\n",
    "mlmodel_torch.input_description[\"input_1\"] = \"Input image to be classified\"\n",
    "mlmodel_torch.output_description[\"classLabel\"] = \"Most likely image category\"\n",
    "mlmodel_torch.output_description[\"Identity\"] = \"Probability of each image category\"\n",
    "mlmodel_torch.short_description = \"MobileNet v2 converted from PyTorch\"\n",
    "mlmodel_torch.version = \"1.0\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "id": "dying-estate",
   "metadata": {},
   "outputs": [],
   "source": [
    "mlmodel_torch.save(\"MobileNetV2_Torch.mlmodel\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "id": "british-maine",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "input {\n",
      "  name: \"input_1\"\n",
      "  shortDescription: \"Input image to be classified\"\n",
      "  type {\n",
      "    imageType {\n",
      "      width: 224\n",
      "      height: 224\n",
      "      colorSpace: RGB\n",
      "    }\n",
      "  }\n",
      "}\n",
      "output {\n",
      "  name: \"Identity\"\n",
      "  shortDescription: \"Probability of each image category\"\n",
      "  type {\n",
      "    dictionaryType {\n",
      "      stringKeyType {\n",
      "      }\n",
      "    }\n",
      "  }\n",
      "}\n",
      "output {\n",
      "  name: \"classLabel\"\n",
      "  shortDescription: \"Most likely image category\"\n",
      "  type {\n",
      "    stringType {\n",
      "    }\n",
      "  }\n",
      "}\n",
      "predictedFeatureName: \"classLabel\"\n",
      "predictedProbabilitiesName: \"Identity\"\n",
      "metadata {\n",
      "  shortDescription: \"MobileNet v2 converted from PyTorch\"\n",
      "  versionString: \"1.0\"\n",
      "  userDefined {\n",
      "    key: \"com.github.apple.coremltools.source\"\n",
      "    value: \"torch==1.6.0\"\n",
      "  }\n",
      "  userDefined {\n",
      "    key: \"com.github.apple.coremltools.version\"\n",
      "    value: \"4.1\"\n",
      "  }\n",
      "}\n",
      "\n"
     ]
    }
   ],
   "source": [
    "print(mlmodel_torch)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "utility-pointer",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}