{"cells":[{"metadata":{},"cell_type":"markdown","source":"# Bengali.AI: Self-supervised pretraining with Context Encoders (fastai2)\n\nThis notebook implements the paper [Context Encoders: Feature Learning by Inpainting](https://arxiv.org/abs/1604.07379) by Deepak Pathak, Philipp Krahenbuhl, Jeff Donahue, Trevor Darrell and Alexei A. Efros.\n\nMy hypothesis is that a model that can successfully fill in the blanks with Bengali handwritten characters, has a pretty darn good understanding of the Bengali language and thus can be repurposed to classify constituent elements of the characters.\n\n\n## 0. Table of contents\n\n1. Context Encoders overview\n2. Libraries and hyperparams\n3. Generating input and targets\n4. Building the model\n5. Training\n6. Visualising the results"},{"metadata":{},"cell_type":"markdown","source":"## 1. Context Encoders overview\n\nThe Context Encoders paper describes a simple pre-training task: remove stuff from images and have the model try to predict what was removed.\n\n[![image.png](https://i.postimg.cc/xTdGCDZC/image.png)](https://postimg.cc/K4d3qVxS)\n\nThe paper describes a number of different techniques for removing stuff. The simplest is to just extract a region from the centre of the image which is what I'm doing in this kernel. In future, I might try the other techniques - I am a little concerned that the task may be a little too easy.\n\nI have made a slight modification to the paper by using EfficientNet-B0 instead of AlexNet for the encoder. Also, I'm only concerning myself with Reconstruction Loss (L2) as I don't mind blurry reconstructions just that the weights can be transferred to the classification task.\n\nIt uses the development version of fastai2."},{"metadata":{},"cell_type":"markdown","source":"## 2. Libraries and hyperparams"},{"metadata":{},"cell_type":"markdown","source":"The library makes use of the development version of [fastai2](https://github.com/fastai/fastai2). Since it isn't available in kernels yet, I'll install it alongside [EfficientNet-PyTorch](https://github.com/zhoudaxia233/EfficientUnet-PyTorch)."},{"metadata":{"trusted":true},"cell_type":"code","source":"!pip install git+https://github.com/fastai/fastai2 > /dev/null\n!pip install efficientnet-pytorch > /dev/null","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"from pathlib import Path\n\nimport pandas as pd\n\nimport torch\nfrom efficientnet_pytorch import EfficientNet\nfrom torch.utils import model_zoo\n\nfrom fastai2.basics import *\nfrom fastai2.data.all import *\nfrom fastai2.callback.all import *\nfrom fastai2.vision.all import *","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"DATA_PATH = Path('/kaggle/input/bengaliai-cv19')\nIMAGE_DATA_PATH = Path('/kaggle/input/grapheme-imgs-128x128')\nOUTPUT_PATH = Path('/kaggle/working')\n\nVALID_PCT = 0.2\nSEED = 420\nBATCH_SIZE = 64\nCROP_SIZE = 32\nIMG_SIZE = 128","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train_df = pd.read_csv(DATA_PATH/'train.csv')","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## 3. Generating input and targets"},{"metadata":{},"cell_type":"markdown","source":"I'm using 2 transforms: one to remove the centre of an image (the `X`) and another to return just the centre (the `y`)."},{"metadata":{"trusted":true},"cell_type":"code","source":"class ImageWithCenterRemoved(Transform):\n    \"\"\"Transform that removes the center part of an image.\"\"\"\n    \n    order = 6\n\n    def __init__(self, crop_size=CROP_SIZE):\n        self.crop_size = crop_size\n\n    def encodes(self, x:PILImageBW) -> PILImageBW:\n        x = array(x)\n    \n        start_height = tuple(IMG_SIZE // 2 - (CROP_SIZE // 2))\n        start_width = tuple(IMG_SIZE // 2 - (CROP_SIZE // 2)) \n        \n        x[\n            ...,\n            start_height:start_height+self.crop_size,\n            start_width:start_width+self.crop_size\n        ] = 0\n    \n        return PILImageBW(Image.fromarray(x))\n    \n    def encodes(self, x:TensorImage):\n        start_height = IMG_SIZE // 2 - (CROP_SIZE // 2)\n        start_width = IMG_SIZE // 2 - (CROP_SIZE // 2)\n        \n        x[\n            ...,\n            start_height:start_height+self.crop_size,\n            start_width:start_width+self.crop_size\n        ] = 0\n        \n        return TensorImage(x)\n    \n    \nclass ImageWithOnlyCenter(Transform):\n    \"\"\"Transform that keeps only the center part of an image.\"\"\"\n    \n    order = 6\n    \n    def __init__(self, crop_size=CROP_SIZE):\n        self.crop_size = crop_size\n\n    def encodes(self, x:TensorImage) -> PILImageBW:\n        start_height = IMG_SIZE // 2 - (CROP_SIZE // 2)\n        start_width = IMG_SIZE // 2 - (CROP_SIZE // 2)\n        \n        output = x[\n            ...,\n            start_height:start_height + self.crop_size,\n            start_width:start_width + self.crop_size\n        ]\n\n        return TensorImage(output)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"items = get_image_files(IMAGE_DATA_PATH)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Next I create a `Datasets` instance splitting the data into a 80/20 train/val split."},{"metadata":{"trusted":true},"cell_type":"code","source":"x_tfms = [PILImageBW.create, ToTensor, ImageWithCenterRemoved()]\ny_tfms = [PILImageBW.create, ToTensor, ImageWithOnlyCenter()]\ntfms = [x_tfms, y_tfms]\n\nsplitter = RandomSplitter(VALID_PCT, seed=SEED)\n\ntds = Datasets(items, tfms, splits=splitter(items))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"imagenet_stats","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Lastly, I use the ImageNet stats to normalise the data, since I will be using a pretrained EfficientNet-B0 model."},{"metadata":{"trusted":true},"cell_type":"code","source":"dl_tfms = [IntToFloatTensor,  Normalize(mean=0.485, std=0.229)]\n\ntrain_dl = TfmdDL(tds.train, bs=BATCH_SIZE, after_batch=dl_tfms)\nvalid_dl = TfmdDL(tds.valid, bs=BATCH_SIZE, after_batch=dl_tfms)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"As you can see, we're now successfully cutting the centre out of an image and have the centre crop as the label."},{"metadata":{"trusted":true},"cell_type":"code","source":"train_dl.show_batch()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Lastly, I'll put the data into a `DataLoaders` class (formally `DataBunch`)."},{"metadata":{"trusted":true},"cell_type":"code","source":"data = DataLoaders(train_dl, valid_dl)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## 4. Building the model"},{"metadata":{},"cell_type":"markdown","source":"The model described below is very similar to a [Unet](https://arxiv.org/abs/1505.04597) model. In that it has an encoder which is responsible for generating a series of downsampled features, then a decoder which upsamples the generates features using a series of fractionally-strided convolution operations.\n\n[![image.png](https://i.postimg.cc/BZ2DnZN4/image.png)](https://postimg.cc/t7C7rjLM)\n\nNote that I am omitting the Adversarial Loss, which is more concerned with real looking results - I only want transferability.\n\nFor the encoder, I simply repurpose the [EfficientNet-PyTorch](https://github.com/lukemelas/EfficientNet-PyTorch) model by removing the final classification layers of the model."},{"metadata":{"trusted":true},"cell_type":"code","source":"class EfficientNetEncoder(EfficientNet):\n    def __init__(self, *args, **kwargs):\n        super().__init__(*args, **kwargs)\n        \n        # the initial layer to convolve into 3 channels\n        # idea from https://www.kaggle.com/aleksandradeis/bengali-ai-efficientnet-pytorch-starter\n        self.input_conv = nn.Conv2d(in_channels=1, out_channels=3, kernel_size=1)\n\n    def forward(self, inputs):\n        x = self.input_conv(inputs)\n        return self.extract_features(x)\n    \n    @classmethod\n    def load_pretrained(cls):\n        model_name = 'efficientnet-b0'\n        model = cls.from_name(model_name, override_params={'num_classes': 1})\n        model_dict = model.state_dict()\n\n        state_dict = model_zoo.load_url('https://publicmodels.blob.core.windows.net/container/aa/efficientnet-b0-355c32eb.pth')\n        state_dict_no_fc = {k: v for k, v in state_dict.items() if not k.startswith('_fc')}\n        model_dict.update(state_dict_no_fc)\n        \n        model.load_state_dict(model_dict)\n\n        return model","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"For the decoder, I create a series of fractional convolutions that upsamples the image to the size of the crop. The paper uses BatchNorm and ReLU in the decoder, so I'm doing the same here."},{"metadata":{"trusted":true},"cell_type":"code","source":"def up_conv(in_channels, out_channels):\n    return nn.Sequential(\n        nn.ConvTranspose2d(in_channels, out_channels, kernel_size=2, stride=2),\n        nn.BatchNorm2d(out_channels),\n        nn.ReLU(inplace=True)\n    )","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class Decoder(nn.Module):\n\n    def __init__(self, encoder, n_channels, out_channels=1):\n        super().__init__()\n\n        self.encoder = encoder\n\n        self.up_conv1 = up_conv(n_channels, 256)    \n        self.up_conv2 = up_conv(256, 128)    # 8x8\n        self.up_conv3 = up_conv(128, 64)    # 16x16\n        self.final_conv = nn.Conv2d(64, out_channels, kernel_size=1)\n    \n    def forward(self, x):\n        x = self.encoder(x)     # input: 1x128x128, output: 1280x4x4\n        x = self.up_conv1(x)    # input: 1280x4x4, output: 256x8x8\n        x = self.up_conv2(x)    # input: 256x8x8, output: 128x16x16\n        x = self.up_conv3(x)    # input: 128x16x16, output: 64x32x32\n        x = self.final_conv(x)  # input: 64x32x32, output: 1x32x32\n        \n        return x","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"encoder = EfficientNetEncoder.load_pretrained()\nmodel = Decoder(encoder, n_channels=1280)  # 1280: EfficientNet b0 output. To do: don't hardcode this.","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## 5. Training"},{"metadata":{},"cell_type":"markdown","source":"With everything in place, the model can just be trained like we normally would in Fast.ai. I'm using standard [OneCycle](https://mc.ai/finding-good-learning-rate-and-the-one-cycle-policy/) training as is familiar to people who have done the Fast.ai course or used the library."},{"metadata":{"trusted":true},"cell_type":"code","source":"if torch.cuda.is_available():\n    print('Cuda available')\n    model = model.cuda()\n    data = data.cuda()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"learner = Learner(data, model, loss_func=nn.MSELoss())","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# learner.lr_find()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"learner.fit_one_cycle(4, 1e-3)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"learner.recorder.plot_loss()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"learner.validate()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## 6. Visualising results"},{"metadata":{"trusted":true},"cell_type":"code","source":"learner.show_results(ds_idx=1)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"That looks pretty good.\n\nIn an upcoming kernel, I'll do a mini-ableation study to see if these weights are useful for transfer learning on the competition's multilabel classification problem."},{"metadata":{"trusted":true},"cell_type":"code","source":"learner.save('model_cycle_1')","execution_count":null,"outputs":[]}],"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":1}