{"cells":[{"metadata":{},"cell_type":"markdown","source":"As far as we know, I/O operations are most expensive in training Neural Networks.<br> In this kernel I wrote pytorch Dataset loader that loads and preprocess all images once and stores them into RAM."},{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"_kg_hide-input":true},"cell_type":"code","source":"import numpy as np\nimport pandas as pd\nimport os\nfrom os.path import join\nimport matplotlib.pyplot as plt\n\nimport torch\nfrom torch import nn, optim\nimport torch.nn.functional as F\nimport torchvision\nfrom torchvision import datasets, transforms\nfrom torchvision.utils import save_image\nfrom torch.utils.data import Dataset, DataLoader\n\nfrom PIL import Image\nimport xml.etree.ElementTree as ET\n\nfrom tqdm import tqdm_notebook as tqdm","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Parameters of GAN"},{"metadata":{"trusted":true},"cell_type":"code","source":"batch_size = 16\nlrG = 0.001\nlrD = 0.001\nbeta1 = 0.5\nepochs = 300\n\nreal_label = 0.5\nfake_label = 0\nnz = 128\n\ndevice = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Pytorch Dataset and DataLoader"},{"metadata":{"trusted":true},"cell_type":"code","source":"class DogDataset(Dataset):\n    def __init__(self, img_dir, annotations_dir, transform1=None, transform2=None):\n\n        self.img_dir = img_dir\n        self.img_names = os.listdir(img_dir)\n        self.transform1 = transform1\n        self.transform2 = transform2\n        \n        self.imgs = []\n        for img_name in self.img_names:\n            path = join(img_dir, img_name)\n            img = torchvision.datasets.folder.default_loader(path)\n    \n            # Crop image\n            annotation_basename = os.path.splitext(os.path.basename(path))[0]\n            annotation_dirname = next(dirname for dirname in os.listdir(annotations_dir) if dirname.startswith(annotation_basename.split('_')[0]))\n            annotation_filename = os.path.join(annotations_dir, annotation_dirname, annotation_basename)\n            tree = ET.parse(annotation_filename)\n            root = tree.getroot()\n            objects = root.findall('object')\n            \n            for o in objects:\n                bndbox = o.find('bndbox')\n                xmin = int(bndbox.find('xmin').text)\n                ymin = int(bndbox.find('ymin').text)\n                xmax = int(bndbox.find('xmax').text)\n                ymax = int(bndbox.find('ymax').text)\n                bbox = (xmin, ymin, xmax, ymax)\n                img_ = img.crop(bbox)\n                # Some crop's are black. if crop is black then don't crop\n                if np.mean(img_) != 0:\n                    img = img_\n\n                if self.transform1 is not None:\n                    img = self.transform1(img)\n\n                self.imgs.append(img)\n\n    def __getitem__(self, index):\n        img = self.imgs[index]\n        \n        if self.transform2 is not None:\n            img = self.transform2(img)\n        \n        return img\n\n    def __len__(self):\n        return len(self.imgs)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Benchmark"},{"metadata":{},"cell_type":"markdown","source":"## Simple Dataloader"},{"metadata":{"trusted":true},"cell_type":"code","source":"%%time\nrandom_transforms = [transforms.RandomRotation(degrees=5)]\ntransform = transforms.Compose([transforms.Resize(64),\n                                transforms.CenterCrop(64),\n                                transforms.RandomHorizontalFlip(p=0.5),\n                                transforms.RandomApply(random_transforms, p=0.2),\n                                transforms.ToTensor(),\n                                transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])\n\ntrain_data = datasets.ImageFolder('../input/all-dogs/', transform=transform)\ntrain_loader = torch.utils.data.DataLoader(train_data, shuffle=True,\n                                           batch_size=batch_size)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"%%time\nfor _ in train_loader:\n    continue","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Simple Dataloader with num_workers=2"},{"metadata":{"trusted":true},"cell_type":"code","source":"%%time\nrandom_transforms = [transforms.RandomRotation(degrees=5)]\ntransform = transforms.Compose([transforms.Resize(64),\n                                transforms.CenterCrop(64),\n                                transforms.RandomHorizontalFlip(p=0.5),\n                                transforms.RandomApply(random_transforms, p=0.2),\n                                transforms.ToTensor(),\n                                transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])\n\ntrain_data = datasets.ImageFolder('../input/all-dogs/', transform=transform)\ntrain_loader = torch.utils.data.DataLoader(train_data, shuffle=True,\n                                           batch_size=batch_size,\n                                          num_workers=2)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"%%time\nfor _ in train_loader:\n    continue","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## RAM Dataloader\nFirst step is **once** download data"},{"metadata":{"trusted":true},"cell_type":"code","source":"%%time\n# First preprocessing of data\ntransform1 = transforms.Compose([transforms.Resize(64),\n                                transforms.CenterCrop(64)])\n\n# Data augmentation and converting to tensors\nrandom_transforms = [transforms.RandomRotation(degrees=5)]\ntransform2 = transforms.Compose([transforms.RandomHorizontalFlip(p=0.5),\n                                 transforms.RandomApply(random_transforms, p=0.3), \n                                 transforms.ToTensor(),\n                                 transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])\n                                 \ntrain_dataset = DogDataset(img_dir='../input/all-dogs/all-dogs/',\n                           annotations_dir='../input/annotation/Annotation/',\n                           transform1=transform1,\n                           transform2=transform2)\n\ntrain_loader = DataLoader(dataset=train_dataset,\n                          batch_size=batch_size,\n                          shuffle=True,\n                          num_workers=4)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Second step: iterate over downloaded images"},{"metadata":{"trusted":true},"cell_type":"code","source":"%%time\nfor _ in train_loader:\n    continue","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Let's calculate how much time you need to train 100 epochs your GAN (forward and backward pass of NNs doesn't counted)\n - Simple Dataloader: ~90sec per epoch * 100 epochs = 9000sec = 2.5h\n - Simple Dataloader with 2 workers: ~ 65sec per epoch * 100 epochs = 1h 50min\n - RAM Dataloader: ~95sec download + 4.5sec per epoch * 100 epochs = 9 min<br>"},{"metadata":{},"cell_type":"markdown","source":"With this dataloader you can make much more experiments and epochs!"},{"metadata":{},"cell_type":"markdown","source":"## Examples of dogs"},{"metadata":{"trusted":true},"cell_type":"code","source":"x = next(iter(train_loader))\n\nfig = plt.figure(figsize=(25, 16))\nfor ii, img in enumerate(x):\n    ax = fig.add_subplot(4, 8, ii + 1, xticks=[], yticks=[])\n    \n    img = img.numpy().transpose(1, 2, 0)\n    plt.imshow((img+1)/2)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Model and training loop taken from https://www.kaggle.com/speedwagon/ralsgan-dogs"},{"metadata":{"trusted":true,"_kg_hide-input":true},"cell_type":"code","source":"class Generator(nn.Module):\n    def __init__(self, nz=128, channels=3):\n        super(Generator, self).__init__()\n        \n        self.nz = nz\n        self.channels = channels\n        \n        def convlayer(n_input, n_output, k_size=4, stride=2, padding=0):\n            block = [\n                nn.ConvTranspose2d(n_input, n_output, kernel_size=k_size, stride=stride, padding=padding, bias=False),\n                nn.BatchNorm2d(n_output),\n                nn.LeakyReLU(inplace=True),\n            ]\n            return block\n\n        self.model = nn.Sequential(\n            *convlayer(self.nz, 512, 4, 1, 0), # Fully connected layer via convolution.\n            *convlayer(512, 256, 4, 2, 1),\n            *convlayer(256, 128, 4, 2, 1),\n            *convlayer(128, 64, 4, 2, 1),\n            *convlayer(64, 32, 4, 2, 1),\n            nn.ConvTranspose2d(32, self.channels, 3, 1, 1),\n            nn.Tanh()\n        )\n\n\n    def forward(self, z):\n        z = z.view(-1, self.nz, 1, 1)\n        img = self.model(z)\n        return img\n\nclass Discriminator(nn.Module):\n    def __init__(self, channels=3):\n        super(Discriminator, self).__init__()\n        \n        self.channels = channels\n\n        def convlayer(n_input, n_output, k_size=4, stride=2, padding=0, bn=False):\n            block = [nn.Conv2d(n_input, n_output, kernel_size=k_size, stride=stride, padding=padding, bias=False)]\n            if bn:\n                block.append(nn.BatchNorm2d(n_output))\n            block.append(nn.LeakyReLU(0.2, inplace=True))\n            return block\n\n        self.model = nn.Sequential(\n            *convlayer(self.channels, 32, 4, 2, 1),\n            *convlayer(32, 64, 4, 2, 1),\n            *convlayer(64, 128, 4, 2, 1, bn=True),\n            *convlayer(128, 256, 4, 2, 1, bn=True),\n            nn.Conv2d(256, 1, 4, 1, 0, bias=False),  # FC with Conv.\n        )\n\n    def forward(self, imgs):\n        out = self.model(imgs)\n        return out.view(-1, 1)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"netG = Generator(nz).to(device)\nnetD = Discriminator().to(device)\n\ncriterion = nn.BCELoss()\n\noptimizerD = optim.Adam(netD.parameters(), lr=lrD, betas=(beta1, 0.999))\noptimizerG = optim.Adam(netG.parameters(), lr=lrG, betas=(beta1, 0.999))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true,"_kg_hide-input":true},"cell_type":"code","source":"def show_generated_img():\n    noise = torch.randn(1, nz, 1, 1, device=device)\n    gen_image = (netG(noise).to(\"cpu\").clone().detach().squeeze(0) + 1) / 2\n    gen_image = gen_image.numpy().transpose(1, 2, 0)\n    plt.imshow(gen_image)\n    plt.show()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true,"_kg_hide-output":true},"cell_type":"code","source":"for epoch in range(epochs):\n    \n    for ii, real_images in tqdm(enumerate(train_loader), total=len(train_loader)):\n        ############################\n        # (1) Update D network\n        ###########################\n        netD.zero_grad()\n        real_images = real_images.to(device)\n        batch_size = real_images.size(0)\n        labels = torch.full((batch_size, 1), real_label, device=device)\n        outputR = netD(real_images)\n        noise = torch.randn(batch_size, nz, 1, 1, device=device)\n        fake = netG(noise)\n        outputF = netD(fake.detach())\n        errD = (torch.mean((outputR - torch.mean(outputF) - labels) ** 2) + \n                torch.mean((outputF - torch.mean(outputR) + labels) ** 2))/2\n        errD.backward(retain_graph=True)\n        optimizerD.step()\n        ############################\n        # (2) Update G network\n        ###########################\n        netG.zero_grad()\n        outputF = netD(fake)   \n        errG = (torch.mean((outputR - torch.mean(outputF) + labels) ** 2) +\n                torch.mean((outputF - torch.mean(outputR) - labels) ** 2))/2\n        errG.backward()\n        optimizerG.step()\n        \n        if (ii+1) % (len(train_loader)//2) == 0:\n            print('[%d/%d][%d/%d] Loss_D: %.4f Loss_G: %.4f'\n                  % (epoch + 1, epochs, ii+1, len(train_loader),\n                     errD.item(), errG.item()))\n\n    show_generated_img()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Let's visualise generated results"},{"metadata":{"trusted":true},"cell_type":"code","source":"gen_z = torch.randn(32, nz, 1, 1, device=device)\ngen_images = (netG(gen_z).to(\"cpu\").clone().detach() + 1)/2\ngen_images = gen_images.numpy().transpose(0, 2, 3, 1)\n\nfig = plt.figure(figsize=(25, 16))\nfor ii, img in enumerate(gen_images):\n    ax = fig.add_subplot(4, 8, ii + 1, xticks=[], yticks=[])\n    plt.imshow(img)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Make predictions and submit"},{"metadata":{"trusted":true},"cell_type":"code","source":"if not os.path.exists('../output_images'):\n    os.mkdir('../output_images')\nim_batch_size = 50\nn_images=10000\nfor i_batch in range(0, n_images, im_batch_size):\n    gen_z = torch.randn(im_batch_size, nz, 1, 1, device=device)\n    gen_images = (netG(gen_z)+1)/2\n    images = gen_images.to(\"cpu\").clone().detach()\n    images = images.numpy().transpose(0, 2, 3, 1)\n    for i_image in range(gen_images.size(0)):\n        save_image(gen_images[i_image, :, :, :], os.path.join('../output_images', f'image_{i_batch+i_image:05d}.png'))\n\n\nimport shutil\nshutil.make_archive('images', 'zip', '../output_images')","execution_count":null,"outputs":[]}],"metadata":{"kernelspec":{"display_name":"Python 3","language":"python","name":"python3"},"language_info":{"name":"python","version":"3.6.4","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"}},"nbformat":4,"nbformat_minor":1}