{"cells":[{"metadata":{},"cell_type":"markdown","source":"Based on:\n\n* [RaLSGAN dogs 000213](https://www.kaggle.com/titericz/ralsgan-dogs-vladislav-bakhteev-kernel)\n\n* [aLSGAN Explained](https://www.kaggle.com/c/generative-dog-images/discussion/99485#latest-580430)\n\n* [Error in Starter Code](https://www.kaggle.com/c/generative-dog-images/discussion/99613#latest-574385)\n\n* [RaLSGAN dogs cropping random](https://www.kaggle.com/sakami/ralsgan-dogs-cropping-random)\n\nChange:\n\n* spectral norm, self-attention\n\n* loss function\n"},{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"import numpy as np\nimport pandas as pd\nimport os\nimport matplotlib.pyplot as plt\nfrom multiprocessing import Pool\nimport xml.etree.ElementTree as ET\nfrom pathlib import Path\nimport torch\nfrom torch import nn, optim\nimport torch.nn.functional as F\nfrom torchvision import datasets, transforms\nfrom torchvision.utils import save_image\nfrom torchvision.datasets.folder import default_loader\nimport random\nfrom tqdm import tqdm_notebook as tqdm    ","execution_count":null,"outputs":[]},{"metadata":{"_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","trusted":true},"cell_type":"code","source":"class Self_Attn(nn.Module):\n    \"\"\" Self attention Layer\"\"\"\n    def __init__(self,in_dim,activation = 'relu'):\n        super(Self_Attn,self).__init__()\n        self.chanel_in = in_dim\n        self.activation = activation\n        \n        self.query_conv = nn.Conv2d(in_channels = in_dim , out_channels = in_dim//8 , kernel_size= 1)\n        self.key_conv = nn.Conv2d(in_channels = in_dim , out_channels = in_dim//8 , kernel_size= 1)\n        self.value_conv = nn.Conv2d(in_channels = in_dim , out_channels = in_dim , kernel_size= 1)\n        self.gamma = nn.Parameter(torch.zeros(1))\n\n        self.softmax  = nn.Softmax(dim=-1) #\n    def forward(self,x):\n        \"\"\"\n            inputs :\n                x : input feature maps( B X C X W X H)\n            returns :\n                out : self attention value + input feature \n                attention: B X N X N (N is Width*Height)\n        \"\"\"\n        m_batchsize,C,width ,height = x.size()\n        proj_query  = self.query_conv(x).view(m_batchsize,-1,width*height).permute(0,2,1) # B X CX(N)\n        proj_key =  self.key_conv(x).view(m_batchsize,-1,width*height) # B X C x (*W*H)\n        energy =  torch.bmm(proj_query,proj_key) # transpose check\n        attention = self.softmax(energy) # BX (N) X (N) \n        proj_value = self.value_conv(x).view(m_batchsize,-1,width*height) # B X C X N\n\n        out = torch.bmm(proj_value,attention.permute(0,2,1) )\n        out = out.view(m_batchsize,C,width,height)\n        \n        out = self.gamma*out + x\n        return out  #,attention\n    \nclass SELayer(nn.Module):\n    def __init__(self, channel, reduction=16):\n        super(SELayer, self).__init__()\n        self.avg_pool = nn.AdaptiveAvgPool2d(1)\n        self.fc = nn.Sequential(\n            nn.Linear(channel, channel // reduction, bias=False),\n            nn.ReLU(inplace=True),\n            nn.Linear(channel // reduction, channel, bias=False),\n            nn.Sigmoid()\n        )\n\n    def forward(self, x):\n        b, c, _, _ = x.size()\n        y = self.avg_pool(x).view(b, c)\n        y = self.fc(y).view(b, c, 1, 1)\n        return x * y.expand_as(x)\n\nclass SeparableConv2d(nn.Module):\n    def __init__(self,in_channels,out_channels,kernel_size=1,stride=1,padding=0,dilation=1,bias=False):\n        super(SeparableConv2d,self).__init__()\n\n        self.conv1 = nn.utils.spectral_norm(nn.Conv2d(in_channels,in_channels,kernel_size,stride,padding,dilation,groups=in_channels,bias=bias))\n        self.pointwise = nn.utils.spectral_norm(nn.Conv2d(in_channels,out_channels,1,1,0,1,1,bias=bias))\n    \n    def forward(self,x):\n        x = self.conv1(x)\n        x = self.pointwise(x)\n        return x\n    \nclass pixelwise_norm_layer(nn.Module):\n    def __init__(self):\n        super(pixelwise_norm_layer, self).__init__()\n        self.eps = 1e-8\n\n    def forward(self, x):\n        return x / (torch.mean(x**2, dim=1, keepdim=True) + self.eps) ** 0.5\n\nclass ResidualBlock(nn.Module):\n    def __init__(self, in_features):\n        super(ResidualBlock, self).__init__()\n\n        conv_block = [\n            nn.Conv2d(in_features, in_features, 3, stride=1, padding=1, bias=False),\n            nn.InstanceNorm2d(in_features, affine=True, track_running_stats=True),\n            nn.ReLU(inplace=True),\n            nn.Conv2d(in_features, in_features, 3, stride=1, padding=1, bias=False),\n            nn.InstanceNorm2d(in_features, affine=True, track_running_stats=True),\n        ]\n\n        self.conv_block = nn.Sequential(*conv_block)\n\n    def forward(self, x):\n        return x + self.conv_block(x)","execution_count":null,"outputs":[]},{"metadata":{"trusted":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.utils.spectral_norm(nn.ConvTranspose2d(n_input, n_output, kernel_size=k_size, stride=stride, padding=padding, bias=False)),\n                pixelwise_norm_layer(),\n                nn.LeakyReLU(0.01),\n            ]\n            return block\n\n        self.model = nn.Sequential(\n            *convlayer(self.nz, 1024, 4, 1, 0), # Fully connected layer via convolution.\n            Self_Attn(1024),\n            *convlayer(1024, 512, 4, 2, 1),\n            *convlayer(512, 256, 4, 2, 1),\n            *convlayer(256, 128, 4, 2, 1),\n            Self_Attn(128),\n            *convlayer(128, 64, 4, 2, 1),\n            nn.ConvTranspose2d(64, self.channels, 3, 1, 1),\n            nn.Tanh()\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.utils.spectral_norm(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.head1 = nn.Sequential(\n            *convlayer(self.channels, 32, 4, 2, 1),\n            nn.LeakyReLU(0.2, inplace=False),\n        )\n        self.head2 = nn.Sequential(\n            *convlayer(32, 64, 4, 2, 1),\n            nn.LeakyReLU(0.2, inplace=False),\n            *convlayer(64, 128, 4, 2, 1, bn=True),\n        )\n        \n        self.model = nn.Sequential(\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        imgs = self.head1(imgs)\n        imgs = self.head2(imgs)\n        logits = self.model(imgs)\n        # out = torch.sigmoid(logits)\n        return logits.view(-1, 1)\n    \n    def feature(self, imgs):\n        feature_1 = self.head1(imgs)\n        feature_2 = self.head2(feature_1)\n        return feature_1, feature_2","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"batch_size = 32\ng_lr = 0.0005\nd_lr = 0.0005\nbeta1 = 0.5\nepochs = 350\n\nreal_label = 0.95\nfake_label = 0\nnz = 256\ndevice = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class DogsDataset(torch.utils.data.Dataset):\n    \n    def __init__(self, root, annotation_root, transform=None,\n                 target_transform=None, loader=default_loader, n_process=4):\n        self.root = Path(root)\n        self.annotation_root = Path(annotation_root)\n        self.transform = transform\n        self.target_transform = target_transform\n        self.loader = loader\n        self.imgs = self.cut_out_dogs(n_process)\n\n    def _get_annotation_path(self, img_path):\n        dog = Path(img_path).stem\n        breed = dog.split('_')[0]\n        breed_dir = next(self.annotation_root.glob(f'{breed}-*'))\n        return breed_dir / dog\n    \n    @staticmethod\n    def _get_dog_box(annotation_path):\n        tree = ET.parse(annotation_path)\n        root = tree.getroot()\n        objects = root.findall('object')\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            yield (xmin, ymin, xmax, ymax)\n            \n    def crop_dog(self, path):\n        imgs = []\n        annotation_path = self._get_annotation_path(path)\n        for bndbox in self._get_dog_box(annotation_path):\n            img = self.loader(path)\n            img_ = img.crop(bndbox)\n            if np.sum(img_) != 0:\n                img = img_\n            imgs.append(img)\n        return imgs\n    \n    def cut_out_dogs(self, n_process):\n        with Pool(n_process) as p:\n            imgs = p.map(self.crop_dog, self.root.iterdir())\n        return imgs\n    \n    def __getitem__(self, index):\n        samples = random.choice(self.imgs[index])\n        if self.transform is not None:\n            samples = self.transform(samples)\n        return samples\n    \n    def __len__(self):\n        return len(self.imgs)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"random_transforms = [\n    #transforms.ColorJitter(brightness=0.75, contrast=0.75, saturation=0.75, hue=0.51), \n    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.3),\n                                transforms.ToTensor(),\n                                transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])\nTRAIN_DIR = Path('../input/all-dogs/')\nANNOTATION_DIR = Path('../input/annotation/Annotation/')\n#train_data = datasets.ImageFolder('../input/all-dogs/', transform=transform)\ntrain_data = DogsDataset(TRAIN_DIR / 'all-dogs/', ANNOTATION_DIR, transform=transform)\ntrain_loader = torch.utils.data.DataLoader(train_data, shuffle=True,\n                                           batch_size=batch_size, num_workers=4)\n                                           \nimgs = next(iter(train_loader))\nimgs = imgs.numpy().transpose(0, 2, 3, 1)\nplt.imshow(imgs[0])\nplt.show()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"netG = Generator(nz).to(device)\nnetD = Discriminator().to(device)\n\ncriterion = nn.BCEWithLogitsLoss()\n\noptimizerD = optim.Adam(netD.parameters(), lr=g_lr, betas=(beta1, 0.999))\noptimizerG = optim.Adam(netG.parameters(), lr=d_lr, betas=(beta1, 0.999))\n\n#fixed_noise = torch.randn(25, nz, 1, 1, device=device)\ndef show_generated_img(netG):\n    noise = torch.randn(1, nz, 1, 1, device=device)\n    gen_image = netG(noise).to(\"cpu\").clone().detach().squeeze(0)\n    gen_image = (gen_image.numpy().transpose(1, 2, 0) + 1)/2.\n    plt.imshow(gen_image)\n    plt.show()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"for epoch in tqdm(range(epochs)):\n    avgErrD = 0.0\n    avgErrG = 0.0\n    avgD_x, avgD_G_z1, avgD_G_z2 = 0.0, 0.0, 0.0\n    count = 0\n    for ii, real_images in enumerate(train_loader):\n        ############################\n        # (1) Update D network: maximize log(D(x)) + log(1 - D(G(z)))\n        ###########################\n        # train with real\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\n        outputR = netD(real_images)\n        errD_real = criterion(outputR, labels)\n        errD_real.backward()\n        D_x = outputR.mean().item()\n\n        # train with fake\n        noise = torch.randn(batch_size, nz, 1, 1, device=device)\n        fake = netG(noise)\n        labels.fill_(fake_label)\n        outputF = netD(fake.detach())\n        errD_fake = criterion(outputF, labels)\n        errD_fake.backward()\n        D_G_z1 = outputF.mean().item()\n        errD = errD_real + errD_fake\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: maximize log(D(G(z)))\n        ###########################\n        netG.zero_grad()\n        labels.fill_(1.0)  # fake labels are real for generator cost\n        outputF = netD(fake)\n        featureR1,featureR2 = netD.feature(real_images)\n        featureF1, featureF2 = netD.feature(fake)\n        errG = torch.mean((outputF - torch.mean(outputR.detach()) - labels) ** 2)\n        errG += 0.01* torch.mean((featureF1 - torch.mean(featureR1.detach(),0))**2)\n        errG += 0.01* torch.mean((featureF2 - torch.mean(featureR2.detach(),0))**2)\n        errG.backward()\n        D_G_z2 = outputF.mean().item()\n        optimizerG.step()\n        \n        avgErrD += errD.item()\n        avgErrG += errG.item()\n        avgD_x += D_x\n        avgD_G_z1 += D_G_z1 \n        avgD_G_z2 += D_G_z2\n        count += 1\n    print('[%d/%d]Loss_D: %.4f Loss_G: %.4f D(x): %.4f D(G(z)): %.4f / %.4f'% (epoch + 1, epochs, avgErrD/count, avgErrG/count, avgD_x/count, avgD_G_z1/count, avgD_G_z2/count))\n    if((epoch+1)%5 == 0):\n        #noise = torch.randn(32, nz, 1, 1, device=device)\n        #show_images(netG(noise),'dogs')\n        show_generated_img(netG)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import zipfile\n\nz = zipfile.PyZipFile('images.zip', mode='w')\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)\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        f = 'image_{}_{}.png'.format(i_batch, i_image)\n        save_image((gen_images[i_image, :, :, :]+1)/2., f);z.write(f); os.remove(f)\nz.close()","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}