{"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":"#### ⚠️⚠️⚠️ Note that this notebook is still **WORK IN PROGRESS**, any constructive feedbacks are welcome!\n\n---\n### **🎯 OBJECTIVE**\n- Demonstrate how you could utilise Flax/JAX to build ⚡high performance model training pipeline for Kaggle competition\n- Demonstrate how you could apply transfer learning on different CV backbones (e.g. ResNet) with Flax/JAX\n\n### **🚧 TODO**\n- [Basic] Why TPU is not significantly faster than GPU??\n- [Basic] In inference, determine when to assign new individual\n- [Basic] Add validation set into training loop\n- [Basic] Find the best learning rate\n- [Basic] Varying learning rate over epochs\n- [Basic] Train on bigger image (224x224/ 512x512)\n- [Basic] Add K-fold validation split\n- [Advanced] Change PyTorch dataloader to more performant data loader\n- [Advanced] Add W&B for experiment logging\n- [Advanced] Pad remainder batch as last iteration instead of skipping it\n- [Advanced] Different head, backbones, loss function... etc.\n\n### **📝 LOG**\n- [20/02/2022] First draft released\n- [22/02/2022] Resolve bug in cross entropy loss\n- [23/02/2022] Work in Kaggle TPU\n- [24/02/2022] (1) Speed up training by pre-resized dataset, (2) enable to work in Kaggle TPU, (3) freeze backbone from param update, (4) get the loss to converge\n- [25/02/2022] add LB submission pipeline\n\n### **📚 REFERENCES**\n- [Flax & JAX training loop example from huggingface](https://github.com/huggingface/transformers/blob/master/examples/flax/image-captioning/run_image_captioning_flax.py)\n- [Flax & JAX training loop example from Flax repo](https://github.com/google/flax/blob/d068512a932da3e05b822790a591bac391aeab36/examples/nlp_seq/train.py#L206-L207)\n- [Sentiment CLF - JAX/FLAX on TPUs + 🤗 + W&B 🚀](https://www.kaggle.com/heyytanay/sentiment-clf-jax-flax-on-tpus-w-b)\n- [[Pytorch|0.374] 🐳 EffB5 KFold with FocalLoss](https://www.kaggle.com/snoop2head/pytorch-0-374-effb5-kfold-with-focalloss)\n- [Train EffNet in TF TPU](https://www.kaggle.com/dschettler8845/baseline-solution-train-indiv-model#dataset_preparation)\n- [Resize Images of Happy Whales and Dolphins](https://www.kaggle.com/gpreda/resize-images-of-happy-whales-and-dolphins)\n- [How to freeze layer or detach a layer in training stage](https://github.com/google/flax/issues/825)\n- [How to checkpoint Flax model](https://github.com/google/flax/discussions/1876)\n\n---\n","metadata":{}},{"cell_type":"markdown","source":"### 0. Pre-setup","metadata":{}},{"cell_type":"code","source":"%%capture\n# upgrade jax, jaxlib, flax is essential, otherwise it may fail to load pretrained model\nimport os\nimport requests\nimport torch\n\n# Kaggle mode\nif os.path.isdir(\"/kaggle/input\"):\n\n    # TPU mode\n    if \"TPU_NAME\" in os.environ:\n        print(\"TPU & Kaggle mode: Upgrade JAX and set up TPU\")\n        !pip install --upgrade jax\n        !pip install --upgrade jaxlib\n        !pip install git+https://github.com/deepmind/optax.git\n        !pip install flax\n        \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 TPU:', config.FLAGS.jax_backend_target)\n        \n    # GPU mode\n    elif torch.cuda.is_available():\n        print(\"GPU & Kaggle mode: Upgrade JAX specific to CUDA and CuDNN\")\n        !pip install --upgrade \"jax[cuda11_cudnn805]\" -f https://storage.googleapis.com/jax-releases/jax_releases.html\n        !pip install --upgrade jaxlib\n        !pip install git+https://github.com/deepmind/optax.git\n        !pip install flax\n            \n    # CPU mode\n    else:\n        raise NotImplementedError(\"CPU & Kaggle mode is not supported yet\")\n#         print(\"CPU & Kaggle mode: Upgrade JAX\")\n#         !pip install jax\n#         !pip install --upgrade jaxlib\n#         !pip install git+https://github.com/deepmind/optax.git\n#         !pip install flax\n        \n# Colab mode\nelse:\n    !pip install kaggle\n    \n    # TPU mode\n    if \"TPU_NAME\" in os.environ:\n        print(\"TPU & Colab mode: Set up TPU devices\")\n        import jax.tools.colab_tpu\n        jax.tools.colab_tpu.setup_tpu()\n\n!pip install flax\n!pip install jax-resnet","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","execution":{"iopub.status.busy":"2022-02-25T15:02:13.576091Z","iopub.execute_input":"2022-02-25T15:02:13.576576Z","iopub.status.idle":"2022-02-25T15:02:59.475078Z","shell.execute_reply.started":"2022-02-25T15:02:13.576538Z","shell.execute_reply":"2022-02-25T15:02:59.473943Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import os\nimport joblib\nfrom typing import *\nfrom functools import partial\n\nfrom tqdm import tqdm\nimport pandas as pd\nimport numpy as np\nfrom PIL import Image\nfrom sklearn.preprocessing import LabelEncoder\n\n# for data loading\nimport torch\nfrom torch.utils.data import Dataset, DataLoader\n# for augmentation\nfrom albumentations import *\n\nimport optax\nimport flax\nfrom flax import linen as nn\nfrom flax.core import frozen_dict\nfrom flax.core import FrozenDict\nfrom flax.jax_utils import replicate, unreplicate\nfrom flax.training import train_state, checkpoints\nfrom flax.training.common_utils import shard, shard_prng_key\n\nimport jaxlib\nimport jax\nimport jax.numpy as jnp\nfrom jax import lax\nfrom jax_resnet import pretrained_resnet, slice_variables, Sequential\n\nprint(f\"Flax version: {flax.__version__}\")\nprint(f\"Optax version: {optax.__version__}\")\nprint(f\"JAX version: {jax.__version__}\")\nprint(f\"JAX lib version: {jaxlib.__version__}\")\nprint(f\"JAX device: {jax.devices()}\")","metadata":{"execution":{"iopub.status.busy":"2022-02-25T15:02:59.477980Z","iopub.execute_input":"2022-02-25T15:02:59.478269Z","iopub.status.idle":"2022-02-25T15:02:59.492114Z","shell.execute_reply.started":"2022-02-25T15:02:59.478231Z","shell.execute_reply":"2022-02-25T15:02:59.491328Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### **1. Configuration**","metadata":{}},{"cell_type":"code","source":"EPOCH_N = 20\nFIRST_N_ITER = None  # early stop on iterations for testing\nLEARNING_RATE = 0.05\n# GPU: (128, 128) -> 64, (256, 256) -> 16\n# TPU: (128, 128) -> 64*8*2\nBATCH_SIZE = 64*8*2\nNUM_WORKERS = 2\nPRE_RESIZE = True # if so, no need resizing during dataloading\nFREEZE_BACKBONE = True\nMOCK_DATASET = False  # control dataloading for speed test\nROOT_DIR = '/kaggle/input/happy-whale-and-dolphin'\nTRAIN_CSV = os.path.join(ROOT_DIR, 'train.csv')\nTEST_CSV = os.path.join(ROOT_DIR, 'sample_submission.csv')\nTRAIN_DIR = '/kaggle/input/resize-images-of-happy-whales-and-dolphins/train_images_128'\nTEST_DIR = '/kaggle/input/resize-images-of-happy-whales-and-dolphins/test_images_128'\n\nIMAGE_SIZE: Optional[Tuple[int, int]] = (128, 128)\n# ref: https://pytorch.org/vision/stable/models.html\nIMAGENET_MEAN = [0.485, 0.456, 0.406]\nIMAGENET_STD = [0.229, 0.224, 0.225]\n\nOUTPUT_N = 15587\nMODEL_ARCH = 'resnet18'\nMODEL_SAVE_DIR = 'resnet_checkpoints'\n\n# validator\nassert all([IMAGE_SIZE[0] == s for s in IMAGE_SIZE]), \"must be equal\"","metadata":{"execution":{"iopub.status.busy":"2022-02-25T15:16:24.319203Z","iopub.execute_input":"2022-02-25T15:16:24.319467Z","iopub.status.idle":"2022-02-25T15:16:24.328145Z","shell.execute_reply.started":"2022-02-25T15:16:24.319438Z","shell.execute_reply":"2022-02-25T15:16:24.327190Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### **2. Data Preproc and Loading**","metadata":{}},{"cell_type":"code","source":"class HappyWhaleDataset(Dataset):\n    \"\"\"\n    borrow from:\n    - https://www.kaggle.com/snoop2head/pytorch-0-374-effb5-kfold-with-focalloss\n    \"\"\"\n    def __init__(self, df, transforms=None, mock_dataset=False, inference=False):\n        if mock_dataset:\n            print(f\"Mocking the dataset for testing...\")\n        \n        self.mock_dataset = mock_dataset\n        self.transforms = transforms\n        self.df = df\n        self.file_names = df.file_path.values\n        if inference:\n            self.species = None\n            self.labels = None\n        else:\n            self.species = df.species.values\n            self.labels = df.individual_id.values\n\n    def __getitem__(self, index):\n        getter = self.__get_mock_item if self.mock_dataset else self.__get_item\n        return getter(index)\n    \n    def __get_mock_item(self, index):\n        mock_img = np.ones((*IMAGE_SIZE, 3))\n        if self.labels is None and self.species is None:\n            return mock_img\n        mock_label = np.random.randint(0, OUTPUT_N)\n        return mock_img, mock_label\n    \n    def __get_item(self, index):\n        image_path = self.file_names[index]\n        image = Image.open(image_path)\n        image = np.array(image)\n        \n        # take care of grayscale image by adding channel\n        if len(image.shape) == 2: \n            image = np.dstack((image,)*3)        \n        if self.transforms:\n            image= self.transforms(image=image)['image']\n        \n        if self.labels is None and self.species is None:\n            return image\n        \n        label = self.labels[index]\n        return image, label\n\n    def __len__(self):\n        return len(self.df)\n\n    def set_transform(self, transform):\n        self.transform = transform\n\n\ndef get_train_file_path(filename, image_dir):\n    return os.path.join(image_dir, filename)","metadata":{"execution":{"iopub.status.busy":"2022-02-25T15:16:24.623089Z","iopub.execute_input":"2022-02-25T15:16:24.623433Z","iopub.status.idle":"2022-02-25T15:16:24.637726Z","shell.execute_reply.started":"2022-02-25T15:16:24.623400Z","shell.execute_reply":"2022-02-25T15:16:24.635978Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"df = pd.read_csv(TRAIN_CSV)\ndf['file_path'] = df['image'].apply(\n    partial(get_train_file_path, image_dir=TRAIN_DIR)\n)\nprint(f\"No. of unique labels: {len(df.individual_id.unique())}\")\n\n# encode label from str to integer\nlabel_encoder = LabelEncoder()\nlabel_encoder.fit(df.individual_id)\ndf.individual_id = label_encoder.transform(df.individual_id)\nwith open(\"label_encoder.pkl\", \"wb\") as p:\n    joblib.dump(label_encoder, p)\n\nassert len(df.individual_id.unique()) == OUTPUT_N\n    \ndf.head(3)","metadata":{"execution":{"iopub.status.busy":"2022-02-25T15:16:24.791704Z","iopub.execute_input":"2022-02-25T15:16:24.792215Z","iopub.status.idle":"2022-02-25T15:16:24.996861Z","shell.execute_reply.started":"2022-02-25T15:16:24.792185Z","shell.execute_reply":"2022-02-25T15:16:24.996117Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# define transformations applied on dataset\nif not PRE_RESIZE:\n    assert IMAGE_SIZE is not None\ntransforms_ls = [Resize(*IMAGE_SIZE, p=1.0)] if not PRE_RESIZE else []\ntransforms_ls += [Normalize(mean=IMAGENET_MEAN, std=IMAGENET_STD, max_pixel_value=255.0, p=1.0)]\ntransforms = Compose(transforms_ls)\n\n# set dataset and dataloader\ntrain_dataset = HappyWhaleDataset(\n    df, transforms=transforms,\n    mock_dataset=MOCK_DATASET,\n    inference=False\n)\n# default collate_func converts numpy to tensor in dataloading\n# reference: https://discuss.pytorch.org/t/torch-dataloader-gives-torch-tensor-as-ouput/31084/6\ntrain_loader = DataLoader(\n    train_dataset, batch_size=BATCH_SIZE,\n    num_workers=NUM_WORKERS, shuffle=True\n)\n\n# sanity check\nx, _ = train_dataset[0]\nbatch, labels = next(iter(train_loader))\n\nassert x.shape[:2] == IMAGE_SIZE and batch.shape[0] == BATCH_SIZE\nprint(f\"Shape for one batch: {batch.shape}\")","metadata":{"execution":{"iopub.status.busy":"2022-02-25T15:16:24.998701Z","iopub.execute_input":"2022-02-25T15:16:24.999317Z","iopub.status.idle":"2022-02-25T15:16:29.497276Z","shell.execute_reply.started":"2022-02-25T15:16:24.999275Z","shell.execute_reply":"2022-02-25T15:16:29.495523Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### **3. Model Construction**","metadata":{}},{"cell_type":"code","source":"# @TODO fill in later\nclass MarginLayer(nn.Module):\n    \"\"\"\n    reference:\n    - https://www.kaggle.com/ragnar123/unsupervised-baseline-arcface\n    - https://arxiv.org/pdf/1801.07698.pdf\n    \"\"\"\n    @nn.compact\n    def __call__(self, inputs):\n        raise NotImplementedError\n\n\nclass Head(nn.Module):\n    \"\"\"\n    references:\n    - fastai\n    - https://pytorch.org/docs/stable/generated/torch.nn.BatchNorm2d.html\n    \"\"\"\n    batch_norm_cls: partial = partial(nn.BatchNorm, momentum=0.9)\n        \n    @nn.compact\n    def __call__(self, inputs, train: bool):\n        \"\"\" param train determines dropout and batchnorm \"\"\"\n        output_n = inputs.shape[-1]\n        x = self.batch_norm_cls(use_running_average=not train)(inputs)\n        x = nn.Dropout(rate=0.25)(x, deterministic=not train)\n        x = nn.Dense(features=output_n)(x)\n        x = nn.relu(x)\n        x = self.batch_norm_cls(use_running_average=not train)(x)\n        x = nn.Dropout(rate=0.5)(x, deterministic=not train)\n        x = nn.Dense(features=OUTPUT_N)(x)\n        return x\n\n    \nclass HappyWhaleModel(nn.Module):\n    backbone: Sequential\n    head: Head\n        \n    def __call__(self, inputs, train: bool):\n        x = self.backbone(inputs)\n        # average pool layer\n        x = jnp.mean(x, axis=(1, 2))\n        x = self.head(x, train)\n        return x\n\n    \ndef _get_backbone_and_params(model_arch: str) -> Tuple[nn.Module, FrozenDict]:\n    # get model & param structure for pretrained model\n    if model_arch == 'resnet18':\n        resnet_tmpl, params = pretrained_resnet(18)\n        model = resnet_tmpl()\n    else:\n        raise NotImplementedError\n        \n    # get model & param structure for backbone\n    start, end = 0, len(model.layers) - 2\n    backbone = Sequential(model.layers[start:end])\n    backbone_params = slice_variables(params, start, end)\n    return backbone, backbone_params\n\n\ndef get_model_and_variables(model_arch: str, head_init_key: int) -> Tuple[nn.Module, FrozenDict]:\n    \"\"\"\n    variables is a composition of params and batch state\n    \"\"\"\n    inputs = jnp.ones((1, *IMAGE_SIZE, 3), jnp.float32)\n    key = jax.random.PRNGKey(head_init_key)\n    # get backbone\n    backbone, backbone_params = _get_backbone_and_params(model_arch)\n    \n    # determine input size for head model\n    backbone_output = backbone.apply(backbone_params, inputs, mutable=False)\n    head_inputs = jnp.ones((1, backbone_output.shape[-1]), jnp.float32)\n    # get head\n    head = Head()\n    head_params = head.init(key, head_inputs, train=False)\n    \n    # get final model\n    model = HappyWhaleModel(backbone, head)\n    # combine params from backbone and head\n    variables = FrozenDict({\n        'params': {\n            'backbone': backbone_params['params'],\n            'head': head_params['params']\n        },\n        'batch_stats': {\n            'backbone': backbone_params['batch_stats'],\n            'head': head_params['batch_stats']\n        }\n    })\n    return model, variables","metadata":{"execution":{"iopub.status.busy":"2022-02-25T15:16:29.499795Z","iopub.execute_input":"2022-02-25T15:16:29.500122Z","iopub.status.idle":"2022-02-25T15:16:29.520713Z","shell.execute_reply.started":"2022-02-25T15:16:29.500079Z","shell.execute_reply":"2022-02-25T15:16:29.519647Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"%%capture\n# sanity check\nmodel, variables = get_model_and_variables('resnet18', 0)\ninputs = jnp.ones((1, *IMAGE_SIZE, 3), jnp.float32)\nkey = jax.random.PRNGKey(0)\no = model.apply(variables, inputs, train=False, mutable=False)\n\nassert o.shape[0] == inputs.shape[0] and o.shape[-1] == OUTPUT_N","metadata":{"execution":{"iopub.status.busy":"2022-02-25T15:16:29.522279Z","iopub.execute_input":"2022-02-25T15:16:29.522550Z","iopub.status.idle":"2022-02-25T15:16:30.912862Z","shell.execute_reply.started":"2022-02-25T15:16:29.522514Z","shell.execute_reply":"2022-02-25T15:16:30.912115Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### **4. Loss Function and Metrics**","metadata":{}},{"cell_type":"code","source":"def topk_accuracy(logits, labels, top_k = (1,)) -> List[jnp.DeviceArray]:\n    \"\"\"\n    reference:\n    - https://www.kaggle.com/snoop2head/pytorch-0-374-effb5-kfold-with-focalloss\n    \n    logits: (bs, class_n)\n    labels: (bs, )\n    \"\"\"\n    max_k = max(top_k)\n    batch_size = labels.shape[0]\n    \n    _, preds = jax.lax.top_k(logits, k = max_k)\n    # (max_k, bs)\n    preds = preds.transpose()\n    # convert labels to have shape: (max_k, bs)\n    correct = (preds == jnp.repeat(labels[None, :], repeats=preds.shape[0], axis = 0) * 1.)\n    \n    accuracies = []\n    for k in top_k:\n        _accuracy = correct[:k].sum() * 100. / batch_size\n        accuracies.append(_accuracy)\n    return accuracies\n\n\n# inputs: (logits, labels)\nloss_fn = optax.softmax_cross_entropy\n# inputs: (logits, labels)\neval_fn = partial(topk_accuracy, top_k = (1, 3))","metadata":{"execution":{"iopub.status.busy":"2022-02-25T15:16:30.915122Z","iopub.execute_input":"2022-02-25T15:16:30.915360Z","iopub.status.idle":"2022-02-25T15:16:30.923534Z","shell.execute_reply.started":"2022-02-25T15:16:30.915325Z","shell.execute_reply":"2022-02-25T15:16:30.922681Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### **5. Optimizer and TrainState**","metadata":{}},{"cell_type":"code","source":"class TrainState(train_state.TrainState):\n    batch_stats: FrozenDict\n    loss_fn: Callable = flax.struct.field(pytree_node=False)\n    eval_fn: Callable = flax.struct.field(pytree_node=False)\n\n\ndef create_mask(params, label_fn):\n    \"\"\"\n    reference:\n    - https://colab.research.google.com/drive/16wcmLt0AIKzMmLPrliuBMfmBvM8VEc4p#scrollTo=TqDvTL_tIQCH\n    \"\"\"\n    def _map(params, mask, label_fn):\n        for k in params:\n            if label_fn(k):\n                mask[k] = 'zero'\n            else:\n                if isinstance(params[k], FrozenDict):\n                    mask[k] = {}\n                    _map(params[k], mask[k], label_fn)\n                else:\n                    mask[k] = 'adam'\n    mask = {}\n    _map(params, mask, label_fn)\n    return frozen_dict.freeze(mask)\n\n\ndef zero_grads():\n    \"\"\"\n    reference:\n    - https://github.com/deepmind/optax/issues/159#issuecomment-896459491\n    \"\"\"\n    def init_fn(_): \n        return ()\n    def update_fn(updates, state, params=None):\n        return jax.tree_map(jnp.zeros_like, updates), ()\n    return optax.GradientTransformation(init_fn, update_fn)","metadata":{"execution":{"iopub.status.busy":"2022-02-25T15:16:30.925847Z","iopub.execute_input":"2022-02-25T15:16:30.926557Z","iopub.status.idle":"2022-02-25T15:16:30.938077Z","shell.execute_reply.started":"2022-02-25T15:16:30.926518Z","shell.execute_reply":"2022-02-25T15:16:30.937217Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"adamw = optax.adamw(\n    learning_rate=LEARNING_RATE,\n    b1=0.9, b2=0.999, \n    eps=1e-6, weight_decay=1e-2\n)\n\n# freeze backbone by masking target params\nif FREEZE_BACKBONE:\n    # label as zero if params belong to \"backbone\" key\n    optimizer = optax.multi_transform(\n        {'adam': adamw, 'zero': zero_grads()},\n        create_mask(variables['params'], lambda s: s.startswith('backbone'))\n    )\n# update all params\nelse:\n    optimizer = adamw\n\nstate = TrainState.create(\n    apply_fn = model.apply,\n    params = variables['params'],\n    tx = optimizer,\n    batch_stats = variables['batch_stats'],\n    loss_fn = loss_fn,\n    eval_fn = eval_fn\n)","metadata":{"execution":{"iopub.status.busy":"2022-02-25T15:16:30.939505Z","iopub.execute_input":"2022-02-25T15:16:30.939804Z","iopub.status.idle":"2022-02-25T15:16:30.962163Z","shell.execute_reply.started":"2022-02-25T15:16:30.939767Z","shell.execute_reply":"2022-02-25T15:16:30.961523Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### **6. Training Loop**","metadata":{}},{"cell_type":"code","source":"def train_step(state: TrainState, batch, labels, dropout_rng) -> Tuple[TrainState, dict, jnp.DeviceArray]:\n    dropout_rng, new_dropout_rng = jax.random.split(dropout_rng)\n    \n    # params as input because we differentiate wrt it \n    def loss_function(params):\n        # if you set state.params, then params can't be backpropagated through!\n        variables = {'params': params, 'batch_stats': state.batch_stats}\n        # return mutated states if mutable is specified\n        logits, new_batch_stats = state.apply_fn(\n            variables, batch, train=True, \n            mutable=['batch_stats'],\n            rngs={'dropout': dropout_rng}\n        )\n        # logits: (BS, OUTPUT_N), one_hot: (BS, OUTPUT_N)\n        one_hot = jax.nn.one_hot(labels, OUTPUT_N)\n        loss = state.loss_fn(logits, one_hot).mean()\n        return loss, (logits, new_batch_stats)\n    \n    # if you wanna vary lr per step\n    #lr = learning_rate_fn(state.step)\n    \n    # backpropagation and update params & batch_stats \n    grad_fn = jax.value_and_grad(loss_function, has_aux=True)\n    (loss, aux), grads = grad_fn(state.params)\n    logits, new_batch_stats = aux\n    grads = lax.pmean(grads, axis_name='batch')\n    new_state = state.apply_gradients(\n        grads=grads, batch_stats=new_batch_stats['batch_stats']\n    )\n    \n    # evaluation metrics\n    accuracy = state.eval_fn(logits, labels)\n    \n    # store metadata\n    metadata = jax.lax.pmean(\n        {'loss': loss, 'accuracy': accuracy},\n        axis_name='batch'\n    )\n    return new_state, metadata, new_dropout_rng\n\n\ndef val_step(state: TrainState, batch, labels):\n    variables = {'params': state.params, 'batch_stats': state.batch_stats}\n    logits = state.apply_fn(variables, batch, train=False)\n    return state.eval_fn(logits, labels)","metadata":{"execution":{"iopub.status.busy":"2022-02-25T15:16:30.963506Z","iopub.execute_input":"2022-02-25T15:16:30.963733Z","iopub.status.idle":"2022-02-25T15:16:30.977032Z","shell.execute_reply.started":"2022-02-25T15:16:30.963702Z","shell.execute_reply":"2022-02-25T15:16:30.976286Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"parallel_train_step = jax.pmap(train_step, axis_name='batch', donate_argnums=(0,))\nparallel_val_step = jax.pmap(val_step, axis_name='batch', donate_argnums=(0,))\n\n# required for parallelism\nstate = replicate(state)\n\n# control randomness on dropout and update inside train_step\nrng = jax.random.PRNGKey(0)\ndropout_rng = jax.random.split(rng, jax.local_device_count())  # for parallelism","metadata":{"execution":{"iopub.status.busy":"2022-02-25T15:16:30.978150Z","iopub.execute_input":"2022-02-25T15:16:30.978640Z","iopub.status.idle":"2022-02-25T15:16:31.043793Z","shell.execute_reply.started":"2022-02-25T15:16:30.978604Z","shell.execute_reply":"2022-02-25T15:16:31.043148Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"for epoch_i in tqdm(range(EPOCH_N), desc=f\"{EPOCH_N} epochs\", position=0, leave=True):\n    # training set\n    train_loss, train_accuracy = [], []\n    iter_n = len(train_dataset)//BATCH_SIZE\n    \n    # @TODO: print out from tqdm can be hard to read, find an alternative\n    with tqdm(total=iter_n, desc=f\"{iter_n} iterations\", leave=False) as progress_bar:\n        for iter_i, (batch, labels) in enumerate(train_loader):\n            \n            # unavoid reminder cant be evenly distributed across tpu devices\n            if iter_i + 1 > iter_n:\n                break\n            if (epoch_i == 0) and (FIRST_N_ITER is not None) and (iter_i >= FIRST_N_ITER):\n                print(f\"Early stop at iteration: {iter_i+1}\")\n                break\n                \n            # shard to enable parallelism\n            batch, labels = shard(batch), shard(labels)\n            batch = jnp.array(batch, dtype=jnp.float32)\n            labels = jnp.array(labels, dtype=jnp.float32)\n            \n            # backprop and update param & batch stats\n            state, train_metadata, dropout_rng = parallel_train_step(state, batch, labels, dropout_rng)\n            train_metadata = unreplicate(train_metadata)\n            \n            # update train statistics\n            _train_loss, _train_top1_acc, _train_top3_acc = map(float, [train_metadata['loss'], *train_metadata['accuracy']])\n            train_loss.append(_train_loss)\n            train_accuracy.append(_train_top1_acc)\n            progress_bar.update(1)\n    \n    avg_train_loss = sum(train_loss)/len(train_loss)\n    avg_train_acc = sum(train_accuracy)/len(train_accuracy)\n    print(f\"[{epoch_i+1}/{EPOCH_N}] Train Loss: {avg_train_loss:.03} | Train Accuracy: {avg_train_acc:.03}\")\n    \n    # validation set\n    pass","metadata":{"execution":{"iopub.status.busy":"2022-02-25T15:16:31.046593Z","iopub.execute_input":"2022-02-25T15:16:31.047039Z","iopub.status.idle":"2022-02-25T15:17:43.119946Z","shell.execute_reply.started":"2022-02-25T15:16:31.047012Z","shell.execute_reply":"2022-02-25T15:17:43.117875Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"state = unreplicate(state)\n\n# checkpoint the model\nos.makedirs(MODEL_SAVE_DIR, exist_ok=True)\ncheckpoints.save_checkpoint(ckpt_dir=MODEL_SAVE_DIR, target=state,\n                            step=int(state.step), overwrite=True)\n#restored_state = checkpoints.restore_checkpoint(ckpt_dir=MODEL_SAVE_DIR, target=state)\nprint(f\"Model checkpoint saved: {MODEL_SAVE_DIR}\")","metadata":{"execution":{"iopub.status.busy":"2022-02-25T15:17:43.121374Z","iopub.status.idle":"2022-02-25T15:17:43.121926Z","shell.execute_reply.started":"2022-02-25T15:17:43.121663Z","shell.execute_reply":"2022-02-25T15:17:43.121691Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### 5. Submission","metadata":{}},{"cell_type":"code","source":"def infer_step(state: TrainState, batch):\n    \"\"\"\n    return top5 predicted index for each sample\n    \"\"\"\n    variables = {'params': state.params, 'batch_stats': state.batch_stats}\n    logits = state.apply_fn(variables, batch, train=False)\n    _, top_preds = jax.lax.top_k(logits, k = 5)\n    return top_preds\n\n\n# load csv for submission\ntest_df = pd.read_csv(TEST_CSV)\ntest_df['file_path'] = test_df['image'].apply(\n    partial(get_train_file_path, image_dir=TEST_DIR)\n)\nprint(f\"No. of Submission Entries: {len(test_df)}\")\n\n# setup dataloader for submission\ntest_dataset = HappyWhaleDataset(\n    test_df, transforms=transforms, \n    mock_dataset=MOCK_DATASET,\n    inference=True\n)\ntest_loader = DataLoader(\n    test_dataset, batch_size=BATCH_SIZE,\n    num_workers=NUM_WORKERS, shuffle=False\n)\n\n# ready for running batch\nparallel_infer_step = jax.pmap(infer_step, axis_name='batch', donate_argnums=(0,))\nstate = replicate(state)","metadata":{"execution":{"iopub.status.busy":"2022-02-25T15:17:43.123476Z","iopub.status.idle":"2022-02-25T15:17:43.129098Z","shell.execute_reply.started":"2022-02-25T15:17:43.128856Z","shell.execute_reply":"2022-02-25T15:17:43.128886Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# batch processing samples for submission\ntop5_idxs = []\niter_n = len(test_dataset)//BATCH_SIZE\nfor iter_i, batch in tqdm(enumerate(test_loader)):\n    if iter_i + 1 > iter_n:\n        state = unreplicate(state)\n        _batch_top5_idxs = infer_step(state, batch)\n        # make this cell block re-runnable\n        state = replicate(state)\n    else:\n        batch = shard(batch)\n        batch = jnp.array(batch, dtype=jnp.float32)\n        _batch_top5_idxs = parallel_infer_step(state, batch)\n        _batch_top5_idxs = _batch_top5_idxs.reshape(-1, 5)\n    top5_idxs.append(_batch_top5_idxs)\n\ntop5_idxs = jnp.concatenate(top5_idxs)\nassert top5_idxs.shape[0] == test_df.shape[0]","metadata":{"execution":{"iopub.status.busy":"2022-02-25T15:17:43.129832Z","iopub.status.idle":"2022-02-25T15:17:43.130158Z","shell.execute_reply.started":"2022-02-25T15:17:43.129972Z","shell.execute_reply":"2022-02-25T15:17:43.129998Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# label decoding\nprediction_parser = lambda idxs: ' '.join(label_encoder.inverse_transform(idxs)).strip()\ntop5_labels = list(map(prediction_parser, top5_idxs))\ntest_df['predictions'] = top5_labels\n\n# write the same csv\ntest_df.drop(columns=['file_path'], inplace=True)\ntest_df.to_csv('submission.csv', index=False)\nprint(f\"Submission csv written!!\")","metadata":{"execution":{"iopub.status.busy":"2022-02-25T15:17:43.130846Z","iopub.status.idle":"2022-02-25T15:17:43.131168Z","shell.execute_reply.started":"2022-02-25T15:17:43.130995Z","shell.execute_reply":"2022-02-25T15:17:43.131013Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}]}