{"cells":[{"metadata":{},"cell_type":"markdown","source":"### reference: "},{"metadata":{"trusted":true},"cell_type":"code","source":"from functools import partial\nfrom multiprocessing import Pool\nimport os\nfrom pathlib import Path\nimport random\nimport shutil\nimport time\nimport xml.etree.ElementTree as ET\n\nfrom matplotlib import pyplot as plt\nimport numpy as np\nfrom PIL import Image\nfrom scipy.stats import truncnorm\nimport torch\nfrom torch import nn, optim\nfrom torch.optim.optimizer import Optimizer\nimport torch.nn.functional as F\nfrom torch.utils.data import Dataset, DataLoader\nfrom torchvision import datasets, transforms\nfrom torchvision.datasets.folder import default_loader\nfrom torchvision.utils import save_image\nfrom torch.nn.utils import spectral_norm\nfrom tqdm import tqdm\nfrom torch.autograd import Variable\nimport sys","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"start_time = time.time()\n\nbatch_size = 32\nepochs = 210\nseed = 1029\n\nTRAIN_DIR = Path('../input/all-dogs/')\nANNOTATION_DIR = Path('../input/annotation/Annotation/')\nDOG_DIR = Path('../dogs/dogs/')\nOUT_DIR = Path('../output_images/')\nDOG_DIR.mkdir(parents=True, exist_ok=True)\nOUT_DIR.mkdir(exist_ok=True)\n\ndevice = torch.device('cuda')\n\nlr = 0.0005\nbeta1 = 0.5\nnz = 256\n\nreal_label = 0.95\nfake_label = 0\n\nrandom.seed(seed)\nos.environ['PYTHONHASHSEED'] = str(seed)\nnp.random.seed(seed)\ntorch.manual_seed(seed)\ntorch.cuda.manual_seed(seed)\ntorch.backends.cudnn.deterministic = True","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"mylist = os.listdir(ANNOTATION_DIR)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"label_map = {}\nfor listname in mylist:\n    label_map[listname.split('-')[0]]=listname.split('-')[1]","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"label_number={}\nfor index, label in enumerate(label_map):\n    label_number[label]=index","execution_count":null,"outputs":[]},{"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        \"\"\"\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","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class ConditionalBatchNorm2d(nn.BatchNorm2d):\n\n    \"\"\"Conditional Batch Normalization\"\"\"\n\n    def __init__(self, num_features, eps=1e-05, momentum=0.1,\n                 affine=False, track_running_stats=True):\n        super(ConditionalBatchNorm2d, self).__init__(\n            num_features, eps, momentum, affine, track_running_stats\n        )\n\n    def forward(self, input, weight, bias, **kwargs):\n        self._check_input_dim(input)\n\n        exponential_average_factor = 0.0\n\n        if self.training and self.track_running_stats:\n            self.num_batches_tracked += 1\n            if self.momentum is None:  # use cumulative moving average\n                exponential_average_factor = 1.0 / self.num_batches_tracked.item()\n            else:  # use exponential moving average\n                exponential_average_factor = self.momentum\n\n        output = F.batch_norm(input, self.running_mean, self.running_var,\n                              self.weight, self.bias,\n                              self.training or not self.track_running_stats,\n                              exponential_average_factor, self.eps)\n        if weight.dim() == 1:\n            weight = weight.unsqueeze(0)\n        if bias.dim() == 1:\n            bias = bias.unsqueeze(0)\n        size = output.size()\n        weight = weight.unsqueeze(-1).unsqueeze(-1).expand(size)\n        bias = bias.unsqueeze(-1).unsqueeze(-1).expand(size)\n        return weight * output + bias\n    \nclass CategoricalConditionalBatchNorm2d(ConditionalBatchNorm2d):\n\n    def __init__(self, num_classes, num_features, eps=1e-5, momentum=0.1,\n                 affine=False, track_running_stats=True):\n        super(CategoricalConditionalBatchNorm2d, self).__init__(\n            num_features, eps, momentum, affine, track_running_stats\n        )\n        self.weights = nn.Embedding(num_classes, num_features)\n        self.biases = nn.Embedding(num_classes, num_features)\n\n        self._initialize()\n\n    def _initialize(self):\n        torch.nn.init.ones_(self.weights.weight.data)\n        torch.nn.init.zeros_(self.biases.weight.data)\n\n    def forward(self, input, c, **kwargs):\n        weight = self.weights(c)\n        bias = self.biases(c)\n\n        return super(CategoricalConditionalBatchNorm2d, self).forward(input, weight, bias)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class ResBlockGenerator(nn.Module):\n\n    def __init__(self, in_channels, out_channels, stride=1):\n        super(ResBlockGenerator, self).__init__()\n\n        self.conv1 = nn.Conv2d(in_channels, out_channels, 3, 1, padding=1)\n        self.conv2 = nn.Conv2d(out_channels, out_channels, 3, 1, padding=1)\n        #nn.init.xavier_uniform(self.conv1.weight.data, 1.)\n        #nn.init.xavier_uniform(self.conv2.weight.data, 1.)\n\n        self.model = nn.Sequential(\n            nn.BatchNorm2d(in_channels),\n            nn.ReLU(),\n            nn.Upsample(scale_factor=2),\n            self.conv1,\n            nn.BatchNorm2d(out_channels),\n            nn.ReLU(),\n            self.conv2\n            )\n        self.bypass = nn.Sequential()\n        if stride != 1:\n            if in_channels != out_channels:\n                self.bypass = nn.Sequential(\n                    nn.Upsample(scale_factor=2),\n                    nn.Conv2d(in_channels,out_channels, 1, 1, padding=0)\n                )\n            else:\n                self.bypass = nn.Upsample(scale_factor=2)\n\n    def forward(self, x):\n        #print (x.shape)\n        #print (self.model(x).shape)\n        #print (self.bypass(x).shape)\n        return self.model(x) + self.bypass(x)\n\n\nclass ResBlockDiscriminator(nn.Module):\n\n    def __init__(self, in_channels, out_channels, stride=1):\n        super(ResBlockDiscriminator, self).__init__()\n\n        self.conv1 = nn.Conv2d(in_channels, out_channels, 3, 1, padding=1)\n        self.conv2 = nn.Conv2d(out_channels, out_channels, 3, 1, padding=1)\n        nn.init.xavier_uniform(self.conv1.weight.data, 1.)\n        nn.init.xavier_uniform(self.conv2.weight.data, 1.)\n\n        if stride == 1:\n            self.model = nn.Sequential(\n                nn.ReLU(),\n                spectral_norm(self.conv1),\n                nn.ReLU(),\n                spectral_norm(self.conv2)\n                )\n        else:\n            self.model = nn.Sequential(\n                nn.ReLU(),\n                spectral_norm(self.conv1),\n                nn.ReLU(),\n                spectral_norm(self.conv2),\n                nn.AvgPool2d(2, stride=stride, padding=0)\n                )\n        self.bypass = nn.Sequential()\n        if stride != 1:\n\n            self.bypass_conv = nn.Conv2d(in_channels,out_channels, 1, 1, padding=0)\n            nn.init.xavier_uniform(self.bypass_conv.weight.data, np.sqrt(2))\n\n            self.bypass = nn.Sequential(\n                spectral_norm(self.bypass_conv),\n                nn.AvgPool2d(2, stride=stride, padding=0)\n            )\n            \n            #if in_channels == out_channels:\n            #    self.bypass = nn.AvgPool2d(2, stride=stride, padding=0)\n            #else:\n            #    self.bypass = nn.Sequential(\n            #        spectral_norm(nn.Conv2d(in_channels,out_channels, 1, 1, padding=0)),\n            #        nn.AvgPool2d(2, stride=stride, padding=0)\n            #     )\n\n\n    def forward(self, x):\n        return self.model(x) + self.bypass(x)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class Generator(nn.Module):\n    \n    def __init__(self, nz=128, channels=3):\n        super(Generator, self).__init__()\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                spectral_norm(nn.ConvTranspose2d(n_input, n_output, kernel_size=k_size, stride=stride, padding=padding, bias=False)),\n                nn.ReLU(inplace=True),\n            ]\n            return block\n\n        self.layer1 = nn.Sequential(\n            *convlayer(self.nz, 64, 4, 1, 0),\n            #*convlayer(512, 64, 4, 2, 1)\n        )\n        \n        #self.CBN1 = CategoricalConditionalBatchNorm2d(120,512)\n        #self.layer2 = nn.Sequential(*convlayer(512, 256, 4, 2, 1))\n        self.layer2 = nn.Sequential(ResBlockGenerator(64, 64, 2))\n        self.CBN2 = CategoricalConditionalBatchNorm2d(120,64)\n        self.layer3 = nn.Sequential(ResBlockGenerator(64, 64, 2))\n        self.CBN3 = CategoricalConditionalBatchNorm2d(120,64)\n        self.layer4 = nn.Sequential(ResBlockGenerator(64, 64, 2))\n        self.CBN4 = CategoricalConditionalBatchNorm2d(120,64)\n        self.layer5 = nn.Sequential(ResBlockGenerator(64, 64, 2))\n        self.CBN5 = CategoricalConditionalBatchNorm2d(120,64)\n        self.outlay = spectral_norm(nn.Conv2d(64, 3, 3, 1, 1))\n        self.Tanh = nn.Tanh()\n        \n\n    def forward(self, z, y):\n        #label = label.view(-1, 120,1,1).type(torch.cuda.FloatTensor)\n        #label = self.label_conv1(label)\n        z = z.view(-1, self.nz, 1, 1)\n        #z = self.first_layer(z)\n        z = self.layer1(z)\n        #z = self.CBN1(z,y)\n        z = self.layer2(z)\n        z = self.CBN2(z,y)\n        z = self.layer3(z)\n        z = self.CBN3(z,y)\n        z = self.layer4(z)\n        z = self.CBN4(z,y)\n        z = self.layer5(z)\n        z = self.CBN5(z,y)\n        z = self.outlay(z)\n        img = self.Tanh(z)\n        #print (img.shape)\n        return img\n\n\nclass Discriminator(nn.Module):\n    \n    def __init__(self, channels=3):\n        super(Discriminator, self).__init__()\n        self.channels = channels\n\n        def convlayer(n_input, n_output, k_size=4, stride=2, padding=0, bn=False):\n            block =[]\n            if bn:\n                block.append(spectral_norm(nn.Conv2d(n_input, n_output, kernel_size=k_size, stride=stride, padding=padding, bias=False)))\n                block.append(nn.BatchNorm2d(n_output))\n            else:\n                block.append(nn.Conv2d(n_input, n_output, kernel_size=k_size, stride=stride, padding=padding, bias=False))\n            #if bn:\n            #    block.append(spectral_norm())\n            block.append(nn.LeakyReLU(0.2, inplace=True))\n            return block\n        \n        #self.label_conv2 = nn.Conv2d(120, 128, 4, 2, 1, bias=False)\n        #self.first_layer2 = nn.Sequential(*convlayer(self.channels, 128, 4, 2, 1))\n\n        self.model = nn.Sequential(\n            *convlayer(self.channels, 64, 4, 2, 1),\n            ResBlockDiscriminator(64,128,2),\n            ResBlockDiscriminator(128,256,2),\n            ResBlockDiscriminator(256,728,2),\n            ResBlockDiscriminator(728,728,1),\n            \n            #*convlayer(128, 256, 4, 2, 1),\n            #*convlayer(256, 512, 4, 2, 1, bn=True),\n            #*convlayer(512, 1024, 4, 2, 1, bn=True),\n        )\n        self.logic_out = nn.Conv2d(728, 1, 4, 1, 0, bias=False)\n        self.aux_out = nn.Conv2d(728, 120, 4, 1, 0, bias=False)\n\n    def forward(self, imgs, label):\n        #label = self.label_conv2(label)\n        #imgs = self.first_layer2(imgs)\n        out = self.model(imgs)\n        logic_output = self.logic_out(out)\n        #print (logic_output.shape)\n        aux_output = self.aux_out(out)\n        #print (aux_output.shape)\n        return logic_output.view(-1, 1),aux_output.view(-1,120)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class DogsDataset(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        dog = Path(path).stem\n        label = label_number[dog.split('_')[0]]\n        #label = np.eye(120)[label]\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,label])\n        return imgs\n    \n    def label_dog(self, path):\n        label = []\n        dog = Path(path).stem\n        label.append(label_number[dog.split('_')[0]])\n        return label\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            #labels = p.map(self.label_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(self.imgs[index][0][0])\n        lables = self.imgs[index][0][1]\n        return samples,lables\n    \n    def __len__(self):\n        return len(self.imgs)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class ParamScheduler(object):\n    \n    def __init__(self, optimizer, scale_fn, step_size):\n        if not isinstance(optimizer, Optimizer):\n            raise TypeError('{} is not an Optimizer'.format(\n                type(optimizer).__name__))\n        \n        self.optimizer = optimizer\n        self.scale_fn = scale_fn\n        self.step_size = step_size\n        self.last_batch_iteration = 0\n        \n    def batch_step(self):\n        for param_group in self.optimizer.param_groups:\n            param_group['lr'] = self.scale_fn(self.last_batch_iteration / self.step_size)\n        \n        self.last_batch_iteration += 1\n\n\ndef combine_scale_functions(scale_fns, phases=None):\n    if phases is None:\n        phases = [1. / len(scale_fns)] * len(scale_fns)\n    phases = [phase / sum(phases) for phase in phases]\n    phases = torch.tensor([0] + phases)\n    phases = torch.cumsum(phases, 0)\n    \n    def _inner(x):\n        idx = (x >= phases).nonzero().max()\n        actual_x = (x - phases[idx]) / (phases[idx + 1] - phases[idx])\n        return scale_fns[idx](actual_x)\n        \n    return _inner\n\n\ndef scale_cos(start, end, x):\n    return start + (1 + np.cos(np.pi * (1 - x))) * (end - start) / 2","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"random_transforms = [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))])\n\ntrain_data = DogsDataset(TRAIN_DIR / 'all-dogs/', ANNOTATION_DIR, transform=transform)\ntrain_loader = DataLoader(train_data, shuffle=True, batch_size=batch_size, num_workers=4)\n\n#imgs = next(iter(train_loader))\n#imgs = imgs.numpy().transpose(0, 2, 3, 1)\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"imgs,labels = next(iter(train_loader))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"net_g = Generator(nz).to(device)\nnet_d = Discriminator().to(device)\n\ncriterion = nn.BCELoss()\nauxiliary_loss = torch.nn.CrossEntropyLoss()\nscale_fn = combine_scale_functions(\n    [partial(scale_cos, 1e-5, 5e-4), partial(scale_cos, 5e-4, 1e-4)], [0.2, 0.8])\n\noptimizer_g = optim.Adam(net_g.parameters(), lr=lr, betas=(beta1, 0.999))\noptimizer_d = optim.Adam(net_d.parameters(), lr=lr, betas=(beta1, 0.999))\n\n#step_size = 2000\n#base_lr, max_lr = lr, lr*10\nscheduler_g = ParamScheduler(optimizer_g, scale_fn, epochs * len(train_loader))\nscheduler_d = ParamScheduler(optimizer_d, scale_fn, epochs * len(train_loader))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"for epoch in range(epochs):\n    for i, (real_images,label) in enumerate(train_loader):\n        # --------------------------------------\n        # Update Discriminator network: maximize log(D(x)) + log(1 - D(G(z)))\n        # --------------------------------------\n        net_d.zero_grad()\n        real_images = real_images.to(device)\n        label = label.type(torch.LongTensor).to(device)\n        #label_g = torch.eye(120)[label.tolist()].to(device)\n        batch_size = real_images.size(0)\n        labels = torch.full((batch_size, 1), real_label, device=device)\n        \n        #label_d = fill[label.tolist()].to(device)\n        \n        scheduler_d.batch_step()\n        output_real,aux_real = net_d(real_images,label)\n        \n        noise = torch.randn(batch_size, nz, 1, 1, device=device)\n        #y_noise = (torch.rand(batch_size, 1) * 120).type(torch.LongTensor).squeeze()\n        #label_g_fake = torch.eye(120)[y_noise.tolist()].to(device)\n        #label_d_fake = fill[y_noise.tolist()].to(device)\n        \n        fake = net_g(noise,label)\n        \n        output_fake,aux_fake = net_d(fake.detach(),label)\n        \n        loss_aux_real = auxiliary_loss(aux_real,label)\n        loss_aux_fake = auxiliary_loss(aux_fake,label)\n        \n        err_d = (torch.mean((output_real - torch.mean(output_fake) - labels) ** 2) + \n                 torch.mean((output_fake - torch.mean(output_real) + labels) ** 2)) / 2\n        \n        err_d = err_d*0.7 + loss_aux_real*0.15 + loss_aux_fake*0.15\n        err_d.backward(retain_graph=True)\n        optimizer_d.step()\n        \n        # --------------------------------------\n        # Update Generator network: maximize log(D(G(z)))\n        # --------------------------------------\n        net_g.zero_grad()\n        scheduler_g.batch_step()\n        output_fake, aux_fake = net_d(fake,label)\n        \n        loss_aux_fake_g = auxiliary_loss(aux_fake,label)\n        err_g = (torch.mean((output_real - torch.mean(output_fake) + labels) ** 2) +\n                 torch.mean((output_fake - torch.mean(output_real) - labels) ** 2)) / 2\n        err_g = err_g*0.85 + loss_aux_fake_g*0.15\n        err_g.backward()\n        optimizer_g.step()\n        \n    print(f'[{epoch + 1}/{epochs}] Loss_d: {err_d.item():.4f} Loss_g: {err_g.item():.4f}')","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","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"im_batch_size = 50\nn_images = 10000\n\nfor i_batch in range(0, n_images, im_batch_size):\n    z = truncated_normal((im_batch_size, nz, 1, 1), threshold=1)\n    gen_z = torch.from_numpy(z).float().to(device)\n    gen_labels = Variable(torch.LongTensor(np.random.randint(0, 120, im_batch_size))).to(device)\n    #y_noise = (torch.rand(im_batch_size, 1) * 120).type(torch.LongTensor).squeeze()\n    #label_g_fake = torch.eye(120)[y_noise.tolist()].to(device)\n    gen_images = (net_g(gen_z,gen_labels) + 1) / 2\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, :, :, :], OUT_DIR / f'image_{i_batch + i_image:05d}.png')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"fig = plt.figure(figsize=(25, 16))\nfor i, j in enumerate(images[:32]):\n    ax = fig.add_subplot(4, 8, i + 1, xticks=[], yticks=[])\n    plt.imshow(j)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"shutil.make_archive('images', 'zip', OUT_DIR)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"elapsed_time = time.time() - start_time\nprint(f'All process done in {int(elapsed_time // 3600)} hours {int(elapsed_time % 3600 // 60)} min.')","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}