{"cells":[{"metadata":{},"cell_type":"markdown","source":"This code is a fork of the [ignite/VAE.ipynb](https://github.com/pytorch/ignite/blob/master/examples/notebooks/VAE.ipynb) (pytorch-ignite)  guide for training a variational autoencoder on the MNIST dataset. That code is primarily based on the [official PyTorch example](https://github.com/pytorch/examples/tree/master/vae). \n\nTheir goal was to replicate the goal as to replicate  [Auto-Encoding Variational Bayes](https://arxiv.org/abs/1312.6114) by Kingma and Welling.  This paper uses an encoder-decoder architecture to encode images to a vector and then reconstruct the images. Here I am going to run a similar experiment with BengaliAi handwritten characters. The ignite event handlers make debugging training bottle necks easier and the visualisation of the relative difficulty of resolving the finer details of graphemes.\n \nIt must be noted that I **have not** used the target classes (grapheme_root, consonant_diacritic, vowel_diacritic). For those looking for more information about how variational encoders work check out this great guide posted by Louis Tia: \n\nhttps://tiao.io/post/tutorial-on-variational-autoencoders-with-a-concise-keras-implementation/. \n\n> Like all autoencoders, the variational autoencoder is primarily used for unsupervised learning of hidden representations"},{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"# This Python 3 environment comes with many helpful analytics libraries installed\n# It is defined by the kaggle/python docker image: https://github.com/kaggle/docker-python\n# For example, here's several helpful packages to load in \n\nimport numpy as np # linear algebra\nimport pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)\n\nimport matplotlib.pyplot as plt\n\n%matplotlib inline\n\nimport numpy as np\nimport pandas as pd\nimport datetime as dt\nfrom mpl_toolkits.axes_grid1 import ImageGrid\nfrom os import listdir, makedirs\nfrom os.path import join, exists, expanduser\nfrom tqdm import tqdm\nfrom keras.preprocessing import image\n\nimport PIL\n\nimport torch\nfrom torch.utils.data import DataLoader\nfrom torch import nn, optim\nfrom torch.nn import functional as F\n\nimport cv2\nfrom torchvision import datasets, transforms\nimport torchvision\nfrom torchvision.utils import save_image, make_grid\n# from torchvision.transforms import Compose, ToTensor\nfrom ignite.engine import Engine, Events\nfrom ignite.metrics import MeanSquaredError, Loss, RunningAverage\n","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true},"cell_type":"code","source":"import torch\nimport random\nfrom sklearn.model_selection import KFold\n","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Load Raw Data"},{"metadata":{"trusted":true},"cell_type":"code","source":"%%time\nimg_data = pd.read_parquet('../input/bengaliai-cv19/train_image_data_0.parquet').set_index('image_id')\n","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Image Processing Functions\n\nGoal here was to crop out as much background as possible and down sizing image to match the MNIST data (28x28 pixels)"},{"metadata":{"trusted":true},"cell_type":"code","source":"def expand2square(pil_img, background_color=0):\n    width, height = pil_img.size\n    if width == height:\n        return pil_img\n    elif width > height:\n        result = PIL.Image.new(pil_img.mode, (width, width), background_color)\n        result.paste(pil_img, (0, (width - height) // 2))\n        return result\n    else:\n        result = PIL.Image.new(pil_img.mode, (height, height), background_color)\n        result.paste(pil_img, ((height - width) // 2, 0))\n        return result","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"HEIGHT = 137\nWIDTH = 236\n#\n# aiming for the same size as the MNIST data set\nIMGSIZE=28\ndef process_image(img_):\n#     \n    img = img_[5:-5,3:-7]\n    img_inv = 255-img\n    img_invb = cv2.GaussianBlur(img_inv, (7,7) , 0)\n    img_inv = cv2.addWeighted(img_inv, 1.0 + 4.0, img_invb, -4.0, 0) # im1 = im + 4.0*(im - im_blurred)\n\n    mask = cv2.threshold( img_inv, 120, 255, cv2.THRESH_BINARY+cv2.THRESH_OTSU)[1]\n    img_inv=mask\n    img_inv = img_inv[np.ix_(mask.any(1), mask.any(0))]\n\n    kernel = np.ones((1,1),np.uint8)\n    close = cv2.morphologyEx(img_inv, cv2.MORPH_CLOSE, kernel)\n    dilation = cv2.dilate(close,kernel,iterations = 2)\n    aug = transforms.Compose([transforms.ToPILImage(),\n                              expand2square,\n                              transforms.Pad(5),\n                              transforms.Resize((int(IMGSIZE),int(IMGSIZE)),PIL.Image.ANTIALIAS),\n                              transforms.ToTensor(),\n                             ])\n    image = aug(dilation)\n    return image","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### torchvision transforms:\n* **ToPILImage** convert numpy array to PIL Image.\n* **Pad** avoid clipping after futher augmentation.\n* **Resize** to resize the image to target size.\n* **ToTensor** convert to PIL Image to tensor.\n\nThe high scale factor makes some feature impreceptible so this function also implements some morphological transforms(dialation) to increase line thickness and back ground cropping and thresholding in using cv2."},{"metadata":{"trusted":true},"cell_type":"code","source":"from tqdm import tqdm_notebook\ntqdm_notebook().pandas()\nimage_tensor =torch.stack(img_data.progress_apply(lambda x:process_image(x.values.reshape(HEIGHT,WIDTH)),axis=1).values.tolist())","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# LABELS\n\nThese are not used in this notebook(we will be looking at the reconstruction loss and KL divergence) but are loaded into the dataset / dataloader and required by ignite during training and inference."},{"metadata":{"trusted":true},"cell_type":"code","source":"LABELS = '../input/bengaliai-cv19/train.csv'\ndf = pd.read_csv(LABELS).set_index('image_id')\nnunique = list(df.nunique())[1:-1]\nprint(nunique)\ndf.head()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# dataframe to tensors to train/val split pytorch dataset"},{"metadata":{"trusted":true},"cell_type":"code","source":"label_tensor= torch.from_numpy( df.loc[img_data.index,['grapheme_root','vowel_diacritic','consonant_diacritic']].values)\ndataset=torch.utils.data.TensorDataset(image_tensor,label_tensor)\ntrain_data,val_data = torch.utils.data.random_split(dataset,lengths=[4*len(dataset)//5,len(dataset)//5])\n\nprint ('len(train_data) : ', len(train_data))\nprint ('len(val_data) : ', len(val_data))\nprint ('image.shape : ', image_tensor.shape)\nprint ('label.item() : ', label_tensor[0])\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"test_image = make_grid(image_tensor[:5].detach().cpu(), nrow=5)\nplt.figure(figsize=(25, 10));\nplt.imshow(test_image.permute(1, 2, 0));\n","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"These downsampled images are missing finer detail but will be sufficient for our purposes."},{"metadata":{},"cell_type":"markdown","source":"# Load DataLoader"},{"metadata":{"trusted":true},"cell_type":"code","source":"SEED = 42\n\ntorch.manual_seed(SEED)\ndevice = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n\nEPOCHS = 40\nINIT_LR = 5*1e-3\nBS = int(128*6)\n\n# CUDA memory\nkwargs = {'num_workers': 1, 'pin_memory': True} if device == 'cuda' else {}\ntrain_loader = torch.utils.data.DataLoader(train_data, batch_size=BS, shuffle=True, **kwargs)\nval_loader = torch.utils.data.DataLoader(val_data, batch_size=BS, shuffle=True, **kwargs)\n\nfor batch in train_loader:\n    x, y = batch\n    break\n\nprint ('x.shape : ', x.shape)\nprint ('y.shape : ', y.shape)\nfixed_images = x.to(device)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"HEIGHT = 137\nWIDTH = 236\n\nclass VAE(nn.Module):\n    def __init__(self,input_size=IMGSIZE*IMGSIZE):\n        self.input_size=input_size\n        super(VAE, self).__init__()\n        \n        fc1out= int(2*input_size/3)\n        embeddim =int(np.sqrt(fc1out)) \n        \n#         input layer\n        self.fc1 = nn.Linear(input_size, fc1out)\n\n#         embedded vector    \n        self.fc21 = nn.Linear(fc1out, embeddim)\n        self.fc22 = nn.Linear(fc1out, embeddim)\n        self.fc3 = nn.Linear(embeddim, fc1out)\n        \n#         output layer\n        self.fc4 = nn.Linear(fc1out,input_size)\n\n    def encode(self, x):\n        h1 = F.relu(self.fc1(x))\n        return self.fc21(h1), self.fc22(h1)\n\n    def reparameterize(self, mu, logvar):\n        std = torch.exp(0.5*logvar)\n        eps = torch.randn_like(std)\n        return mu + eps*std\n\n    def decode(self, z):\n        h3 = F.relu(self.fc3(z))\n        return torch.sigmoid(self.fc4(h3))\n\n    def forward(self, x):\n        mu, logvar = self.encode(x.view(-1,self.input_size))\n        z = self.reparameterize(mu, logvar)\n        return self.decode(z), mu, logvar\n\n\ndef kld_loss(x_pred, x, mu, logvar):\n    # see Appendix B from VAE paper:\n    # Kingma and Welling. Auto-Encoding Variational Bayes. ICLR, 2014\n    # https://arxiv.org/abs/1312.6114\n    # 0.5 * sum(1 + log(sigma^2) - mu^2 - sigma^2)\n    return -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp())\n\nbce_loss = nn.BCELoss(reduction='sum')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def print_num_params(model, display_all_modules=False):\n    total_num_params = 0\n    for n, p in model.named_parameters():\n        num_params = 1\n        for s in p.shape:\n            num_params *= s\n        if display_all_modules: print(\"{}: {}\".format(n, num_params))\n        total_num_params += num_params\n    print(\"-\" * 50)\n    print(\"Total number of parameters for {}: {:.2e}\".format(model.__class__,total_num_params))\n    \n\n\nmodel = torchvision.models.resnext50_32x4d(pretrained=False)\nprint_num_params(model)\nmodel = torchvision.models.resnet50(pretrained=False)\nprint_num_params(model)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"model = VAE(IMGSIZE*IMGSIZE).to(device)\noptimizer = optim.Adam(model.parameters(), lr=6e-3)\nprint_num_params(model)\nmodel","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Reconstruction + KL divergence losses summed over all elements and batch\ndef loss_function(recon_x, x, mu, logvar):\n    BCE = F.binary_cross_entropy(recon_x, x.view(-1, np.prod(fixed_images.shape[-2:])), reduction='sum')\n\n    # see Appendix B from VAE paper:\n    # Kingma and Welling. Auto-Encoding Variational Bayes. ICLR, 2014\n    # https://arxiv.org/abs/1312.6114\n    # 0.5 * sum(1 + log(sigma^2) - mu^2 - sigma^2)\n    KLD = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp())\n\n    return BCE + KLD","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"\ndef process_function(engine, batch):\n    model.train()\n    # clear the gradients from previous run\n    optimizer.zero_grad()\n    x, _ = batch\n    x = x.to(device)\n    x = x.view(-1,  np.prod(fixed_images.shape[-2:]))\n    # get predictions    \n    x_pred, mu, logvar = model(x)\n    # calculate loss     \n    BCE = bce_loss(x_pred, x)\n    KLD = kld_loss(x_pred, x, mu, logvar)\n    loss = BCE + KLD\n    # apply gradients     \n    loss.backward()\n    optimizer.step()\n    return loss.item(), BCE.item(), KLD.item()\n\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def evaluate_function(engine, batch):\n    model.eval()\n    with torch.no_grad():\n        x, _ = batch\n        x = x.to(device)\n        x = x.view(-1,  np.prod(fixed_images.shape[-2:]))\n        x_pred, mu, logvar = model(x)\n        kwargs = {'mu': mu, 'logvar': logvar}\n        return x_pred, x, kwargs","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"trainer = Engine(process_function)\nevaluator = Engine(evaluate_function)\ntraining_history = {'bce': [], 'kld': [], 'mse': []}\nvalidation_history = {'bce': [], 'kld': [], 'mse': []}","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"\nRunningAverage(output_transform=lambda x: x[0]).attach(trainer, 'loss')\nRunningAverage(output_transform=lambda x: x[1]).attach(trainer, 'bce')\nRunningAverage(output_transform=lambda x: x[2]).attach(trainer, 'kld')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"\nMeanSquaredError(output_transform=lambda x: [x[0], x[1]]).attach(evaluator, 'mse')\nLoss(bce_loss, output_transform=lambda x: [x[0], x[1]]).attach(evaluator, 'bce')\nLoss(kld_loss).attach(evaluator, 'kld')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"\ndef print_logs(engine, dataloader, mode, history_dict):\n    evaluator.run(dataloader, max_epochs=1)\n    metrics = evaluator.state.metrics\n    avg_mse = metrics['mse']\n    avg_bce = metrics['bce']\n    avg_kld = metrics['kld']\n    avg_loss =  avg_bce + avg_kld\n    print(\n        mode + \" Results - Epoch {} - Avg mse: {:.2f} Avg loss: {:.2f} Avg bce: {:.2f} Avg kld: {:.2f}\"\n        .format(engine.state.epoch, avg_mse, avg_loss, avg_bce, avg_kld))\n    for key in evaluator.state.metrics.keys():\n        history_dict[key].append(evaluator.state.metrics[key])\n\ntrainer.add_event_handler(Events.EPOCH_COMPLETED, print_logs, train_loader, 'Training', training_history)\ntrainer.add_event_handler(Events.EPOCH_COMPLETED, print_logs, val_loader, 'Validation', validation_history)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"\ndef compare_images(engine, save_img=False):\n    epoch = engine.state.epoch\n    reconstructed_images = model(fixed_images.view(-1, np.prod(fixed_images.shape[-2:])))[0].view(-1,1,*tuple(fixed_images.shape[-2:]))\n    comparison = torch.cat([fixed_images, reconstructed_images])\n    if save_img:\n        save_image(comparison.detach().cpu(), 'reconstructed_epoch_' + str(epoch) + '.png', nrow=32)\n    comparison_image = make_grid(comparison.detach().cpu(), nrow=32)\n    fig = plt.figure(figsize=(25, 10));\n    output = plt.imshow(comparison_image.permute(1, 2, 0));\n    plt.title('Epoch ' + str(epoch));\n    plt.show();\n    \ntrainer.add_event_handler(Events.STARTED(every=3), compare_images, save_img=False)\ntrainer.add_event_handler(Events.EPOCH_COMPLETED(every=8), compare_images, save_img=False)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"This is an example of what the network starts with."},{"metadata":{"trusted":true},"cell_type":"code","source":"\nreconstructed_images = model(fixed_images.view(-1, np.prod(fixed_images.shape[-2:])))[0].view(-1,1,*tuple(fixed_images.shape[-2:]))\ncomparison = torch.cat([fixed_images, reconstructed_images])\ncomparison_image = make_grid(comparison.detach().cpu(), nrow=32)\nfig = plt.figure(figsize=(25, 18));\noutput = plt.imshow(comparison_image.permute(1, 2, 0));\nplt.title('Epoch ' + str(-1));\nplt.show();","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# let the training begin!"},{"metadata":{"trusted":true},"cell_type":"code","source":"BS,EPOCHS","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"e = trainer.run(train_loader, max_epochs=EPOCHS)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"plt.plot(range(EPOCHS), training_history['bce'], 'dodgerblue', label='training')\nplt.plot(range(EPOCHS), validation_history['bce'], 'orange', label='validation')\nplt.xlim(0, EPOCHS);\nplt.xlabel('Epoch')\nplt.ylabel('BCE')\nplt.title('Binary Cross Entropy on Training/Validation Set')\nplt.legend();","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"plt.plot(range(EPOCHS), training_history['kld'], 'dodgerblue', label='training')\nplt.plot(range(EPOCHS), validation_history['kld'], 'orange', label='validation')\nplt.xlim(0, EPOCHS);\nplt.xlabel('Epoch')\nplt.ylabel('KLD')\nplt.title('KL Divergence on Training/Validation Set')\nplt.legend();\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"","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":4}