{"cells":[{"metadata":{},"cell_type":"markdown","source":"# CGAN with historical weight averaging\nThe historical averaging trick from Ian Goodfellow's paper - https://arxiv.org/pdf/1606.03498.pdf did the trick for me. I couldn't go below 47 though. "},{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"import matplotlib.pyplot as plt\nimport numpy as np\nimport torch\nfrom torch import nn, optim\nimport torch.nn.functional as F\nfrom torchvision import datasets, transforms\nfrom torchvision.utils import save_image\nimport os\nfrom torch import nn, optim\nimport torch.nn.functional as F\nfrom torchvision import datasets, transforms\nfrom torchvision.utils import save_image\nfrom torch.utils.data import Dataset, DataLoader\nimport os\nfrom time import time\nfrom PIL import Image\nimport torch.nn.parallel\nimport torch.optim as optim\nimport torch.utils.data\nimport torchvision\nimport torchvision.datasets as dset\nimport torchvision.transforms as transforms\nimport torchvision.utils as vutils\nfrom torch.autograd import Variable\nfrom torch import nn, optim\nimport torch.nn.functional as F\nfrom torchvision import datasets, transforms\nfrom torchvision.utils import save_image\nimport matplotlib.image as mpimg\nimport torchvision\nimport torchvision.datasets as dset\nfrom torch.utils.data import Dataset, DataLoader\nimport torchvision.transforms as transforms\nimport torchvision.utils as vutils\nfrom torch.autograd import Variable\nimport xml.etree.ElementTree as ET\nimport random\nfrom torch.nn.utils import spectral_norm\nfrom scipy.stats import truncnorm\nfrom tqdm import tqdm_notebook as tqdm\n\nbatch_size = 32\nstart = time()\n\ndef 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\nseed_everything()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def truncated_normal(size, threshold=1):\n    values = truncnorm.rvs(-threshold, threshold, size=size)\n    return values\n\ndef mse(imageA, imageB):\n        err = np.sum((imageA.astype(\"float\") - imageB.astype(\"float\")) ** 2)\n        err /= float(imageA.shape[0] * imageA.shape[1])\n        return err\n    \ndef make_answer():\n    good_breeds = analyse_generated_by_class()\n    create_submit(good_breeds)\n    \ndef analyse_generated_by_class(n_images=5):\n    good_breeds = []\n    for l in range(len(decoded_dog_labels)):\n        sample = []\n        for _ in range(n_images):\n            noise = torch.randn(1, nz, 1, 1, device=device)\n            dog_label = torch.full((1,) , l, device=device, dtype=torch.long)\n            gen_image = netG(noise, dog_label).to(\"cpu\").clone().detach().squeeze(0)\n            gen_image = gen_image.numpy().transpose(1, 2, 0)\n            sample.append(gen_image)\n        \n        d = np.round(np.sum([mse(sample[k], sample[k+1]) for k in range(len(sample)-1)])/n_images, 1)\n        if d < 1.0: continue  # had mode colapse(discard)\n            \n        print(f\"Generated breed({d}): \", decoded_dog_labels[l])\n        figure, axes = plt.subplots(1, len(sample), figsize=(64, 64))\n        for index, axis in enumerate(axes):\n            axis.axis('off')\n            image_array = (sample[index] + 1.) / 2.\n            axis.imshow(image_array)\n        plt.show()\n        \n        good_breeds.append(l)\n    return good_breeds\n\n\ndef create_submit(good_breeds):\n    print(\"Creating submit\")\n    os.makedirs('../output_images', exist_ok=True)\n    im_batch_size = 100\n    n_images = 10000\n    \n    all_dog_labels = np.random.choice(good_breeds, size=n_images, replace=True)\n    for i_batch in range(0, n_images, im_batch_size):\n        noise = torch.randn(im_batch_size, nz, 1, 1, device=device)\n        dog_labels = torch.from_numpy(all_dog_labels[i_batch: (i_batch+im_batch_size)]).to(device)\n        gen_images = netG(noise, dog_labels)\n        gen_images = (gen_images.to(\"cpu\").clone().detach() + 1.) / 2.\n        for ii, img in enumerate(gen_images):\n            save_image(gen_images[ii, :, :, :], os.path.join('../output_images', f'image_{i_batch + ii:05d}.png'))\n            \n    import shutil\n    shutil.make_archive('images', 'zip', '../output_images')  ","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class DataGenerator(Dataset):\n    def __init__(self, directory, transform=None, n_samples=np.inf, crop_dogs=True):\n        self.directory = directory\n        self.transform = transform\n        self.n_samples = n_samples        \n        self.samples, self.labels = self.load_dogs_data(directory, crop_dogs)\n\n    def load_dogs_data(self, directory, crop_dogs):\n        required_transforms = torchvision.transforms.Compose([\n                torchvision.transforms.Resize(64),\n                torchvision.transforms.CenterCrop(64),\n        ])\n\n        imgs = []\n        labels = []\n        paths = []\n        for root, _, fnames in sorted(os.walk(directory)):\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            # Load image\n            try: img = dset.folder.default_loader(path)\n            except: continue\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                \n            if crop_dogs:\n                tree = ET.parse(os.path.join('../input/annotation/Annotation/',\n                                             annotation_dirname, annotation_basename))\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                    object_img = required_transforms(img.crop((xmin, ymin, xmax, ymax)))\n                    imgs.append(object_img)\n                    labels.append(annotation_dirname.split('-')[1].lower())\n\n            else:\n                object_img = required_transforms(img)\n                imgs.append(object_img)\n                labels.append(annotation_dirname.split('-')[1].lower())\n            \n        return imgs, labels\n    \n    \n    def __getitem__(self, index):\n        sample = self.samples[index]\n        label = self.labels[index]\n        \n        if self.transform is not None: \n            sample = self.transform(sample)\n        return np.asarray(sample), label\n\n    \n    def __len__(self):\n        return len(self.samples)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"%%time\ndatabase = '../input/all-dogs/all-dogs/'\n\nrandom_transforms = [transforms.ColorJitter(brightness=(1,1.3), contrast=(1,1.3), saturation=0, hue=0), transforms.RandomRotation(degrees=5)]\ntransform = transforms.Compose([transforms.RandomHorizontalFlip(p=0.5),\n                                transforms.RandomApply(random_transforms, p=0.5),\n                                transforms.ToTensor(),\n                                transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])\n\ntrain_data = DataGenerator(database, transform=transform, n_samples=25000, crop_dogs=True)\n\ndecoded_dog_labels = {i:breed for i, breed in enumerate(sorted(set(train_data.labels)))}\nencoded_dog_labels = {breed:i for i, breed in enumerate(sorted(set(train_data.labels)))}\ntrain_data.labels = [encoded_dog_labels[l] for l in train_data.labels] # encode dog labels in the data generator\n\n\n\ntrain_loader = torch.utils.data.DataLoader(train_data, shuffle=True,\n                                           batch_size=batch_size, num_workers=4)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Gan Helpers"},{"metadata":{"trusted":true},"cell_type":"code","source":"class PixelwiseNorm(nn.Module):\n    def __init__(self):\n        super(PixelwiseNorm, self).__init__()\n\n    def forward(self, x, alpha=1e-8):\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    \ndef show_generated_img_all(n_images=5, nz=128):\n    sample = []\n    for _ in range(n_images):\n        noise = torch.randn(1, nz, 1, 1, device=device)\n        dog_label = torch.randint(0, len(encoded_dog_labels), (1, ), device=device)\n        gen_image = netG(noise, dog_label).to(\"cpu\").clone().detach().squeeze(0)\n        gen_image = gen_image.numpy().transpose(1, 2, 0)\n        sample.append(gen_image)\n        \n    figure, axes = plt.subplots(1, len(sample), figsize=(64, 64))\n    for index, axis in enumerate(axes):\n        axis.axis('off')\n        image_array = (sample[index] + 1.) / 2.\n        axis.imshow(image_array)\n    plt.show()\n        \n# 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)\n#     gen_image = gen_image.numpy().transpose(1, 2, 0)\n#     gen_image = ((gen_image+1.0)/2.0)\n#     plt.imshow(gen_image)\n#     plt.show()  \n    \ndef show_generated_img(n_images=5, nz=128):\n    sample = []\n    for _ in range(n_images):\n        noise = torch.randn(1, nz, 1, 1, device=device)\n        dog_label = torch.randint(0, len(encoded_dog_labels), (1, ), device=device)\n        gen_image = netG(noise, dog_label).to(\"cpu\").clone().detach().squeeze(0)\n        gen_image = gen_image.numpy().transpose(1, 2, 0)\n        sample.append(gen_image)\n        \n    figure, axes = plt.subplots(1, len(sample), figsize=(64, 64))\n    for index, axis in enumerate(axes):\n        axis.axis('off')\n        image_array = (sample[index] + 1.) / 2.\n        axis.imshow(image_array)\n    plt.show()    \n    \nclass MinibatchStdDev(nn.Module):\n    def __init__(self):\n        super(MinibatchStdDev, self).__init__()\n    def forward(self, x, alpha=1e-8):\n        batch_size, _, height, width = x.shape\n        y = x - x.mean(dim=0, keepdim=True)\n        y = torch.sqrt(y.pow(2.).mean(dim=0, keepdim=False) + alpha)\n        y = y.mean().view(1, 1, 1, 1)\n        y = y.repeat(batch_size,1, height, width)\n        y = torch.cat([x, y], 1)\n        return y    \n    \n    \ndef snconv2d(in_channels, out_channels, kernel_size, stride=1, padding=0, dilation=1, groups=1, bias=False):\n    return spectral_norm(nn.Conv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=kernel_size,\n                                   stride=stride, padding=padding, dilation=dilation, groups=groups, bias=bias))  \n\ndef sndeconv2d(in_channels, out_channels, kernel_size, stride=1, padding=0, dilation=1, groups=1, bias=False):\n    return spectral_norm(nn.ConvTranspose2d(in_channels=in_channels, out_channels=out_channels, kernel_size=kernel_size,\n                                   stride=stride, padding=padding, dilation=dilation, groups=groups, bias=bias))  \n\ndef snlinear(in_features, out_features):\n    return spectral_norm(nn.Linear(in_features=in_features, out_features=out_features))\n\ndef sn_embedding(num_embeddings, embedding_dim):\n    return spectral_norm(nn.Embedding(num_embeddings=num_embeddings, embedding_dim=embedding_dim))\n\nclass ConditionalBatchNorm2d(nn.Module):\n    # https://github.com/pytorch/pytorch/issues/8985#issuecomment-405080775\n    def __init__(self, num_features, num_classes):\n        super().__init__()\n        self.num_features = num_features\n        self.bn = nn.BatchNorm2d(num_features, momentum=0.001, affine=False)\n        self.embed = nn.Embedding(num_classes, num_features * 2)\n        # self.embed.weight.data[:, :num_features].normal_(1, 0.02)  # Initialise scale at N(1, 0.02)\n        self.embed.weight.data[:, :num_features].fill_(1.)  # Initialize scale to 1\n        self.embed.weight.data[:, num_features:].zero_()  # Initialize bias at 0\n\n    def forward(self, x, y):\n        out = self.bn(x)\n        gamma, beta = self.embed(y).chunk(2, 1)\n        out = gamma.view(-1, self.num_features, 1, 1) * out + beta.view(-1, self.num_features, 1, 1)\n        return out","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class GenBlock(nn.Module):\n    def __init__(self, in_channels, out_channels, num_classes):\n        super(GenBlock, self).__init__()\n        self.cond_bn1 = ConditionalBatchNorm2d(in_channels, num_classes)\n        self.relu = nn.ReLU(inplace=True)\n        self.snconv2d1 = snconv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=3, stride=1, padding=1)\n        self.cond_bn2 = ConditionalBatchNorm2d(out_channels, num_classes)\n        self.snconv2d2 = snconv2d(in_channels=out_channels, out_channels=out_channels, kernel_size=3, stride=1, padding=1)\n        self.snconv2d0 = snconv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=1, stride=1, padding=0)\n\n    def forward(self, x, labels):\n        x0 = x\n\n        x = self.cond_bn1(x, labels)\n        x = self.relu(x)\n        x = F.interpolate(x, scale_factor=2, mode='nearest') # upsample\n        x = self.snconv2d1(x)\n        x = self.cond_bn2(x, labels)\n        x = self.relu(x)\n        x = self.snconv2d2(x)\n\n        x0 = F.interpolate(x0, scale_factor=2, mode='nearest') # upsample\n        x0 = self.snconv2d0(x0)\n\n        out = x + x0\n        return out    \n    \nclass DiscBlock(nn.Module):\n    def __init__(self, in_channels, out_channels):\n        super(DiscBlock, self).__init__()\n        self.relu = nn.ReLU()\n        self.snconv2d1 = snconv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=3, stride=1, padding=1)\n        self.snconv2d2 = snconv2d(in_channels=out_channels, out_channels=out_channels, kernel_size=3, stride=1, padding=1)\n        self.downsample = nn.AvgPool2d(2)\n        self.ch_mismatch = False\n        if in_channels != out_channels:\n            self.ch_mismatch = True\n        self.snconv2d0 = snconv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=1, stride=1, padding=0)\n\n    def forward(self, x, downsample=True):\n        x0 = x\n\n        x = self.relu(x)\n        x = self.snconv2d1(x)\n        x = self.relu(x)\n        x = self.snconv2d2(x)\n        if downsample:\n            x = self.downsample(x)\n\n        if downsample or self.ch_mismatch:\n            x0 = self.snconv2d0(x0)\n            if downsample:\n                x0 = self.downsample(x0)\n\n        out = x + x0\n        return out        ","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Generator and Discriminator"},{"metadata":{"trusted":true},"cell_type":"code","source":"class Generator(nn.Module):\n    def __init__(self, nz, nfeats, nchannels, num_classes):\n        super(Generator, self).__init__()\n\n        # input is Z, going into a convolution\n        self.conv1 = spectral_norm(nn.ConvTranspose2d(nz, nfeats * 8, 4, 1, 0, bias=False))\n        #self.bn1 = nn.BatchNorm2d(nfeats * 8)\n        # state size. (nfeats*8) x 4 x 4\n        self.block1 = GenBlock(nfeats * 8, nfeats * 8, num_classes)\n        # state size. (nfeats*8) x 8 x 8\n        self.block2 = GenBlock(nfeats * 8, nfeats * 4, num_classes)\n        # state size. (nfeats*4) x 16 x 16\n        self.block3 = GenBlock(nfeats * 4, nfeats * 2, num_classes)        \n        # state size. (nfeats*2) x 32 x 32\n        self.block4 = GenBlock(nfeats * 2, nfeats, num_classes)\n        self.bn = nn.BatchNorm2d(nfeats, eps=1e-5, momentum=0.0001, affine=True)\n        # state size. nfeats x 64 x 64\n        self.conv6 = spectral_norm(nn.ConvTranspose2d(nfeats, nchannels, 3, 1, 1, bias=False))\n        # state size. nfeats x 64 x 64\n        self.pixnorm = PixelwiseNorm()\n        self.relu = nn.ReLU()\n        \n    def forward(self, x, labels):\n        #x = F.leaky_relu(self.bn1(self.conv1(x)))\n        #x = F.leaky_relu(self.bn2(self.conv2(x)))\n        #x = F.leaky_relu(self.bn3(self.conv3(x)))\n        #x = F.leaky_relu(self.bn4(self.conv4(x)))\n        #x = F.leaky_relu(self.bn5(self.conv5(x)))\n        x = self.conv1(x)\n        x = self.block1(x, labels)\n        x = self.block2(x, labels)\n        x = self.block3(x, labels)\n        x = self.block4(x, labels)\n        x = self.bn(x)\n        x = self.relu(x)\n#         x = self.pixnorm(x)\n        x = self.conv6(x)\n        x = torch.tanh(x)\n        return x\n\n\n\n    \n\n\n# class Discriminator(nn.Module):\n#     def __init__(self, nchannels, nfeats, num_classes):\n#         super(Discriminator, self).__init__()\n\n#         # input is (nchannels) x 64 x 64\n#         self.conv1 = nn.Conv2d(nchannels, nfeats, 4, 2, 1, bias=False)\n#         # state size. (nfeats) x 32 x 32\n#         self.block1 = DiscBlock(nfeats, nfeats*2)\n#         self.bn1 = nn.BatchNorm2d(nfeats * 2)\n#         # state size. (nfeats*2) x 16 x 16\n#         self.block2 = DiscBlock(nfeats*2, nfeats*4)\n# #         self.bn2 = nn.BatchNorm2d(nfeats * 4)\n#         # state size. (nfeats*4) x 8 x 8\n#         self.block3 = DiscBlock(nfeats*4, nfeats*8)\n#         # state size. (nfeats*8) x 4 x 4\n#         self.block4 = DiscBlock(nfeats*8, nfeats*8)\n#         # state size. (nfeats*8) x 4 x 4 \n#         self.snlinear1 = snlinear(in_features=nfeats*8, out_features=1)\n#         self.sn_embedding1 = sn_embedding(num_classes, nfeats*8)\n        \n#     def forward(self, x, labels):\n#         x = self.conv1(x)\n#         x = self.block1(x)\n#         x = self.bn1(x)\n#         x = self.block2(x)\n# #         x = self.bn2(x)\n#         x = self.block3(x)\n#         x = self.block4(x, downsample=False)\n#         x = F.relu(x)\n#         x = torch.sum(x, dim=[2,3]) # n x (nfeats*8)\n#         output1 = torch.squeeze(self.snlinear1(x)) # n x 1\n        \n#         # Projection\n#         h_labels = self.sn_embedding1(labels)\n#         proj = torch.mul(x, h_labels) \n#         output2 = torch.sum(proj, dim=[1])\n        \n#         output = output1 + output2\n# #         return output, torch.sigmoid(output), \n#         return torch.sigmoid(output).view(-1, 1)\n\nclass Discriminator(nn.Module):\n    def __init__(self, nchannels, nfeats, num_classes):\n        super(Discriminator, self).__init__()\n        self.label_emb = sn_embedding(num_classes, 64*64)\n        # input is (nchannels) x 64 x 64\n        self.conv1 = nn.Conv2d(nchannels+1, nfeats, 4, 2, 1, bias=False)\n        # state size. (nfeats) x 32 x 32\n        self.block1 = DiscBlock(nfeats, nfeats*2)\n        self.bn1 = nn.BatchNorm2d(nfeats * 2)\n#         self.self_attn = Self_Attn(nfeats*2)\n        # state size. (nfeats*2) x 16 x 16\n        self.block2 = DiscBlock(nfeats*2, nfeats*4)\n#         self.bn2 = nn.BatchNorm2d(nfeats * 4)\n        # state size. (nfeats*4) x 8 x 8\n        self.block3 = DiscBlock(nfeats*4, nfeats*8)\n        # state size. (nfeats*8) x 4 x 4\n        self.block4 = DiscBlock(nfeats*8, nfeats*8)\n#         self.downscale = nn.MaxPool2d(2)\n        # state size. (nfeats*8) x 2 x 2\n#         self.batch_discriminator = MinibatchStdDev()\n        # state size. (nfeats*8+1) x 2 x 2\n        \n        self.conv5 = spectral_norm(nn.Conv2d(nfeats * 8, 1, 2, 1, 0, bias=False))\n        # state size. 1 x 1 x 1\n        \n    def forward(self, imgs, labels):\n        enc = self.label_emb(labels).view((-1, 1, 64, 64))\n        enc = F.normalize(enc, p=2, dim=1)\n        x = torch.cat((imgs, enc), 1)\n        \n        \n        x = self.conv1(x)\n        x = self.block1(x)\n        x = self.bn1(x)\n#         x = self.self_attn(x)\n        x = self.block2(x)\n#         x = self.bn2(x)\n        x = self.block3(x)\n        x = self.block4(x)\n#         x = self.downscale(x)\n        x = F.leaky_relu(x, 0.2)\n#         x = self.batch_discriminator(x)\n        x = torch.sigmoid(self.conv5(x))\n        return x.view(-1, 1)\n     ","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Hyperparameters"},{"metadata":{"trusted":true},"cell_type":"code","source":"device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n\n# TTUR\nlr_d = 0.0006\nlr_g = 0.0003\nbeta1 = 0.5\nepochs = 300\nnum_classes = len(encoded_dog_labels)\nnetG = Generator(128, 32, 3, num_classes).to(device)\nnetD = Discriminator(3, 48, num_classes).to(device)\n\ncriterion = nn.BCELoss()\ncriterionH = nn.MSELoss()\n\noptimizerD = optim.Adam(netD.parameters(), lr=lr_d, betas=(beta1, 0.999))\noptimizerG = optim.Adam(netG.parameters(), lr=lr_g, betas=(beta1, 0.999))\n# lr_schedulerG = torch.optim.lr_scheduler.CosineAnnealingWarmRestarts(optimizerG,\n#                                                                      T_0=epochs//200, eta_min=0.00005)\n# lr_schedulerD = torch.optim.lr_scheduler.CosineAnnealingWarmRestarts(optimizerD,\n#                                                                      T_0=epochs//200, eta_min=0.00005)\n\nnz = 128\nfixed_noise = torch.randn(25, nz, 1, 1, device=device)\n\nreal_label = 0.7\nfake_label = 0.0\nbatch_size = train_loader.batch_size\n\ndef get_model_weights(net):\n    average = {}\n    params = dict(net.named_parameters())\n    for p in params:\n        average[p] = params[p].detach()    \n    return average    \n\nprint(sum(p.numel() for p in netG.parameters()))\nprint(sum(p.numel() for p in netD.parameters()))\n\naverageD = False\naverageG = False\nhist_average_cutoff = 96 #Emperically found this good\n# hist_average_cutoff = -1","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Train"},{"metadata":{"trusted":true},"cell_type":"code","source":"fids = {}\nstep = 1\nfor epoch in range(epochs):\n    if (time() - start) > 310 : #Change to 31000  \n        break\n   \n        \n    for ii, (real_images, dog_labels) in tqdm(enumerate(train_loader), total=len(train_loader)):\n        if (time() - start) > 310 : #Change to 31000\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        \n#         dog_labels = torch.tensor(dog_labels, device=device)\n        dog_labels = dog_labels.to(device)\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, dog_labels)\n        errD_real = criterion(output, labels)\n        errD_real.backward()\n        D_x = output.mean().item()\n        \n        # Historical averaging weights\n        err_hD = 0\n        if epoch > hist_average_cutoff:\n            if not averageD:\n                print(\"Starting historical weight averaging for discriminator\")\n                averageD = get_model_weights(netD)\n            paramsD = dict(netD.named_parameters())\n            for p in paramsD:\n                err_hD += criterionH(paramsD[p], averageD[p])\n                averageD[p] = (averageD[p] * (step-1) + paramsD[p].detach())/step\n            err_hD.backward()\n\n        # train with fake\n        noise = torch.randn(batch_size, nz, 1, 1, device=device)\n\n#         print(noise.shape)\n        fake = netG(noise, dog_labels)\n        labels.fill_(fake_label) + np.random.uniform(0, 0.2)\n        output = netD(fake.detach(), dog_labels)\n        errD_fake = criterion(output, labels)\n        errD_fake.backward()\n        D_G_z1 = output.mean().item()\n        errD_final = errD_real + errD_fake + err_hD\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, dog_labels)\n        errG = criterion(output, labels)\n        errG.backward()\n        D_G_z2 = output.mean().item()\n                \n        err_hG = 0\n        if epoch > hist_average_cutoff:\n            if not averageG:\n                print(\"Starting historical weight averaging for generator\")\n                averageG = get_model_weights(netG)\n            paramsG = dict(netG.named_parameters())\n            for p in paramsG:\n                err_hG += criterionH(paramsG[p], averageG[p])\n                averageG[p] = (averageG[p] * (step-1) + paramsG[p].detach())/step\n            err_hG.backward()\n            step += 1\n        \n        errG_final = errG + err_hG\n        \n        optimizerG.step()\n        \n        if ii % 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_final.item(), errG_final.item(), D_x, D_G_z1, D_G_z2))\n            \n#         lr_schedulerG.step(epoch)\n#         lr_schedulerD.step(epoch)\n\n#     if epoch % 5 == 0:\n#         show_generated_img(6)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# show_generated_img_all()\nmake_answer()","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":1}