{"cells":[{"metadata":{},"cell_type":"markdown","source":"## SAGAN with residual connections\n\ncode is based on https://github.com/brain-research/self-attention-gan"},{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"import numpy as np # linear algebra\nimport pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)\nimport xml.etree.ElementTree as ET \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/\"))\nfrom pathlib import Path\nimport enum\nimport functools\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\nimport matplotlib.pyplot as plt\n\nImage.open(\"../input/all-dogs/all-dogs/n02110627_13662.jpg\")","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"from typing import List, Tuple\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.init as init\nimport torch.nn.functional as F\nfrom torch.nn.utils import spectral_norm\nimport torch.utils.data\n\nimport torchvision\nfrom torchvision import transforms, datasets\nfrom torchvision.utils import save_image\n\nfrom tqdm import tqdm_notebook as tqdm","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"EPOCHS = 270\nBATCH_SIZE = 64\nNORMALIZE_Z = True\n\nZ_SIZE = 128\nN_FEATURES = 32\n\n\nCONDITIONAL = True\n\nSELF_ATTENTION = True\n    \nSPECTRAL_IN_G = True\nSPECTRAL_IN_D = True\nBN_IN_D = False\nRESIDUAL_IN_D = True\n\nGENERATOR_ACTIVATION = F.relu  # functools.partial(F.leaky_relu, negative_slope=0.02)\n\n## Discriminator\nDISCRIMINATOR_ACTIVATION = F.relu  # functools.partial(F.leaky_relu, negative_slope=0.02)\n\n\nclass LossType(enum.Enum):\n    LS = 1\n    RaLS = 2\n    Hinge = 3\n    RaHinge = 4\n    \n    def requires_real_for_generator(self) -> bool:\n        return self == LossType.RaLS\n\n\nLOSS = LossType.Hinge\nBETA1 = 0.0\nBETA2 = 0.999\nLR_GENERATOR = 0.0001\nLR_DISCRIMINATOR = 0.0004\n\nCROP_PROB = 1.0\n\n\n## Debugging\nSHOW_FREQ = None\nSHOW_FREQ_EPOCH = EPOCHS // 10","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def sample(size):\n    z = np.random.normal(loc=0, scale=1, size=(size, Z_SIZE)).astype(np.float32)\n    if NORMALIZE_Z:\n        z /= np.sqrt(np.sum(z ** 2, axis=1, keepdims=True))\n    return torch.as_tensor(z)\n\ntorch.manual_seed(42)\nFIXED_NOISES = sample(16)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def extract_class(path: str) -> str:\n    name = Path(path).name\n    idx = name.index(\"_\")\n    return name[:idx]\n\n\ndef gather_image_and_annots(annot_path: str = \"../input/annotation/Annotation/\", image_path: str = \"../input/all-dogs/all-dogs/\") -> List[Tuple[str, str]]:\n    results: List[Tuple[str, str]] = []\n    breeds = os.listdir(annot_path)\n    for breed in breeds:\n        for dog in os.listdir(os.path.join(annot_path, breed)):\n            annot = os.path.join(annot_path, breed, dog)\n            image = os.path.join(image_path, dog + \".jpg\")\n            if os.path.exists(image):\n                results.append((image, annot))\n    return results\n        \n\ndef get_classes(path: str):\n    classes = set()\n    for filepath in Path(path).glob(\"*.jpg\"):\n        classes.add(extract_class(filepath))\n    return sorted(classes)\n\n\ndef filter_small_boxes(boxes: List[List[int]]) -> List[List[int]]:\n    return [x for x in boxes if min(x[2] - x[0], x[3] - x[1]) >= 64]\n\n\ndef get_bounding_boxes(path: str, filter_small: bool = True) -> List[List[int]]:\n    tree = ET.parse(path)\n    root = tree.getroot()\n    objects = root.findall('object')\n    boxes: List[List[int]] = []\n    for o in objects:\n        bndbox = o.find('bndbox') # reading bound box\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        boxes.append([xmin, ymin, xmax, ymax])\n    if filter_small:\n        boxes = filter_small_boxes(boxes)\n    return boxes\n\nBoxes = List[List[int]]\n\ndef append_bounding_boxes(pairs: List[Tuple[str, str]], q: float = 0.99) -> List[Tuple[str, Boxes]]:\n    n = len(pairs)\n    sizes = []\n    for img, _ in tqdm(pairs):\n        sizes.append(os.stat(img).st_size)\n    pairs = np.array(pairs)\n    sizes = np.array(sizes)\n    indices = sizes < np.quantile(sizes, q)\n    pairs = pairs[indices]\n    results = []\n    for i, (img_path, annot) in enumerate(tqdm(pairs)):\n        img = Image.open(img_path)\n        if min(*img.size) < 64:\n            continue\n        boxes = get_bounding_boxes(annot, filter_small=True)\n        if not boxes:\n            continue\n        results.append((img_path, boxes))\n    print(\"Valid sized images: {} / {}\".format(len(results), n))\n    return results\n\nCLASSES = get_classes(\"../input/all-dogs/all-dogs/\")\nPAIRS = append_bounding_boxes(gather_image_and_annots())\nIMAGES, BOXES = [p[0] for p in PAIRS], [p[1] for p in PAIRS]\nlen(CLASSES), len(PAIRS)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class Dataset(torch.utils.data.Dataset):\n    def __init__(self, image_files, boxes, transform, crop_prob: float = 0.5, jitter_ratio: float = 0.05):\n        super().__init__()\n        self.transform = transform\n        self.image_files = image_files\n        self.boxes = boxes\n        self.crop_prob = crop_prob\n        self.jitter_ratio = jitter_ratio\n        \n    def __len__(self) -> int:\n        return len(self.image_files)\n    \n    def __getitem__(self, idx):\n        filepath = self.image_files[idx]\n        image = Image.open(str(filepath)).convert(\"RGB\")\n        if self.crop_prob >= 1.0 or np.random.random() < self.crop_prob:\n            boxes = self.boxes[idx]\n            image = self._crop_bounding_boxes(image, boxes)\n        if CONDITIONAL:\n            class_id = CLASSES.index(extract_class(filepath))\n        else:\n            class_id = 0  # dummy\n        image_tensor = self.transform(image)\n        return image_tensor, torch.as_tensor(class_id)\n    \n    def _crop_bounding_boxes(self, image: Image.Image, boxes: List[List[int]]) -> Image.Image:\n        i = np.random.randint(0, len(boxes))\n        xmin, ymin, xmax, ymax = boxes[i]\n        if min(xmax - xmin, ymax - ymin) < 64:\n            return image\n        x_center = (xmin + xmax) // 2\n        y_center = (ymin + ymax) // 2\n        size = max(xmax - xmin, ymax - ymin)\n        if int(size * self.jitter_ratio):\n            x_jitter = np.random.randint(int(-size * self.jitter_ratio), int(size * self.jitter_ratio))\n            x_center += x_jitter\n            y_jitter = np.random.randint(int(-size * self.jitter_ratio), int(size * self.jitter_ratio))\n            y_center += y_jitter\n        size_jitter = np.random.randint(0, int(size * self.jitter_ratio))\n        xmin = max(0, x_center - (size + size_jitter) // 2)\n        ymin = max(0, y_center - (size + size_jitter) // 2)\n        xmax = min(image.size[0], x_center + (size + size_jitter) // 2)\n        ymax = min(image.size[1], y_center + (size + size_jitter) // 2)\n        return image.crop((xmin, ymin, xmax, ymax))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"transform = transforms.Compose([\n    transforms.Resize(64),\n    transforms.RandomCrop(64),\n    transforms.RandomHorizontalFlip(p=0.5),\n    transforms.ToTensor(),\n    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),\n])\ndataset = Dataset(IMAGES, BOXES, transform, crop_prob=CROP_PROB)\ndataloader = torch.utils.data.DataLoader(dataset, shuffle=True, batch_size=BATCH_SIZE, num_workers=4)\nprint(len(dataset), len(dataloader))\n_ = next(iter(dataloader))","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### ops"},{"metadata":{"trusted":true},"cell_type":"code","source":"def init_linear(linear):\n    init.xavier_uniform_(linear.weight)\n    linear.bias.data.zero_()\n\n\ndef init_conv(conv, glu=True):\n    init.xavier_uniform_(conv.weight)\n    if conv.bias is not None:\n        conv.bias.data.zero_()\n\n\ndef spectral_init(module, use_norm: bool, gain=1):\n    init.kaiming_uniform_(module.weight, gain)\n    if module.bias is not None:\n        module.bias.data.zero_()\n    if use_norm:\n        return spectral_norm(module)\n    else:\n        return module\n\n\ndef leaky_relu(input):\n    return F.leaky_relu(input, negative_slope=0.2)\n\n\nclass SelfAttention(nn.Module):\n    def __init__(self, in_channel: int, use_norm: bool, gain=1):\n        super().__init__()\n\n        self.query = spectral_init(nn.Conv1d(in_channel, in_channel // 8, 1),\n                                   use_norm=use_norm,\n                                   gain=gain)\n        self.key = spectral_init(nn.Conv1d(in_channel, in_channel // 8, 1),\n                                 use_norm=use_norm,\n                                 gain=gain)\n        self.value = spectral_init(nn.Conv1d(in_channel, in_channel, 1),\n                                   use_norm=use_norm,\n                                   gain=gain)\n\n        self.gamma = nn.Parameter(torch.tensor(0.0))\n\n    def forward(self, input):\n        shape = input.shape\n        flatten = input.view(shape[0], shape[1], -1)\n        query = self.query(flatten).permute(0, 2, 1)\n        key = self.key(flatten)\n        value = self.value(flatten)\n        query_key = torch.bmm(query, key)\n        attn = F.softmax(query_key, 1)\n        attn = torch.bmm(value, attn)\n        attn = attn.view(*shape)\n        out = self.gamma * attn + input\n\n        return out\n\n\nclass ConditionalNorm(nn.Module):\n    def __init__(self, in_channel, n_class):\n        super().__init__()\n\n        self.bn = nn.BatchNorm2d(in_channel, affine=False)\n        self.embed = nn.Embedding(n_class, in_channel * 2)\n        self.embed.weight.data[:, :in_channel] = 1\n        self.embed.weight.data[:, in_channel:] = 0\n\n    def forward(self, input, class_id):\n        out = self.bn(input)\n        embed = self.embed(class_id)\n        gamma, beta = embed.chunk(2, 1)\n        gamma = gamma.unsqueeze(2).unsqueeze(3)\n        beta = beta.unsqueeze(2).unsqueeze(3)\n        out = gamma * out + beta\n\n        return out\n\ndef snconv2d(\n    in_features: int, \n    out_features: int, \n    kernel_size: int, \n    stride: int, \n    padding: int,\n    use_norm: bool,\n) -> nn.Module:\n    conv = nn.Conv2d(in_features, out_features, kernel_size, stride=stride, padding=padding, bias=True)\n    init.xavier_uniform_(conv.weight)\n    conv.bias.data.zero_()\n    if use_norm:\n        return spectral_norm(conv)\n    else:\n        return conv\n\n\ndef snlinear(in_features: int, out_features: int, use_norm: bool) -> nn.Module:\n    linear = nn.Linear(in_features, out_features, bias=True)\n    init.xavier_uniform_(linear.weight)\n    linear.bias.data.zero_()\n    if use_norm:\n        return spectral_norm(linear)\n    else:\n        return linear\n\n\ndef sn_embedding(num_classes: int, embedding_dim: int, use_norm: bool) -> nn.Module:\n    emb = nn.Embedding(num_classes, embedding_dim)\n    init.xavier_uniform_(emb.weight)\n    if use_norm:\n        return spectral_norm(emb)\n    else:\n        return emb","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Generator"},{"metadata":{"trusted":true},"cell_type":"code","source":"def upsample_conv(in_features: int, out_features: int, kernel_size: int, stride: int, padding: int) -> nn.Module:\n    return nn.Sequential(\n        nn.Upsample(scale_factor=2),\n        snconv2d(in_features, out_features, kernel_size=kernel_size, stride=stride, padding=padding, use_norm=SPECTRAL_IN_G),\n    )\n\n\nclass GeneratorBlock(nn.Module):\n    def __init__(self, in_features: int, out_features: int, num_classes: int) -> None:\n        super().__init__()\n        self.in_features = in_features\n        if CONDITIONAL:\n            self.bn0 = ConditionalNorm(in_features, num_classes)\n            self.bn1 = ConditionalNorm(out_features, num_classes)\n        else:\n            self.bn0 = nn.BatchNorm2d(in_features)\n            self.bn1 = nn.BatchNorm2d(out_features)\n        self.upconv = upsample_conv(in_features, out_features, kernel_size=3, stride=1, padding=1)\n        self.conv1 = snconv2d(out_features, out_features, kernel_size=3, stride=1, padding=1, use_norm=SPECTRAL_IN_G)\n        self.skip_upconv = upsample_conv(in_features, out_features, kernel_size=1, stride=1, padding=0)\n        self.activation = GENERATOR_ACTIVATION\n        \n    def forward(self, x, class_id):\n        x0 = x\n        if CONDITIONAL:\n            x = self.activation(self.bn0(x, class_id))\n        else:\n            x = self.activation(self.bn0(x))\n        x = self.upconv(x)\n        if CONDITIONAL:\n            x = self.activation(self.bn1(x, class_id))\n        else:\n            x = self.activation(self.bn1(x))\n        x = self.conv1(x)\n        \n        x0 = self.skip_upconv(x0)\n        return x + x0\n    \n\nclass Generator(nn.Module):\n    def __init__(self, z_size: int = Z_SIZE, num_classes: int = len(CLASSES), features: int = N_FEATURES) -> None:\n        super().__init__()\n        self.linear0 = snlinear(Z_SIZE, features * 4 * 4 * 8, use_norm=SPECTRAL_IN_G)\n        # 8C x 4 x 4\n        self.block0 = GeneratorBlock(features * 8, features * 8, num_classes=num_classes) \n        # 8C x 8 x 8\n        self.block1 = GeneratorBlock(features * 8, features * 4, num_classes=num_classes) \n        # 4C x 16 x 16\n        self.block2 = GeneratorBlock(features * 4, features * 2, num_classes=num_classes)\n        # 2C x 32 x 32\n        if SELF_ATTENTION:\n            self.attn = SelfAttention(features * 2, use_norm=SPECTRAL_IN_G)\n        # 2C x 32 x 32\n        self.block3 = GeneratorBlock(features * 2, features * 1, num_classes=num_classes)\n        # C x 64 x 64\n        self.bn = nn.BatchNorm2d(features * 1)\n        # 3 x 64 x 64\n        self.last_conv = snconv2d(features, 3, kernel_size=3, stride=1, padding=1, use_norm=SPECTRAL_IN_G)\n        self.activation = GENERATOR_ACTIVATION\n        \n    def forward(self, z, class_id):\n        x = self.linear0(z)\n        x = x.view(x.size(0), -1, 4, 4)\n        x = self.block0(x, class_id)\n        x = self.block1(x, class_id)\n        x = self.block2(x, class_id)\n        if SELF_ATTENTION:\n            x = self.attn(x)\n        x = self.block3(x, class_id)\n        x = self.activation(self.bn(x))\n        x = self.last_conv(x)\n        x = torch.tanh(x)\n        return x\n    \n    \ngenerator = Generator()\nprint(generator(torch.randn(1, Z_SIZE), torch.zeros(1, dtype=torch.long)).shape)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Discriminator"},{"metadata":{"trusted":true},"cell_type":"code","source":"def downsample(x):\n    return F.avg_pool2d(x, kernel_size=2, stride=2)\n\n\nclass DiscriminatorBlock(nn.Module):\n    def __init__(self, in_features: int, out_features: int, down=True, activation=DISCRIMINATOR_ACTIVATION):\n        super().__init__()\n        self.down = down\n        self.activation = activation\n        if BN_IN_D:\n            self.bn0 = nn.BatchNorm2d(in_features)\n            self.bn1 = nn.BatchNorm2d(out_features)\n        self.conv0 = snconv2d(in_features,  out_features, kernel_size=3, stride=1, padding=1, use_norm=SPECTRAL_IN_D)\n        self.conv1 = snconv2d(out_features, out_features, kernel_size=3, stride=1, padding=1, use_norm=SPECTRAL_IN_D)\n        if RESIDUAL_IN_D and (down or in_features != out_features):\n            self.skip_conv = snconv2d(in_features, out_features, kernel_size=1, stride=1, padding=0, use_norm=SPECTRAL_IN_D)\n        else:\n            self.skip_conv = None\n            \n    def forward(self, x):\n        x0 = x\n        if BN_IN_D:\n            x = self.bn0(x)\n        x = self.activation(x)\n        x = self.conv0(x)\n        if BN_IN_D:\n            x = self.bn1(x)\n        x = self.activation(x)\n        x = self.conv1(x)\n        if self.down:\n            x = downsample(x)\n        if self.skip_conv is not None:\n            x0 = self.skip_conv(x0)\n            if self.down:\n                x0 = downsample(x0)\n        if RESIDUAL_IN_D:\n            x = x0 + x\n        return x\n\n\nclass DiscriminatorDownOptimizedBlock(nn.Module):\n    def __init__(self, in_features: int, out_features: int, activation=DISCRIMINATOR_ACTIVATION) -> None:\n        super().__init__()\n        self.conv1 = snconv2d(in_features, out_features, kernel_size=3, stride=1, padding=1, use_norm=SPECTRAL_IN_D)\n        self.activation = activation\n        self.conv2 = snconv2d(out_features, out_features, kernel_size=3, stride=1, padding=1, use_norm=SPECTRAL_IN_D)\n        if RESIDUAL_IN_D:\n            self.skip_conv = snconv2d(in_features, out_features, kernel_size=1, stride=1, padding=0, use_norm=SPECTRAL_IN_D)\n        \n    def forward(self, x):\n        x0 = x\n        x = self.conv1(x)\n        x = self.activation(x)\n        x = self.conv2(x)\n        x = downsample(x)\n        if RESIDUAL_IN_D:\n            x0 = downsample(x0)\n            x0 = self.skip_conv(x0)\n            x = x0 + x\n        return x\n\n    \nclass Discriminator(nn.Module):\n    def __init__(self, num_classes: int = len(CLASSES), features: int = N_FEATURES, activation=DISCRIMINATOR_ACTIVATION) -> None:\n        super().__init__()\n        self.activation = activation\n        self.block0 = DiscriminatorDownOptimizedBlock(3, features * 1)\n        if SELF_ATTENTION:\n            self.attn = SelfAttention(features * 1, use_norm=SPECTRAL_IN_D)\n        self.block1 = DiscriminatorBlock(features * 1, features * 2)\n        self.block2 = DiscriminatorBlock(features * 2, features * 4)\n        self.block3 = DiscriminatorBlock(features * 4, features * 8)\n        self.block4 = DiscriminatorBlock(features * 8, features * 8, down=False)\n        self.linear = snlinear(features * 8, 1, use_norm=SPECTRAL_IN_D)\n        if CONDITIONAL:\n            self.emb = sn_embedding(num_classes, features * 8, use_norm=SPECTRAL_IN_D)\n        \n    def forward(self, x, class_id):\n        x = self.block0(x)\n        if SELF_ATTENTION:\n            x = self.attn(x)\n        x = self.block1(x)\n        x = self.block2(x)\n        x = self.block3(x)\n        x = self.block4(x)\n        h = self.activation(x)\n        h = h.sum(3).sum(2)\n        output = self.linear(h).squeeze(1)\n        if CONDITIONAL:\n            h_labels = self.emb(class_id)\n            output = output + (h * h_labels).sum(1)\n        return output\n    \nDiscriminator()(torch.randn(4, 3, 64, 64), torch.zeros(4, dtype=torch.long)).shape","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"@torch.no_grad()\ndef inverse_transformation(tensor):\n    mean=[0.5, 0.5, 0.5]\n    std=[0.5, 0.5, 0.5]\n    for i in range(3):\n        m = mean[i]\n        s = std[i]\n        tensor[:, i].mul_(s).add_(m)\n    return tensor\n\n@torch.no_grad()\ndef generate_images(generator, noises, labels=None):\n    if labels is None:\n        labels = torch.multinomial(torch.ones(len(CLASSES), dtype=torch.float32) / len(CLASSES), noises.size(0)).to(noises.device)\n    images = generator(noises, labels)\n    return inverse_transformation(images)\n\n\ndef show_images(tensor):\n    grid = torchvision.utils.make_grid(tensor)\n    grid_numpy = grid.permute(1, 2, 0).cpu().numpy()\n    plt.imshow(grid_numpy)\n    plt.show()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class DiscriminatorCriterion(nn.Module):\n    def forward(self, fake_images_output, real_images_output):\n        if LOSS == LossType.RaLS:\n            return torch.mean((real_images_output - fake_images_output.mean() - 1) ** 2) + \\\n                   torch.mean((fake_images_output - real_images_output.mean() + 1) ** 2)\n        elif LOSS == LossType.LS:\n            return torch.mean(real_images_output ** 2) + torch.mean((fake_images_output - 1) ** 2)\n        elif LOSS == LossType.Hinge:\n            return F.relu(1.0 - real_images_output).mean() + F.relu(1 + fake_images_output).mean()\n        else:\n            raise Exception(\"unimplemented: {}\".format(LOSS))\n    \n    \nclass GeneratorCriterion(nn.Module):\n    def forward(self, fake_images_output, real_images_output):\n        if LOSS == LossType.RaLS:\n            return torch.mean((real_images_output - fake_images_output.mean() + 1) ** 2) + \\\n                   torch.mean((fake_images_output - real_images_output.mean() - 1) ** 2)\n        elif LOSS == LossType.LS:\n            return torch.mean(fake_images_output ** 2)\n        elif LOSS == LossType.Hinge:\n            return -fake_images_output.mean()\n        else:\n            raise Exception(\"unimplemented: {}\".format(LOSS))\n\n\ndef requires_grad(model, flag=True):\n    for p in model.parameters():\n        p.requires_grad = flag\n        \ndef warmup_lr_scheduler(optimizer, warmup_iters, warmup_factor):\n    def f(x):\n        if x >= warmup_iters:\n            return 1\n        alpha = float(x) / warmup_iters\n        return warmup_factor * (1 - alpha) + alpha\n    return torch.optim.lr_scheduler.LambdaLR(optimizer, f)\n    \n\ndef train_one_epoch(epoch, generator, discriminator, generator_optimizer, discriminator_optimizer, dataloader, device, show_freq=None):\n    generator.train()\n    discriminator.train()\n    disc_criterion = DiscriminatorCriterion()\n    gen_criterion = GeneratorCriterion()\n    \n    lr_scheduler_gen = None\n    lr_scheduler_disc = None\n    if False and epoch == 0:\n        lr_scheduler_gen = warmup_lr_scheduler(generator_optimizer, len(dataloader), warmup_factor=1. / 1000)\n        lr_scheduler_disc = warmup_lr_scheduler(discriminator_optimizer, len(dataloader), warmup_factor=1. / 1000)\n    \n    loss_discriminators = []\n    loss_generators = []\n    for step, (real_images, real_labels) in enumerate(tqdm(dataloader, total=len(dataloader))):\n        # 1. Update D\n        # 1.1. train with real images\n        discriminator.zero_grad()\n        real_images = real_images.to(device)\n        real_labels = real_labels.to(device)\n        batch_size = real_images.size(0)\n        \n        real_images_output = discriminator(real_images, real_labels)\n\n        # 1.2. train with fake images\n        noise = sample(batch_size).to(device)\n        # fake_labels = torch.multinomial(torch.ones(len(CLASSES), dtype=torch.float32) / len(CLASSES), batch_size).to(device)\n        fake_images = generator(noise, real_labels)\n        fake_images_output = discriminator(fake_images.detach(), real_labels)\n\n        loss_discriminator = disc_criterion(fake_images_output, real_images_output)\n        loss_discriminator.backward()\n        discriminator_optimizer.step()\n        loss_discriminators.append(loss_discriminator.detach().item())\n        \n        # 2. Update G\n        generator.zero_grad()\n        noise = sample(batch_size).to(device)\n        fake_images = generator(noise, real_labels)\n        fake_images_output = discriminator(fake_images, real_labels)\n        if LOSS.requires_real_for_generator():\n            real_images_output = discriminator(real_images, real_labels).detach()\n        else:\n            real_images_output = None\n        loss_generator = gen_criterion(fake_images_output, real_images_output)\n        loss_generator.backward()\n        generator_optimizer.step()\n        loss_generators.append(loss_generator.detach().item())\n        \n        if show_freq and step % show_freq == 0:\n            generator.eval()\n            fake_images = generate_images(generator, FIXED_NOISES.to(device))\n            show_images(fake_images)\n            generator.train()\n            \n        if lr_scheduler_gen is not None:\n            lr_scheduler_gen.step()\n        if lr_scheduler_disc is not None:\n            lr_scheduler_disc.step()\n    print('[%d/%d] Loss_D: %.4f Loss_G: %.4f'\n      % (epoch + 1, EPOCHS, \n         np.mean(loss_discriminators), np.mean(loss_generators)))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import gc\ngc.collect()\ntorch.cuda.empty_cache()\n\nimport datetime\nprint(datetime.datetime.now(tz=datetime.timezone(datetime.timedelta(hours=9))).strftime(\"%Y/%m/%d %H:%M:%S\"))\n\ndevice = torch.device(\"cuda\") if torch.cuda.is_available() else torch.device(\"cpu\")\n\ngenerator = Generator().to(device)\ndiscriminator = Discriminator().to(device)\n\ngenerator_optimizer = torch.optim.Adam(generator.parameters(), lr=LR_GENERATOR, betas=(BETA1, BETA2))\ndiscriminator_optimizer = torch.optim.Adam(discriminator.parameters(), lr=LR_DISCRIMINATOR, betas=(BETA1, BETA2))\n\nif not os.path.exists(\"./tmp\"):\n    os.mkdir(\"./tmp\")\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"\nfor epoch in range(EPOCHS):\n    print(f\"[Epoch {epoch}]\")\n    train_one_epoch(epoch, generator, discriminator, generator_optimizer, discriminator_optimizer, dataloader, device, show_freq=SHOW_FREQ)\n    generator.eval()\n    if (epoch + 1) % SHOW_FREQ_EPOCH == 0:\n        fake_images = generate_images(generator, FIXED_NOISES.to(device))\n        show_images(fake_images)\n        if CONDITIONAL:\n            fake_images = generate_images(generator, FIXED_NOISES.to(device))\n            show_images(fake_images)\n        save_image(fake_images, \"./tmp/epoch{}.png\".format(epoch + 1))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"print(datetime.datetime.now(tz=datetime.timezone(datetime.timedelta(hours=9))).strftime(\"%Y/%m/%d %H:%M:%S\"))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"from torchvision.utils import save_image\n\nif not os.path.exists('../output_images'):\n    os.mkdir('../output_images')\nim_batch_size = 1\nn_images = 10000\ngenerator.eval()\nfor i_batch in tqdm(range(0, n_images), total=n_images):\n    gen_z = sample(im_batch_size).to(device)\n    gen_images = generate_images(generator, gen_z)\n    images = gen_images.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')\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"","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}