nb_inference.general.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# default_exp inference.general"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# General inference methods\n",
    "\n",
    "> Various general methods to be used during inference"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# hide\n",
    "from nbdev.showdoc import *\n",
    "from fastcore.test import *\n",
    "\n",
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export \n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "from pytorch_lightning.callbacks import BasePredictionWriter\n",
    "import torch\n",
    "import os\n",
    "from pathlib import Path"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "def convert_test_df(df):\n",
    "    '''Converts a df designed for testing to a train/val df as this easier to use for fastai batch inference'''\n",
    "    dummy_train = df.copy()\n",
    "    val = df.copy()\n",
    "    \n",
    "    dummy_train[\"is_valid\"] = False\n",
    "    val[\"is_valid\"] = True\n",
    "    \n",
    "    return val.append(dummy_train, ignore_index=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "class PredictionWriter(BasePredictionWriter):\n",
    "    '''Blabla\n",
    "    \n",
    "    Parameters\n",
    "    ----------\n",
    "    \n",
    "    splits : dict; optional\n",
    "        Contains \n",
    "    '''\n",
    "    def __init__(self, output_dir:str, output_file:str, dataset_names=list(), write_interval=\"epoch\"):\n",
    "        super().__init__(write_interval)\n",
    "        self.output_dir = output_dir\n",
    "        self.output_file = output_file\n",
    "        self.dataset_names = dataset_names\n",
    "        Path(self.output_dir).mkdir(parents=True, exist_ok=True)\n",
    "        \n",
    "    def write_on_epoch_end(self, trainer, pl_module, predictions, batch_indices):   \n",
    "        \n",
    "        dataset_names = self.dataset_names\n",
    "        if len(dataset_names)==0:\n",
    "            dataset_names = np.arange(len(predictions))\n",
    "        \n",
    "        preds = dict()\n",
    "        for dataset_name, prediction in zip(dataset_names, predictions):\n",
    "            probs, gts = zip(*prediction)\n",
    "            probs, gts = torch.concat(probs, dim=0), torch.concat(gts, dim=0)\n",
    "            preds[dataset_name] = (probs, gts)\n",
    "            \n",
    "        torch.save(preds, os.path.join(self.output_dir, f\"{self.output_file}.pt\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}