{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.7.12","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# Pix2Struct Fine-tuning for Graph data extraction + AMP + W&B logging 🚀\n\nThis notebook (read scratch pad) shows how you can fine-tune Google's Pix2Struct base model for DocVQA (Document Visual Question Answering) on the competition's data.\n\nCurrently, the loss is quite terrible (reason being I can't adjust the hyperparameters too much without running out of memory) but IT TRAINS. Sticking on little hope there :')\n\nIf you wish to play with hyperparameters (might cause brain cell loss in certain cases due to OOM), please fork this notebook and be my guest.\nI really wish to see some of you extend my work and do something cool with this!\n\nYou can find my current selection of hyperparameters in the `Config` dictionary.\n\nHope you all find this useful!\n\nP.S: Huge thanks to [@nbroad's](https://www.kaggle.com/nbroad) donut training [notebook](https://www.kaggle.com/code/nbroad/donut-train-benetech).","metadata":{}},{"cell_type":"markdown","source":"
\n\n
","metadata":{}},{"cell_type":"code","source":"%%capture\n! pip install git+https://github.com/huggingface/transformers.git","metadata":{"_kg_hide-output":true,"_kg_hide-input":true,"execution":{"iopub.status.busy":"2023-04-27T18:34:09.410212Z","iopub.execute_input":"2023-04-27T18:34:09.410721Z","iopub.status.idle":"2023-04-27T18:34:53.465358Z","shell.execute_reply.started":"2023-04-27T18:34:09.410675Z","shell.execute_reply":"2023-04-27T18:34:53.464017Z"},"trusted":true},"execution_count":1,"outputs":[]},{"cell_type":"code","source":"import os\nimport sys\nimport cv2\nimport json\nimport glob\nimport numpy as np\nimport pandas as pd\nfrom tqdm.auto import tqdm\nimport matplotlib.pyplot as plt\n\nfrom random import shuffle\n\nimport wandb\nimport torch\nimport torch.nn as nn\nfrom torch.utils.data import DataLoader, Dataset\nfrom torch.cuda.amp import GradScaler, autocast\n\nfrom transformers import (\n AutoProcessor,\n Pix2StructConfig,\n Pix2StructForConditionalGeneration,\n get_linear_schedule_with_warmup,\n)\n\nimport albumentations as A\nfrom albumentations.pytorch import ToTensorV2\n\nimport warnings\nwarnings.simplefilter(\"ignore\")","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","execution":{"iopub.status.busy":"2023-04-27T18:34:53.46782Z","iopub.execute_input":"2023-04-27T18:34:53.468593Z","iopub.status.idle":"2023-04-27T18:35:05.073897Z","shell.execute_reply.started":"2023-04-27T18:34:53.468528Z","shell.execute_reply":"2023-04-27T18:35:05.072833Z"},"trusted":true},"execution_count":2,"outputs":[]},{"cell_type":"code","source":"Config = {\n 'IMAGE_DIR': '/kaggle/input/benetech-making-graphs-accessible/train/images/',\n 'MAX_PATCHES': 1024,\n 'MODEL_NAME': 'ybelkada/pix2struct-base',\n 'IMG_SIZE': (256, 256),\n 'MAX_LEN': 256,\n 'LR': 3e-5,\n 'NB_EPOCHS': 2,\n 'TRAIN_BS': 2,\n 'VALID_BS': 2,\n 'ALL_SAMPLES': int(1e+100),\n# '_wandb_kernel': 'tanaym',\n}","metadata":{"execution":{"iopub.status.busy":"2023-04-27T18:35:05.075349Z","iopub.execute_input":"2023-04-27T18:35:05.077177Z","iopub.status.idle":"2023-04-27T18:35:05.084174Z","shell.execute_reply.started":"2023-04-27T18:35:05.077128Z","shell.execute_reply":"2023-04-27T18:35:05.082954Z"},"trusted":true},"execution_count":3,"outputs":[]},{"cell_type":"markdown","source":"### About W&B:\n
\"Weights

\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":"def wandb_log(**kwargs):\n for k, v in kwargs.items():\n wandb.log({k: v})\n\n# Start W&B logging\n# W&B Login\nfrom kaggle_secrets import UserSecretsClient\nuser_secrets = UserSecretsClient()\nwb_key = user_secrets.get_secret(\"WANDB_API_KEY\")\n\nwandb.login(key=wb_key)\n\nrun = wandb.init(\n project='pytorch',\n config=Config,\n group='multi_modal',\n job_type='train',\n)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Let's add chart types as special tokens and a special BOS token\nBOS_TOKEN = \"<|BOS|>\"\nX_START = \"\"\nX_END = \"\"\nY_START = \"\"\nY_END = \"\"\n\nnew_tokens = [\n \"\",\n \"\",\n \"\",\n \"\",\n \"\",\n X_START,\n X_END,\n Y_START,\n Y_END,\n BOS_TOKEN,\n]","metadata":{"execution":{"iopub.status.busy":"2023-04-27T18:35:05.087975Z","iopub.execute_input":"2023-04-27T18:35:05.088353Z","iopub.status.idle":"2023-04-27T18:35:05.104289Z","shell.execute_reply.started":"2023-04-27T18:35:05.088315Z","shell.execute_reply":"2023-04-27T18:35:05.103115Z"},"trusted":true},"execution_count":4,"outputs":[]},{"cell_type":"markdown","source":"Just resizing and graph image normalization as augments for now!","metadata":{}},{"cell_type":"code","source":"def augments():\n return A.Compose([\n A.Resize(width=Config['IMG_SIZE'][0], height=Config['IMG_SIZE'][1]),\n A.Normalize(\n mean=[0, 0, 0],\n std=[1, 1, 1],\n max_pixel_value=255,\n ),\n ToTensorV2(),\n ])","metadata":{"execution":{"iopub.status.busy":"2023-04-27T18:35:05.106014Z","iopub.execute_input":"2023-04-27T18:35:05.106376Z","iopub.status.idle":"2023-04-27T18:35:05.114705Z","shell.execute_reply.started":"2023-04-27T18:35:05.106339Z","shell.execute_reply":"2023-04-27T18:35:05.113755Z"},"trusted":true},"execution_count":5,"outputs":[]},{"cell_type":"code","source":"class BeneTechDataset(Dataset):\n def __init__(self, dataset, processor, augments=None):\n self.dataset = dataset\n self.processor = processor\n self.augments = augments\n\n def __len__(self):\n return len(self.dataset)\n\n def __getitem__(self, idx):\n item = self.dataset[idx]\n image = cv2.imread(item['image'])\n if self.augments:\n image = self.augments(image=image)['image']\n encoding = self.processor(\n images=image,\n return_tensors=\"pt\", \n add_special_tokens=True, \n max_patches=Config['MAX_PATCHES']\n )\n \n encoding = {k:v.squeeze() for k,v in encoding.items()}\n encoding[\"text\"] = item[\"label\"]\n return encoding","metadata":{"execution":{"iopub.status.busy":"2023-04-27T18:35:05.116389Z","iopub.execute_input":"2023-04-27T18:35:05.116774Z","iopub.status.idle":"2023-04-27T18:35:05.127312Z","shell.execute_reply.started":"2023-04-27T18:35:05.116733Z","shell.execute_reply":"2023-04-27T18:35:05.126284Z"},"trusted":true},"execution_count":6,"outputs":[]},{"cell_type":"markdown","source":"If you make changes to the vocab or anything else in the below cell, please don't forget to resize model token embeddings (as done below).","metadata":{}},{"cell_type":"code","source":"def get_model(extra_tokens=new_tokens):\n processor = AutoProcessor.from_pretrained(Config['MODEL_NAME'])\n model = Pix2StructForConditionalGeneration.from_pretrained(Config['MODEL_NAME'])\n processor.image_processor.size = {\n \"height\": Config['IMG_SIZE'][0],\n \"width\": Config['IMG_SIZE'][1],\n }\n\n processor.tokenizer.add_tokens(extra_tokens)\n model.resize_token_embeddings(len(processor.tokenizer))\n return processor, model","metadata":{"execution":{"iopub.status.busy":"2023-04-27T18:35:05.128993Z","iopub.execute_input":"2023-04-27T18:35:05.129409Z","iopub.status.idle":"2023-04-27T18:35:05.138707Z","shell.execute_reply.started":"2023-04-27T18:35:05.129372Z","shell.execute_reply":"2023-04-27T18:35:05.137739Z"},"trusted":true},"execution_count":7,"outputs":[]},{"cell_type":"code","source":"def collator(batch):\n new_batch = {\"flattened_patches\":[], \"attention_mask\":[]}\n texts = [item[\"text\"] for item in batch]\n text_inputs = processor(\n text=texts, \n padding=\"max_length\", \n truncation=True, \n return_tensors=\"pt\", \n add_special_tokens=True, \n max_length=Config['MAX_LEN']\n )\n new_batch[\"labels\"] = text_inputs.input_ids\n for item in batch:\n new_batch[\"flattened_patches\"].append(item[\"flattened_patches\"])\n new_batch[\"attention_mask\"].append(item[\"attention_mask\"])\n new_batch[\"flattened_patches\"] = torch.stack(new_batch[\"flattened_patches\"])\n new_batch[\"attention_mask\"] = torch.stack(new_batch[\"attention_mask\"])\n\n return new_batch","metadata":{"execution":{"iopub.status.busy":"2023-04-27T18:35:05.14068Z","iopub.execute_input":"2023-04-27T18:35:05.141172Z","iopub.status.idle":"2023-04-27T18:35:05.150185Z","shell.execute_reply.started":"2023-04-27T18:35:05.141135Z","shell.execute_reply":"2023-04-27T18:35:05.149228Z"},"trusted":true},"execution_count":8,"outputs":[]},{"cell_type":"code","source":"def train_one_epoch(model, processor, train_loader, optimizer, scaler):\n \"\"\"\n Trains the model on all batches for one epoch with NVIDIA's AMP\n \"\"\"\n model.train()\n avg_loss = 0\n with autocast():\n prog_bar = tqdm(enumerate(train_loader), total=len(train_loader))\n for idx, batch in prog_bar:\n labels = batch.pop(\"labels\").to('cuda')\n flattened_patches = batch.pop(\"flattened_patches\").to('cuda')\n attention_mask = batch.pop(\"attention_mask\").to('cuda')\n\n outputs = model(\n flattened_patches=flattened_patches,\n attention_mask=attention_mask,\n labels=labels\n )\n\n loss = outputs.loss\n scaler.scale(loss).backward()\n scaler.step(optimizer)\n scaler.update()\n optimizer.zero_grad(set_to_none=True)\n prog_bar.set_description(f\"loss: {loss.item():.4f}\")\n avg_loss += loss.item()\n \n avg_loss = avg_loss / len(train_loader)\n print(f\"Average training loss: {avg_loss:.4f}\")\n return avg_loss\n\n@torch.no_grad()\ndef valid_one_epoch(model, processor, valid_loader):\n \"\"\"\n Validates the model on all batches (in val set) for one epoch\n \"\"\"\n model.eval()\n avg_loss = 0\n prog_bar = tqdm(enumerate(valid_loader), total=len(valid_loader))\n for idx, batch in prog_bar:\n labels = batch.pop(\"labels\").to('cuda')\n flattened_patches = batch.pop(\"flattened_patches\").to('cuda')\n attention_mask = batch.pop(\"attention_mask\").to('cuda')\n \n outputs = model(\n flattened_patches=flattened_patches,\n attention_mask=attention_mask,\n labels=labels\n )\n \n loss = outputs.loss\n prog_bar.set_description(f\"loss: {loss.item():.4f}\")\n avg_loss += loss.item()\n \n avg_loss = avg_loss / len(valid_loader)\n print(f\"Average validation loss: {avg_loss:.4f}\")\n return avg_loss","metadata":{"execution":{"iopub.status.busy":"2023-04-27T18:35:05.151984Z","iopub.execute_input":"2023-04-27T18:35:05.152438Z","iopub.status.idle":"2023-04-27T18:35:05.165213Z","shell.execute_reply.started":"2023-04-27T18:35:05.1524Z","shell.execute_reply":"2023-04-27T18:35:05.164413Z"},"trusted":true},"execution_count":9,"outputs":[]},{"cell_type":"code","source":"def fit(model, processor, train_loader, valid_loader, optimizer, scaler):\n \"\"\"\n A nice function that binds it all together and reminds me of Keras days from 2018 :)\n \"\"\"\n best_val_loss = int(1e+5)\n for epoch in range(2):\n print(f\"{'='*20} Epoch: {epoch+1} / {Config['NB_EPOCHS']} {'='*20}\")\n _ = train_one_epoch(model, processor, train_loader, optimizer, scaler)\n val_avg_loss = valid_one_epoch(model, processor, valid_loader)\n \n if val_avg_loss < best_val_loss:\n best_val_loss = val_avg_loss\n print(f\"Saving best model so far with loss: {best_val_loss:.4f}\")\n torch.save(model.state_dict(), f\"pix2struct_base_benetech.pt\")\n print(f\"Best model with val_loss: {best_val_loss:.4f}\")","metadata":{"execution":{"iopub.status.busy":"2023-04-27T18:35:05.169309Z","iopub.execute_input":"2023-04-27T18:35:05.169596Z","iopub.status.idle":"2023-04-27T18:35:05.178724Z","shell.execute_reply.started":"2023-04-27T18:35:05.169551Z","shell.execute_reply":"2023-04-27T18:35:05.177739Z"},"trusted":true},"execution_count":10,"outputs":[]},{"cell_type":"code","source":"# Training cell\n# if __name__ == \"__main__\":\n# Read the processed JSON file\nwith open(\"/kaggle/input/benetech-processed-data-json/data.json\", \"r\") as fl:\n dataset = json.load(fl)['data']\n\n# Shuffle the dataset and select however samples you want for training\nshuffle(dataset)\ndataset = dataset[:Config['ALL_SAMPLES']]\n\n# We are splitting the data naively for now\nsplit = 0.90\ntrain_samples = int(len(dataset) * split)\ntrain_ds = dataset[:train_samples+1]\nvalid_ds = dataset[train_samples:]\n\n# Yeah all that\nprocessor, model = get_model()\nmodel.to('cuda')\n# wandb.watch(model)\noptimizer = torch.optim.Adam(params=model.parameters(), lr=Config['LR'])\n\n# Load the data into Datasets and then make DataLoaders for training\ntrain_dataset = BeneTechDataset(train_ds, processor, augments=augments())\ntrain_dataloader = DataLoader(train_dataset, shuffle=True, batch_size=Config['TRAIN_BS'], collate_fn=collator)\n\nvalid_dataset = BeneTechDataset(valid_ds, processor, augments=augments())\nvalid_dataloader = DataLoader(valid_dataset, shuffle=False, batch_size=Config['VALID_BS'], collate_fn=collator)\n\nnb_train_steps = int(train_samples / Config['TRAIN_BS'] * Config['NB_EPOCHS'])\n\n# Print out the data sizes we are training on\nprint(f\"Training on {len(train_ds)} samples, Validating on {len(valid_ds)} samples\")\n\n# Train the model now\nfit(\n model=model,\n processor=processor,\n train_loader=train_dataloader,\n valid_loader=valid_dataloader,\n optimizer=optimizer,\n scaler=GradScaler(),\n)","metadata":{"execution":{"iopub.status.busy":"2023-04-27T18:35:05.180491Z","iopub.execute_input":"2023-04-27T18:35:05.181573Z","iopub.status.idle":"2023-04-28T03:35:14.463666Z","shell.execute_reply.started":"2023-04-27T18:35:05.181536Z","shell.execute_reply":"2023-04-28T03:35:14.462151Z"},"trusted":true},"execution_count":11,"outputs":[{"output_type":"display_data","data":{"text/plain":"Downloading (…)rocessor_config.json: 0%| | 0.00/231 [00:00\u001b[0;34m\u001b[0m\n\u001b[0;32m----> 1\u001b[0;31m \u001b[0mval_avg_loss\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mvalid_one_epoch\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mmodel\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mprocessor\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mvalid_loader\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m","\u001b[0;31mNameError\u001b[0m: name 'valid_loader' is not defined"],"ename":"NameError","evalue":"name 'valid_loader' is not defined","output_type":"error"}]},{"cell_type":"code","source":"valid_loader","metadata":{"execution":{"iopub.status.busy":"2023-04-27T11:34:41.142702Z","iopub.execute_input":"2023-04-27T11:34:41.143421Z","iopub.status.idle":"2023-04-27T11:34:41.230516Z","shell.execute_reply.started":"2023-04-27T11:34:41.143379Z","shell.execute_reply":"2023-04-27T11:34:41.228571Z"},"trusted":true},"execution_count":1,"outputs":[{"traceback":["\u001b[0;31m---------------------------------------------------------------------------\u001b[0m","\u001b[0;31mNameError\u001b[0m Traceback (most recent call last)","\u001b[0;32m/tmp/ipykernel_23/3582571110.py\u001b[0m in \u001b[0;36m\u001b[0;34m\u001b[0m\n\u001b[0;32m----> 1\u001b[0;31m \u001b[0mvalid_loader\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m","\u001b[0;31mNameError\u001b[0m: name 'valid_loader' is not defined"],"ename":"NameError","evalue":"name 'valid_loader' is not defined","output_type":"error"}]},{"cell_type":"markdown","source":"
\n\n
","metadata":{}}]}