{"cells":[{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"##### This Python 3 environment comes with many helpful analytics libraries installed\n# It is defined by the kaggle/python docker image: https://github.com/kaggle/docker-python\n# For example, here's several helpful packages to load in \n\nimport numpy as np # linear algebra\nimport pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)\n\n# Input data files are available in the \"../input/\" directory.\n# For example, running this (by clicking run or pressing Shift+Enter) will list the files in the input directory\n\nimport os\nprint(os.listdir(\"../input\"))\n\n# Any results you write to the current directory are saved as output.","execution_count":null,"outputs":[]},{"metadata":{"_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","trusted":true},"cell_type":"code","source":"from PIL import Image\n#from PIL import ImageFile\n#ImageFile.LOAD_TRUNCATED_IMAGES = True\nimport time\nimport random\nimport numpy as np\nimport torch.nn as nn\nimport torch.optim as optim\nimport torch.utils.data\nfrom torch.utils.data import Dataset, DataLoader\nimport torchvision\nfrom torch.autograd import Variable\nimport torch.nn.functional as F\nimport torchvision.datasets as dset\nimport torchvision.transforms as transforms\nimport torchvision.utils as vutils\nimport matplotlib.pyplot as plt\nfrom torchvision.utils import save_image\nimport math\nfrom torch.nn.utils import spectral_norm\n\nimport glob\n\nimport xml.etree.ElementTree as ET\n\nfrom tqdm import tqdm_notebook as tqdm\n\nimport tensorboardX as tbx\nwriter = tbx.SummaryWriter(\".\")","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"start_time = time.time()\nrandom.seed(0)\nnp.random.seed(0)\ntorch.manual_seed(0)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"if not os.path.exists('../class_images'):\n    os.mkdir('../class_images')\nif not os.path.exists('../class_images/train'):\n    os.mkdir('../class_images/train')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# This loader will use the underlying loader plus crop the image based on the annotation\ndef doggo_loader(path):\n    img = torchvision.datasets.folder.default_loader(path) # default loader\n    \n    # Get bounding box\n    annotation_basename = os.path.splitext(os.path.basename(path))[0]\n    annotation_dirname = next(dirname for dirname in os.listdir('../input/annotation/Annotation/') if dirname.startswith(annotation_basename.split('_')[0]))\n    annotation_filename = os.path.join('../input/annotation/Annotation', annotation_dirname, annotation_basename)\n    tree = ET.parse(annotation_filename)\n    root = tree.getroot()\n    objects = root.findall('object')\n    for idx, o in enumerate(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        w = xmax - xmin\n        h = ymax - ymin\n        l = max([w,h])\n\n        if w > h:\n            ymin = ymin - int((l-h)/2)\n            #ymin = max(ymin - int((l-h)/2), 0)\n            #ymax = min(ymin - int((l-h)/2) + l, img.height)\n        else:\n            xmin = xmin - int((l-w)/2)\n            #xmin = max(xmin - int((l-w)/2), 0)\n            #xmax = min(xmin - int((l-w)/2) + l, img.width)\n        ##bbox = (xmin, ymin , xmax, ymax)\n        bbox = (xmin, ymin , xmin+l, ymin+l)\n\n        c_img = img.crop(bbox).copy()\n        #img = img.resize((100, 100))\n        c_img.thumbnail((100, 100), Image.ANTIALIAS)\n\n        c_img.save('../class_images/train/image_{}_{}.png'.format(annotation_basename, idx))\n\n    return img\n\n\n# The dataset (example)\ndataset = torchvision.datasets.ImageFolder(\n    '../input/all-dogs/',\n    loader=doggo_loader, # THE CUSTOM LOADER\n    transform=torchvision.transforms.Compose([\n        transforms.RandomResizedCrop(64, scale=(1.0, 1.0), ratio=(1., 1.)),\n\n        torchvision.transforms.ToTensor(),\n    ]) # some transformations, add your data preprocessing here\n)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"if not os.path.exists('../class_images'):\n    os.mkdir('../class_images')\nif not os.path.exists('../class_images/train'):\n    os.mkdir('../class_images/train')\ndataloader = torch.utils.data.DataLoader(dataset, batch_size=50,\n                                         shuffle=False, num_workers=8)\nfor i, data in enumerate(tqdm(dataloader)):\n    pass","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Initial_setting\nworkers = 8\nbatch_size=32\nnz = 128\nnch_g = 64\nnch_d = 64\nn_epoch = 20000000   \nlr = 0.0002\nbeta1 = 0.5\noutf = './result_lsgan'\ndisplay_interval = 100\nsave_fake_image_interval = 1500\nplt.rcParams['figure.figsize'] = 10, 6\n \ntry:\n    os.makedirs(outf, exist_ok=True)\nexcept OSError as error: \n    print(error)\n    pass\n \nrandom.seed(0)\nnp.random.seed(0)\ntorch.manual_seed(0)\n \ndef weights_init(m):\n    classname = m.__class__.__name__\n    if classname.find('Conv') != -1:            \n        m.weight.data.normal_(0.0, 0.02)\n        m.bias.data.fill_(0)\n    elif classname.find('Linear') != -1:        \n        m.weight.data.normal_(0.0, 0.02)\n        m.bias.data.fill_(0)\n    elif classname.find('BatchNorm') != -1:    \n        m.weight.data.normal_(1.0, 0.02)\n        m.bias.data.fill_(0)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"'''\nclass Generator(nn.Module):\n    def __init__(self, nz=100, nch_g=64, nch=3):\n        super(Generator, self).__init__()\n        self.layers = nn.ModuleDict({\n            'layer0': nn.Sequential(\n                nn.ConvTranspose2d(nz, nch_g * 8, 4, 1, 0),     \n                nn.BatchNorm2d(nch_g * 8),\n                nn.ReLU()                                       \n            ),  # (100, 1, 1) -> (512, 4, 4)\n            'layer1': nn.Sequential(\n                nn.ConvTranspose2d(nch_g * 8, nch_g * 4, 4, 2, 1),\n                nn.BatchNorm2d(nch_g * 4),\n                nn.ReLU(),\n                #nn.Dropout(p=0.5, inplace=False)\n            ),  # (512, 4, 4) -> (256, 8, 8)\n            'layer2': nn.Sequential(\n                nn.ConvTranspose2d(nch_g * 4, nch_g * 2, 4, 2, 1),\n                nn.BatchNorm2d(nch_g * 2),\n                nn.ReLU(),\n                #nn.Dropout(p=0.5, inplace=False)\n            ),  # (256, 8, 8) -> (128, 16, 16)\n            'layer3': nn.Sequential(\n                nn.ConvTranspose2d(nch_g * 2, nch_g, 4, 2, 1),\n                nn.BatchNorm2d(nch_g),\n                nn.ReLU()\n            ),  # (128, 16, 16) -> (64, 32, 32)\n            # 'layer3-5': nn.Sequential(\n            #    nn.ConvTranspose2d(nch_g, nch_g, 4, 2, 1),\n            #    nn.BatchNorm2d(nch_g),\n            #    nn.ReLU()\n            #),  # (64, 32, 32) -> (64, 64, 64)\n            'layer4': nn.Sequential(\n                #nn.ConvTranspose2d(nch_g, nch, 3, 1, 0),\n                nn.ConvTranspose2d(nch_g, nch, 4, 2, 1),\n                nn.Tanh()\n            )   # (64, 32, 32) -> (3, 64, 64)\n        })\n \n    def forward(self, z):\n        for layer in self.layers.values():  \n            z = layer(z)\n        return z\n    \n    def denorm(self, z):\n        for layer in self.layers.values():  \n            z = layer(z)   \n        z = z * 0.5 + 0.5\n        return z\n \n \nclass Discriminator(nn.Module):\n    def __init__(self, nch=3, nch_d=64):\n        super(Discriminator, self).__init__()\n        self.layers = nn.ModuleDict({\n            'layer0': nn.Sequential(\n                nn.Conv2d(nch, nch_d, 4, 2, 1),     \n                nn.LeakyReLU(negative_slope=0.2)    \n            ),  # (3, 64, 64) -> (64, 32, 32)\n            'layer1': nn.Sequential(\n                nn.Conv2d(nch_d, nch_d * 2, 4, 2, 1),\n                nn.BatchNorm2d(nch_d * 2),\n                nn.LeakyReLU(negative_slope=0.2)\n            ),  # (64, 32, 32) -> (128, 16, 16)\n            'layer2': nn.Sequential(\n                nn.Conv2d(nch_d * 2, nch_d * 4, 4, 2, 1),\n                nn.BatchNorm2d(nch_d * 4),\n                nn.LeakyReLU(negative_slope=0.2)\n            ),  # (128, 16, 16) -> (256, 8, 8)\n            'layer3': nn.Sequential(\n                nn.Conv2d(nch_d * 4, nch_d * 8, 4, 2, 1),\n                nn.BatchNorm2d(nch_d * 8),\n                nn.LeakyReLU(negative_slope=0.2)\n            ),  # (256, 8, 8) -> (512, 4, 4)\n            'layer4': nn.Conv2d(nch_d * 8, 1, 4, 1, 0)\n            # (512, 4, 4) -> (1, 1, 1)\n        })\n \n    def forward(self, x):\n        for layer in self.layers.values():  \n            x = layer(x)\n        return x.squeeze()  \n'''","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# ----------------------------------------------------------------------------\n# Pixelwise feature vector normalization.\n# reference: https://github.com/tkarras/progressive_growing_of_gans/blob/master/networks.py#L120\n# ----------------------------------------------------------------------------\nclass PixelwiseNorm(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    \nclass 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\n\nclass Generator(nn.Module):\n    def __init__(self, nz, nfeats, nchannels):\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        \n        self.conv2 = spectral_norm(nn.ConvTranspose2d(nfeats * 8, nfeats * 8, 4, 2, 1, bias=False))\n        #self.bn2 = nn.BatchNorm2d(nfeats * 8)\n        # state size. (nfeats*8) x 8 x 8\n        \n        self.conv3 = spectral_norm(nn.ConvTranspose2d(nfeats * 8, nfeats * 4, 4, 2, 1, bias=False))\n        #self.bn3 = nn.BatchNorm2d(nfeats * 4)\n        # state size. (nfeats*4) x 16 x 16\n        \n        self.conv4 = spectral_norm(nn.ConvTranspose2d(nfeats * 4, nfeats * 2, 4, 2, 1, bias=False))\n        #self.bn4 = nn.BatchNorm2d(nfeats * 2)\n        # state size. (nfeats * 2) x 32 x 32\n        \n        self.conv5 = spectral_norm(nn.ConvTranspose2d(nfeats * 2, nfeats, 4, 2, 1, bias=False))\n        #self.bn5 = nn.BatchNorm2d(nfeats)\n        # state size. (nfeats) x 64 x 64\n        \n        self.conv6 = spectral_norm(nn.ConvTranspose2d(nfeats, nchannels, 3, 1, 1, bias=False))\n        # state size. (nchannels) x 64 x 64\n        self.pixnorm = PixelwiseNorm()\n    def forward(self, x):\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 = F.leaky_relu(self.conv1(x))\n        x = F.leaky_relu(self.conv2(x))\n        x = self.pixnorm(x)\n        x = F.leaky_relu(self.conv3(x))\n        x = self.pixnorm(x)\n        x = F.leaky_relu(self.conv4(x))\n        x = self.pixnorm(x)\n        x = F.leaky_relu(self.conv5(x))\n        x = self.pixnorm(x)\n        x = torch.tanh(self.conv6(x))\n        \n        return x\n    def denorm(self, z):\n        #for layer in self.layers.values():  \n        z = self.forward(z)   \n        z = z * 0.5 + 0.5\n        return z\n\n\n\nclass Discriminator(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 = nn.Conv2d(nchannels, nfeats, 4, 2, 1, bias=False)\n        # state size. (nfeats) x 32 x 32\n        \n        self.conv2 = spectral_norm(nn.Conv2d(nfeats, nfeats * 2, 4, 2, 1, bias=False))\n        self.bn2 = nn.BatchNorm2d(nfeats * 2)\n        # state size. (nfeats*2) x 16 x 16\n        \n        self.conv3 = spectral_norm(nn.Conv2d(nfeats * 2, nfeats * 4, 4, 2, 1, bias=False))\n        self.bn3 = nn.BatchNorm2d(nfeats * 4)\n        # state size. (nfeats*4) x 8 x 8\n       \n        self.conv4 = spectral_norm(nn.Conv2d(nfeats * 4, nfeats * 8, 4, 2, 1, bias=False))\n        self.bn4 = nn.MaxPool2d(2)\n        # state size. (nfeats*8) x 4 x 4\n        self.batch_discriminator = MinibatchStdDev()\n        self.pixnorm = PixelwiseNorm()\n        self.conv5 = spectral_norm(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 = F.leaky_relu(self.conv1(x), 0.2)\n        x = F.leaky_relu(self.bn2(self.conv2(x)), 0.2)\n       # x = self.pixnorm(x)\n        x = F.leaky_relu(self.bn3(self.conv3(x)), 0.2)\n       # x = self.pixnorm(x)\n        x = F.leaky_relu(self.bn4(self.conv4(x)), 0.2)\n       # x = self.pixnorm(x)\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":{"trusted":true},"cell_type":"code","source":"class DogDataset(Dataset):\n    \"\"\"Dog Dataset.\"\"\"\n    \n    def __init__(self, data_path, transform=None):\n        \"\"\"\n        Args:\n            data_path (string): data path(glob_pattern) for dataset images\n            transform (callable, optional): Optional transform to be applied on a sample.\n        \"\"\"\n        self.data_list = sorted(glob.glob(data_path, recursive=True))\n        self.transform = transform\n        #self.data_name_list = sorted(glob.glob(data_path, recursive=True))\n        #self.transform = transform\n        #self.data_list = [transforms.ToTensor()(Image.open(img_name)) for img_name in self.data_name_list]\n        \n    def __len__(self):\n        return len(self.data_list)\n        #return len(self.data_name_list)\n    \n    def __getitem__(self, idx):\n        image = Image.open(self.data_list[idx])\n        #image = transforms.ToPILImage()(self.data_list[idx])\n        \n        if self.transform:\n            image = self.transform(image)\n        return image","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"dataset = DogDataset(data_path='../class_images/train/*.png',\n                              transform=transforms.Compose([\n                              transforms.RandomResizedCrop(64, scale=(0.7, 1.0), ratio=(1., 1.)),\n                              transforms.RandomHorizontalFlip(),\n                              #transforms.ColorJitter(brightness=0.05, contrast=0.05, saturation=0.05, hue=0.05),\n                              transforms.ToTensor(),\n                              transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),\n                          ]))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"dataloader = torch.utils.data.DataLoader(dataset, batch_size=batch_size,\n                                         shuffle=True, num_workers=int(workers))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"\nx = next(iter(dataloader))\n\nfig = plt.figure(figsize=(25, 16))\nfor ii, img in enumerate(x[:32]):\n    ax = fig.add_subplot(4, 8, ii + 1, xticks=[], yticks=[])\n    \n    img = img.numpy().transpose(1, 2, 0)\n    plt.imshow(img)\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\nprint('device:', device)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"#netG = Generator(nz=nz, nch_g=nch_g).to(device)\nnetG = Generator(nz=nz, nfeats=nch_g, nchannels=3).to(device)\n#netG.apply(weights_init)    \nprint(netG)\n    \n#netD = Discriminator(nch_d=nch_d).to(device)\nnetD = Discriminator(nchannels=3, nfeats=nch_d).to(device)\n#netD = nn.Sequential(netD, nn.Sigmoid()) # For DRAGAN\n#netD.apply(weights_init)\nprint(netD)\n    \ncriterion = nn.MSELoss()\n#criterion = nn.BCELoss() # For DRAGAN\n \noptimizerD = optim.Adam(netD.parameters(), lr=lr, betas=(beta1, 0.999), weight_decay=1e-5)  \noptimizerG = optim.Adam(netG.parameters(), lr=lr, betas=(beta1, 0.999), weight_decay=1e-5)  \n\n#optimizerD = optim.Adam(netD.parameters(), lr=lr, betas=(beta1, 0.999))  # For DRAGAN \n#optimizerG = optim.Adam(netG.parameters(), lr=lr, betas=(beta1, 0.999))  # For DRAGAN\n    \n    \nfixed_noise = torch.randn(batch_size, nz, 1, 1, device=device)  # for save_fake_image\nLoss_D_list, Loss_G_list = [], []\n    \nsave_fake_image_count = 1","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def save_fig(epoch):\n    if not os.path.exists('../output_images_{}'.format(epoch)):\n        os.mkdir('../output_images_{}'.format(epoch))\n    im_batch_size = 50\n    n_images=10000\n\n    for 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.denorm(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            save_image(gen_images[i_image, :, :, :], os.path.join('../output_images_{}'.format(epoch), f'image_{i_batch+i_image:05d}.png'))\n\n\n    import shutil\n    shutil.make_archive('images_{}'.format(epoch), 'zip', '../output_images_{}'.format(epoch))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def cosine_annealing(optimizer, start_lr, cur_steps, num_cycle):\n    t_cur = cur_steps % num_cycle\n    lr = 0.5 * start_lr * (math.cos(math.pi * t_cur / num_cycle) + 1)\n    for param_group in optimizer.param_groups:\n        param_group[\"lr\"] = lr\n    return lr","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# For DRAGAN\ndef compute_gradient_penalty(D, X):\n    \"\"\"Calculates the gradient penalty loss for DRAGAN\"\"\"\n    # Loss weight for gradient penalty\n    lambda_gp = 10\n    # Random weight term for interpolation\n    alpha = torch.Tensor(np.random.random(size=X.shape)).cuda()\n\n    interpolates = alpha * X + ((1 - alpha) * (X + 0.5 * X.std() * torch.rand(X.size()).cuda()))\n    interpolates = Variable(interpolates, requires_grad=True).cuda()\n\n    d_interpolates = D(interpolates)\n\n    fake = Variable(torch.Tensor(X.shape[0], ).fill_(1.0), requires_grad=False).cuda()\n\n    # Get gradient w.r.t. interpolates\n    gradients = torch.autograd.grad(\n        outputs=d_interpolates,\n        inputs=interpolates,\n        grad_outputs=fake,\n        create_graph=True,\n        retain_graph=True,\n        only_inputs=True,\n    )[0]\n\n    gradient_penalty = lambda_gp * ((gradients.norm(2, dim=1) - 1) ** 2).mean()\n    return gradient_penalty","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"rap_time = start_time\nfor epoch in range(n_epoch):\n    #random.seed(0)\n    #np.random.seed(0)\n    #torch.manual_seed(0)\n    for itr, data in tqdm(enumerate(dataloader), total = len(dataloader)):\n        real_image = data.to(device)   # Real Images\n        sample_size = real_image.size(0)  # The number of images\n        noise = torch.randn(sample_size, nz, 1, 1, device=device)   # generate input noise \n\n        real_target = torch.full((sample_size,), 1., device=device)   # real target\n        #real_target = torch.full((sample_size,), 0.7, device=device) + torch.rand(sample_size, device=device)/2   # real target\n        fake_target = torch.full((sample_size,), 0., device=device)   # fake target\n        #fake_target = torch.full((sample_size,), 0., device=device) + torch.rand(sample_size, device=device)/3.3  # fake target\n        #--------  Update Discriminator  ---------\n        netD.zero_grad()    # initialize gradient\n \n        output = netD(real_image)   # Discriminator output for real image\n        errD_real = criterion(output, real_target)  # MSELoss\n        D_x = output.mean().item()  # for logging\n \n        fake_image = netG(noise)    # fake images\n        \n        output = netD(fake_image.detach())  # Discriminator output for fake image\n        errD_fake = criterion(output, fake_target)  # MSELoss\n        D_G_z1 = output.mean().item()  # for logging\n        \n        #gradient_penalty = compute_gradient_penalty(netD, real_image.data) # For DRAGAN\n \n        errD = errD_real + errD_fake    # Discriminator Loss\n        #errD = errD + gradient_penalty   # For DRAGAN\n        \n        errD.backward()    # backward\n        optimizerD.step()   # Updata Discriminator params\n \n        #---------  Update Generator   ----------\n        netG.zero_grad()    # initialize gradient      \n        output = netD(fake_image)   # Discriminator output for fake image\n        errG = criterion(output, real_target)   # MSELoss\n        errG.backward()     # backward\n        D_G_z2 = output.mean().item()  # for logging\n \n        optimizerG.step()   # Updata Generator params\n        writer.add_scalar(\"errD\", errD.item(), epoch*len(dataloader)+itr)\n        writer.add_scalar(\"errG\", errG.item(), epoch*len(dataloader)+itr)\n        writer.add_scalar(\"D_G_z1\", D_G_z1, epoch*len(dataloader)+itr)\n        writer.add_scalar(\"D_G_z2\", D_G_z2, epoch*len(dataloader)+itr)\n    writer.add_images('fake_images', fake_image[:10] * 0.5 + 0.5, global_step=epoch*len(dataloader)+itr, walltime=None, dataformats='NCHW')\n\n    cosine_annealing(optimizerG, 0.0008, time.time() - start_time, 60*60*10)\n    cosine_annealing(optimizerD, 0.0002, time.time() - start_time, 60*60*10)\n\n    if time.time() - start_time > 60*60*8.95:\n        break","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"netG.eval()\n#netG.train()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"from scipy.stats import truncnorm\n\ndef truncated_normal_(tensor, mean=0, std=1):\n    size = tensor.shape\n    tmp = tensor.new_empty(size + (4,)).normal_()\n    valid = (tmp < 2) & (tmp > -2)\n    ind = valid.max(-1, keepdim=True)[1]\n    tensor.data.copy_(tmp.gather(-1, ind).squeeze(-1))\n    tensor.data.mul_(std).add_(mean)\n    return tensor","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 torch.from_numpy(values).float()","execution_count":null,"outputs":[]},{"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\n\nfor i_batch in range(0, n_images, im_batch_size):\n    gen_z = truncated_normal((im_batch_size, nz, 1, 1), threshold=0.5).to(device)\n    gen_images = netG.denorm(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        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":{"trusted":true},"cell_type":"code","source":"samples = truncated_normal((32, nz, 1, 1), threshold=0.5).to(device)\nsamples = netG.denorm(samples).detach().cpu().numpy().transpose(0, 2, 3, 1)\n\nfig = plt.figure(figsize=(25, 16))\nfor ii, img in enumerate(samples):\n    ax = fig.add_subplot(4, 8, ii + 1, xticks=[], yticks=[])\n    plt.imshow(img)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"torch.save(netG.state_dict(), 'netG.pth')\ntorch.save(netD.state_dict(), 'netD.pth')","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}