{"cells":[{"metadata":{},"cell_type":"markdown","source":"# Generative Dog Images\n## Best Public Score: 90.03206 (Version 7)\n<hr>"},{"metadata":{},"cell_type":"markdown","source":"## Generator and Discriminator Networks"},{"metadata":{"trusted":true},"cell_type":"code","source":"# imports\nimport torch.nn as nn\n\n\n# weights initialization\ndef weights_init(m):\n    classname = m.__class__.__name__\n    if classname.find('Conv') != -1:\n        nn.init.normal_(m.weight.data, 0.0, 0.02)\n    elif classname.find('BatchNorm') != -1:\n        nn.init.normal_(m.weight.data, 1.0, 0.02)\n        nn.init.constant_(m.bias.data, 0)\n\n\n# Generator network\nclass Generator(nn.Module):\n    \n    def __init__(self, nz=100, nc=3, ngf=64, init_weights=False):\n        super(Generator, self).__init__()\n        self.network = nn.Sequential(\n            # input shape: nz x 1 x 1 -> output shape: (ngf*8) x 4 x 4\n            nn.ConvTranspose2d(nz, ngf*8, 4, 1, 0, bias=False),\n            nn.BatchNorm2d(ngf*8),\n            nn.ReLU(inplace=True),\n            \n            # input shape: (ngf*8) x 4 x 4 -> output shape: (ngf*4) x 8 x 8\n            nn.ConvTranspose2d(ngf*8, ngf*4, 4, 2, 1, bias=False),\n            nn.BatchNorm2d(ngf*4),\n            nn.ReLU(inplace=True),\n            \n            # input shape: (ngf*4) x 8 x 8 -> output shape: (ngf*2) x 16 x 16\n            nn.ConvTranspose2d(ngf*4, ngf*2, 4, 2, 1, bias=False),\n            nn.BatchNorm2d(ngf*2),\n            nn.ReLU(inplace=True),\n            \n            # input shape: (ngf*2) x 16 x 16 -> output shape: (ngf) x 32 x 32\n            nn.ConvTranspose2d(ngf*2, ngf, 4, 2, 1, bias=False),\n            nn.BatchNorm2d(ngf),\n            nn.ReLU(inplace=True),\n            \n            # input shape: (ngf) x 32 x 32 -> output shape: (nc) x 64 x 64\n            nn.ConvTranspose2d(ngf, nc, 4, 2, 1, bias=False),\n            nn.Tanh()\n        )\n        if init_weights:\n            self.apply(weights_init)\n    \n    def forward(self, x):\n        return self.network(x)\n\n\n# Discriminator network\nclass Discriminator(nn.Module):\n    \n    def __init__(self, nc=3, ndf=64, init_weights=False):\n        super(Discriminator, self).__init__()\n        self.network = nn.Sequential(\n            # input shape: (nc) x 64 x 64 -> output shape: (ndf) x 32 x 32\n            nn.Conv2d(nc, ndf, 4, 2, 1, bias=False),\n            nn.LeakyReLU(0.2, inplace=True),\n            \n            # input shape: (ndf) x 32 x 32 -> output shape: (ndf*2) x 16 x 16\n            nn.Conv2d(ndf, ndf*2, 4, 2, 1, bias=False),\n            nn.BatchNorm2d(ndf*2),\n            nn.LeakyReLU(0.2, inplace=True),\n            \n            # input shape: (ndf*2) x 16 x 16 -> output shape: (ndf*4) x 8 x 8\n            nn.Conv2d(ndf*2, ndf*4, 4, 2, 1, bias=False),\n            nn.BatchNorm2d(ndf*4),\n            nn.LeakyReLU(0.2, inplace=True),\n            \n            # input shape: (ndf*4) x 8 x 8 -> output shape: (ndf*8) x 4 x 4\n            nn.Conv2d(ndf*4, ndf*8, 4, 2, 1, bias=False),\n            nn.BatchNorm2d(ndf*8),\n            nn.LeakyReLU(0.2, inplace=True),\n            \n            # input shape: (ndf*8) x 4 x 4 -> output shape: 1 x 1 x 1\n            nn.Conv2d(ndf*8, 1, 4, 1, 0, bias=False),\n            nn.Sigmoid()\n        )\n        if init_weights:\n            self.apply(weights_init)\n    \n    def forward(self, x):\n        return self.network(x)\n","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Dependencies"},{"metadata":{"trusted":true},"cell_type":"code","source":"# imports\nimport glob\nimport matplotlib.pyplot as plt\nimport numpy as np\nimport os\nimport pandas as pd\nimport shutil\nimport torch\nimport torchvision\n\nfrom PIL import Image\nfrom torch.utils.data import Dataset, DataLoader\nfrom xml.etree import ElementTree\n#from dcgan.models import Generator, Discriminator\n","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Configurations"},{"metadata":{"trusted":true},"cell_type":"code","source":"# configurations\nRANDOMSEED = None\nIMAGE_ROOT = '../input/all-dogs/all-dogs/'\nANNOT_ROOT = '../input/annotation/Annotation/'\nIMAGE_SIZE = 64\nN_CHANNELS = 3\nLATENT_DIM = 100\nBATCH_SIZE = 64\nNUM_EPOCHS = 10\nLEARN_RATE = 0.0002\nUSE_N_GPUS = 1\n","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Custom Dataset"},{"metadata":{"trusted":true},"cell_type":"code","source":"# custom dataset\nclass CustomDataset(Dataset):\n    \n    def __init__(self, image_dir, annotation_dir, transform=None):\n        self._imgbboxes = {}\n        self._transform = transform if not transform is None else torchvision.transforms.ToTensor()\n        breed_map = {os.path.basename(breed_dir).split('-')[0]: breed_dir \\\n                     for breed_dir in glob.glob(f'{annotation_dir}/*')}\n        index_img = 0\n        for image_path in glob.glob(f'{image_dir}/*.*'):\n            breed, index = os.path.splitext(os.path.basename(image_path))[0].split('_')\n            for obj in ElementTree.parse(f'{breed_map[breed]}/{breed}_{index}').getroot().findall('object'):\n                bbox = obj.find('bndbox')\n                xmin = int(bbox.find('xmin').text)\n                ymin = int(bbox.find('ymin').text)\n                xmax = int(bbox.find('xmax').text)\n                ymax = int(bbox.find('ymax').text)\n                self._imgbboxes[index_img] = (image_path, (xmin, ymin, xmax, ymax))\n                index_img += 1\n    \n    def __len__(self):\n        return len(self._imgbboxes)\n    \n    def __getitem__(self, index):\n        image_path, bbox = self._imgbboxes[index]\n        image = Image.open(image_path).crop(bbox)\n        image = self._transform(image)\n        return {'image': image, 'label': 1}\n","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Visualization Function"},{"metadata":{"trusted":true},"cell_type":"code","source":"# show image tensor as grid\ndef show(x, title=None):\n    grid_image = torchvision.utils.make_grid(x, normalize=True).numpy()\n    if title:\n        plt.title(title)\n    plt.axis('off')\n    plt.imshow(np.transpose(grid_image, (1, 2, 0)))\n    plt.show()\n","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Setup Output Directory Structure, Seed and Device"},{"metadata":{"trusted":true},"cell_type":"code","source":"# setup output directory structure\nfor directory in ['fake', 'ckpt']:\n    if not os.path.isdir(f'output/{directory}/'):\n        os.makedirs(f'output/{directory}/')\n\n\n# setup seed for reproducible results\nif RANDOMSEED:\n    torch.manual_seed(RANDOMSEED)\n    torch.backends.cudnn.deterministic = True\n    torch.backends.cudnn.benchmark = False\n\n\n# setup device\ndevice = torch.device('cuda:0' if torch.cuda.is_available() and USE_N_GPUS > 0 else 'cpu')\ndevice_info = 'CPU' if device.type.lower() == 'cpu' else f'{torch.cuda.get_device_name(device)} [CUDA]'\nprint(f'[INFO] Using device: {device_info}')\n","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## ETL Pipeline"},{"metadata":{"trusted":true},"cell_type":"code","source":"# create transform, dataset and dataloader\nif N_CHANNELS == 1:\n    transform = torchvision.transforms.Compose([\n        torchvision.transforms.Grayscale(),\n        torchvision.transforms.Resize(IMAGE_SIZE),\n        torchvision.transforms.CenterCrop(IMAGE_SIZE),\n        torchvision.transforms.ToTensor(),\n        torchvision.transforms.Normalize((0.5,), (0.5,))\n    ])\nelse:\n    transform = torchvision.transforms.Compose([\n        torchvision.transforms.Resize(IMAGE_SIZE),\n        torchvision.transforms.CenterCrop(IMAGE_SIZE),\n        torchvision.transforms.ToTensor(),\n        torchvision.transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))\n    ])\n\ndataset = CustomDataset(IMAGE_ROOT, ANNOT_ROOT, transform=transform)\ndataloader = DataLoader(dataset, batch_size=BATCH_SIZE, shuffle=True)\n","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Dataset Samples"},{"metadata":{"trusted":true},"cell_type":"code","source":"# example of real samples\nreal_batch = next(iter(dataloader))['image'][:64]\ntorchvision.utils.save_image(real_batch, 'output/real.jpg', normalize=True)\nshow(real_batch, 'Real samples')\n","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Initialize Generator and Discriminator Networks"},{"metadata":{"trusted":true},"cell_type":"code","source":"# create generator and discriminator networks\nnetG = Generator(LATENT_DIM, N_CHANNELS, init_weights=True).to(device)\nnetD = Discriminator(N_CHANNELS, init_weights=True).to(device)\n","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Initialize Criterion (Loss) and Optimizers"},{"metadata":{"trusted":true},"cell_type":"code","source":"# create criterion and optimizers\ncriterion = torch.nn.BCELoss()\noptimizerG = torch.optim.Adam(netG.parameters(), lr=LEARN_RATE, betas=(0.5, 0.999))\noptimizerD = torch.optim.Adam(netD.parameters(), lr=LEARN_RATE, betas=(0.5, 0.999))\n","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Training"},{"metadata":{"trusted":true},"cell_type":"code","source":"# initialize variables\nreal_label = 1\nfake_label = 0\nseed_noise = torch.randn(64, LATENT_DIM, 1, 1, device=device)\nloop_width = len(str(NUM_EPOCHS))\nstep_width = len(str(len(dataloader)))\nhist_lossG = []\nhist_lossD = []\nhist_ckptG = []\nhist_ckptD = []\n\n\n# begin epoch\nfor epoch in range(NUM_EPOCHS):\n    \n    \n    # iterate over data\n    for i, batch in enumerate(dataloader):\n        \n        \n        # ---- train discriminator: maximize log(D(x)) + log(1 - D(G(z))) ----\n        netD.zero_grad()\n        \n        # pass one batch of real images\n        real_x = batch['image'].to(device)\n        labels = torch.full((real_x.size(0),), real_label, device=device)\n        output = netD(real_x).view(-1)\n        lossD_real = criterion(output, labels)\n        lossD_real.backward()\n        D_x = output.mean().item()\n        \n        # pass one batch of fake images\n        z = torch.randn(real_x.size(0), LATENT_DIM, 1, 1, device=device)\n        fake_x = netG(z)\n        labels.fill_(fake_label)\n        output = netD(fake_x.detach()).view(-1)\n        lossD_fake = criterion(output, labels)\n        lossD_fake.backward()\n        D_Gz_1 = output.mean().item()\n        \n        # estimate total loss over both batches\n        lossD = lossD_real + lossD_fake\n        \n        # update discriminator\n        optimizerD.step()\n        \n        \n        # ---- train generator: maximize log(D(G(z))) ----\n        netG.zero_grad()\n        \n        # update generator\n        labels.fill_(real_label)\n        output = netD(fake_x).view(-1)\n        lossG = criterion(output, labels)\n        lossG.backward()\n        D_Gz_2 = output.mean().item()\n        \n        optimizerG.step()\n        \n        \n        # ---- record statistics ----\n        hist_lossG.append(lossG.item())\n        hist_lossD.append(lossD.item())\n        step = epoch * len(dataloader) + i\n        if step % 500 == 0 or step == NUM_EPOCHS * len(dataloader) - 1:\n            print(f'[Epoch {epoch+1:{loop_width}d}/{NUM_EPOCHS}]',\n                  f'[Batch {i+1:{step_width}d}/{len(dataloader)}]',\n                  f'- Loss_G: {lossG.item():7.4f}',\n                  f'- Loss_D: {lossD.item():7.4f}',\n                  f'- D(x): {D_x:7.4f}',\n                  f'- D(G(z)): {D_Gz_1:7.4f} -> {D_Gz_2:7.4f}')\n            with torch.no_grad():\n                fake_batch = netG(seed_noise).detach().cpu()\n            #torchvision.utils.save_image(fake_batch, f'output/fake/step-{step}.jpg', normalize=True)\n            show(fake_batch, f'Fake samples [step: {step}]')\n            torch.save(netG.state_dict(), f'output/ckpt/netG-{step}.pt')\n            torch.save(netD.state_dict(), f'output/ckpt/netD-{step}.pt')\n            hist_ckptG.append(f'output/ckpt/netG-{step}.pt')\n            hist_ckptD.append(f'output/ckpt/netD-{step}.pt')\n            if len(hist_ckptG) > 5:\n                os.remove(hist_ckptG.pop(0))\n            if len(hist_ckptD) > 5:\n                os.remove(hist_ckptD.pop(0))\n\n\n# save statistics\ndf = pd.DataFrame()\ndf['lossG'] = hist_lossG\ndf['lossD'] = hist_lossD\ndf.to_csv('output/statistics.csv', index=False)\n\n\n# plot statistics\nplt.figure(figsize=(10, 5))\nplt.plot(hist_lossG, label='G Loss')\nplt.plot(hist_lossD, label='D Loss')\nplt.xlabel('Iterations')\nplt.ylabel('Loss')\nplt.title('Generator and Discriminator Loss')\nplt.legend()\nplt.savefig('output/statistics.jpg')\n","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Submission"},{"metadata":{"trusted":true},"cell_type":"code","source":"# prepare submission for Kaggle\nprint('[INFO] Creating archive for submission... ', end='')\nif not os.path.isdir('output/kaggle/'):\n    os.makedirs('output/kaggle/')\n\nfor i in range(0, 10000, 50):\n    z = torch.randn(50, LATENT_DIM, 1, 1, device=device)\n    with torch.no_grad():\n        fake_x = netG(z).detach().cpu()\n    for j in range(50):\n        filename = f'output/kaggle/{str(i+j).zfill(4)}.png'\n        torchvision.utils.save_image(fake_x[j, :, :, :], filename, normalize=True)\n\nshutil.make_archive('images', 'zip', 'output/kaggle/')\nshutil.rmtree('output/kaggle/')\nprint('done')\n","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}