{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 64,
   "id": "b8574bee-e4d5-4640-babc-86178f6003fe",
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import ResNetModel, ResNetForImageClassification\n",
    "import torch\n",
    "import coremltools as ct"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 65,
   "id": "0eb4ebac-7d9c-44ca-9bd6-8014313b08cf",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = ResNetForImageClassification.from_pretrained(\"kwooleryorg/finetuned-style-resnet-50\", return_dict=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 77,
   "id": "f91aebb7-5d2b-4057-ba38-b1f9e1191b7e",
   "metadata": {},
   "outputs": [],
   "source": [
    "model.eval();\n",
    "custom_labels=['Boho',\n",
    " 'Coastal',\n",
    " 'Contemporary',\n",
    " 'Farmhouse',\n",
    " 'Glam',\n",
    " 'Industrial',\n",
    " 'Mid-Century',\n",
    " 'Modern',\n",
    " 'Scandinavian',\n",
    " 'Traditional']\n",
    "labels = [f\"label{i}\" for i in range(1000)]\n",
    "for i in range(len(custom_labels)):\n",
    "    labels[i] = custom_labels[i]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 67,
   "id": "03dd37f0-8114-430a-9065-81adb3986006",
   "metadata": {},
   "outputs": [],
   "source": [
    "random_tokens = torch.rand(1, 3, 224, 224)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 68,
   "id": "df2a0262-d381-413c-829e-311ae73669d7",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/kevin/anaconda3/lib/python3.9/site-packages/transformers/models/resnet/modeling_resnet.py:95: TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!\n",
      "  if num_channels != self.num_channels:\n"
     ]
    }
   ],
   "source": [
    "traced_model = torch.jit.trace(model, random_tokens)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 69,
   "id": "67b1a9a5-d997-4bbd-87b5-5329f39ba1cb",
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.jit.save(traced_model, \"./traced_resnet50.pt\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 82,
   "id": "2377f748-985a-4bc8-bfdd-a0659afdb1ab",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Support for converting Torch Script Models is experimental. If possible you should use a traced model for conversion.\n",
      "Tuple detected at graph output. This will be flattened in the converted model.\n",
      "Converting PyTorch Frontend ==> MIL Ops: 100%|█████████▉| 402/404 [00:00<00:00, 2081.87 ops/s]\n",
      "Running MIL frontend_pytorch pipeline: 100%|██████████| 5/5 [00:00<00:00, 156.82 passes/s]\n",
      "Running MIL default pipeline: 100%|██████████| 56/56 [00:03<00:00, 16.78 passes/s] \n",
      "Running MIL backend_neuralnetwork pipeline: 100%|██████████| 8/8 [00:00<00:00, 469.32 passes/s]\n",
      "Translating MIL ==> NeuralNetwork Ops: 100%|██████████| 451/451 [00:04<00:00, 95.32 ops/s]  \n"
     ]
    }
   ],
   "source": [
    "input_shape = ct.Shape(shape=(1, 3, 224, 224))\n",
    "input_tensor = ct.TensorType(name=\"image\", shape=input_shape)\n",
    "classifier_config = ct.ClassifierConfig(labels)\n",
    "outputs=ct.TensorType(name=\"logits\")\n",
    "mlmodel = ct.convert(\"./traced_resnet50.pt\", inputs=[input_tensor], outputs=[outputs], classifier_config=classifier_config, source=\"pytorch\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 83,
   "id": "1ed63292-a3a6-485d-80c6-e6c8b751afd5",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "input {\n",
       "  name: \"image\"\n",
       "  type {\n",
       "    multiArrayType {\n",
       "      shape: 1\n",
       "      shape: 3\n",
       "      shape: 224\n",
       "      shape: 224\n",
       "      dataType: FLOAT32\n",
       "    }\n",
       "  }\n",
       "}\n",
       "output {\n",
       "  name: \"logits\"\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: \"logits\"\n",
       "metadata {\n",
       "  userDefined {\n",
       "    key: \"com.github.apple.coremltools.source\"\n",
       "    value: \"torch==1.13.1+cu117\"\n",
       "  }\n",
       "  userDefined {\n",
       "    key: \"com.github.apple.coremltools.version\"\n",
       "    value: \"6.3.0\"\n",
       "  }\n",
       "}"
      ]
     },
     "execution_count": 83,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "mlmodel"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 84,
   "id": "26a64763-64da-452f-9747-556a0e931262",
   "metadata": {},
   "outputs": [],
   "source": [
    "mlmodel.save(\"./resnet50-trained.mlmodel\")"
   ]
  }
 ],
 "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.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}