{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%reload_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from fastai import *\n",
    "from fastai.vision import *\n",
    "from fastai.vision.models.darknet import Darknet"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "PATH = Path('../data/cifar10/')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ds_tfms = ([*rand_pad(4, 32), flip_lr(p=0.5)], [])\n",
    "data = image_data_from_folder(PATH, valid='test', ds_tfms=ds_tfms, tfms=cifar_norm, bs=64)\n",
    "learn = Learner(data, Darknet([1,1,1], 10), metrics=accuracy)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Example of loss in two parts"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This is a toy example. We pretend that our final loss comes from two different losses loss1 and loss2."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class CombinedLoss(nn.Module):\n",
    "    \n",
    "    def forward(self, output, target):\n",
    "        pct = uniform(0,1)\n",
    "        loss = F.cross_entropy(output, target)\n",
    "        self.loss1 = pct * loss\n",
    "        self.loss2 = (1-pct) * loss\n",
    "        return loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class HandleDualLoss(LearnerCallback):\n",
    "    _order = -20 #Needs to run before the recorder\n",
    "    \n",
    "    def on_train_begin(self, **kwargs):\n",
    "        self.learn.recorder.add_metric_names(['loss1', 'loss2'])\n",
    "    \n",
    "    def on_epoch_begin(self, **kwargs):\n",
    "        self.avg1, self.avg2, self.nums = 0., 0., 0\n",
    "    \n",
    "    def on_batch_end(self, last_target, train, **kwargs):\n",
    "        if not train:\n",
    "            bs = last_target.size(0)\n",
    "            self.avg1 += bs * learn.loss_fn.loss1.detach()\n",
    "            self.avg2 += bs * learn.loss_fn.loss2.detach()\n",
    "            self.nums += bs\n",
    "    \n",
    "    def on_epoch_end(self, **kwargs):\n",
    "        self.learn.recorder.add_metrics([self.avg1/self.nums, self.avg2/self.nums])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "learn.loss_fn = CombinedLoss()\n",
    "learn.callback_fns.append(HandleDualLoss)\n",
    "learn.fit_one_cycle(2, 3e-3, wd=0.4, div_factor=10, pct_start=0.5)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Handle metrics that aren't an average over batches."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Compute the precision for the first class."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Precision(LearnerCallback):\n",
    "    _order = -20 #Needs to run before the recorder\n",
    "    \n",
    "    def on_train_begin(self, **kwargs):\n",
    "        self.learn.recorder.add_metric_names(['precision'])\n",
    "    \n",
    "    def on_epoch_begin(self, **kwargs):\n",
    "        self.correct, self.total = 0, 0\n",
    "    \n",
    "    def on_batch_end(self, last_output, last_target, train, **kwargs):\n",
    "        if not train:\n",
    "            preds = last_output.argmax(1)\n",
    "            pdb.set_trace()\n",
    "            self.correct += ((preds==0) * (last_target==0)).float().sum()\n",
    "            self.total += (preds==0).long().sum()\n",
    "    \n",
    "    def on_epoch_end(self, **kwargs):\n",
    "        self.learn.recorder.add_metrics([self.correct/self.total])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "learn = Learner(data, Darknet([1,1,1], 10), metrics=accuracy)\n",
    "learn.callback_fns.append(Precision)\n",
    "learn.fit_one_cycle(2, 3e-3, wd=0.4, div_factor=10, pct_start=0.5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pdb"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}