{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# default_exp utils"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Utils\n",
    "\n",
    "> Basic utilities with few dependencies."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#hide\n",
    "from nbdev.showdoc import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "from collections.abc import Iterable, Mapping\n",
    "from functools import partial\n",
    "import inspect\n",
    "import numpy as np\n",
    "import torch\n",
    "\n",
    "from htools import flatten"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "At training time, we will typically want to put the model and the current mini batch on the GPU. When developing on a CPU, a GPU isn't available, so we define a variable that will automatically find the right device. This goes in utils rather than core to avoid circular imports with the callbacks module."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "device(type='cpu')"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "DEVICE"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "def hasarg(func, arg):\n",
    "    \"\"\"Checks if a function has a given argument.\n",
    "    Works with args and kwargs as well if you exclude the\n",
    "    stars. See example below.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    func: function\n",
    "    arg: str\n",
    "        Name of argument to look for.\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    bool\n",
    "\n",
    "    Example\n",
    "    -------\n",
    "    def foo(a, b=6, *args):\n",
    "        return\n",
    "\n",
    "    >>> hasarg(foo, 'b')\n",
    "    True\n",
    "\n",
    "    >>> hasarg(foo, 'args')\n",
    "    True\n",
    "\n",
    "    >>> hasarg(foo, 'c')\n",
    "    False\n",
    "    \"\"\"\n",
    "    return arg in inspect.signature(func).parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "def quick_stats(x, digits=3):\n",
    "    \"\"\"Quick wrapper to get mean and standard deviation of a tensor.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    x: torch.Tensor\n",
    "    digits: int\n",
    "        Number of digits to round mean and standard deviation to.\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    tuple[float]\n",
    "    \"\"\"\n",
    "    return round(x.mean().item(), digits), round(x.std().item(), digits)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "def concat(*args, dim=-1):\n",
    "    \"\"\"Wrapper to torch.cat which accepts tensors as non-keyword\n",
    "    arguments rather than requiring them to be wrapped in a list.\n",
    "    This can be useful if we've built some generalized functionality\n",
    "    where parameters must be passed in a consistent manner.\n",
    "    \n",
    "    Parameters\n",
    "    ----------\n",
    "    args: torch.tensor\n",
    "        Multiple tensors to concatenate.\n",
    "    dim: int\n",
    "        Dimension to concatenate on (last dimension by default).\n",
    "        \n",
    "    Returns\n",
    "    -------\n",
    "    torch.tensor\n",
    "    \"\"\"\n",
    "    return torch.cat(args, dim=dim)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "def weighted_avg(*args, weights):\n",
    "    \"\"\"Compute a weighted average of multiple tensors.\n",
    "    \n",
    "    Parameters\n",
    "    ----------\n",
    "    args: torch.tensor\n",
    "        Multiple tensors with the same dtype and shape that you want to average.\n",
    "    weights: list\n",
    "        Ints or floats to weight each input tensor. The length of this list must\n",
    "        match the number of tensors passed in: the first weight will be multiplied \n",
    "        by the first tensor, the second weight by the second tensor, etc. If your\n",
    "        weights don't sum to 1, they will be normalized automatically.\n",
    "        \n",
    "    Returns\n",
    "    -------\n",
    "    torch.tensor: Same dtype and shape as each of the input tensors.\n",
    "    \"\"\"\n",
    "    weights = torch.tensor(weights)\n",
    "    total = weights.sum().float()\n",
    "    if total != 1: weights = weights / total\n",
    "    res = torch.stack(args)\n",
    "    weights_shape = [-1 if i == 0 else 1 for i, _ in enumerate(range(res.ndim))]\n",
    "    return torch.mean(res * weights.view(*weights_shape), axis=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "def identity(x):\n",
    "    \"\"\"Temporarily copied from htools.\n",
    "    \n",
    "    Returns the input argument. Sometimes it is convenient to have this if\n",
    "    we sometimes apply a function to an item: rather than defining a None\n",
    "    variable, sometimes setting it to a function, then checking if it's None\n",
    "    every time we're about to call it, we can set the default as identity and\n",
    "    safely call it without checking.\n",
    "    \n",
    "    Parameters\n",
    "    ----------\n",
    "    x: any\n",
    "    \n",
    "    Returns\n",
    "    -------\n",
    "    x: Unchanged input.\n",
    "    \"\"\"\n",
    "    return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "def tensor_dict_diffs(d1, d2):\n",
    "    \"\"\"Compare two dictionaries of tensors. The two dicts must have the \n",
    "    same keys.\n",
    "    \n",
    "    Parameters\n",
    "    ----------\n",
    "    d1: dict[any: torch.Tensor]\n",
    "    d2: dict[any: torch.Tensor]\n",
    "    \n",
    "    Returns\n",
    "    -------\n",
    "    list: Returns the keys where tensors differ for d1 and d2. \n",
    "    \"\"\"\n",
    "    assert d1.keys() == d2.keys()\n",
    "    res = []\n",
    "    for k, v in d1.items():\n",
    "        if not torch.eq(v, d2[k]).all():\n",
    "            res.append(k)\n",
    "    return res"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "def find_tensors(gpu_only=True):\n",
    "    \"\"\"Prints a list of the Tensors being tracked by the garbage collector.\n",
    "    From\n",
    "    https://forums.fast.ai/t/gpu-memory-not-being-freed-after-training-is-over/10265/8\n",
    "    with some minor reformatting.\n",
    "    \n",
    "    Parameters\n",
    "    ----------\n",
    "    gpu_only: bool\n",
    "        If True, only find tensors that are on the GPU.\n",
    "        \n",
    "    Returns\n",
    "    -------\n",
    "    None: Output is printed to stdout.\n",
    "    \"\"\"\n",
    "    def pretty_size(size):\n",
    "        \"\"\"Pretty prints a torch.Size object.\"\"\"\n",
    "        assert(isinstance(size, torch.Size))\n",
    "        return \" × \".join(map(str, size))\n",
    "\n",
    "    total_size = 0\n",
    "    for obj in gc.get_objects():\n",
    "        try:\n",
    "            if torch.is_tensor(obj) and (not gpu_only or obj.is_cuda):\n",
    "                print(f'{type(obj).__name__}: {\"GPU\" if obj.is_cuda else \"\"}',\n",
    "                      \"pinned\" if obj.is_pinned else \"\",\n",
    "                      pretty_size(obj.size()))\n",
    "                total_size += obj.numel()\n",
    "            elif hasattr(obj, \"data\") and torch.is_tensor(obj.data):\n",
    "                if not gpu_only or obj.is_cuda:\n",
    "                    print(\n",
    "                        f'{type(obj).__name__} → {type(obj.data).__name__}',\n",
    "                        'GPU' if obj.is_cuda else '',\n",
    "                        'pinned' if obj.data.is_pinned else '',\n",
    "                        'grad' if obj.requires_grad else '',\n",
    "                        'volatile' if obj.volatile else '',\n",
    "                        pretty_size(obj.data.size())\n",
    "                    )\n",
    "                    total_size += obj.data.numel()\n",
    "        except Exception:\n",
    "            pass\n",
    "    print(\"Total size:\", total_size)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "def inverse_sigmoid(y):\n",
    "    \"\"\"Use to determine the bias initializer for the final linear layer of\n",
    "    model.\n",
    "    \n",
    "    Parameters\n",
    "    ----------\n",
    "    x: float\n",
    "        Value between 0 and 1 (e.g. the proportion of the training data that\n",
    "        are postives).\n",
    "        \n",
    "    Returns\n",
    "    -------\n",
    "    float: Inverse sigmoid of input.\n",
    "        I.E. if y=sigmoid(x), inverse_sigmoid(y)=x.\n",
    "    \"\"\"\n",
    "    return np.log(y / (1-y))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "def init_bias_constant_(layer, value=None, target_pct=None):\n",
    "    \"\"\"Helper to initialize a layer's bias term to a constant. This is\n",
    "    particularly useful for the final layer of a binary classifier where it's\n",
    "    often helpful to intialize it to the value that, when passed through a \n",
    "    sigmoid activation, is equal to the percent of your dataset belonging to\n",
    "    the majority class. This reduces the chance that the first epoch or so\n",
    "    will be spent simply learning a bias term, and has two benefits:\n",
    "    \n",
    "    1. May reduce training time slightly.\n",
    "    2. The beginning of training can be deceptively important - a messy first\n",
    "    couple epochs can have long-lasting repercussions. This is often hard to\n",
    "    identify without in-depth digging into model weights so it often goes\n",
    "    unnoticed. \n",
    "    \n",
    "    Parameters\n",
    "    ----------\n",
    "    layer: nn.Module\n",
    "        The layer to initialize a bias for. This will often be the last layer\n",
    "        of our network - we rarely need to initialize a constant bias \n",
    "        otherwise.\n",
    "    value: float or None\n",
    "        If provided, the bias will be initialized to this value.\n",
    "    target_pct: float or None\n",
    "        If provided, must be a float between 0 and 1.\n",
    "        \n",
    "    Returns\n",
    "    -------\n",
    "    None: The layer is updated in place.\n",
    "    \"\"\"\n",
    "    assert bool(value is None) + bool(target_pct is None) == 1, \\\n",
    "        'Specify either `value` OR `target_pct`.'\n",
    "    if target_pct is not None:\n",
    "        assert 0 < target_pct < 1, \\\n",
    "        '`target_pct` must be in range (0, 1) (exclusive).'\n",
    "\n",
    "    value = value or inverse_sigmoid(target_pct)\n",
    "    with torch.no_grad():\n",
    "        layer.bias.fill_(value)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "def is_builtin(x, drop_callables=True):\n",
    "    \"\"\"Check if an object belongs to the Python standard library.\n",
    "    \n",
    "    Parameters\n",
    "    ----------\n",
    "    drop_callables: bool\n",
    "        If True, we won't consider callables (classes/functions) to be builtin.\n",
    "        Classes have class `type` and functions have class \n",
    "        `builtin_function_or_method`, both of which are builtins - however, \n",
    "        this is often not what we mean when we want to know if something is\n",
    "        built in. Note: knowing the class alone is not enough to determine if\n",
    "        the objects it creates are built-in; this may depend on the kwargs\n",
    "        passed to its constructor. This will NOT check if a class was defined\n",
    "        in the standard library.\n",
    "    \n",
    "    Returns\n",
    "    -------\n",
    "    bool: True if the object is built-in. If the object is list-like, each\n",
    "    item will be checked as well the container. If the object is dict-like,\n",
    "    each key AND value will be checked (you can always pass in d.keys() or \n",
    "    d.values() for more limited checking). Again, the container itself will\n",
    "    be checked as well.\n",
    "    \"\"\"\n",
    "    def _builtin(x, drop_callables):\n",
    "        if callable(x) and drop_callables:\n",
    "            return False\n",
    "        return x.__class__.__module__ == 'builtins'\n",
    "    \n",
    "    builtin = partial(_builtin, drop_callables=drop_callables)\n",
    "    # Check mapping before iterable because mappings are iterable.\n",
    "    if isinstance(x, Mapping):\n",
    "        return builtin(x) and all(builtin(o) for o in flatten(x.items()))\n",
    "    elif isinstance(x, Iterable):\n",
    "        return builtin(x) and all(builtin(o) for o in flatten(x))\n",
    "    return builtin(x)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.7.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
