{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-08-07T20:39:13.881585Z",
     "start_time": "2021-08-07T20:39:13.868461Z"
    }
   },
   "outputs": [],
   "source": [
    "# default_exp utils"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Utils\n",
    "\n",
    "> Basic utilities with few dependencies."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-08-07T20:39:15.620796Z",
     "start_time": "2021-08-07T20:39:14.202820Z"
    }
   },
   "outputs": [],
   "source": [
    "#hide\n",
    "from nbdev.showdoc import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-08-07T20:39:15.647042Z",
     "start_time": "2021-08-07T20:39:15.622681Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The autoreload extension is already loaded. To reload it, use:\n",
      "  %reload_ext autoreload\n"
     ]
    }
   ],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-08-07T20:39:17.559071Z",
     "start_time": "2021-08-07T20:39:15.651141Z"
    }
   },
   "outputs": [],
   "source": [
    "# export\n",
    "from collections.abc import Iterable, Mapping\n",
    "from functools import partial\n",
    "import io\n",
    "import inspect\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import os\n",
    "from pathlib import Path\n",
    "import random\n",
    "import requests\n",
    "import torch\n",
    "import warnings\n",
    "\n",
    "from htools import flatten, ifnone, save, load"
   ]
  },
  {
   "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 reproducible(seed=1, verbose=True):\n",
    "    if verbose: print('Setting seeds for reproducible training.')\n",
    "    os.environ['PYTHONHASHSEED'] = str(seed)\n",
    "    random.seed(seed)\n",
    "    np.random.seed(seed)\n",
    "    torch.manual_seed(seed)\n",
    "    torch.backends.cudnn.deterministic = True\n",
    "    torch.backends.cudnn.benchmark = False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "def gpu_setup(make_reproducible=True, seed=1, verbose=1):\n",
    "    if make_reproducible: reproducible(seed, verbose)\n",
    "    if not torch.cuda.is_available(): warnings.warn('Cuda not available.')\n",
    "    if DEVICE.type != 'cuda': \n",
    "        print(DEVICE)\n",
    "        warnings.warn('Incendio device is not cuda.')"
   ]
  },
  {
   "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": 2,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-08-09T00:49:34.328960Z",
     "start_time": "2021-08-09T00:49:34.313474Z"
    }
   },
   "outputs": [],
   "source": [
    "# export\n",
    "def pool_stds(ns, means, stds):\n",
    "    \"\"\"Pool standard deviations from multiple groups. We cannot simply take a\n",
    "    weighted average like for means. \n",
    "    \n",
    "    Formulas and some notation is from:\n",
    "    http://www.obg.cuhk.edu.hk/ResearchSupport/StatTools/CombineMeansSDs_Pgm.php\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    ns: list[int] or torch.LongTensor\n",
    "        Size of each group. Must be in same order as other args.\n",
    "    means: list[float] or torch.FloatTensor\n",
    "        Mean of each group. Must be in same order as other args.\n",
    "    stds: list[float] or torch.FloatTensor\n",
    "        Standard deviation of each group. Must be in same order as other args.\n",
    "        \n",
    "    Returns\n",
    "    -------\n",
    "    torch.FloatTensor: pooled standard deviation across all groups.\n",
    "    \"\"\"\n",
    "    ns, means, stds = map(torch.tensor, [ns, means, stds])\n",
    "    n_total = ns.sum()\n",
    "    totals = means * ns\n",
    "    tx = totals.sum()\n",
    "    txx = torch.sum(stds.pow(2).mul(ns - 1) + totals.pow(2).div(ns))\n",
    "    return torch.sqrt((txx - tx.pow(2).div(n_total)) / (n_total - 1))"
   ]
  },
  {
   "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 \n",
    "        average.\n",
    "    weights: list\n",
    "        Ints or floats to weight each input tensor. The length of this list\n",
    "        must match the number of tensors passed in: the first weight will be\n",
    "        multiplied by the first tensor, the second weight by the second \n",
    "        tensor, etc. If your weights don't sum to 1, they will be normalized \n",
    "        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": 6,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-10-22T05:09:38.544627Z",
     "start_time": "2020-10-22T05:09:38.500868Z"
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/hmamin/anaconda3/lib/python3.7/site-packages/ipykernel/ipkernel.py:287: DeprecationWarning: `should_run_async` will not call `transform_cell` automatically in the future. Please pass the result to `transformed_cell` argument and any exception that happen during thetransform in `preprocessing_exc_tuple` in IPython 7.17 and above.\n",
      "  and should_run_async(code)\n"
     ]
    }
   ],
   "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": null,
   "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": null,
   "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)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-07-24T22:54:22.054689Z",
     "start_time": "2021-07-24T22:54:22.037347Z"
    }
   },
   "outputs": [],
   "source": [
    "# export\n",
    "def corr(x, y):\n",
    "    \"\"\"Pearson correlation coefficient.\n",
    "    \n",
    "    Parameters\n",
    "    ----------\n",
    "    x: torch.Tensor\n",
    "        Same shape as y.\n",
    "    y: torch.Tensor\n",
    "        Same shape as x.\n",
    "    \"\"\"\n",
    "    x_diff, y_diff = x - x.mean(), y - y.mean()\n",
    "    return (x_diff @ y_diff) / (x_diff.pow(2).sum().sqrt()\n",
    "                                * y_diff.pow(2).sum().sqrt())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-10-20T06:21:22.416562Z",
     "start_time": "2020-10-20T06:21:22.353363Z"
    }
   },
   "outputs": [],
   "source": [
    "# export\n",
    "def out_features(model): \n",
    "    \"\"\"Try to extract number of output features from the last layer of a \n",
    "    model. This is often useful when building encoder-decoder models or \n",
    "    stacking encoders and classification heads. Not sure how airtight the \n",
    "    logic is here so use with caution.\n",
    "    \n",
    "    Parameters\n",
    "    ----------\n",
    "    model: nn.Module\n",
    "        Model to examine.\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    int: Number of output features.\n",
    "    \"\"\"\n",
    "    return [c.out_features for c in model.children() \n",
    "            if hasattr(c, 'out_features')][-1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-08-07T20:39:30.244389Z",
     "start_time": "2021-08-07T20:39:30.207060Z"
    }
   },
   "outputs": [],
   "source": [
    "# export\n",
    "def get_einops_image(path=None, as_tensor=True, channels_first=True, \n",
    "                     color=True):\n",
    "    \"\"\"Easily load the sample image einops uses for tutorials. Nice way to\n",
    "    play around with tensor reshaping/transformation to make sure you're doing\n",
    "    what you think you are. The file will be stored in a dir called \n",
    "    ~/.incendio.\n",
    "    \n",
    "    Parameters\n",
    "    ----------\n",
    "    path: None or str or Path\n",
    "        If not None, this is the location to try to load the file from and to\n",
    "        save it to if it's not already there. Otherwise, this defaults to\n",
    "        ~/.incendio/einops.zip.\n",
    "    as_tensor: bool\n",
    "        Return tensor if True, np array if False.\n",
    "    channels_first: bool\n",
    "        Return with shape (bs, c, h, w) if True, (bs, h, w, c) if False. \n",
    "        Ignored if color=False.\n",
    "        \n",
    "    Returns\n",
    "    -------\n",
    "    torch.Tensor or np.ndarray\n",
    "    \"\"\"\n",
    "    path = Path(ifnone(path, '~/.incendio/einops.zip')).expanduser()\n",
    "    if not path.is_file():\n",
    "        r = requests.get('https://github.com/arogozhnikov/einops/raw/master/'\n",
    "                         'docs/resources/test_images.npy')\n",
    "        r.raise_for_status()\n",
    "        save(np.load(io.BytesIO(r.content)).astype(np.float32), path)\n",
    "    array = load(path).astype(np.float32)\n",
    "    if not color: array = array[..., 0]\n",
    "    elif channels_first: array = array.transpose(0, 3, 1, 2)\n",
    "    if as_tensor: array = torch.tensor(array)\n",
    "    return array"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-08-07T20:39:31.822077Z",
     "start_time": "2021-08-07T20:39:30.859635Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Writing data to /Users/harrisonmamin/.incendio/einops.zip.\n",
      "Object loaded from /Users/harrisonmamin/.incendio/einops.zip.\n",
      "Object loaded from /Users/harrisonmamin/.incendio/einops.zip.\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "(torch.Size([6, 3, 96, 96]), torch.Size([6, 96, 96]))"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "img = get_einops_image()\n",
    "img_bw = get_einops_image(color=False)\n",
    "img.shape, img_bw.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-08-07T20:44:05.699367Z",
     "start_time": "2021-08-07T20:44:05.660152Z"
    }
   },
   "outputs": [],
   "source": [
    "# export\n",
    "def plot_img_batch(x, channels_first=True):\n",
    "    \"\"\"Plot batch of image tensors. Helpful for assessing operations on\n",
    "    tensors, particularly when reshaping things with einops or equivalent ops\n",
    "    (view, transpose, etc.).\n",
    "    \n",
    "    Parameters\n",
    "    ----------\n",
    "    x: torch.Tensor or np.array\n",
    "        Shape (bs, h, w), (bs, c, h, w), or (bs, h, w, c). \n",
    "        Example: the tensor loaded by `get_einops_image`.\n",
    "    channels_first: bool\n",
    "        Specify if your input x has channels before the spatial dimensions,\n",
    "        e.g. True if x has shape (bs, c, h, w) and False if it has shape\n",
    "        (bs, h, w, c). (Technically the batch dimension is always first, this\n",
    "        just means first relative to the spatial dimensions.)\n",
    "    \"\"\"\n",
    "    ndim = x.ndim\n",
    "    assert ndim in {3, 4}, \\\n",
    "        'x ndim must be 3 or 4: (bs, h, w), (bs, c, h, w) or (bs, h, w, c)'\n",
    "\n",
    "    if isinstance(x, np.ndarray): x = torch.tensor(x)\n",
    "    if getattr(x, 'requires_grad', False): x = x.detach()\n",
    "    fig, ax = plt.subplots(len(x), 1, figsize=(4, len(x)*2))\n",
    "    cmap = 'gray' if ndim == 3 else None\n",
    "    for i, (xi, axi) in enumerate(zip(x, ax.flatten())):\n",
    "        if ndim == 4 and channels_first: xi = xi.permute(1, 2, 0)\n",
    "        axi.imshow(xi, cmap=cmap)\n",
    "        axi.set_title(f'Batch {i}')\n",
    "        axi.set_xticks([])\n",
    "        axi.set_yticks([])\n",
    "    plt.tight_layout()\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-08-07T20:44:06.848686Z",
     "start_time": "2021-08-07T20:44:06.588084Z"
    }
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 28060 characters]\n",
      "text/plain": [
       "<Figure size 288x864 with 6 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "plot_img_batch(img)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-08-07T20:44:15.933271Z",
     "start_time": "2021-08-07T20:44:15.655929Z"
    }
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 25580 characters]\n",
      "text/plain": [
       "<Figure size 288x864 with 6 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Works on numpy arrays and single channel images too.\n",
    "plot_img_batch(img_bw.numpy())"
   ]
  }
 ],
 "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.3"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": true,
   "toc_window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
