{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/root/miniconda3/envs/torch/lib/python3.9/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n",
      "/root/miniconda3/envs/torch/lib/python3.9/site-packages/diffusers/models/cross_attention.py:30: FutureWarning: Importing from cross_attention is deprecated. Please import from diffusers.models.attention_processor instead.\n",
      "  deprecate(\n"
     ]
    }
   ],
   "source": [
    "from transformers import pipeline\n",
    "import torch\n",
    "from optimum.onnxruntime import ORTModelForImageClassification\n",
    "from transformers import AutoImageProcessor\n",
    "\n",
    "from models import ResNetFN"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Could not find image processor class in the image processor config or the model config. Loading based on pattern matching with the model's feature extractor configuration.\n",
      "Some weights of the model checkpoint at microsoft/resnet-50 were not used when initializing ResNetModel: ['classifier.1.bias', 'classifier.1.weight']\n",
      "- This IS expected if you are initializing ResNetModel 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 ResNetModel from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<All keys matched successfully>"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "processor = AutoImageProcessor.from_pretrained(\"microsoft/resnet-50\")\n",
    "model = ResNetFN()\n",
    "state_dict = torch.load('/opt/data/private/graduation/autopilot/outputs/checkpoint-6321/pytorch_model.bin')\n",
    "model.load_state_dict(state_dict)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Could not find image processor class in the image processor config or the model config. Loading based on pattern matching with the model's feature extractor configuration.\n"
     ]
    }
   ],
   "source": [
    "from datasets import load_dataset\n",
    "from transformers import AutoImageProcessor\n",
    "from torch.utils.data import DataLoader\n",
    "import evaluate\n",
    "from tqdm import tqdm\n",
    "from torchvision.transforms import (\n",
    "    CenterCrop,\n",
    "    Compose,\n",
    "    Normalize,\n",
    "    RandomHorizontalFlip,\n",
    "    RandomResizedCrop,\n",
    "    Resize,\n",
    "    ToTensor,\n",
    ")\n",
    "\n",
    "image_processor = AutoImageProcessor.from_pretrained('microsoft/resnet-50')\n",
    "size = image_processor.size[\"shortest_edge\"]\n",
    "normalize = Normalize(mean=image_processor.image_mean, std=image_processor.image_std)\n",
    "val_transforms = Compose(\n",
    "  [\n",
    "    Resize(size),\n",
    "    CenterCrop(size),\n",
    "    ToTensor(),\n",
    "    normalize,\n",
    "  ]\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Resolving data files: 100%|██████████| 1805/1805 [00:00<00:00, 11224.96it/s]\n",
      "Found cached dataset imagefolder (/root/.cache/huggingface/datasets/imagefolder/default-5ead7e559c62c602/0.0.0/37fbb85cc714a338bea574ac6c7d0b5be5aff46c1862c1989b20e0771199e93f)\n"
     ]
    }
   ],
   "source": [
    "def collate_fn(examples):\n",
    "  pixel_values = torch.stack([example[\"pixel_values\"] for example in examples])\n",
    "  labels = torch.tensor([example[\"label\"] for example in examples])\n",
    "  return {\"pixel_values\": pixel_values, \"labels\": labels}\n",
    "\n",
    "def preprocess_val(example_batch):\n",
    "  \"\"\"Apply _val_transforms across a batch.\"\"\"\n",
    "  example_batch[\"pixel_values\"] = [val_transforms(image.convert(\"RGB\")) for image in example_batch[\"image\"]]\n",
    "  return example_batch\n",
    "\n",
    "eval_dataset = load_dataset('imagefolder', data_dir='/opt/data/private/graduation/autopilot/data/test', split='train').with_transform(preprocess_val)\n",
    "eval_dataloader = DataLoader(eval_dataset, collate_fn=collate_fn, batch_size=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "metric = evaluate.load(\"accuracy\")\n",
    "\n",
    "model = model.cuda()\n",
    "model = model.cpu()\n",
    "model.eval()\n",
    "for step, batch in tqdm(enumerate(eval_dataloader)):\n",
    "  # for key in batch.keys():\n",
    "  #   batch[key] = batch[key].cuda()\n",
    "  with torch.no_grad():\n",
    "    outputs = model(**batch)\n",
    "  predictions = outputs.logits.argmax(dim=-1)\n",
    "  references = batch[\"labels\"]\n",
    "  metric.add_batch(\n",
    "    predictions=predictions,\n",
    "    references=references,\n",
    "  )\n",
    "eval_metric = metric.compute()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/root/miniconda3/envs/torch/lib/python3.9/site-packages/transformers/models/convnext/feature_extraction_convnext.py:28: FutureWarning: The class ConvNextFeatureExtractor is deprecated and will be removed in version 5 of Transformers. Please use ConvNextImageProcessor instead.\n",
      "  warnings.warn(\n"
     ]
    }
   ],
   "source": [
    "from optimum.onnxruntime import ORTModelForImageClassification\n",
    "\n",
    "save_directory = '/opt/data/private/graduation/autopilot/quantization/onnx'\n",
    "model_onnx = ORTModelForImageClassification.from_pretrained(save_directory, file_name=\"model_quantized.onnx\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "metric = evaluate.load(\"accuracy\")\n",
    "model_onnx = model.cuda()\n",
    "# model_onnx.eval()\n",
    "\n",
    "for step, batch in tqdm(enumerate(eval_dataloader)):\n",
    "  for key in batch.keys():\n",
    "    batch[key] = batch[key].cuda()\n",
    "  with torch.no_grad():\n",
    "    outputs = model_onnx(**batch)\n",
    "  predictions = outputs.logits.argmax(dim=-1)\n",
    "  references = batch[\"labels\"]\n",
    "  metric.add_batch(\n",
    "    predictions=predictions,\n",
    "    references=references,\n",
    "  )\n",
    "eval_metric_onnx = metric.compute()"
   ]
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "6f64be793538f7fe230f350828c9baf03d97c4df0981f52e8388f53f367f4a42"
  },
  "kernelspec": {
   "display_name": "Python 3.9.16 ('torch')",
   "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.16"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}