{"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":" # TPU Flower Classification with TFRecord & SPMD JAX training\n ### A comprehensive pipeline for JAX+multi-core TPU\nHi Kagglers! Here is a small piece of contribution to Kaggle's ongoing [Google Open Source Software Experts](https://www.kaggle.com/google-oss-expert-prize-winners) Challenge. \n\nUntil coming across Google's TPU and JAX, I have been a long-time Pytorch+GPU user. While Pytorch's **imperative** programming paradigm makes manipulating tensors easier at the cost of making analysis difficult, JAX's **functional** programming paradigm makes magic such as [auto-vectorization](https://jax.readthedocs.io/en/latest/jax-101/03-vectorization.html), [easy manipulation of higher-order derivatives](https://jax.readthedocs.io/en/latest/jax-101/04-advanced-autodiff.html), native integration of [JIT](https://jax.readthedocs.io/en/latest/jax-101/02-jitting.html), and [seamless composition between these functionalities](https://jax.readthedocs.io/en/latest/notebooks/quickstart.html) accessible.\n\n![Image](https://pbs.twimg.com/media/EobMOrBVEAAvrSf?format=jpg&name=large)\n*Image Credit: James Bradbury's Twitter two years ago [here](!https://twitter.com/jekbradbury/status/1335001804732395522)\n\nJAX and TPU go *particularly* well together. JAX uses the XLA compiler (which underlies TPU usage) and is by design choice easier to analyze (and optimize). The combination has driven cut-edge research, as evidenced by [this paper](https://arxiv.org/pdf/2011.03641) and [deepmind's blog post](https://www.deepmind.com/blog/using-jax-to-accelerate-our-research).   \nBuilding on wonderful introductions to JAX by [Nilay Chauhan](https://www.kaggle.com/discussions/getting-started/315696) (general introduction), [aakashnain](https://www.kaggle.com/code/aakashnain/building-models-in-jax-part1-stax/notebook) (STAX library) and [Sanyam Bhutani](https://www.kaggle.com/getting-started/308753) (JAX 201), \nthis notebook hopes to provide those (like myself) who have experience with Pytorch+GPU with a start in the next generation of deep learning technology.\n\\\nDo you still remember the panic and chaos that is pytorch DDP / XLA (and worse so, both) - and did I forget to mention that **switching between single/multiple GPU/TPU** for a JAX pipeline is **as easy as changing the Kaggle runtime**? Try it out!\n\nThis introduction series is the product of the past two weeks I spent with JAX in hopes that it'll make this wonderful library accessible to everyone\n### This introduction series features\n#### 1. How to perform [common flax model surgery](https://www.kaggle.com/roguekk007/flax-model-surgery/)\n#### 2. How to define an SOTA computer-vision backbone in `flax`.\n#### 3. Port pretrained weights from pytorch implementation for [EfficientnetV2](https://www.kaggle.com/code/roguekk007/efficientnetv2-jax) (the first, to the best of my knowledge).\n#### 4. (This notebook) A highly-efficient general-purpose TPU training pipeline in JAX\n\n![image.png](https://storage.googleapis.com/kaggle-competitions/kaggle/21154/logos/header.png?t=2020-06-04-00-33-35\")\n\n## Contents\n#### 1. Set up kagggle TPU / JAX runtime\n#### 2. TFRecords: Don't make data the bottleneck\n#### 3. Flax+Optax: [Model surgery](https://www.kaggle.com/roguekk007/flax-model-surgery/) and optimization\n#### 4. Single-Process Multi-Data (SPMD) TPU training using JAX\n#### 5. Inference\n\nWe will be using the beautiful [Flower Classification on TPU](https://www.kaggle.com/competitions/tpu-getting-started) dataset to exemplify a general CV classification pipeline. Feel free to create wonders with it - and don't forget to upvote!","metadata":{}},{"cell_type":"markdown","source":"# 1. Set up Kagggle TPU / JAX Runtime\n\nAs mentioned [here](https://www.kaggle.com/discussions/getting-started/315696), JAX is not natively supported by Kaggle TPU Runtime (yet. Kaggle what are you waiting for?) so we'll do some setup - credits to [this kernel](https://www.kaggle.com/code/alexlwh/happywhale-flax-jax-tpu-gpu-resnet-baseline). We do not have to do this on cloud TPU in general.","metadata":{}},{"cell_type":"code","source":"import os\nimport warnings\nimport tensorflow as tf\nos.environ['TF_XLA_FLAGS'] = '--tf_xla_enable_xla_devices'\n\n# Seems like the new one will break jax\n!pip install --upgrade jax jaxlib==0.3.5 git+https://github.com/deepmind/optax.git flax optax -q\n    \nif 'TPU_NAME' in os.environ and 'KAGGLE_DATA_PROXY_TOKEN' in os.environ:\n    use_tpu = True\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    # Enforce bfloat16 multiplication\n    config.update('jax_default_matmul_precision', 'bfloat16')\n    print('Registered (Kaggle) TPU:', config.FLAGS.jax_backend_target)\nelse:\n    use_tpu = False","metadata":{"execution":{"iopub.status.busy":"2022-04-21T22:41:39.48498Z","iopub.execute_input":"2022-04-21T22:41:39.485462Z","iopub.status.idle":"2022-04-21T22:42:40.794621Z","shell.execute_reply.started":"2022-04-21T22:41:39.485365Z","shell.execute_reply":"2022-04-21T22:42:40.793503Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import functools\nfrom tqdm.notebook import tqdm\n\nimport jax\nimport jax.numpy as jnp\nfrom jax import lax \n\nimport flax\nimport flax.linen as nn\nfrom flax.core import freeze, unfreeze\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\nfrom jax_efficientnetv2 import efficientnet_v2_pretrained\n\n# Apologies for shamelessly oppressing the warnings here\nimport warnings\nwarnings.filterwarnings(\"ignore\")","metadata":{"execution":{"iopub.status.busy":"2022-04-21T22:42:40.79655Z","iopub.execute_input":"2022-04-21T22:42:40.796822Z","iopub.status.idle":"2022-04-21T22:42:41.08605Z","shell.execute_reply.started":"2022-04-21T22:42:40.796795Z","shell.execute_reply":"2022-04-21T22:42:41.085087Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Specify Training Arguments\n\nNote we specify `batch_size` now for *a single core* and the effective batch size for training is `device_count * batch_size`. ","metadata":{}},{"cell_type":"code","source":"args = {\n    'experiment_name': 'starter',\n    # Efficientnetv2-m\n    'model': 'm',\n    'batch_size': 8, \n    'epochs': 16,\n    'base_lr': 7e-5, # should directly correspond to `batch_size`\n    \n    # Data / augmentation \n    'img_size': 512, # 192, 224, 331, 512\n    # The actual label number is 104\n    'num_labels': 104 if not use_tpu else 128,\n    'saving_dir': '/kaggle/working/',\n    'device_count' : jax.device_count(),\n    \n    # Debugging-purposes\n    'sanity_check': False,\n}\n\n# The effective lr should linearly scale with batch size\nargs['lr'] = args['base_lr'] * args['device_count']\nif args['sanity_check']:\n    args['epochs'] = 1\n\n# Data-specific, should change for each dataset \nargs['data_dir'] = '/kaggle/input/tpu-getting-started/tfrecords-jpeg-'\\\n                    + str(args['img_size'])+'x'+str(args['img_size'])\nargs['train_dir'] = os.path.join(args['data_dir'], 'train')\nargs['val_dir'] = os.path.join(args['data_dir'], 'val')\nargs['test_dir'] = os.path.join(args['data_dir'], 'test')\n\nprint('Running on', args['device_count'], 'processors')\nprint(jax.devices())","metadata":{"execution":{"iopub.status.busy":"2022-04-21T22:42:41.091101Z","iopub.execute_input":"2022-04-21T22:42:41.091415Z","iopub.status.idle":"2022-04-21T22:43:05.48399Z","shell.execute_reply.started":"2022-04-21T22:42:41.091375Z","shell.execute_reply":"2022-04-21T22:43:05.483003Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# 2. Dataloading using TFRecords\nTo be honest, Pytorch Dataloader -> TFRecords is a hard but worthwhile switch. For those used to freely playing around with tensors and numpy arrays (at the cost of huge uneliminatable CPU overhead), `tf.data` provides a more efficient pipeline. Here is an [excellent introduction](https://towardsdatascience.com/a-practical-guide-to-tfrecords-584536bc786c). Please refer to this [notebook on this dataset](https://www.kaggle.com/code/ryanholbrook/tfrecords-basics/notebook#Parsing-Serialized-Examples); also refer to [this link](https://www.kaggle.com/code/yihdarshieh?scriptVersionId=41561359&cellId=29) for tips on `tf.data.TFRecordDataset`. The code below are integrated from these sources.\n\n### Note on TFRecord: TFRecord addresses data bottleneck by:\n* Parallelize I/O by sharding data across files.\n* I/O pre-fetching on large files\n* Rule of thumb from [the official docs](https://www.tensorflow.org/tutorials/load_data/tfrecord): create at least 10 * N files for N hosts, and each file should be ideally 100MB+\n\n### Different data pipelines\n* `PIL / OpenCV image` -> `Pytorch Dataset` -> `DataLoader`\n* `tfrecord` - (specify decoding) -> `TFRecordDataset` - (batching specifications) -> `DataLoader`\n\nThe following code block defines `read_labeled_tfrecord` and `read_unlabeled_tfrecord` to specify encoding. They are later used as mapping functions on `TFRecordsDataset` to parse the data.","metadata":{}},{"cell_type":"code","source":"# Given an `tf.string`, returns a legible tensor\ndef decode_image(image_data):\n    # image is of type `tf.uint8`\n    image = tf.image.decode_jpeg(image_data, channels=3)\n    image = tf.reshape(image, [args['img_size'], args['img_size'], 3]) \n    return image\n\n# parses general type example, returns data sample (tuple of (image : ), (label : ))\ndef read_labeled_tfrecord(example):\n    # Note how we are defining the example structure here\n    labeled_format = {\n        \"image\": tf.io.FixedLenFeature([], tf.string), # tf.string means bytestring\n        \"class\": tf.io.FixedLenFeature([], tf.int64),  # shape [] means single element\n    }\n    parsed_example = tf.io.parse_single_example(example, labeled_format)\n    image = decode_image(parsed_example['image'])\n    label = tf.cast(parsed_example['class'], tf.int32)\n    return {'image': image, 'label': label} \n\ndef read_unlabeled_tfrecord(example):   \n    unlabeled_format = {\n        \"image\": tf.io.FixedLenFeature([], tf.string), \n        \"id\": tf.io.FixedLenFeature([], tf.string),  \n    }\n    parsed_example = tf.io.parse_single_example(example, unlabeled_format)\n    image = decode_image(parsed_example['image'])\n    idnum = parsed_example['id']\n    return {'image': image, 'id': idnum}\n\ntrain_filenames = tf.io.gfile.glob(args['train_dir']+'/*.tfrec')\nval_filenames = tf.io.gfile.glob(args['val_dir']+'/*.tfrec')\ntest_filenames = tf.io.gfile.glob(args['test_dir']+'/*.tfrec')","metadata":{"execution":{"iopub.status.busy":"2022-04-21T22:43:05.485524Z","iopub.execute_input":"2022-04-21T22:43:05.486423Z","iopub.status.idle":"2022-04-21T22:43:05.532005Z","shell.execute_reply.started":"2022-04-21T22:43:05.486375Z","shell.execute_reply":"2022-04-21T22:43:05.530904Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Batching and Augmentation\nBelow are batching operations such as `.batch`, `.prefetch`. Use `to_jax` to convert `tf.Tensor` data to jax datatype and `shard` them across computational cores in the last step. Here is a visible walkthrough, also see comments below. Augmentation code adopted from [here](https://gist.github.com/sayakpaul/e0024bae08afcd3d75b6d52fda191025).\n\n* `to_jax`: Applied at the last step of data pipeline to convert TF tensor to jax tensors\n* `tf_randaugment`, `normalize`: image processing / augmentation. `tf_randaugment` exemplifies how integrate *any* python function into the `tf.data` pipeline via `tf.py_function` (also see the [helpful doc](https://www.tensorflow.org/api_docs/python/tf/py_function))\n* `load_dataset`: Define the data pipeline, see comments\n\n**Remarks:** `dataset.map` is the key to flexibility! With familiarity I find `tf.data` API no less flexible than Pytorch","metadata":{}},{"cell_type":"code","source":"import tensorflow_datasets as tfds\n\ninput_dtype = jnp.bfloat16 if use_tpu else jnp.float32\nlabel_dtype = jnp.int16\n    \n# Convert a datasample to JAX-pipeline-compatible format. last step\ndef to_jax(sample):\n    sample['image'] = jnp.array(sample['image'], dtype=input_dtype)\n    sample['label'] = jnp.array(sample['label'], dtype=label_dtype)\n    # Convert labels to one_hot\n    sample['label'] = jax.nn.one_hot(sample['label'], args['num_labels'], dtype=label_dtype, axis=-1)\n    return shard(sample)\n\n# Augmentation\nfrom imgaug import augmenters as iaa\naug = iaa.RandAugment(n=2, m=15)\ndef tf_randaugment(sample):\n    augment_fn = lambda img: aug(images=img.numpy())\n    im_shape = sample['image'].shape\n    [sample['image'],] = tf.py_function(augment_fn, [sample['image']], [tf.float32])\n    sample['image'].set_shape(im_shape)\n    return sample\n\ndef normalize_and_resize(sample):\n    sample['image'] = tf.cast(sample['image'], tf.float32) / 128. - 1.\n    return sample\n    \ndef load_dataset(filenames, labeled=True, ordered=False, shuffle_buffer_size=1, drop_remainder=False,\\\n                augment=True):\n    AUTO = tf.data.experimental.AUTOTUNE\n    # tf.data runtime will optimize this parameter\n    dataset = tf.data.TFRecordDataset(filenames, num_parallel_reads=AUTO)\n\n    options = tf.data.Options()\n    if not ordered:\n        options.experimental_deterministic = False\n        \n    # Step 1: Read in the data, shuffle and batching\n    dataset = dataset.with_options(options)\\\n                .map(read_labeled_tfrecord if labeled else read_unlabeled_tfrecord,\n                    num_parallel_calls=tf.data.experimental.AUTOTUNE)\\\n                .shuffle(shuffle_buffer_size)\\\n                .batch(args['batch_size'] * args['device_count'], drop_remainder=drop_remainder)\n    \n    # We exemplify augmentation using RandAugment\n    if augment:\n        dataset = dataset.map(tf_randaugment, num_parallel_calls=AUTO)\n    # Add `prefetch` at the last step to parallize as much as possible!\n    dataset = dataset.map(normalize_and_resize).prefetch(AUTO)\n    # Finally, apply to_jax transformation\n    return map(to_jax, tfds.as_numpy(dataset))","metadata":{"execution":{"iopub.status.busy":"2022-04-21T22:43:05.533802Z","iopub.execute_input":"2022-04-21T22:43:05.534505Z","iopub.status.idle":"2022-04-21T22:43:07.279047Z","shell.execute_reply.started":"2022-04-21T22:43:05.534458Z","shell.execute_reply":"2022-04-21T22:43:07.277967Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# 3. Flax & Optax: Model Surgery and Optimization\n\nSince JAX **strictly decouples model state (parameters) and the pass-forward function**, performing surgery on pretrained models is a big challenge. I find [flax](https://flax.readthedocs.io/en/latest/) as flexible for model-building after getting used to its design, and [optax](https://optax.readthedocs.io/en/latest/) is a trivial extension in most usage cases.\n\n### Model Surgery: Customize your Backbone\nPlease refer to my **[Flax Model Surgery](https://www.kaggle.com/roguekk007/flax-model-surgery/edit)** kernel for more resources and examples on this section. The following topics are addressed in that kernel, they are used in the code block below\n1. Defining using `nn.Module`\n2. Instantiating model parameters separately using `model.init`\n3. Stochastic / Mutable behaviors in `nn.Dropout` and `nn.BatchNorm`\n4. Common model surgery\n\nIn [this utility script](https://www.kaggle.com/code/roguekk007/jax-efficientnetv2), I implemented EfficientnetV2 in flax - go check it out and compare to pytorch implementations! Since JAX is still young and many pretrained weights are in pytorch, I spent some time to port pretrained EfficientnetV2 weights to JAX/flax model - check out [**this notebook on porting pretrained weights**](https://www.kaggle.com/code/roguekk007/efficientnetv2-jax/notebook)\n\nIn this notebook, we make the (reasonable) assumption that we need no stochasticity when there are no mutable states (see `__call__` below)\n\n### 🔪 The Sharp Bits: Models on TPU 🔪 ###\n**Very Important**: Whenever possible, **make *everything* a multiple of 128** (at least 8) when working with TPU's! \n\nTPU's underlying architecture does matrix multiplication in these sizes. \n\nIf you don't manually pad the tensors, XLA will pad them for you and there can be *absolutely horrendous* paddings\n* As an example, invoking a `512 x 81313` `NormalizedLinear` layer with our current model will cause TPU memory overflow, which is totally unreasonable because such a layer is comparatively little computation. This is because during norm calculation, TPU somehow ends up with a `[Very very big size] * 2` matrix and wants to pad it to `[Very very big size] * 8`. Simply maually padding our layer to `512 x 81408` will fix this problem. This problem this took me 3 hours to debug","metadata":{}},{"cell_type":"code","source":"from functools import partial\nfrom typing import Any\n\n# A very minimal head, really\n# Retrieve output, perform average pooling, then output\nclass AddHeadtoBackbone(nn.Module):\n    backbone: Any\n    num_features : int\n    dtype : Any = input_dtype\n    \n    @nn.compact \n    def __call__(self, x):\n        mutable = self.is_mutable_collection('batch_stats')\n        x = self.backbone(x)\n        x = jax.nn.swish(jnp.mean(x, axis=(1, 2)))\n        return nn.Dense(self.num_features, use_bias=False, param_dtype=self.dtype, dtype=self.dtype)(x)\n\n# The random_key will be \"consumed\" by this function so pass a subkey\ndef get_model(random_key):\n    pretrained_weights_path = f'/kaggle/input/efficientnetv2-jax/efficientnetv2-{args[\"model\"]}.msgpack'\n    # Go check the utility script out!\n    backbone, backbone_params = efficientnet_v2_pretrained(args['model'], pretrained_weights_path, \\\n                                                           input_size=args['img_size'], dtype=input_dtype, verbose=True)\n    # Add our own head and instantiate model parameters\n    model = AddHeadtoBackbone(backbone=backbone, num_features=args['num_labels'])\n    dummy_inputs = jnp.ones((1, args['img_size'], args['img_size'], 3), input_dtype)  \n    random_key, key1, key2 = jax.random.split(random_key, 3)  \n    params = unfreeze(model.init({'params': key1, 'dropout': key2}, dummy_inputs))\n    \n    # Replace the backbone portion of `params` with pretrained weights\n    params['params']['backbone'] = backbone_params['params']\n    params['batch_stats']['backbone'] = backbone_params['batch_stats']\n    return model, params","metadata":{"execution":{"iopub.status.busy":"2022-04-21T22:43:07.280714Z","iopub.execute_input":"2022-04-21T22:43:07.281048Z","iopub.status.idle":"2022-04-21T22:43:07.298719Z","shell.execute_reply.started":"2022-04-21T22:43:07.281003Z","shell.execute_reply":"2022-04-21T22:43:07.297557Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Initialize Model and Optimizer\nOptax optimizer is fairly straightforward: \n1. Initialize with parameters we wish to optimize **(don't pass in `batch_stats`)**. \n2. Use `train_state` to aggregate model and optimizer parameter states\n3. Use `flax.jax_utils.replicate` to broadcast `train_state` to each process\n\nThe most relevant usage of `Optax.optimizer`:\n```\n### Define model and `optimizer`\nop_states = optimizer.init(model_params)\n\n### calculate `grads`\nupdates, op_states = optimizer.update(grads, op_states, model_params)\nnew_model_states = optax.apply_updates(model_states, updates)\n```","metadata":{}},{"cell_type":"code","source":"random_key = jax.random.PRNGKey(0)\nrandom_key, subkey = jax.random.split(random_key)\n\n# This subkey is \"consumed\" by `get_model`\nmodel, model_params = get_model(subkey)\n# Cast to training dtype\nmodel_params = jax.tree_map(lambda x : x.astype(input_dtype), model_params)\n\nscheduler = optax.constant_schedule(args['lr'])\noptimizer = optax.chain(\n  optax.clip(1.0),\n  optax.adamw(learning_rate=scheduler, weight_decay=1e-4))\n\n# Don't throw the whole model_params in there!! else we'll be optimizing running statistics\nop_params = optimizer.init(model_params['params'])\ntrain_state = {'model': model_params, 'op': op_params}\n# Broadcast the train state across cores (analagously done for data via \"shard\")\n# Consider: override the name to save memory\npl_train_state = flax.jax_utils.replicate(train_state)","metadata":{"execution":{"iopub.status.busy":"2022-04-21T22:43:07.300291Z","iopub.execute_input":"2022-04-21T22:43:07.300603Z","iopub.status.idle":"2022-04-21T22:45:44.642796Z","shell.execute_reply.started":"2022-04-21T22:43:07.300565Z","shell.execute_reply":"2022-04-21T22:45:44.641632Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# 4. A JAX Training Pipeline\n\n### What's at stake?\nA JAX training pipeline like this offers:\n* Seamless switch between XLA devices (CPU / GPU / TPU).\n* Easy device parallelism - due to JAX's functional design.  While we need to take more care in writing JAX programs, the reward is running the *same code* on 1 GPU or 8 TPUs. \n* No more! are the days of illegible pytorch DDP and cryptic yet irresolvable TPU hanging - the final straw which motivated this project.\n\n## 🔪 [The Sharp bits](https://jax.readthedocs.io/en/latest/notebooks/Common_Gotchas_in_JAX.html)  🔪\n#### An (incomprehensive) list of remarks (and bugs which took a long time):\n1. Stochastic processes (drop-out, initialization) always consume random seeds - provide them\n    * At each step, *split* your key and feed one to `training_step` via `model.apply(...,rngs={'dropout':subkey})`\n    * `model.apply` *will not* require rng in eval mode - provide and check for this!\n2. Stateful changes with each forward pass: `BatchNorm` has never been such a pain\n    * The initialized parameters for each mode contains two fields: `params` and `batch_stats`\n    * Do NOT use the gradients for `batch_stats`. They're not meant to be optimized.\n3. Multi-process computation: sharding and syncing\n    * Use `shard` to scatter a data batch (multiple of device count) across processes, use `shard_prng` for rng keys\n    * `replicate` and `unreplicate` helps you manage model states\n    * Use `lax.pmean` to sync data (loss, gradients, batch stats) across processes - but using `pmean` in a function means it can only be used when composed with `pmap`\n4. `jax.tree_map` is your best friend: broadcast a function on a parameters, dictionary, lists - you name it\n5. Carefully differentiate the behavior during training and testing (validation):\n    * Don't be sloppy in managing `mutable=False` during validation, don't provide RNG unless you *have to* - they'll come and bite you in the back\n    * `FrozenDicts` from flax throws an error whenever you try to change it, `unfreeze/freeze` the state immediately before/after training for good practice","metadata":{}},{"cell_type":"code","source":"from pipeline_utils_jax import accuracy_fn\nfrom jax.scipy.special import logsumexp\n\ncross_entropy = lambda logits, labels : -(jax.nn.log_softmax(logits * 16) * labels).sum(-1).mean()\nmetric_fn = accuracy_fn\n\nmetric_fn = accuracy_fn\ncriterion = cross_entropy\n\ndef training_step(apply_fn, update_fn, train_state, batch, subkey):\n    def calculate_loss(model_params, apply_fn, batch):\n        # Here the model *is* mutable!! set mutable='batch_stats' to enable dropout behavior\n        # The mutable state is {'batch_stats': ...} (this bug took me two hours)\n        logits, mutable_states = apply_fn(model_params, batch['image'], \\\n                                    mutable='batch_stats', rngs={'dropout': subkey})\n        model_params['batch_stats'] = mutable_states['batch_stats']\n        loss = criterion(logits, batch['label'])\n        # Sync batch_stats and loss across devices\n        model_params['batch_stats'] = jax.tree_map(functools.partial(lax.pmean, axis_name='devices'), \\\n                                                         model_params['batch_stats'])\n        loss = lax.pmean(criterion(logits, batch['label']), 'devices')\n        return loss, (logits, model_params)\n    \n    grad_fn = jax.value_and_grad(calculate_loss, has_aux=True)\n    (loss, (logits, train_state['model'])), grads = grad_fn(train_state['model'], apply_fn, batch)\n    # Only take the gradients for parameters, and sync across devices\n    grads = lax.pmean(grads['params'], 'devices')\n    \n    updates, train_state['op'] = update_fn(grads, train_state['op'], train_state['model']['params'])\n    train_state['model']['params'] = optax.apply_updates(train_state['model']['params'], updates)\n    \n    metric_value = lax.pmean(metric_fn(logits, batch['label']), 'devices')\n    return {'train_state': train_state, 'loss': loss, 'metric':metric_value}\n\ndef val_step(apply_fn, model_params, batch):\n    # Do not enable batch statistic update (mutable=False) ^or^ stochastic behavior\n    # We do not need a rng on this forward pass because Dropout is not stochastic anymore\n    logits = apply_fn(model_params, batch['image'], mutable=False)\n    loss = lax.pmean(criterion(logits, batch['label']), 'devices')\n    # Calculate metric\n    metric_value = lax.pmean(metric_fn(logits, batch['label']), 'devices')\n    return {'loss': loss, 'metric':metric_value}\n\ntrain_step_parallel = jax.pmap(functools.partial(training_step, apply_fn=model.apply, update_fn=optimizer.update),\\\n         axis_name='devices')\nval_step_parallel = jax.pmap(functools.partial(val_step, apply_fn=model.apply),\\\n         axis_name='devices')","metadata":{"execution":{"iopub.status.busy":"2022-04-21T22:45:44.648949Z","iopub.execute_input":"2022-04-21T22:45:44.649305Z","iopub.status.idle":"2022-04-21T22:45:44.693999Z","shell.execute_reply.started":"2022-04-21T22:45:44.649263Z","shell.execute_reply":"2022-04-21T22:45:44.692758Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### A note on `lax.pmap`\nThis **is** the magic function which makes parallelization trivial. \n\nWe wrote `train_step(apply_fn, update_fn, train_state, batch, subkey)` and `val_step` *exactly as we would have written for single processor*. The only changes are:\n* Use `lax.pmean` for syncing necessary information - the cores host separate models and data, they can easily run off on their own\n* `apply_fn, update_fn` are arguments which do not vary across processes, so we wrap them in `functools.partial`\n* The rest of the arguments **must all be sharded / replicated** else there'll be error. Wrap process-unspecific elements into `partial` or manually shard them.\n* Name the pmap-ped axes (I like to call it `devices`) and use this name for `pmean(..., axis_name='devices')` for cross-device syncing","metadata":{}},{"cell_type":"code","source":"def save_checkpoint(log_path, logs, verbose=True):\n    with open(log_path, \"wb\") as outfile:\n        outfile.write(msgpack_serialize(to_state_dict(logs)))\n    if verbose:\n        print(\"Checkpoint written to:\", log_path)\n        \ndef load_checkpoint(log_path, init_log, verbose=True):\n    # Read msgpack file\n    with open(log_path, \"rb\") as data_file:\n        byte_data = data_file.read()\n    if verbose:\n        print('Checkpoint retrieved from:', log_path)\n    return from_bytes(init_log, byte_data)","metadata":{"execution":{"iopub.status.busy":"2022-04-21T22:45:44.697423Z","iopub.execute_input":"2022-04-21T22:45:44.697896Z","iopub.status.idle":"2022-04-21T22:45:44.706778Z","shell.execute_reply.started":"2022-04-21T22:45:44.697854Z","shell.execute_reply":"2022-04-21T22:45:44.705737Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Boilerplate Pipeline Code\nI copied this line-for-line from my pytorch pipeline, the only differences are:\n* Create a new random key every step\n* Generate dataloaders anew at every epoch\n* Manage `freeze` / `unfreeze` meticulously\n\n### And that's it! All there's needed for easy parallelism to work in JAX\n#### Several notes on congratulating ourselves:\n1. We're running *very efficient training* on TPU via this pipeline!! We're effectively running a model **the size of `efficientnet-b7`** on **batch size of `64`** with an **image size of `512x512`** (!!!) TFRecords is a necessity to endure this throughput. Try changing the kernel runtime to GPU and see what P100 can endure.\n2. By changing `input_dtype`, we can easily see how `bfloat16` effectively saves the RAM withou the pains of `fp16` mixed precision - mixed precision is always another layer of wraparound in Pytorch (and GPU in general) and TPU's `bfloat16` has totally solved this problem.","metadata":{}},{"cell_type":"code","source":"logs = {\n    'train_loss': [], 'train_metric': [], 'val_loss': [], 'val_metric': [], \n    'train_state': None, 'args': args}\n\nfor epoch in range(args['epochs']):\n    ### Training ###\n    trainloader = load_dataset(train_filenames, labeled=True, augment=True,\\\n                    ordered=False, shuffle_buffer_size=4*args['batch_size'], drop_remainder=True)\n    # This \"total\" is specific to batch size!\n    counter = tqdm(trainloader, total=199, leave=False)\n    pl_train_state['model'] = unfreeze(pl_train_state['model'])\n    train_loss, train_metric, train_counter = 0, 0, 0\n    for batch in counter:\n        random_key, subkey = jax.random.split(random_key)\n        out = train_step_parallel(train_state=pl_train_state, batch=batch, subkey=shard_prng_key(subkey))\n        pl_train_state = out['train_state']\n        loss_value, metric_value = out['loss'][0].item(), out['metric'][0].item()\n        \n        counter.set_postfix({'loss':loss_value, 'metric':metric_value})\n        train_loss += loss_value\n        train_metric += metric_value\n        train_counter += 1\n        if args['sanity_check']:\n            break\n    pl_train_state['model'] = freeze(pl_train_state['model'])\n    train_loss, train_metric = train_loss / train_counter, train_metric / train_counter\n    \n    \n    ### Validation ###\n    valloader = load_dataset(val_filenames, labeled=True, augment=False,\\\n                    ordered=False, shuffle_buffer_size=4*args['batch_size'], drop_remainder=True)\n    \n    val_steps, val_loss, val_metric = 0, 0, 0\n    counter = tqdm(valloader, total=58, leave=False)\n    for batch in counter:\n        out = val_step_parallel(model_params=pl_train_state['model'], batch=batch)\n        loss_value, metric_value = out['loss'][0].item(), out['metric'][0].item()\n        val_loss, val_metric, val_steps = val_loss + loss_value, val_metric + metric_value,\\\n                                        val_steps + 1\n        counter.set_postfix({'val_loss':loss_value, 'val_metric':metric_value})\n        if args['sanity_check']:\n            break\n    val_metric, val_loss = val_metric / val_steps, val_loss / val_steps\n    \n    # Logging\n    logs['train_loss'].append(train_loss)\n    logs['train_metric'].append(train_metric)\n    logs['val_loss'].append(val_loss)\n    logs['val_metric'].append(val_metric)\n    \n    # Checkpointing\n    if (val_metric == max(logs['val_metric'])) or (val_loss == min(logs['val_loss'])):\n        print(f'Epoch {epoch}, loss:{train_loss:.4f} acc:{train_metric:.4f} val_loss:{val_loss:.4f} val_metric:{val_metric:.4f}')\n        out_path = os.path.join(args['saving_dir'], 'checkpoint.msgpack')\n        logs['train_state'] = jax.tree_map(lambda x : x.astype(jnp.float16), \\\n                    flax.jax_utils.unreplicate(pl_train_state))\n        save_checkpoint(out_path, logs, verbose=False)\n    # load_checkpoint(out_path, logs)","metadata":{"execution":{"iopub.status.busy":"2022-04-21T22:45:44.708552Z","iopub.execute_input":"2022-04-21T22:45:44.709117Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}]}