derive_channel_stats.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import jax.numpy as jnp\n",
    "from functools import partial\n",
    "import jax\n",
    "from jax import jit, vmap"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "take a sample of the training set"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(2700, 64, 64, 13)"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import tensorflow_datasets as tfds\n",
    "dataset = tfds.load('eurosat/all', split='train[:10%]', as_supervised=True)\n",
    "for tfe_imgs, tfe_labels in dataset.batch(100_000):\n",
    "    break\n",
    "imgs = jnp.array(tfe_imgs)\n",
    "del tfe_imgs\n",
    "imgs.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "construct a flattened version of the array, we only care about per channel stats"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "flat_imgs = imgs.reshape(-1, 13)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "what is the max and 99.9th percentile per channel?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DeviceArray([10388., 27661., 28000., 28000., 20748., 23288., 25600.,\n",
       "             28002.,  5209.,   183., 16716., 16337., 27534.],            dtype=float32)"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "jnp.max(flat_imgs, axis=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DeviceArray([2612., 3492., 3582., 3970., 3709., 4407., 5678., 5604.,\n",
       "             2264.,   52., 4848., 3876., 6124.], dtype=float32)"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "channel_p999s = jnp.percentile(flat_imgs, 99.9, axis=0)\n",
    "channel_p999s"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "definitely looks like the ~28000 values are sensor noise (?)\n",
    "let's clip at the 99.9th percentile"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "@partial(vmap, in_axes=(1, 0), out_axes=1)\n",
    "def clip_per_channel(x, a_max):    \n",
    "    return jnp.clip(x, a_min=0, a_max=a_max)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "clipped_flat_imgs = clip_per_channel(flat_imgs, channel_p999s)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DeviceArray([2612., 3492., 3582., 3970., 3709., 4407., 5678., 5604.,\n",
       "             2264.,   52., 4848., 3876., 6124.], dtype=float32)"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "jnp.max(clipped_flat_imgs, axis=0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "and now calc channel mean/std for standardisation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(DeviceArray([1354.7904  , 1117.1971  , 1041.1869  ,  946.0516  ,\n",
       "              1198.119   , 2003.2725  , 2375.2615  , 2302.0972  ,\n",
       "               730.6099  ,   12.077799, 1821.995   , 1119.2013  ,\n",
       "              2602.028   ], dtype=float32),\n",
       " DeviceArray([ 242.14961 ,  324.46646 ,  386.99976 ,  587.74664 ,\n",
       "               565.23846 ,  859.7307  , 1086.1215  , 1116.8077  ,\n",
       "               404.7259  ,    4.397278,  998.8627  ,  756.0413  ,\n",
       "              1231.3727  ], dtype=float32))"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "@jit\n",
    "def channel_means_stds(x):    \n",
    "    means = jnp.mean(x, axis=0)\n",
    "    stds = jnp.std(x, axis=0)\n",
    "    return means, stds\n",
    "\n",
    "channel_means, channel_stds = channel_means_stds(clipped_flat_imgs)\n",
    "channel_means, channel_stds"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "can now roll the entire thing into a preprocess; clip -> standardise"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "@partial(vmap, in_axes=(1, 0, 0), out_axes=1)\n",
    "def standardise_per_channel(x, mean, std):        \n",
    "    return (x - mean) / std      \n",
    "\n",
    "@jit\n",
    "def preprocess(x): #, channel_p999s, channel_means, channel_stds):\n",
    "    orig_shape = x.shape\n",
    "    x = x.reshape(-1, 13)\n",
    "    x = clip_per_channel(x, channel_p999s)\n",
    "    x = standardise_per_channel(x, channel_means, channel_stds)    \n",
    "    return x.reshape(orig_shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(5400, 64, 64, 13)"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "std_clipped_imgs = preprocess(imgs) #, channel_p999s, channel_means, channel_stds)\n",
    "std_clipped_imgs.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(DeviceArray([-5.15690566e-08,  1.16206984e-07,  5.12158449e-09,\n",
       "               1.40137146e-07,  1.01107140e-07,  3.64692134e-08,\n",
       "               5.69776262e-08, -1.21019511e-07,  1.09849154e-07,\n",
       "               6.64039916e-08,  1.87202733e-08,  1.57091350e-07,\n",
       "               8.04883484e-08], dtype=float32),\n",
       " DeviceArray([1.        , 1.0000001 , 1.        , 0.99999994, 1.        ,\n",
       "              1.        , 1.        , 0.99999994, 1.        , 1.0000001 ,\n",
       "              0.99999994, 0.9999999 , 1.        ], dtype=float32))"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "channel_means_stds(std_clipped_imgs.reshape(-1, 13))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.6.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}