{"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":"
\n

🍀 Herbarium Jax/Flax Training - 🤗 + KFolds + W&B Tracking 📊

\n
\n\n

\n JAX/Flax training notebook that fine-tunes Vision Transformer from 🤗 transformers also with Weights and Biases tracking.
Do keep in mind, this *is* a proper training notebook but I have set it to train only for 1 epoch only for the first fold since there is so much data, it would be otherwise quite hard to train for all folds with more than one epochs.
\n

\n\n

If you liked this notebook, kindly leave an upvote ⬆️

","metadata":{}},{"cell_type":"markdown","source":"#### Attribution\nThis notebook takes a lot of functions and major inspiration from HuggingFace's Flax Vision example [here](https://github.com/huggingface/transformers/blob/master/examples/flax/vision/run_image_classification.py)","metadata":{}},{"cell_type":"markdown","source":"

1. Installation & Imports 📩

","metadata":{}},{"cell_type":"code","source":"%%sh\npip install -q wandb\npip install -q transformers","metadata":{"_uuid":"187c2962-56ea-45fc-b039-597d1e254c19","_cell_guid":"c6a81ca6-7f26-4574-9e80-14973339b706","collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2022-02-20T15:09:09.794227Z","iopub.execute_input":"2022-02-20T15:09:09.794563Z","iopub.status.idle":"2022-02-20T15:09:24.936818Z","shell.execute_reply.started":"2022-02-20T15:09:09.794484Z","shell.execute_reply":"2022-02-20T15:09:24.936025Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import json\nimport os\nfrom typing import Callable\n\nimport wandb\nimport numpy as np\nimport pandas as pd\nfrom PIL import Image\nimport logging\n\nimport torch\nimport torchvision.transforms as transforms\nfrom tqdm.notebook import tqdm\nfrom sklearn.model_selection import StratifiedKFold\n\nimport jax\nimport jax.numpy as jnp\nimport optax\nimport transformers\nfrom flax import jax_utils\nfrom flax.jax_utils import unreplicate\nfrom flax.training import train_state\nfrom flax.training.common_utils import get_metrics, onehot, shard, shard_prng_key\nfrom transformers import AutoConfig, FlaxAutoModelForImageClassification\n\n# To keep out those nasty warnings\nimport warnings\nwarnings.simplefilter('ignore')\ntransformers.utils.logging.set_verbosity(transformers.logging.ERROR)\nos.environ['TF_CPP_MIN_LOG_LEVEL'] = '3' ","metadata":{"_uuid":"acbcdddc-1d0a-427a-bafc-e9070f4660ef","_cell_guid":"3d530023-cb70-4e2a-a6bb-1125783ba091","collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2022-02-20T15:09:24.938996Z","iopub.execute_input":"2022-02-20T15:09:24.939266Z","iopub.status.idle":"2022-02-20T15:09:29.398751Z","shell.execute_reply.started":"2022-02-20T15:09:24.939228Z","shell.execute_reply":"2022-02-20T15:09:29.398005Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"Config = {\n 'MODEL_NAME': 'google/vit-base-patch16-224-in21k',\n 'JSON_PATH': \"../input/herbarium-2022-fgvc9/train_metadata.json\",\n 'IMG_PATH': \"../input/herbarium-2022-fgvc9/train_images\",\n 'NUM_LABELS': 15501,\n 'N_SPLITS': 5,\n 'TRAIN_BS': 32,\n 'VALID_BS': 32,\n 'N_EPOCHS': 2,\n 'NUM_WORKERS': 4,\n 'LR': 1e-3,\n 'IMG_SIZE': 224,\n 'infra': \"Kaggle\",\n 'competition': 'herbarium_2022',\n '_wandb_kernel': 'tanaym',\n \"wandb\": True,\n}","metadata":{"_uuid":"73e18a8b-c784-4119-9c19-9c101fd7fb79","_cell_guid":"625ad516-83cc-4cbd-8b58-d2e8b770025e","collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2022-02-20T15:09:29.40112Z","iopub.execute_input":"2022-02-20T15:09:29.401526Z","iopub.status.idle":"2022-02-20T15:09:29.407144Z","shell.execute_reply.started":"2022-02-20T15:09:29.40149Z","shell.execute_reply":"2022-02-20T15:09:29.406203Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"

2. About Weights and Biases 📊


\n

WandB is a developer tool for companies turn deep learning research projects into deployed software by helping teams track their models, visualize model performance and easily automate training and improving models.\nWe will use their tools to log hyperparameters and output metrics from your runs, then visualize and compare results and quickly share findings with your colleagues.

","metadata":{}},{"cell_type":"markdown","source":"To login to W&B, you can use below snippet.\n\n```python\nfrom kaggle_secrets import UserSecretsClient\nuser_secrets = UserSecretsClient()\nwb_key = user_secrets.get_secret(\"WANDB_API_KEY\")\n\nwandb.login(key=wb_key)\n```\nMake sure you have your W&B key stored as `WANDB_API_KEY` under Add-ons -> Secrets\n\nYou can view [this](https://www.kaggle.com/ayuraj/experiment-tracking-with-weights-and-biases) notebook to learn more about W&B tracking.\n\nIf you don't want to login to W&B, the kernel will still work and log everything to W&B in anonymous mode.","metadata":{}},{"cell_type":"code","source":"# Start W&B logging\nif Config['wandb']:\n run = wandb.init(\n project='jax',\n config=Config,\n group='vision',\n job_type='train',\n anonymous='must'\n )","metadata":{"execution":{"iopub.status.busy":"2022-02-20T15:09:29.408329Z","iopub.execute_input":"2022-02-20T15:09:29.408638Z","iopub.status.idle":"2022-02-20T15:09:37.183907Z","shell.execute_reply.started":"2022-02-20T15:09:29.408598Z","shell.execute_reply":"2022-02-20T15:09:37.183172Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"

3. Utility functions ⚒️

","metadata":{}},{"cell_type":"code","source":"# Some utility functions\ndef wandb_log(**kwargs):\n \"\"\"\n Logs a key-value pair to W&B\n \"\"\"\n for k, v in kwargs.items():\n wandb.log({k: v})\n\ndef load_json(json_file):\n \"\"\"\n Loads metadata json file and returns a processed pandas dataframe\n \"\"\"\n with open(json_file, \"r\", encoding=\"ISO-8859-1\") as file:\n train = json.load(file)\n train_img = pd.DataFrame(train['images'])\n train_ann = pd.DataFrame(train['annotations'])\n train_df = train_img.merge(train_ann, on='image_id')\n return train_df\n\ndef create_learning_rate_fn(\n train_ds_size: int, \n train_batch_size: int, \n num_train_epochs: int, \n num_warmup_steps: int, \n learning_rate: float\n) -> Callable[[int], jnp.array]:\n \n \"\"\"Returns a linear warmup, linear_decay learning rate function.\"\"\"\n steps_per_epoch = train_ds_size // train_batch_size\n num_train_steps = steps_per_epoch * num_train_epochs\n warmup_fn = optax.linear_schedule(init_value=0.0, end_value=learning_rate, transition_steps=num_warmup_steps)\n decay_fn = optax.linear_schedule(\n init_value=learning_rate, end_value=0, transition_steps=num_train_steps - num_warmup_steps\n )\n schedule_fn = optax.join_schedules(schedules=[warmup_fn, decay_fn], boundaries=[num_warmup_steps])\n return schedule_fn\n\ndef loss_fn(logits, labels):\n loss = optax.softmax_cross_entropy(logits, onehot(labels, logits.shape[-1]))\n return loss.mean()\n\n# Define gradient update step fn\n# Pretty self explanatory for the most part if you have some exposure in deep learning\ndef train_step(state, batch):\n dropout_rng, new_dropout_rng = jax.random.split(state.dropout_rng)\n\n def compute_loss(params):\n labels = batch.pop(\"labels\")\n logits = state.apply_fn(**batch, params=params, dropout_rng=dropout_rng, train=True)[0]\n loss = loss_fn(logits, labels)\n return loss\n\n grad_fn = jax.value_and_grad(compute_loss)\n loss, grad = grad_fn(state.params)\n grad = jax.lax.pmean(grad, \"batch\")\n\n new_state = state.apply_gradients(grads=grad, dropout_rng=new_dropout_rng)\n\n metrics = {\"loss\": loss, \"learning_rate\": linear_decay_lr_schedule_fn(state.step)}\n metrics = jax.lax.pmean(metrics, axis_name=\"batch\")\n\n return new_state, metrics\n\n# Define validation function\ndef valid_step(params, batch):\n labels = batch.pop(\"labels\")\n logits = model(**batch, params=params, train=False)[0]\n loss = loss_fn(logits, labels)\n\n # summarize metrics\n accuracy = (jnp.argmax(logits, axis=-1) == labels).mean()\n metrics = {\"loss\": loss, \"accuracy\": accuracy}\n metrics = jax.lax.pmean(metrics, axis_name=\"batch\")\n return metrics","metadata":{"_uuid":"80925003-931a-4423-92e5-018fb9319c04","_cell_guid":"5d1f350b-7bc2-4868-a227-52ab8e616b0d","collapsed":false,"_kg_hide-input":true,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2022-02-20T15:09:37.186215Z","iopub.execute_input":"2022-02-20T15:09:37.186418Z","iopub.status.idle":"2022-02-20T15:09:37.20549Z","shell.execute_reply.started":"2022-02-20T15:09:37.186391Z","shell.execute_reply":"2022-02-20T15:09:37.204615Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"

4. Custom dataset class & Augmentations 🖼️

","metadata":{}},{"cell_type":"code","source":"class HerbariumData(torch.utils.data.Dataset):\n \"\"\"\n Custom dataset class for this competition's data\n \"\"\"\n def __init__(self, df, labels=None, augments=True, is_test=False):\n self.df = df\n self.labels = labels\n self.augments = augments\n self.is_test = is_test\n \n def __getitem__(self, idx):\n file_name = self.df['file_name'].values[idx]\n file_path = os.path.join(Config['IMG_PATH'], file_name)\n image = Image.open(file_path)\n \n if self.augments:\n image = self.augments(image)\n \n if not self.is_test:\n labels = self.labels.values[idx]\n return image, labels\n return image\n \n def __len__(self):\n return len(self.df)","metadata":{"_uuid":"4dd661ff-af69-469f-9ddd-65365548a0bd","_cell_guid":"4fb1d9b6-4b6e-4906-ba27-8a8c0e29efb5","collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2022-02-20T15:09:37.207039Z","iopub.execute_input":"2022-02-20T15:09:37.207602Z","iopub.status.idle":"2022-02-20T15:09:37.218647Z","shell.execute_reply.started":"2022-02-20T15:09:37.207565Z","shell.execute_reply":"2022-02-20T15:09:37.217925Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def collate_fn(samples):\n pixel_values = torch.stack([sample[0] for sample in samples])\n labels = torch.tensor([sample[1] for sample in samples])\n\n batch = {\"pixel_values\": pixel_values, \"labels\": labels}\n batch = {k: v.numpy() for k, v in batch.items()}\n\n return batch","metadata":{"_uuid":"4f6790c8-fdfb-4e50-b59b-7b1124cbe62b","_cell_guid":"aeb3c6cb-dd00-4fd8-acf1-d693d7ca28b7","collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2022-02-20T15:09:37.220094Z","iopub.execute_input":"2022-02-20T15:09:37.220367Z","iopub.status.idle":"2022-02-20T15:09:37.2288Z","shell.execute_reply.started":"2022-02-20T15:09:37.220322Z","shell.execute_reply":"2022-02-20T15:09:37.228126Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class Augments:\n def train_augments():\n return transforms.Compose(\n [\n transforms.RandomResizedCrop(Config['IMG_SIZE']),\n transforms.RandomHorizontalFlip(),\n transforms.ToTensor(),\n transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),\n ]\n )\n def valid_augments():\n return transforms.Compose(\n [\n transforms.Resize(Config['IMG_SIZE']),\n transforms.CenterCrop(Config['IMG_SIZE']),\n transforms.ToTensor(),\n transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),\n ]\n )","metadata":{"_uuid":"45f5f941-9ec3-4388-8ae7-0d9d94e23a17","_cell_guid":"6c6c1f57-89e0-4b75-a4b8-32cba3f48d49","collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2022-02-20T15:09:37.230805Z","iopub.execute_input":"2022-02-20T15:09:37.231711Z","iopub.status.idle":"2022-02-20T15:09:37.240435Z","shell.execute_reply.started":"2022-02-20T15:09:37.231685Z","shell.execute_reply":"2022-02-20T15:09:37.239664Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"

5. Training 🚀

","metadata":{}},{"cell_type":"markdown","source":"","metadata":{}},{"cell_type":"code","source":"# Create a model Config\nconfig = AutoConfig.from_pretrained(\n Config['MODEL_NAME'],\n num_labels=Config['NUM_LABELS'],\n image_size=Config['IMG_SIZE'],\n)","metadata":{"_uuid":"df3d5a54-cc7e-4074-a366-f32aa323fd0e","_cell_guid":"1e4e700c-ffe5-4509-8805-79fa473849b2","collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2022-02-20T15:09:37.24213Z","iopub.execute_input":"2022-02-20T15:09:37.243874Z","iopub.status.idle":"2022-02-20T15:09:38.483895Z","shell.execute_reply.started":"2022-02-20T15:09:37.243839Z","shell.execute_reply":"2022-02-20T15:09:38.482905Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"A simple training state for a single optax optimizer.","metadata":{}},{"cell_type":"code","source":"class TrainState(train_state.TrainState):\n dropout_rng: jnp.ndarray\n\n def replicate(self):\n return jax_utils.replicate(self).replace(dropout_rng=shard_prng_key(self.dropout_rng))","metadata":{"_uuid":"638a08b8-9b04-42de-941b-5b4b4c52d6c1","_cell_guid":"b2730e37-583a-4301-becc-3eb12aac589b","collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2022-02-20T15:09:38.485342Z","iopub.execute_input":"2022-02-20T15:09:38.485718Z","iopub.status.idle":"2022-02-20T15:09:38.492706Z","shell.execute_reply.started":"2022-02-20T15:09:38.485683Z","shell.execute_reply":"2022-02-20T15:09:38.491997Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"A quick overview of what I'm doing in the training cell down below;\n1. Reading the JSON file into a pandas dataframe and splitting it into k-folds\n2. Getting the current fold's split of training and validation data\n3. Making a dataset out of the training and validation pandas dataframes and then making dataloaders of them\n4. Defining the model with the proper parameters and the config we defined above\n5. Initializing jax-specific variables such as random key and linear decay learning rate schedule function\n6. Defining Adam optimizer using Optax\n7. Creating parallel versions of the training and validation step functions (for multi-device functionality)\n8. Making a train state and replicating it on each device (for multi-device functionality)\n9. Now running the epochs, inside it, we are running training and validation loops over all batches\n10. Inside a typical logic, we have;\n * We shard (basically, breaking into small subsets) the batch data\n * Pass the current state and the sharded batch data through corresponding step function (train or valid)\n * Getting the metric from the output and printing it along with logging it to WandB","metadata":{}},{"cell_type":"code","source":"# Main trainer\nif __name__ == \"__main__\":\n kf = StratifiedKFold(n_splits=Config['N_SPLITS'])\n train_file = load_json(Config['JSON_PATH'])\n \n for fold_, (train_idx, valid_idx) in enumerate(kf.split(X=train_file, y=train_file['category_id'])):\n # Only training for one fold since the data is huge and I don't want you all to wait an eternity\n if fold_ != 0:\n continue\n print(f\"{'='*40} Fold: {fold_+1} / {Config['N_SPLITS']} {'='*40}\")\n \n train_ = train_file.loc[train_idx]\n valid_ = train_file.loc[valid_idx]\n \n # Create train and validation dataloaders\n train_dataset = HerbariumData(train_, labels=train_['category_id'], augments=Augments.train_augments())\n valid_dataset = HerbariumData(valid_, labels=valid_['category_id'], augments=Augments.valid_augments())\n\n train = torch.utils.data.DataLoader(\n train_dataset,\n batch_size=Config['TRAIN_BS']*jax.device_count(),\n shuffle=True,\n pin_memory=True,\n num_workers=Config['NUM_WORKERS'],\n collate_fn=collate_fn,\n )\n valid = torch.utils.data.DataLoader(\n valid_dataset,\n batch_size=Config['VALID_BS']*jax.device_count(),\n shuffle=False,\n num_workers=Config['NUM_WORKERS'],\n collate_fn=collate_fn,\n )\n \n # Instantiate the model\n model = FlaxAutoModelForImageClassification.from_pretrained(\n Config['MODEL_NAME'], config=config, seed=42, dtype='float32'\n )\n \n # Initialize our training parameters\n rng = jax.random.PRNGKey(42)\n rng, dropout_rng = jax.random.split(rng)\n\n # Create learning rate schedule\n linear_decay_lr_schedule_fn = create_learning_rate_fn(\n len(train_dataset),\n Config['TRAIN_BS']*jax.device_count(),\n Config['N_EPOCHS'],\n 0,\n Config['LR'],\n )\n\n # Create Adam optimizer\n adamw = optax.adam(\n learning_rate=linear_decay_lr_schedule_fn,\n )\n \n # Create parallel version of the train and eval step\n p_train_step = jax.pmap(train_step, \"batch\", donate_argnums=(0,))\n p_valid_step = jax.pmap(valid_step, \"batch\")\n\n # Replicate the train state on each device\n state = TrainState.create(apply_fn=model.__call__, params=model.params, tx=adamw, dropout_rng=dropout_rng)\n state = state.replicate()\n \n for epoch in range(Config['N_EPOCHS']):\n print(f\"{'-'*20} Epoch: {epoch+1} / {Config['N_EPOCHS']} {'-'*20}\")\n rng, input_rng = jax.random.split(rng)\n\n # Training\n train_metrics = []\n steps_per_epoch = len(train_dataset) // (Config['TRAIN_BS']*jax.device_count())\n train_prog = tqdm(train, total=steps_per_epoch)\n\n for batch in train_prog:\n batch = shard(batch)\n state, train_metric = p_train_step(state, batch)\n train_metrics.append(train_metric)\n\n train_prog.set_description(f\"loss: {train_metric['loss'].tolist()[0]:.4f}\")\n \n # Log to wandb\n if Config['wandb']:\n wandb_log(\n train_loss=train_metric['loss'].tolist()[0]\n )\n train_metric = unreplicate(train_metric)\n train_prog.close()\n print(f\"Train loss at Epoch {epoch+1}: {train_metric['loss']}\")\n\n # Evaluating\n # Not doing the valid part, only 2 epochs of training since it takes a lot of time.\n continue\n valid_metrics = []\n steps_per_epoch = len(valid_dataset) // (Config['VALID_BS']*jax.device_count())\n valid_prog = tqdm(valid, total=steps_per_epoch)\n\n for batch in valid_prog:\n batch = shard(batch)\n metric = p_valid_step(state.params, batch)\n valid_metrics.append(metric)\n\n valid_prog.set_description(f\"val_loss: {metric['loss'].tolist()[0]:.4f}\")\n \n # Log to wandb\n if Config['wandb']:\n wandb_log(\n train_loss=metric['loss'].tolist()[0]\n )\n\n # Normalize eval metrics\n valid_metrics = get_metrics(valid_metrics)\n valid_metrics = jax.tree_map(jnp.mean, valid_metrics)\n\n valid_prog.close()\n print(f\"Valid loss at Epoch {epoch+1}: {metric['loss'][0]}\")","metadata":{"execution":{"iopub.status.busy":"2022-02-20T15:12:24.475119Z","iopub.execute_input":"2022-02-20T15:12:24.475648Z","iopub.status.idle":"2022-02-20T15:13:46.611045Z","shell.execute_reply.started":"2022-02-20T15:12:24.475608Z","shell.execute_reply":"2022-02-20T15:13:46.60984Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Code taken from https://www.kaggle.com/ayuraj/interactive-eda-using-w-b-tables\n\n# Finish the logging run\nif Config['wandb']:\n run.finish()","metadata":{"execution":{"iopub.status.busy":"2022-02-20T15:10:42.463281Z","iopub.status.idle":"2022-02-20T15:10:42.464387Z","shell.execute_reply.started":"2022-02-20T15:10:42.4641Z","shell.execute_reply":"2022-02-20T15:10:42.464129Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"
\n\n
","metadata":{}}]}