{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# Supplementary Material to [pytorch-friendly introduction to JAX+TPU](https://www.kaggle.com/roguekk007/pytorch-friendly-comprehensive-jax-tpu-intro/edit)\nLet's dive right in!\n\n####  Set up Kaggle / JAX Runtime (GPU will suffice for this kernel)","metadata":{}},{"cell_type":"code","source":"import os\nimport warnings\nimport tensorflow as tf\n\n!pip install --upgrade jax jaxlib git+https://github.com/deepmind/optax.git flax -q\n    \nif 'TPU_NAME' in os.environ and 'KAGGLE_DATA_PROXY_TOKEN' in os.environ:\n    use_tpu = True\n    !pip install --upgrade jax jaxlib git+https://github.com/deepmind/optax.git flax -q\n    \n    import requests \n    from jax.config import config\n    if 'TPU_DRIVER_MODE' not in globals():\n        url = 'http:' + os.environ['TPU_NAME'].split(':')[1] + ':8475/requestversion/tpu_driver_nightly'\n        resp = requests.post(url)\n        TPU_DRIVER_MODE = 1\n    config.FLAGS.jax_xla_backend = \"tpu_driver\"\n    config.FLAGS.jax_backend_target = os.environ['TPU_NAME']\n    print('Registered (Kaggle) TPU:', config.FLAGS.jax_backend_target)\nelse:\n    use_tpu = False\n!pip install --upgrade git+https://github.com/n2cholas/jax-resnet.git -q","metadata":{"execution":{"iopub.status.busy":"2022-04-07T01:00:12.47714Z","iopub.execute_input":"2022-04-07T01:00:12.47739Z","iopub.status.idle":"2022-04-07T01:00:30.895128Z","shell.execute_reply.started":"2022-04-07T01:00:12.477363Z","shell.execute_reply":"2022-04-07T01:00:30.894299Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import functools\nfrom tqdm.notebook import tqdm\n\nimport jax\nfrom jax import lax \n\nimport flax\nimport flax.linen as nn\nfrom flax.core.frozen_dict import FrozenDict\nfrom flax.training.common_utils import shard, shard_prng_key\nfrom flax.serialization import to_state_dict, from_state_dict,\\\n                        msgpack_serialize, msgpack_restore, from_bytes\n\nimport optax\nimport msgpack\n\nfrom jax_resnet import pretrained_resnet\nfrom jax_resnet.common import slice_variables\n\n# Apologies for shamelessly oppressing the warnings here\nimport warnings\nwarnings.filterwarnings(\"ignore\")","metadata":{"execution":{"iopub.status.busy":"2022-04-07T01:00:30.898725Z","iopub.execute_input":"2022-04-07T01:00:30.898968Z","iopub.status.idle":"2022-04-07T01:00:36.317611Z","shell.execute_reply.started":"2022-04-07T01:00:30.898941Z","shell.execute_reply":"2022-04-07T01:00:36.316936Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Here are some typical model manipulations, and we exemplify their usage in this notebook\n1. Creating a backbone\n2. Remove head from a pretrained backbone & append own head\n\n### To think in the way of JAX / Flax, we must first **decouple a model's architecture (application function) and parameters**. \n\nLet that sink in...\n\nFrom now on, a better (and compiler-tolerant) way to think about `model.forward(x)` is `model_forward_fn(x, params)` where `model_forward_fn` only depends on the model architecture and `params` denotes the model parameters.","metadata":{}},{"cell_type":"code","source":"# Create a dummy input\n# always use the subkey and propagate the original key!\nrandom_key = jax.random.PRNGKey(0)\nrandom_key, subkey = jax.random.split(random_key)\nx = jax.random.normal(subkey, (2, 128, 128, 3))","metadata":{"execution":{"iopub.status.busy":"2022-04-07T01:02:39.325743Z","iopub.execute_input":"2022-04-07T01:02:39.326468Z","iopub.status.idle":"2022-04-07T01:02:39.779104Z","shell.execute_reply.started":"2022-04-07T01:02:39.326428Z","shell.execute_reply":"2022-04-07T01:02:39.77829Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# 1. Create a backbone: \n### Simple example\nWe begin with a simple CNN. \n\nNote that for `nn.Dense` we only specified the `out_channels` in the definition - how would flax know what is the `in_channel`? ","metadata":{}},{"cell_type":"code","source":"import flax\nimport flax.linen as nn\nfrom typing import Sequence\nfrom flax.core import freeze, unfreeze\n\n# Straightforward definition\nclass simpleCNN(nn.Module):\n    # We specify constructing elements here\n    channels : Sequence[int]\n    num_classes : int\n        \n    def setup(self):\n        self.conv_layers = [nn.Conv(channel, kernel_size=(3, 3))\n                            for channel in self.channels]\n        # Note that for linear layer, \n        self.linear_layer = nn.Dense(self.num_classes)\n    \n    def __call__(self, x):\n        for conv_layer in self.conv_layers:\n            x = conv_layer(x)\n            x = nn.relu(x)\n        # Manual average pooling over last 2 dimensions:\n        x = x.mean(-1).mean(-1)\n        x = self.linear_layer(x)\n        return x\n    \nmodel = simpleCNN(channels=[4, 8, 16], num_classes=12)","metadata":{"execution":{"iopub.status.busy":"2022-04-07T01:10:58.352807Z","iopub.execute_input":"2022-04-07T01:10:58.353346Z","iopub.status.idle":"2022-04-07T01:10:58.361564Z","shell.execute_reply.started":"2022-04-07T01:10:58.353311Z","shell.execute_reply":"2022-04-07T01:10:58.360837Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Remarks on Construction\n* Put constructing parameters (`channels`, `num_classes`) with their typing in the class field (in Pytorch this is done in `__init__`)\n* Use `setup` to register submodules, variables, and parameters \n* `__call__` defines the model's behavior, but **to apply the model to data call `model.apply`**\n\n\n## On using the model:\nSince model weights and application are decoupled, we need to:\n* Call `model.init` *with a random key and sample input* to instantiate model parameters\n* Use `apply_fn=model.apply` to do a forward pass on *any* models with the same architecture - we only need to vary `params`. \n* Wait - what about `model.eval(), model.train()` and batch statistics (which should vary with each forward pass?) - see below \"compact\" section","metadata":{}},{"cell_type":"code","source":"random_key, subkey = jax.random.split(random_key)\nparams = model.init(subkey, x)\nprint(\"Parameter:\", params.keys(), type(params))\n\napply_fn = model.apply\nmodel_output = apply_fn(params, x)\nprint(\"Model output:\", model_output.shape)","metadata":{"execution":{"iopub.status.busy":"2022-04-07T01:10:58.65108Z","iopub.execute_input":"2022-04-07T01:10:58.651527Z","iopub.status.idle":"2022-04-07T01:11:00.789765Z","shell.execute_reply.started":"2022-04-07T01:10:58.651493Z","shell.execute_reply":"2022-04-07T01:11:00.788976Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Compact module definition\nThis interface integrates module declaration and application. Replace the **application of a module**  with its **definition**\n\ne.g. Instead of\n```\ndef setup(...\n    self.conv = nn.Conv(16, (3, 3))\n...\ndef __call__(self, x): ...\n    x = self.conv(x)\n``` \ntantamountly do this:\n```\n@nn.compact\ndef __call__(self, x): ...\n    x = nn.Conv(16, (3, 3))(x)\n```\n\nHere is an example:","metadata":{}},{"cell_type":"code","source":"class simpleCNN_compact(nn.Module):\n    # We specify constructing elements here\n    channels : Sequence[int]\n    num_classes : int\n    \n    @nn.compact\n    def __call__(self, x):\n        for i in range(len(self.channels)):\n            x = nn.Conv(self.channels[i], kernel_size=(3, 3))(x)\n            x = nn.relu(x)\n        # Manual average pooling over last 2 dimensions:\n        x = x.mean(-1).mean(-1)\n        x = nn.Dense(self.num_classes)(x)\n        return x\n    \nmodel = simpleCNN_compact(channels=[4, 8, 16], num_classes=12)\nparams = model.init(subkey, x)\nprint(\"Parameter:\", params.keys(), type(params))\n\napply_fn = model.apply\nmodel_output = apply_fn(params, x)\nprint(\"Model output:\", model_output.shape)","metadata":{"execution":{"iopub.status.busy":"2022-04-07T01:11:00.791437Z","iopub.execute_input":"2022-04-07T01:11:00.791854Z","iopub.status.idle":"2022-04-07T01:11:00.967477Z","shell.execute_reply.started":"2022-04-07T01:11:00.791817Z","shell.execute_reply":"2022-04-07T01:11:00.966798Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Manage randomness (dropout) and mutable parameters (batchnorm)\nLet's add `BatchNorm` and `Dropout` to our model: ","metadata":{}},{"cell_type":"code","source":"# Straightforward definition\nclass complexCNN(nn.Module):\n    # We specify constructing elements here\n    channels : Sequence[int]\n    num_classes : int\n    train : bool = True\n    \n    @nn.compact\n    def __call__(self, x):\n        for i in range(len(self.channels)):\n            x = nn.Conv(self.channels[i], kernel_size=(3, 3))(x)\n            x = nn.Dropout(rate=0.5, deterministic=not self.train)(x)\n            x = nn.BatchNorm(use_running_average=not self.train)(x)\n            x = nn.relu(x)\n        # Manual average pooling over last 2 dimensions:\n        x = x.mean(-1).mean(-1)\n        x = nn.Dense(self.num_classes)(x)\n        return x\n    \nmodel = complexCNN(channels=[4, 8, 16], num_classes=12)","metadata":{"execution":{"iopub.status.busy":"2022-04-07T01:11:04.997026Z","iopub.execute_input":"2022-04-07T01:11:04.997716Z","iopub.status.idle":"2022-04-07T01:11:05.006605Z","shell.execute_reply.started":"2022-04-07T01:11:04.997679Z","shell.execute_reply":"2022-04-07T01:11:05.005891Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### 🔪Why `BatchNorm` and `Dropout` are problematic in JAX🔪\n* `Dropout` and `Batchnorm` **behave differently during training & evaluation**\n* `Dropout` introduces stochastic behavior which we wish to control\n* `Batchnorm` introduces running statistics: they change with each forward pass during training. ","metadata":{}},{"cell_type":"code","source":"# Generate a random key for the forward pass\nrandom_key, key1, key2, key3 = jax.random.split(random_key, 4)\n\nparams = model.init({'params': key1, 'dropout': key2}, x)\nprint(\"Parameter:\", params.keys(), type(params))\nparams = unfreeze(params)\n\nx_out, params['batch_stats'] = model.apply(params, x, rngs={'dropout': key3},\n                    mutable='batch_stats')\nparams = freeze(params)\nprint(\"Model output:\", x_out.shape)","metadata":{"execution":{"iopub.status.busy":"2022-04-07T01:12:23.457937Z","iopub.execute_input":"2022-04-07T01:12:23.458234Z","iopub.status.idle":"2022-04-07T01:12:23.707753Z","shell.execute_reply.started":"2022-04-07T01:12:23.458206Z","shell.execute_reply":"2022-04-07T01:12:23.706958Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Wow! That's quite a mouthful, let's break this down. \n\n#### Now our `params` contain a second key field `batch_stats`! \n### Remarks:\n1. For dropout, provide a random key dictionary `{'params': key1, 'dropout': key2}` during `model.init`\n2. Provide a random key dictionary `{'params': key1, 'dropout': key2}` (don't forget to split your keys each time!)\n3. Specify `mutable=batch_stats` during the forward pass to modify the running stats, the `apply` function returns the altered version of `batch_stats`\n4. **Introduce a state `train` which alters the behavior of training / evaluation forward pass** Alternatively, we can assume that we're in `eval` state whenever nothing is mutable (see the comprehensive intro notebook)\n5. Flax manages parameters in `FrozenDicts`, we need to `unfreeze` before forward pass and `freeze` afterwards\n\nAlso refer to [this excellent source](https://colab.research.google.com/github/google/flax/blob/main/docs/notebooks/linen_intro.ipynb#scrollTo=BBrbcEdCnQ4o) for more information.\nLet us test that the forward function behaves as expected: we use the same input across four runs with different rng keys","metadata":{}},{"cell_type":"code","source":"import jax.numpy as jnp\n\nrandom_key, key1, key2, key3, key4 = jax.random.split(random_key, 5)\n# training mode\nmodel.train = True\nparams = unfreeze(params)\nx_out1, params1 = model.apply(params, x, rngs={'dropout':key1}, mutable='batch_stats')\nx_out2, params2 = model.apply(params, x, rngs={'dropout':key2}, mutable='batch_stats')\nparams = freeze(params)\n# Compare the two outputs\nprint('Difference between outputs for training forward pass:', jnp.abs(x_out1 - x_out2).sum().item())\n\n# Evaluation mode:\nmodel.train = False\nparams = unfreeze(params)\nx_out1, params1 = model.apply(params, x, rngs={'dropout':key1}, mutable='batch_stats')\nx_out2, params2 = model.apply(params, x, rngs={'dropout':key2}, mutable='batch_stats')\nparams = freeze(params)\n# Compare the two outputs\nprint('Difference between outputs for evaluation forward pass:', jnp.abs(x_out1 - x_out2).sum().item())","metadata":{"execution":{"iopub.status.busy":"2022-04-07T01:13:07.105851Z","iopub.execute_input":"2022-04-07T01:13:07.10611Z","iopub.status.idle":"2022-04-07T01:13:07.718352Z","shell.execute_reply.started":"2022-04-07T01:13:07.106083Z","shell.execute_reply":"2022-04-07T01:13:07.716745Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"!pip install --upgrade git+https://github.com/n2cholas/jax-resnet.git -q","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Head Surgery\nWe have access to imagenet-pretrained backbone with 3-channel input and 1000-class output, but we want a 6-channel model with 10-class output, what should we do?\n\nThe problem here, again, is that weights are decoupled from application, so we simply declare a module within another as in Pytorch","metadata":{}},{"cell_type":"code","source":"# Let us grab a backbone and truncate it\nfrom jax_resnet import pretrained_resnet, pretrained_resnest\n\nresnet_template, backbone_params = pretrained_resnest(50)\nbackbone = resnet_template()\nbackbone = nn.Sequential(backbone.layers[:-2])\n\n# The number of in_channels and num_classes we need\nin_channels = 6\nnum_classes = 24\n\nrandom_key, k1, k2 = jax.random.split(random_key, 3)\nx_3 = jax.random.normal(k1, (1, 128, 128, 3))\nx_6 = jax.random.normal(k2, (1, 128, 128, 6))","metadata":{"execution":{"iopub.status.busy":"2022-04-07T01:38:29.917744Z","iopub.execute_input":"2022-04-07T01:38:29.918466Z","iopub.status.idle":"2022-04-07T01:38:30.000657Z","shell.execute_reply.started":"2022-04-07T01:38:29.91843Z","shell.execute_reply":"2022-04-07T01:38:29.999932Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Let's start with head surgery\n\nWe begin by creating a model wrapper which \"wraps around the backbone\"","metadata":{}},{"cell_type":"code","source":"class AddHeadtoBackbone(nn.Module):\n    backbone : nn.Sequential\n    num_classes : int\n        \n    def setup(self):\n        self.head = nn.Dense(num_classes)\n    \n    def __call__(self, x):\n        x = self.backbone(x)\n        # Avg pooling\n        x = x.mean(-1).mean(-1)\n        return self.head(x)\n        \nmodel = AddHeadtoBackbone(backbone=backbone, num_classes=num_classes)","metadata":{"execution":{"iopub.status.busy":"2022-04-07T01:38:30.447935Z","iopub.execute_input":"2022-04-07T01:38:30.44879Z","iopub.status.idle":"2022-04-07T01:38:30.455467Z","shell.execute_reply.started":"2022-04-07T01:38:30.44873Z","shell.execute_reply":"2022-04-07T01:38:30.45447Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"The next step in creating a model in Flax is always to initialize it.\n\nNote: by providing a sample input, we do not even need to know how many \nchannels the backbone output","metadata":{}},{"cell_type":"code","source":"random_key, subkey = jax.random.split(random_key)\nparams = unfreeze(model.init(subkey, x_3))\nparams['params'].keys(), params['batch_stats'].keys()","metadata":{"execution":{"iopub.status.busy":"2022-04-07T01:38:31.22145Z","iopub.execute_input":"2022-04-07T01:38:31.221699Z","iopub.status.idle":"2022-04-07T01:38:32.574598Z","shell.execute_reply.started":"2022-04-07T01:38:31.221673Z","shell.execute_reply":"2022-04-07T01:38:32.573721Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"We next load the pretrained weights into `params`, to do this call `unfreeze` first","metadata":{}},{"cell_type":"code","source":"params['params']['backbone'] = backbone_params['params']\nparams['batch_stats']['backbone'] = backbone_params['batch_stats']\nparams = freeze(params)","metadata":{"execution":{"iopub.status.busy":"2022-04-07T01:38:32.576388Z","iopub.execute_input":"2022-04-07T01:38:32.576666Z","iopub.status.idle":"2022-04-07T01:38:32.582802Z","shell.execute_reply.started":"2022-04-07T01:38:32.57663Z","shell.execute_reply":"2022-04-07T01:38:32.581986Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Let's try a forward pass\nx_out = model.apply(params, x_3, mutable=False)\n# Great! we have what we want\nx_out.shape, x_out.dtype","metadata":{"execution":{"iopub.status.busy":"2022-04-07T01:38:43.512652Z","iopub.execute_input":"2022-04-07T01:38:43.513335Z","iopub.status.idle":"2022-04-07T01:38:45.957589Z","shell.execute_reply.started":"2022-04-07T01:38:43.513299Z","shell.execute_reply":"2022-04-07T01:38:45.956917Z"},"trusted":true},"execution_count":null,"outputs":[]}]}