finetune-detr-hf.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Finetune DETR to detect female-ish faces in paintings"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "! pip install --upgrade scipy transformers datasets huggingface_hub pytorch-lightning pycocotools"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!huggingface-cli login"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from huggingface_hub import notebook_login\n",
    "notebook_login()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from Cocordiais import CocordiaisDataset\n",
    "\n",
    "from datasets import load_dataset\n",
    "from torch.utils.data import DataLoader\n",
    "from transformers import AutoModelForObjectDetection, DetrImageProcessor, Trainer, TrainingArguments"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "DETR_MODEL = \"facebook/detr-resnet-50\"\n",
    "HF_DATASET = \"thiagohersan/cordiais-faces\"\n",
    "HF_MODEL= \"thiagohersan/detr-cordiais-autotrain\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "detr_size = { \"shortest_edge\": 800, \"longest_edge\": 800 }\n",
    "detr_processor = DetrImageProcessor.from_pretrained(DETR_MODEL, size=detr_size)\n",
    "\n",
    "hf_dataset = load_dataset(HF_DATASET)\n",
    "hf_dataset = hf_dataset[\"train\"].train_test_split(test_size=0.2, shuffle=True, seed=101010)\n",
    "\n",
    "dataset_train = CocordiaisDataset(hf_dataset[\"train\"], img_processor=detr_processor, train=True)\n",
    "dataset_eval = CocordiaisDataset(hf_dataset[\"test\"], img_processor=detr_processor, train=False)\n",
    "\n",
    "print(\"Number of examples:\\n  Train: %s\\n  Evaluation: %s\" % (len(dataset_train), len(dataset_eval)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "labels = dataset_train.data.features[\"objects\"].feature[\"category\"].names\n",
    "id2label = {i:l for i,l in enumerate(labels)}\n",
    "label2id = {l:i for i,l in id2label.items()}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = AutoModelForObjectDetection.from_pretrained(\n",
    "  DETR_MODEL,\n",
    "  id2label=id2label,\n",
    "  label2id=label2id,\n",
    "  revision=\"no_timm\", \n",
    "  num_labels=len(id2label),\n",
    "  num_queries=16,\n",
    "  ignore_mismatched_sizes=True\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "training_args = TrainingArguments(\n",
    "  output_dir=HF_MODEL,\n",
    "  per_device_train_batch_size=12,\n",
    "  per_device_eval_batch_size=4,\n",
    "  num_train_epochs=100,\n",
    "  fp16=True,\n",
    "  save_strategy=\"epoch\",\n",
    "  save_total_limit=2,\n",
    "  logging_strategy=\"epoch\",\n",
    "  learning_rate=1e-5,\n",
    "  weight_decay=1e-4,\n",
    "  remove_unused_columns=False\n",
    ")\n",
    "\n",
    "trainer = Trainer(\n",
    "  model=model,\n",
    "  args=training_args,\n",
    "  data_collator=dataset_train.collate_batch,\n",
    "  train_dataset=dataset_train.data,\n",
    "  eval_dataset=dataset_eval.data,\n",
    "  tokenizer=detr_processor\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext tensorboard\n",
    "%tensorboard --logdir lightning_logs/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer.train()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer.push_to_hub(HF_MODEL, private=True)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.9.17 ('hf-model')",
   "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.17"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "89e384cab7c47fb35ec95d2248b519cf922ee174880eed636c26cdfb6c4df768"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}