Upload flexible_shapes_repro.ipynb

flexible_shapes_repro.ipynb

@@ -0,0 +1,567 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "8f5b0950",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import coremltools as ct"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "id": "009656b9",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from transformers import AutoTokenizer, AutoModel\n",
+    "import numpy as np\n",
+    "import torch\n",
+    "import torch.nn as nn"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "c0eb4797",
+   "metadata": {},
+   "source": [
+    "## Model Setup"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "id": "6a3b370e",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "model_id = \"bert-base-uncased\""
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "id": "1b4b35d8",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Some weights of the model checkpoint at bert-base-uncased were not used when initializing BertModel: ['cls.seq_relationship.bias', 'cls.predictions.transform.LayerNorm.bias', 'cls.predictions.transform.LayerNorm.weight', 'cls.seq_relationship.weight', 'cls.predictions.transform.dense.bias', 'cls.predictions.transform.dense.weight', 'cls.predictions.decoder.weight', 'cls.predictions.bias']\n",
+      "- This IS expected if you are initializing BertModel from the checkpoint of a model trained on another task or with another architecture (e.g. initializing a BertForSequenceClassification model from a BertForPreTraining model).\n",
+      "- This IS NOT expected if you are initializing BertModel from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).\n"
+     ]
+    }
+   ],
+   "source": [
+    "tokenizer = AutoTokenizer.from_pretrained(model_id)\n",
+    "model = AutoModel.from_pretrained(model_id)\n",
+    "\n",
+    "model = model.eval()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "id": "f3f55386",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "compute_units = ct.ComputeUnit.CPU_ONLY"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "id": "ccbd0617",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "shape = (1, 128)\n",
+    "inputs = {\n",
+    "    \"input_ids\": np.random.randint(0, tokenizer.vocab_size, shape),\n",
+    "    \"attention_mask\": np.ones(shape, dtype=np.int64),\n",
+    "}"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "id": "20ea1402",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "odict_keys(['last_hidden_state', 'pooler_output'])"
+      ]
+     },
+     "execution_count": 7,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "t_inputs = {k: torch.tensor(v, dtype=torch.int32) for k, v in inputs.items()}\n",
+    "outputs = model(**t_inputs)\n",
+    "outputs.keys()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "e512e19b",
+   "metadata": {},
+   "source": [
+    "## JIT"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "id": "ad66c2eb",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "class Wrapper(nn.Module):\n",
+    "    def __init__(self, model):\n",
+    "        super().__init__()\n",
+    "        self.model = model\n",
+    "        \n",
+    "    def forward(self, *args, **kwargs):\n",
+    "        return self.model(return_dict=False, *args, **kwargs)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 9,
+   "id": "efb91bb7",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "to_jit = Wrapper(model)\n",
+    "jit_inputs = list(t_inputs.values())"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 10,
+   "id": "068cb16c",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "jitted_model = torch.jit.trace(to_jit, jit_inputs)\n",
+    "jitted_model.eval();"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 11,
+   "id": "2ae7472a",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "with torch.no_grad():\n",
+    "    output_jit = jitted_model(*jit_inputs)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 12,
+   "id": "f75237f7",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "tensor(0., grad_fn=<MaxBackward1>)"
+      ]
+     },
+     "execution_count": 12,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "(output_jit[0] - outputs[\"last_hidden_state\"]).abs().max()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 13,
+   "id": "820fd659",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "tensor(0., grad_fn=<MaxBackward1>)"
+      ]
+     },
+     "execution_count": 13,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "(output_jit[1] - outputs[\"pooler_output\"]).abs().max()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "8be44765",
+   "metadata": {},
+   "source": [
+    "## Core ML Conversion"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 14,
+   "id": "5e221907",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "input_shape = ct.Shape(shape=(1, ct.RangeDim(lower_bound=1, upper_bound=128, default=1)))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 15,
+   "id": "bb8e96d5",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def _get_coreml_inputs(sample_inputs):\n",
+    "    return [\n",
+    "        ct.TensorType(\n",
+    "            name=k,\n",
+    "#             shape=v.shape,\n",
+    "            shape=input_shape,\n",
+    "            dtype=v.numpy().dtype if isinstance(v, torch.Tensor) else v.dtype,\n",
+    "        ) for k, v in sample_inputs.items()\n",
+    "    ]"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 16,
+   "id": "e9e83c6a",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Tuple detected at graph output. This will be flattened in the converted model.\n",
+      "Converting PyTorch Frontend ==> MIL Ops:   0%|        | 0/630 [00:00<?, ? ops/s]Core ML embedding (gather) layer does not support any inputs besides the weights and indices. Those given will be ignored.\n",
+      "Converting PyTorch Frontend ==> MIL Ops: 100%|▉| 628/630 [00:00<00:00, 3196.37 o\n",
+      "Running MIL Common passes: 100%|███████████| 40/40 [00:00<00:00, 49.98 passes/s]\n",
+      "Running MIL FP16ComputePrecision pass: 100%|█| 1/1 [00:01<00:00,  1.01s/ passes]\n",
+      "Running MIL Clean up passes: 100%|█████████| 11/11 [00:01<00:00,  5.64 passes/s]\n"
+     ]
+    }
+   ],
+   "source": [
+    "coreml_input_types = _get_coreml_inputs(t_inputs)\n",
+    "coreml_output_types = [ct.TensorType(name=name) for name in outputs.keys()]\n",
+    "\n",
+    "coreml_model = ct.convert(\n",
+    "    jitted_model,\n",
+    "    convert_to = \"mlprogram\",\n",
+    "    minimum_deployment_target = ct.target.macOS13,\n",
+    "    inputs = coreml_input_types,\n",
+    "    outputs = coreml_output_types,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "f3263470",
+   "metadata": {},
+   "source": [
+    "Conversion succeeds. Let's run inference."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 17,
+   "id": "378948b4",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "coreml_outputs = coreml_model.predict(t_inputs)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 18,
+   "id": "bb3e90c9",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "last_hidden_state\n",
+      "\tshape: torch.Size([1, 128, 768])\n",
+      "\tmax diff: 0.010896801948547363\n",
+      "pooler_output\n",
+      "\tshape: torch.Size([1, 768])\n",
+      "\tmax diff: 0.004700236022472382\n"
+     ]
+    }
+   ],
+   "source": [
+    "for name in [\"last_hidden_state\", \"pooler_output\"]:\n",
+    "    coreml_tensor = torch.tensor(coreml_outputs[name])\n",
+    "    diff = (coreml_tensor - outputs[name]).abs().max()\n",
+    "    print(f\"{name}\\n\\tshape: {coreml_tensor.shape}\\n\\tmax diff: {diff}\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "6506507b",
+   "metadata": {},
+   "source": [
+    "### Metadata Manipulation"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 19,
+   "id": "0c48e4ab",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "spec = coreml_model._spec"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 20,
+   "id": "7e24db9b",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def set_multiarray_shape(node, shape):\n",
+    "    \"\"\"Change the shape of the specified input or output in the Core ML model.\"\"\"\n",
+    "    del node.type.multiArrayType.shape[:]\n",
+    "    for x in shape:\n",
+    "        node.type.multiArrayType.shape.append(x)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 21,
+   "id": "9937f578",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "for ml_input in coreml_model._spec.description.input:\n",
+    "    set_multiarray_shape(ml_input, (1, 128))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 22,
+   "id": "d4bbcd8b",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from coremltools.models.neural_network import flexible_shape_utils"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 23,
+   "id": "d44bf932",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "flexible_shape_utils.set_multiarray_ndshape_range(\n",
+    "    coreml_model._spec,\n",
+    "    \"input_ids\",\n",
+    "    [1, 1],\n",
+    "    [1, 128],\n",
+    ")\n",
+    "\n",
+    "flexible_shape_utils.set_multiarray_ndshape_range(\n",
+    "    coreml_model._spec,\n",
+    "    \"attention_mask\",\n",
+    "    [1, 1],\n",
+    "    [1, 128],\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "4cb87162",
+   "metadata": {},
+   "source": [
+    "Output shapes"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 24,
+   "id": "fe0c14ed",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "shapes = ((1, 128, 768), (1, 768))\n",
+    "for output, shape in zip(spec.description.output, shapes):\n",
+    "    set_multiarray_shape(output, shape)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 25,
+   "id": "2a9ab275",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "set_multiarray_shape(coreml_model._spec.description.output[1], (1, 768))"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "63751eec",
+   "metadata": {},
+   "source": [
+    "Flexible shapes for `last_hidden_state`"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 26,
+   "id": "c61921b8",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "flexible_shape_utils.set_multiarray_ndshape_range(\n",
+    "    coreml_model._spec,\n",
+    "    \"last_hidden_state\",\n",
+    "    [1, 1, 768],\n",
+    "    [1, 128, 768],\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 27,
+   "id": "7f666ca0",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "[name: \"last_hidden_state\"\n",
+       "type {\n",
+       "  multiArrayType {\n",
+       "    shape: 1\n",
+       "    shape: 128\n",
+       "    shape: 768\n",
+       "    dataType: FLOAT32\n",
+       "    shapeRange {\n",
+       "      sizeRanges {\n",
+       "        lowerBound: 1\n",
+       "        upperBound: 1\n",
+       "      }\n",
+       "      sizeRanges {\n",
+       "        lowerBound: 1\n",
+       "        upperBound: 128\n",
+       "      }\n",
+       "      sizeRanges {\n",
+       "        lowerBound: 768\n",
+       "        upperBound: 768\n",
+       "      }\n",
+       "    }\n",
+       "  }\n",
+       "}\n",
+       ", name: \"pooler_output\"\n",
+       "type {\n",
+       "  multiArrayType {\n",
+       "    shape: 1\n",
+       "    shape: 768\n",
+       "    dataType: FLOAT32\n",
+       "  }\n",
+       "}\n",
+       "]"
+      ]
+     },
+     "execution_count": 27,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "coreml_model._spec.description.output"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 28,
+   "id": "f36b4233",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "/opt/homebrew/Caskroom/miniforge/base/envs/sdcoreml/lib/python3.9/site-packages/coremltools/models/model.py:146: RuntimeWarning: You will not be able to run predict() on this Core ML model. Underlying exception message was: Error compiling model: \"compiler error:  Encountered an error while compiling a neural network model: validator error: Model output 'pooler_output' has a different shape than its corresponding return value to main.\".\n",
+      "  _warnings.warn(\n"
+     ]
+    }
+   ],
+   "source": [
+    "coreml_model = ct.models.MLModel(coreml_model._spec, weights_dir=coreml_model.weights_dir)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "4eef9027",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "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.9.15"
+  },
+  "toc": {
+   "base_numbering": 1,
+   "nav_menu": {},
+   "number_sections": true,
+   "sideBar": true,
+   "skip_h1_title": false,
+   "title_cell": "Table of Contents",
+   "title_sidebar": "Contents",
+   "toc_cell": false,
+   "toc_position": {},
+   "toc_section_display": true,
+   "toc_window_display": false
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}