{"cells":[{"metadata":{},"cell_type":"markdown","source":"![](\"https://i.imgur.com/F4T8Ys2.jpg\")
","execution_count":null},{"metadata":{},"cell_type":"markdown","source":"# Introduction\n\nHello everyone! Welcome to the \"ALASKA2 Image Steganalysis\" competition on Kaggle! In this competition, contestants are challenged to build machine learning models to predict whether a message has been hidden in an image using steganography. An accurate solution to this problem can open up new possibilities in the area of steganalysis and deep learning-based encryption.\n\nIn this kernel, I will demonstrate how one can finetune EfficientNet-B3 to solve this task using PyTorch. I will use PyTorch v1.5 and Kaggle's TPU v3-8 to train the model on 8 folds and 1 epoch in less than three hours.","execution_count":null},{"metadata":{},"cell_type":"markdown","source":"# Acknowledgements\n\n1. [PyTorch XLA ~ by PyTorch](https://pytorch.org/xla/release/1.5/index.html)\n2. [Torchvision Models ~ by PyTorch](https://pytorch.org/docs/stable/torchvision/models.html)\n3. [PANDA / submit test ~ Yasufumi Nakama](https://www.kaggle.com/yasufuminakama/panda-submit-test)\n4. [Super-duper fast pytorch tpu kernel... ~ by Abhishek](https://www.kaggle.com/abhishek/super-duper-fast-pytorch-tpu-kernel)","execution_count":null},{"metadata":{},"cell_type":"markdown","source":"# Contents\n\n* [Preparing the ground](#1)\n * [Set up PyTorch-XLA](#1.1)\n * [Import libraries](#1.2)\n * [Set hyperparameters and paths](#1.3)\n * [Load .csv data](#1.4)\n * [Display few images](#1.5)\n\n \n* [Modeling](#2)\n * [Build PyTorch dataset](#2.1)\n * [Build EfficientNet-B3 model](#2.2)\n * [Split 300, 000 images into 8 folds](#2.3)\n * [Define cross entropy and accuracy](#2.4)\n * [Define helper function for training logs](#2.5)\n * [Train model on all 8 TPU cores in parallel](#2.6)\n\n\n* [Takeaways](#3)","execution_count":null},{"metadata":{},"cell_type":"markdown","source":"# Preparing the ground ","execution_count":null},{"metadata":{},"cell_type":"markdown","source":"## Set up PyTorch-XLA (inspired by Abhishek's kernel :D)","execution_count":null},{"metadata":{"_kg_hide-input":true,"_kg_hide-output":true,"trusted":true},"cell_type":"code","source":"!curl https://raw.githubusercontent.com/pytorch/xla/master/contrib/scripts/env-setup.py -o pytorch-xla-env-setup.py\n!python pytorch-xla-env-setup.py --version nightly --apt-packages libomp5 libopenblas-dev\n!export XLA_USE_BF16=1","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Import libraries (for data loading, processing, and modeling on TPU)","execution_count":null},{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"import os\nimport gc\nimport cv2\nimport sys\nimport time\nimport copy\n\nimport numpy as np\nimport pandas as pd\nfrom colorama import Fore\nfrom tqdm.notebook import tqdm\nimport matplotlib.pyplot as plt\nfrom joblib import Parallel, delayed\nfrom keras.utils import to_categorical\nfrom sklearn.model_selection import KFold\n\nimport torch\nimport torch.nn as nn\nimport torch.utils.data\nfrom torch.hub import load\nfrom torch.optim import Adam\nfrom torch import DoubleTensor, FloatTensor, LongTensor\n\nimport torch_xla.core.xla_model as xm\nimport torch_xla.distributed.parallel_loader as pl\nimport torch_xla.distributed.xla_multiprocessing as xmp\n\nfrom torchvision.transforms import Normalize\nfrom torch.utils.data import Dataset, DataLoader, sampler\nfrom albumentations import VerticalFlip, HorizontalFlip, Compose","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Set hyperparamerters and paths (adjust these to improve CV and LB :D)","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"H = 512\nW = 512\nVF = 0.5\nHF = 0.5\nDELAY = 30\nFRAC = 0.1\nDROP = 0.225\n\nFOLDS = 8\nEPOCHS = 3\nLR = 1e-3, 1e-3\nBATCH_SIZE = 32\nVAL_BATCH_SIZE = 32\nMODEL_NAME = 'efficientnet_b3'\nMODEL = 'rwightman/gen-efficientnet-pytorch'","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true},"cell_type":"code","source":"PATH = '../input/'\nMODEL_PATH = 'efficientnet_model'\nDATA_PATH = PATH + 'alaska2-image-steganalysis/'\nSAMPLE_SUB_PATH = DATA_PATH + 'sample_submission.csv'\n\nTEST_PATH = DATA_PATH + 'Test/'\nUERD_PATH = DATA_PATH + 'UERD/'\nCOVER_PATH = DATA_PATH + 'Cover/'\nJMiPOD_PATH = DATA_PATH + 'JMiPOD/'\nJUNIWARD_PATH = DATA_PATH + 'JUNIWARD/'\nTRAIN_PATHS = [COVER_PATH, JMiPOD_PATH, JUNIWARD_PATH, UERD_PATH]","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Load .csv data (to access image IDs for training and validation)","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"sample_submission = pd.read_csv(SAMPLE_SUB_PATH)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"sample_submission.head()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Display few images (from Test directory)","execution_count":null},{"metadata":{"trusted":true,"_kg_hide-input":false,"_kg_hide-output":false},"cell_type":"code","source":"def display_images(num):\n sq_num = np.sqrt(num)\n assert sq_num == int(sq_num)\n\n sq_num = int(sq_num)\n image_ids = os.listdir(TEST_PATH)\n fig, ax = plt.subplots(nrows=sq_num, ncols=sq_num, figsize=(20, 20))\n\n for i in range(sq_num):\n for j in range(sq_num):\n idx = i*sq_num + j\n img = cv2.imread(TEST_PATH + image_ids[idx])\n ax[i, j].imshow(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))\n ax[i, j].set_title('Image {}'.format(idx), fontsize=12)\n\n plt.show()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true,"_kg_hide-input":true},"cell_type":"code","source":"display_images(36)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Modeling ","execution_count":null},{"metadata":{},"cell_type":"markdown","source":"## Build PyTorch dataset (with image transforms and targets)","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"def get_img(path, aug):\n img = cv2.imread(path)\n img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)/255\n return aug(image=cv2.resize(img, (H, W)))['image']\n\nclass ALASKADataset(Dataset):\n def __init__(self, image_id, is_test, is_val):\n self.is_test = is_test\n self.image_id = image_id\n self.no_aug = is_test or is_val\n\n self.vertical = VerticalFlip(p=VF)\n self.horizontal = HorizontalFlip(p=HF)\n if self.no_aug: self.transform = lambda image: {'image': image}\n else: self.transform = Compose([self.vertical, self.horizontal], p=1)\n\n def __len__(self):\n multiplier = 1 if self.is_test else 4\n return multiplier*len(self.image_id)\n \n def __getitem__(self, idx):\n index = idx%len(self.image_id)\n\n if self.is_test:\n category = None\n path = TEST_PATH + self.image_id[index]\n return FloatTensor(get_img(path, self.transform))\n else:\n target = idx/len(self.image_id)\n category = [int(np.floor(target) > 0)]\n path = TRAIN_PATHS[int(target)] + self.image_id[index]\n return FloatTensor(get_img(path, self.transform)), FloatTensor(category)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Build EfficientNet-B3 model (with a custom Dense head)","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"class ENSModel(nn.Module):\n def __init__(self):\n super(ENSModel, self).__init__()\n self.dropout = nn.Dropout(p=DROP)\n self.dense_output = nn.Linear(1536, 1)\n self.efn = load(MODEL, MODEL_NAME, pretrained=True)\n self.efn = nn.Sequential(*list(self.efn.children())[:-1])\n \n def forward(self, x):\n x = x.reshape(-1, 3, H, W)\n return self.dense_output(self.dropout(self.efn(x).reshape(-1, 1536)))","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Split 300,000 images into 8 folds (for cross-validation)","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"kfolds = KFold(n_splits=FOLDS)\nimage_id = os.listdir(COVER_PATH)\nsplit_indices = kfolds.split(image_id)\n\nval_ids, train_ids = [], []\nfor index in split_indices:\n val_ids.append(np.array(image_id)[index[1]])\n train_ids.append(np.array(image_id)[index[0]])","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Define binary cross entropy and accuracy (for backpropagation)","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"def bce(inp, targ):\n return nn.BCELoss()(nn.Sigmoid()(inp), targ)\n\ndef acc(inp, targ):\n targ_idx = targ.squeeze()\n inp_idx = torch.round(nn.Sigmoid()(inp)).squeeze()\n return (inp_idx == targ_idx).float().sum(axis=0)/len(inp_idx)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Define helper function for training logs (to check training status)","execution_count":null},{"metadata":{"trusted":true,"_kg_hide-input":true},"cell_type":"code","source":"def print_metric(data, fold, start, end, metric, typ):\n n, value = \"Steganalysis\", np.round(data.item(), 3)\n g, c, y, r = Fore.GREEN, Fore.CYAN, Fore.YELLOW, Fore.RESET\n \n tick = g + '\\u2714' + r\n t = typ, n, metric, c, value, r\n time = np.round(end - start, 1)\n time = \"Time: {}{}{} s\".format(y, time, r)\n string = \"FOLD {} \".format(fold + 1) + tick + \" \"\n print(string + \"{} {} {}: {}{}{}\".format(*t) + \" \" + time)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Train model on all 8 TPU cores in parallel (one fold per core)","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"class ImbSamp(sampler.Sampler):\n\n def __len__(self): return self.num_samples\n def __iter__(self): return (self.indices[i] for i in self._get_probs())\n def _get_label(self, dataset, idx): return int(idx/(len(self.dataset)/4) >= 1)\n \n def _get_weight(self, idx, count_dict):\n return 1.0/count_dict[self._get_label(self.dataset, idx)]\n \n def _get_probs(self):\n return torch.multinomial(self.weights, self.num_samples, replacement=True)\n\n def __init__(self, dataset, indices=None, num_samples=None):\n self.indices = list(range(len(dataset))) if indices is None else indices\n self.num_samples = len(self.indices) if num_samples is None else num_samples\n\n count = {}\n self.dataset = dataset\n for idx in self.indices:\n label = self._get_label(dataset, idx)\n if label in count: count[label] += 1\n if label not in count: count[label] = 1\n\n self.weights = DoubleTensor([self._get_weight(idx, count) for idx in self.indices])","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"model = ENSModel()\n\ndef run(fold):\n val = val_ids[fold]\n train = train_ids[fold]\n device = xm.xla_device(fold + 1)\n\n val_set = ALASKADataset(val, False, True)\n train_set = ALASKADataset(train, False, False)\n val_loader = DataLoader(val_set, batch_size=VAL_BATCH_SIZE)\n train_loader = DataLoader(train_set, batch_size=BATCH_SIZE, sampler=ImbSamp(train_set))\n\n network = copy.deepcopy(model).to(device)\n optimizer = Adam([{'params': network.efn.parameters(), 'lr': LR[0]},\n {'params': network.dense_output.parameters(), 'lr': LR[1]}])\n\n start = time.time()\n for epoch in range(EPOCHS):\n\n batch = 1\n for train_batch in train_loader:\n train_img, train_targs = train_batch\n \n network = network.to(device)\n train_img = train_img.to(device)\n train_targs = train_targs.to(device)\n \n network.train()\n train_preds = network.forward(train_img)\n train_loss = bce(train_preds, train_targs)\n\n optimizer.zero_grad()\n train_loss.backward()\n xm.optimizer_step(optimizer, barrier=True)\n\n batch = batch + 1\n if batch >= FRAC*len(train_loader): break\n\n network.eval()\n val_loss, val_acc = 0, 0\n for val_batch in tqdm(val_loader):\n\n img, targ = val_batch\n with torch.no_grad():\n img = img.to(device)\n targ = targ.to(device)\n network = network.to(device)\n pred = network.forward(val_img)\n \n pred = network.forward(img)\n val_acc += acc(pred, targ.squeeze(dim=1)).item()*len(pred)\n val_loss += bce(pred, targ.squeeze(dim=1)).item()*len(pred)\n\n end = time.time()\n time.sleep(DELAY*fold)\n network = network.cpu()\n model_path = MODEL_PATH + \"_{}.pt\"\n \n val_acc /= len(val_set)\n val_loss /= len(val_set)\n print_metric(val_acc, fold, start, end, metric=\"Accuracy\", typ=\"Val\")\n torch.save(network.state_dict(), model_path.format(fold + 1)); del network; gc.collect()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"Parallel(n_jobs=FOLDS, backend=\"threading\")(delayed(run)(i) for i in range(FOLDS))","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Takeaways \n\n1. Using all 8 TPU cores in parallel can dramatically speed up KFold training.\n2. Using complex models (like ResNet-152, DenseNet-201, Efficient-B3, etc) can improve the model's performance.\n3. Training and inference should be done in separate notebooks to avoid confusion and make it easier to iterate fast.","execution_count":null}],"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":4,"nbformat_minor":4}