{"cells":[{"metadata":{},"cell_type":"markdown","source":"# Dogs GAN \n\nNirjhar Roy deserves all of Credit: https://www.kaggle.com/phoenix9032/gan-dogs-starter-24-jul-custom-layers\nIf you upvote this kernel and haven't upvoted his, please do. \n\nThis is basically a more organized version of his code, hopefully this will allow us (and anyone who wants to for it) to better understand his solution, explore and change the model with more ease.\n\n\n\n\n\n\n\n\n## content Table\n\n1. ** Imports **
\n2. ** Conf**
\n3. ** Pytorch initializations**
\n4. **Data generator**
\n5. **Defining Neural Nets**
\n> **5.1 Generator**
\n> **5.2 Discriminator**
\n6. **Functions **
\n> **6.1 Image show functions**
\n> **6.2 Truncate function**
\n7. **Train Module**
\n8. **Define and train the GAN**
\n9. **Make 10000 Images for competiotion**
\n"},{"metadata":{},"cell_type":"markdown","source":"## 1. Imports \n\n"},{"metadata":{"trusted":true},"cell_type":"code","source":"import shutil\nimport os\nimport random\nfrom time import time\nimport matplotlib.pyplot as plt\nimport numpy as np\nimport torch\nfrom torch.nn.utils import spectral_norm\nfrom torch.utils.data import Dataset\nimport torchvision\nimport xml.etree.ElementTree as ET\nfrom tqdm import tqdm_notebook as tqdm\nfrom scipy.stats import truncnorm\n","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## 2. Conf"},{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"# Define RGB images\nCHANNEL = {'RGB': 3, 'BW': 1}\nIMG_CHANNEL = CHANNEL['RGB']\n\n# Defining path to dog images \nDATA_LOCATION = '../input/all-dogs/all-dogs/'\n\n# Define batch size\nBATCH_SIZE = 32\n\n# Each image in the generator starts from LATENT_DIM points generated from a gaussian distribution (variance = 1 sigma = 0)\nLATENT_DIM = 100\n\n# The amount of parameters in the networks scales with those number\nCONV_DEPTHS_G = 32\nCONV_DEPTHS_D = 48\n\n\n# Learning rate functions\n# The learning scheduale \nLEARNING_RATE_G = 0.0003 # was 0.0003 \nLEARNING_RATE_D = 0.0001 # was 0.0001\nBETA_1 = 0.5\nT0_interval = 200\n# T0_interval = 400\n\n\n# Amount of epochs, EPOCHS/T0_interval has to be integer\n\nEPOCHS = 1600\n# EPOCHS = 700 \n\n\n# If archtiecture changed to CGAN then it might be needed later\nNUM_CLASSES = 10\n\n# Define limit for training time\nTIME_FOR_TRAIN = 25000 # 8.3 hours (in seconds)\n# TIME_FOR_TRAIN = 7 * 60 * 60 # hours\n# TIME_FOR_TRAIN = 5 * 60","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## LR scheduale - cosine annealing \n\nRelevant paper: https://arxiv.org/pdf/1608.03983.pdf
\nAlso explained in: https://sidravi1.github.io/blog/2018/04/25/playing-with-sgdr\n\n
\n\\begin{align}\n\\eta_t = \\eta_{min}^i + \\frac{1}{2}(\\eta_{max}^i - \\eta_{min}^i)(1 + \\cos(\\frac{T_{cur}}{T_i} \\pi))\n\\end{align}\n\n"},{"metadata":{},"cell_type":"markdown","source":"## 3. Pytorch initializations"},{"metadata":{"trusted":true},"cell_type":"code","source":"def seed_everything(seed=42):\n random.seed(seed)\n os.environ['PYTHONHASHSEED'] = str(seed)\n np.random.seed(seed)\n torch.manual_seed(seed)\n torch.cuda.manual_seed(seed)\n torch.backends.cudnn.deterministic = True\n \nseed_everything()\ndevice = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n\n\n","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## 4. Data generator\n\nAlongside data augmentation"},{"metadata":{"_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","trusted":true},"cell_type":"code","source":"class DataGenerator(Dataset):\n def __init__(self, directory, transform=None, n_samples=np.inf):\n self.directory = directory\n self.transform = transform\n self.n_samples = n_samples\n self.samples = self._load_subfolders_images(directory)\n if len(self.samples) == 0:\n raise RuntimeError(\"Found 0 files in subfolders of: {}\".format(directory))\n\n def _load_subfolders_images(self, root):\n IMG_EXTENSIONS = (\n '.jpg', '.jpeg', '.png', '.ppm', '.bmp', '.pgm', '.tif', '.tiff', '.webp')\n\n def is_valid_file(x):\n return torchvision.datasets.folder.has_file_allowed_extension(x, IMG_EXTENSIONS)\n\n \n \n required_transforms = torchvision.transforms.Compose([\n torchvision.transforms.Resize(64),\n torchvision.transforms.CenterCrop(64),\n ])\n imgs = []\n paths = []\n for root, _, fnames in sorted(os.walk(root)):\n for fname in sorted(fnames)[:min(self.n_samples, 999999999999999)]:\n path = os.path.join(root, fname)\n paths.append(path)\n\n for path in paths:\n if is_valid_file(path):\n # Load image\n img = torchvision.datasets.folder.default_loader(path)\n\n # Get bounding boxes\n annotation_basename = os.path.splitext(os.path.basename(path))[0]\n annotation_dirname = next(\n dirname for dirname in os.listdir('../input/annotation/Annotation/') if\n dirname.startswith(annotation_basename.split('_')[0]))\n annotation_filename = os.path.join('../input/annotation/Annotation/',\n annotation_dirname, annotation_basename)\n tree = ET.parse(annotation_filename)\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\n w = np.min((xmax - xmin, ymax - ymin))\n bbox = (xmin, ymin, xmin + w, ymin + w)\n object_img = required_transforms(img.crop(bbox))\n # object_img = object_img.resize((64,64), Image.ANTIALIAS)\n imgs.append(object_img)\n return imgs\n\n def __getitem__(self, index):\n sample = self.samples[index]\n\n if self.transform is not None:\n sample = self.transform(sample)\n\n return np.asarray(sample)\n\n def __len__(self):\n return len(self.samples)\n \n\ntransform = torchvision.transforms.Compose([torchvision.transforms.RandomHorizontalFlip(p=0.1),\n torchvision.transforms.ToTensor(),\n torchvision.transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))]) \ndata_gen = DataGenerator(DATA_LOCATION, transform=transform, n_samples=25000)\ntrain_loader = torch.utils.data.DataLoader(data_gen, shuffle=True, batch_size=BATCH_SIZE, num_workers=4)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## 5. Defining Neural Nets"},{"metadata":{},"cell_type":"markdown","source":"## 5.1 Generator"},{"metadata":{"trusted":true},"cell_type":"code","source":"class PixelwiseNorm(torch.nn.Module):\n def __init__(self):\n super(PixelwiseNorm, self).__init__()\n\n def forward(self, x, alpha=1e-8):\n \"\"\"\n forward pass of the module\n :param x: input activations volume\n :param alpha: small number for numerical stability\n :return: y => pixel normalized activations\n \"\"\"\n y = x.pow(2.).mean(dim=1, keepdim=True).add(alpha).sqrt() # [N1HW]\n y = x / y # normalize the input x volume\n return y\n \n \nclass Generator(torch.nn.Module):\n def __init__(self, latent_dim, nfeats, nchannels): #was nz\n super(Generator, self).__init__()\n\n # input is Z, going into a convolution\n self.conv1 = spectral_norm(torch.nn.ConvTranspose2d(latent_dim, nfeats * 8, 4, 1, 0, bias=False)) #was nz\n\n self.conv2 = spectral_norm(torch.nn.ConvTranspose2d(nfeats * 8, nfeats * 8, 4, 2, 1, bias=False))\n\n self.conv3 = spectral_norm(torch.nn.ConvTranspose2d(nfeats * 8, nfeats * 4, 4, 2, 1, bias=False))\n\n self.conv4 = spectral_norm(torch.nn.ConvTranspose2d(nfeats * 4, nfeats * 2, 4, 2, 1, bias=False))\n\n self.conv5 = spectral_norm(torch.nn.ConvTranspose2d(nfeats * 2, nfeats, 4, 2, 1, bias=False))\n\n self.conv6 = spectral_norm(torch.nn.ConvTranspose2d(nfeats, nchannels, 3, 1, 1, bias=False))\n self.pixnorm = PixelwiseNorm()\n\n def forward(self, x):\n\n x = torch.nn.functional.leaky_relu(self.conv1(x))\n x = self.pixnorm(x) # wasnt here\n x = torch.nn.functional.leaky_relu(self.conv2(x))\n x = torch.nn.Dropout(0.05)(x)\n x = self.pixnorm(x)\n x = torch.nn.functional.leaky_relu(self.conv3(x))\n x = self.pixnorm(x) \n x = torch.nn.functional.leaky_relu(self.conv4(x))\n x = torch.nn.Dropout(0.05)(x)\n x = self.pixnorm(x)\n x = torch.nn.functional.leaky_relu(self.conv5(x))\n x = self.pixnorm(x)\n x = torch.tanh(self.conv6(x))\n\n return x\n","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## 5.2 Discriminator"},{"metadata":{"trusted":true},"cell_type":"code","source":"class MinibatchStdDev(torch.nn.Module):\n \"\"\"\n Minibatch standard deviation layer for the discriminator\n \"\"\"\n\n def __init__(self):\n \"\"\"\n derived class constructor\n \"\"\"\n super(MinibatchStdDev, self).__init__()\n\n def forward(self, x, alpha=1e-8):\n \"\"\"\n forward pass of the layer\n :param x: input activation volume\n :param alpha: small number for numerical stability\n :return: y => x appended with standard deviation constant map\n \"\"\"\n batch_size, _, height, width = x.shape\n # [B x C x H x W] Subtract mean over batch.\n y = x - x.mean(dim=0, keepdim=True)\n # [1 x C x H x W] Calc standard deviation over batch\n y = torch.sqrt(y.pow(2.).mean(dim=0, keepdim=False) + alpha)\n\n # [1] Take average over feature_maps and pixels.\n y = y.mean().view(1, 1, 1, 1)\n\n # [B x 1 x H x W] Replicate over group and pixels.\n y = y.repeat(batch_size, 1, height, width)\n\n # [B x C x H x W] Append as new feature_map.\n y = torch.cat([x, y], 1)\n # return the computed values:\n return y","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class Discriminator(torch.nn.Module):\n def __init__(self, nchannels, nfeats):\n super(Discriminator, self).__init__()\n\n # input is (nchannels) x 64 x 64\n self.conv1 = torch.nn.Conv2d(nchannels, nfeats, 4, 2, 1, bias=False)\n # state size. (nfeats) x 32 x 32\n\n self.conv2 = spectral_norm(torch.nn.Conv2d(nfeats, nfeats * 2, 4, 2, 1, bias=False))\n self.bn2 = torch.nn.BatchNorm2d(nfeats * 2)\n # state size. (nfeats*2) x 16 x 16\n\n self.conv3 = spectral_norm(torch.nn.Conv2d(nfeats * 2, nfeats * 4, 4, 2, 1, bias=False))\n self.bn3 = torch.nn.BatchNorm2d(nfeats * 4)\n # state size. (nfeats*4) x 8 x 8\n\n self.conv4 = spectral_norm(torch.nn.Conv2d(nfeats * 4, nfeats * 8, 4, 2, 1, bias=False))\n self.bn4 = torch.nn.MaxPool2d(2)\n # state size. (nfeats*8) x 4 x 4\n self.batch_discriminator = MinibatchStdDev()\n\n self.conv5 = spectral_norm(torch.nn.Conv2d(nfeats * 8 + 1, 1, 2, 1, 0, bias=False))\n # state size. 1 x 1 x 1\n\n def forward(self, x):\n x = torch.nn.functional.leaky_relu(self.conv1(x), 0.2)\n x = torch.nn.functional.leaky_relu(self.bn2(self.conv2(x)), 0.2)\n x = torch.nn.functional.leaky_relu(self.bn3(self.conv3(x)), 0.2)\n x = torch.nn.functional.leaky_relu(self.bn4(self.conv4(x)), 0.2)\n x = self.batch_discriminator(x)\n x = torch.sigmoid(self.conv5(x))\n # x= self.conv5(x)\n return x.view(-1, 1)\n","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## 6. Functions "},{"metadata":{},"cell_type":"markdown","source":"### 6.1 Image show functions"},{"metadata":{"trusted":true},"cell_type":"code","source":"def show_generated_img_all():\n gen_z = torch.randn(32, LATENT_DIM, 1, 1, device=device)\n gen_images = netG(gen_z).to(\"cpu\").clone().detach()\n gen_images = gen_images.numpy().transpose(0, 2, 3, 1)\n gen_images = (gen_images + 1.0) / 2.0\n fig = plt.figure(figsize=(25, 16))\n for ii, img in enumerate(gen_images):\n ax = fig.add_subplot(4, 8, ii + 1, xticks=[], yticks=[])\n plt.imshow(img)\n # plt.savefig(filename)\n\n\ndef show_generated_img():\n row_num = 1\n col_num = 10 \n gen_z = torch.randn(row_num * col_num , LATENT_DIM, 1, 1, device=device)\n gen_images = netG(gen_z).to(\"cpu\").clone().detach()\n gen_images = gen_images.numpy().transpose(0, 2, 3, 1)\n gen_images = (gen_images + 1.0) / 2.0\n fig = plt.figure(figsize=(20, 4))\n for ii, img in enumerate(gen_images):\n ax = fig.add_subplot(row_num, col_num, ii + 1, xticks=[], yticks=[])\n plt.imshow(img)\n plt.show()\n # plt.savefig(filename)\n \n ","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### 6.2 Truncate function"},{"metadata":{"trusted":true},"cell_type":"code","source":"def truncated_normal(size, threshold=1):\n values = truncnorm.rvs(-threshold, threshold, size=size)\n return values","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## 7. Train Module"},{"metadata":{"trusted":true},"cell_type":"code","source":"def train_module(epochs):\n step = 0\n start = time()\n for epoch in range(epochs):\n for ii, (real_images) in enumerate(train_loader):\n end = time()\n if (end - start) > TIME_FOR_TRAIN:\n break\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) + np.random.uniform(-0.1, 0.1)\n\n output = netD(real_images)\n errD_real = criterion(output, labels)\n errD_real.backward()\n D_x = output.mean().item()\n\n # train with fake\n noise = torch.randn(batch_size, LATENT_DIM, 1, 1, device=device)\n fake = netG(noise)\n labels.fill_(fake_label) + np.random.uniform(0, 0.2)\n output = netD(fake.detach())\n errD_fake = criterion(output, labels)\n errD_fake.backward()\n D_G_z1 = output.mean().item()\n errD = errD_real + errD_fake\n optimizerD.step()\n\n ############################\n # (2) Update G network: maximize log(D(G(z)))\n ###########################\n netG.zero_grad()\n labels.fill_(real_label) # fake labels are real for generator cost\n output = netD(fake)\n errG = criterion(output, labels)\n errG.backward()\n D_G_z2 = output.mean().item()\n optimizerG.step()\n\n if step % 500 == 0:\n print('[%d/%d][%d/%d] Loss_D: %.4f Loss_G: %.4f D(x): %.4f D(G(z)): %.4f / %.4f'\n % (epoch + 1, EPOCHS, ii, len(train_loader),\n errD.item(), errG.item(), D_x, D_G_z1, D_G_z2))\n\n valid_image = netG(fixed_noise)\n step += 1\n lr_schedulerG.step(epoch)\n lr_schedulerD.step(epoch)\n \n if epoch % 5 == 0:\n show_generated_img()\n print(end - start)\n if (end - start) > TIME_FOR_TRAIN:\n break","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## 8. Define and train the GAN"},{"metadata":{"trusted":true},"cell_type":"code","source":"netG = Generator(LATENT_DIM, CONV_DEPTHS_G, IMG_CHANNEL).to(device)\nnetD = Discriminator(3, CONV_DEPTHS_D).to(device)\ncriterion = torch.nn.BCELoss()\noptimizerD = torch.optim.Adam(netD.parameters(), lr=LEARNING_RATE_G, betas=(BETA_1, 0.999))\noptimizerG = torch.optim.Adam(netG.parameters(), lr=LEARNING_RATE_D, betas=(BETA_1, 0.999))\nlr_schedulerG = torch.optim.lr_scheduler.CosineAnnealingWarmRestarts(optimizerG,\n T_0=EPOCHS // T0_interval, eta_min=0.00005)\nlr_schedulerD = torch.optim.lr_scheduler.CosineAnnealingWarmRestarts(optimizerD,\n T_0=EPOCHS // T0_interval, eta_min=0.00005)\n\nfixed_noise = torch.randn(25, LATENT_DIM, 1, 1, device=device)\n\nreal_label = 0.8\nfake_label = 0.0\nbatch_size = train_loader.batch_size\n# criterion = nn.MSELoss()\ntrain_module(EPOCHS)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## 9. Make 10000 Images for competiotion"},{"metadata":{"trusted":true},"cell_type":"code","source":"netG.eval()\n\nif 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 z = truncated_normal((im_batch_size, 100, 1, 1), threshold=1)\n gen_z = torch.from_numpy(z).float().to(device)\n # gen_z = torch.randn(im_batch_size, 100, 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 torchvision.utils.save_image((gen_images[i_image, :, :, :] + 1.0) / 2.0,\n os.path.join('../output_images', f'image_{i_batch + i_image:05d}.png'))\n\nshutil.make_archive('images', 'zip', '../output_images')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"show_generated_img_all()","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}