{"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_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# Import Libraries","metadata":{}},{"cell_type":"code","source":"import os\nimport zipfile\nimport glob\n\nimport numpy as np\nimport matplotlib.pyplot as plt\nimport xml.etree.ElementTree as ET\n\nimport torch\nimport torch.nn as nn\nimport torch.optim as optim\nimport torch.nn.functional as F\nimport torchvision\nimport torchvision.transforms as transforms\n\nfrom PIL import Image\nfrom tqdm.notebook import tqdm\nfrom dataclasses import dataclass","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Extract Data","metadata":{}},{"cell_type":"code","source":"for dirname, _, filenames in os.walk('/kaggle/input'):\n    for filename in filenames:\n        print(os.path.join(dirname, filename))","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"with zipfile.ZipFile(\"../input/generative-dog-images/all-dogs.zip\",\"r\") as zip_ref:\n    zip_ref.extractall(\"/kaggle/temp/\")\n    \nwith zipfile.ZipFile(\"../input/generative-dog-images/Annotation.zip\",\"r\") as zip_ref:\n    zip_ref.extractall(\"/kaggle/temp/\")","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Lisiting Extracted Data","metadata":{}},{"cell_type":"code","source":"!ls /kaggle/temp/Annotation/","metadata":{"scrolled":true,"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"!ls /kaggle/temp/all-dogs","metadata":{"scrolled":true,"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Initializing Constants","metadata":{}},{"cell_type":"code","source":"ROOT = '/kaggle/temp/'\n\nANNOT_PATH = ROOT + 'Annotation/'\nIMAGE_PATH = ROOT + 'all-dogs/'\n\nannotations = os.listdir(ANNOT_PATH) \nimages = os.listdir(IMAGE_PATH)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"@dataclass\nclass TrainConfig:\n    num_workers: int = 4\n    epochs: int = 50\n    batch_size: int = 64\n    generate_size: int = 10_000\n    save_epoch: int = 5\n        \n    mean: float = 0.5\n    std: float = 0.5\n        \n    num_channels: int = 3\n    image_size: int = 64\n    feature_size: int = 64\n    noise_size: int = 100\n    embedding_dim: int = 256\n    attention: bool = True\n\n    device: 'typing.Any' = torch.device('cuda' if torch.cuda.is_available() else 'cpu')","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"train_config = TrainConfig()","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Preprocessing\n\nTook the help of https://www.kaggle.com/tikutiku/gan-dogs-starter-biggan","metadata":{}},{"cell_type":"code","source":"print(f\"Number of breeds: {len(annotations)}\")\nprint(f\"Number of images: {len(images)}\")","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"for breed in annotations:\n    annotations += glob.glob(breed+'/*')\nprint(f\"Number of available annotations: {len(annotations)}\")","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"breed_map = {}\nfor annotation in annotations:\n    index, *breed = annotation.split(\"-\")\n    breed_map[index] = index + \"-\" + \"-\".join(breed)\nprint(f\"Number of breeds in breed_map: {len(breed_map)}\")","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"breed_map","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def bounding_box(annot_path, image):\n    bounding_path = annot_path + str(breed_map[image.split(\"_\")[0]]) + \"/\" + str(image.split(\".\")[0])\n    tree = ET.parse(bounding_path)\n    root = tree.getroot()\n    objects = root.findall(\"object\")\n    bboxes = []\n    for o in objects:\n        bound_box = o.find(\"bndbox\")\n        x_min = int(bound_box.find(\"xmin\").text)\n        y_min = int(bound_box.find(\"ymin\").text)\n        x_max = int(bound_box.find(\"xmax\").text)\n        y_max = int(bound_box.find(\"ymax\").text)\n        bboxes.append((x_min, y_min, x_max, y_max))\n    return bboxes","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def bounding_box_ratio(annot_path, image):\n    bounding_path = annot_path + str(breed_map[image.split(\"_\")[0]]) + \"/\" + str(image.split(\".\")[0])\n    tree = ET.parse(bounding_path)\n    root = tree.getroot()\n    objects = root.findall(\"object\")\n    bbox_ratios = []\n    for o in objects:\n        bound_box = o.find(\"bndbox\")\n        x_min = int(bound_box.find(\"xmin\").text)\n        y_min = int(bound_box.find(\"ymin\").text)\n        x_max = int(bound_box.find(\"xmax\").text)\n        y_max = int(bound_box.find(\"ymax\").text)\n        x_len = x_max - x_min\n        y_len = y_max - y_min\n        ratio = y_len / x_len\n        bbox_ratios.append((x_len, y_len, ratio))\n    return bbox_ratios","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"%%time\n\n#threshold for aspect ratio, at the same time idx for each bbx\nimages_th = []\n\nfor image in tqdm(images):\n    bbox_ratios = bounding_box_ratio(ANNOT_PATH, image)\n    for i,(x_len, y_len, ratio) in enumerate(bbox_ratios):\n        if ((ratio > 0.2) & (ratio < 4.0)):\n            images_th.append(image[:-4] + '_' + str(i) + '.jpg')\n\nprint(f\"Original Length: {len(images)}\")\nprint(f\"After Thresholding Length: {len(images_th)}\")","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"#from https://www.kaggle.com/korovai/dogs-images-intruders-extraction\nintruders = [\n    #n02088238-basset\n    'n02088238_10870_0.jpg',\n    \n    #n02088466-bloodhound\n    'n02088466_6901_1.jpg',\n    'n02088466_6963_0.jpg',\n    'n02088466_9167_0.jpg',\n    'n02088466_9167_1.jpg',\n    'n02088466_9167_2.jpg',\n    \n    #n02089867-Walker_hound\n    'n02089867_2221_0.jpg',\n    'n02089867_2227_1.jpg',\n    \n    #n02089973-English_foxhound # No details\n    'n02089973_1132_3.jpg',\n    'n02089973_1352_3.jpg',\n    'n02089973_1458_1.jpg',\n    'n02089973_1799_2.jpg',\n    'n02089973_2791_3.jpg',\n    'n02089973_4055_0.jpg',\n    'n02089973_4185_1.jpg',\n    'n02089973_4185_2.jpg',\n    \n    #n02090379-redbone\n    'n02090379_4673_1.jpg',\n    'n02090379_4875_1.jpg',\n    \n    #n02090622-borzoi # Confusing\n    'n02090622_7705_1.jpg',\n    'n02090622_9358_1.jpg',\n    'n02090622_9883_1.jpg',\n    \n    #n02090721-Irish_wolfhound # very small\n    'n02090721_209_1.jpg',\n    'n02090721_1222_1.jpg',\n    'n02090721_1534_1.jpg',\n    'n02090721_1835_1.jpg',\n    'n02090721_3999_1.jpg',\n    'n02090721_4089_1.jpg',\n    'n02090721_4276_2.jpg',\n    \n    #n02091032-Italian_greyhound\n    'n02091032_722_1.jpg',\n    'n02091032_745_1.jpg',\n    'n02091032_1773_0.jpg',\n    'n02091032_9592_0.jpg',\n    \n    #n02091134-whippet\n    'n02091134_2349_1.jpg',\n    'n02091134_14246_2.jpg',\n    \n    #n02091244-Ibizan_hound\n    'n02091244_583_1.jpg',\n    'n02091244_2407_0.jpg',\n    'n02091244_3438_1.jpg',\n    'n02091244_5639_1.jpg',\n    'n02091244_5639_2.jpg',\n    \n    #n02091467-Norwegian_elkhound\n    'n02091467_473_0.jpg',\n    'n02091467_4386_1.jpg',\n    'n02091467_4427_1.jpg',\n    'n02091467_4558_1.jpg',\n    'n02091467_4560_1.jpg',\n    \n    #n02091635-otterhound\n    'n02091635_1192_1.jpg',\n    'n02091635_4422_0.jpg',\n    \n    #n02091831-Saluki\n    'n02091831_1594_1.jpg',\n    'n02091831_2880_0.jpg',\n    'n02091831_7237_1.jpg',\n    \n    #n02092002-Scottish_deerhound\n    'n02092002_1551_1.jpg',\n    'n02092002_1937_1.jpg',\n    'n02092002_4218_0.jpg',\n    'n02092002_4596_0.jpg',\n    'n02092002_5246_1.jpg',\n    'n02092002_6518_0.jpg',\n    \n    #02093256-Staffordshire_bullterrier\n    'n02093256_1826_1.jpg',\n    'n02093256_4997_0.jpg',\n    'n02093256_14914_0.jpg',\n    \n    #n02093428-American_Staffordshire_terrier\n    'n02093428_5662_0.jpg',\n    'n02093428_6949_1.jpg'\n            ]\n\nprint(f\"Number of intruders: {len(intruders)}\")","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class DogDataset(torch.utils.data.Dataset):\n    \n    def __init__(self, annot_path, image_path, image_list, label_map, transform=None, intruders=[]):\n        self.ANNOT_PATH = annot_path\n        self.IMAGE_PATH = image_path\n        \n        self.image_list = image_list\n        \n        self.transform = transform\n        \n        self.images = []\n        self.labels = []\n        \n        for image_path in self.image_list:\n            if image_path in intruders:\n                continue\n                \n            *image_name, bbox_idx = image_path.split(\"_\")\n            image_path = \"_\".join(image_name) + \".jpg\"\n            image = self._data_preprocessing(image_path, int(bbox_idx.split(\".\")[0]))\n            \n            if self.transform:\n                image = self.transform(image)\n                \n            self.images.append(image)\n            \n            label = label_map[image_name[0]]\n            self.labels.append(label)\n            \n    def _data_preprocessing(self, image_path, bbox_idx):\n        bbox = bounding_box(self.ANNOT_PATH, image_path)[bbox_idx]\n        image  = Image.open(os.path.join(self.IMAGE_PATH, image_path)) # PILImage format\n        cropped_image  = image.crop(bbox)\n        return cropped_image\n        \n    def __len__(self):\n        return len(self.images)\n    \n    def __getitem__(self, idx):\n        image = self.images[idx]\n        label = self.labels[idx]\n        \n        return {\"images\": image, \"labels\": label}","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"%%time\n\nn_classes = len(annotations)\n\nlabel_map = {breed: i for i, breed in enumerate(breed_map.keys())}\n\ntransform = transforms.Compose([transforms.Resize((train_config.image_size, train_config.image_size)), \n                                 transforms.ToTensor(), \n                                 transforms.Normalize(mean=(train_config.mean), std=(train_config.std))\n                                ])\n\ntrain_set = DogDataset(annot_path=ANNOT_PATH, \n                       image_path=IMAGE_PATH, \n                       image_list=images_th, \n                       label_map=label_map, \n                       transform=transform, \n                       intruders=intruders\n                      )\n\ntrain_loader = torch.utils.data.DataLoader(train_set,\n                          shuffle=True, batch_size=train_config.batch_size,\n                          num_workers=train_config.num_workers, pin_memory=True)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def show_grid(image):\n  npimage = image.numpy()\n  plt.imshow(np.transpose(npimage, (1, 2, 0)))\n  plt.show()","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"iter_loader = iter(train_loader)\ndata = iter_loader.next()\nshow_grid(torchvision.utils.make_grid(data[\"images\"], normalize=True))\nprint(data[\"labels\"])","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Creating Model","metadata":{}},{"cell_type":"markdown","source":"## Helper Functions","metadata":{}},{"cell_type":"code","source":"def initialize_weights(m):\n    if isinstance(m, nn.Linear):\n        nn.init.xavier_uniform_(m.weight.data)\n        m.bias.data.fill_(0)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def l2_normalize(v, eps=1e-12):\n    return v / (v.norm() + eps)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Self Attention Layer","metadata":{}},{"cell_type":"code","source":"class SelfAttentionLayer(nn.Module):\n    def __init__(self, in_dim):\n        super().__init__()\n        \n        self.query_conv = nn.Conv2d(in_channels = in_dim, out_channels = in_dim // 2, kernel_size = 1)\n        self.key_conv = nn.Conv2d(in_channels = in_dim, out_channels = in_dim // 2, kernel_size = 1)\n        self.value_conv = nn.Conv2d(in_channels = in_dim, out_channels = in_dim, kernel_size = 1)\n        \n        self.gamma = nn.Parameter(torch.zeros(1))\n        self.softmax = nn.Softmax(dim = -1)\n        \n    def forward(self, x):\n        batch_size, C, width, height = x.size()\n        query = self.query_conv(x).view(batch_size, -1, width * height).permute(0, 2, 1)\n        key = self.key_conv(x).view(batch_size, -1, width * height)\n        energy = torch.bmm(query, key)\n        \n        attn = self.softmax(energy)\n        value = self.value_conv(x).view(batch_size, -1, width * height)\n        out = torch.bmm(value, attn.permute(0, 2, 1))\n        out = out.view(batch_size, C, width, height)\n        \n        out = self.gamma * out + x\n        return out, attn","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Spectral Normalization Layer","metadata":{}},{"cell_type":"code","source":"class SpectralNormLayer(nn.Module):\n    def __init__(self, module, name=\"weight\", power_iterations=1):\n        super().__init__()\n        self.module = module\n        self.name = name\n        self.power_iterations = power_iterations\n        \n        if not self._made_params():\n            self._make_params()\n            \n    def _update_u_v(self):\n        u = getattr(self.module, self.name + \"_u\")\n        v = getattr(self.module, self.name + \"_v\")\n        w = getattr(self.module, self.name + \"_bar\")\n        \n        height = w.data.shape[0]\n        for _ in range(self.power_iterations):\n            v.data = l2_normalize(torch.mv(torch.t(w.view(height, -1).data), u.data))\n            u.data = l2_normalize(torch.mv(w.view(height, -1).data, v.data))\n            \n        sigma = u.dot(w.view(height, -1).mv(v))\n        setattr(self.module, self.name, w / sigma.expand_as(w))\n        \n    def _made_params(self):\n        try:\n            u = getattr(self.module, self.name + \"_u\")\n            v = getattr(self.module, self.name + \"_v\")\n            w = getattr(self.module, self.name + \"_bar\")\n            return True\n        except AttributeError:\n            return False\n    \n    def _make_params(self):\n        w = getattr(self.module, self.name)\n        \n        height = w.data.shape[0]\n        width = w.view(height, -1).data.shape[1]\n        \n        u = nn.Parameter(w.data.new(height).normal_(0, 1), requires_grad = False)\n        v = nn.Parameter(w.data.new(width).normal_(0, 1), requires_grad = False)\n        u.data = l2_normalize(u.data)\n        v.data = l2_normalize(v.data)\n        w_bar = nn.Parameter(w.data)\n        \n        del self.module._parameters[self.name]\n        \n        self.module.register_parameter(self.name + \"_u\", u)\n        self.module.register_parameter(self.name + \"_v\", v)\n        self.module.register_parameter(self.name + \"_bar\", w_bar)\n        \n    def forward(self, *args):\n        self._update_u_v()\n        return self.module.forward(*args)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Generator","metadata":{}},{"cell_type":"code","source":"class Generator(nn.Module):\n    def __init__(self, noise_size, embedding_dim, num_classes, num_channels, image_size, attn = True, feature_size = 64):\n        super().__init__()\n        self.attn = attn\n        self.noise_dim = noise_size\n        \n        assert image_size == 64, \"Cannot produce images other than size 64x64!\"\n        self.image_size = image_size\n        \n        # Embedding Layer\n        self.label_embedding = nn.Embedding(num_classes, embedding_dim)\n        \n        # Linear Layer\n        self.linear_layer = nn.Sequential(\n            nn.Linear(embedding_dim, noise_size),\n            nn.BatchNorm1d(noise_size),\n            nn.ReLU()\n        )\n        \n        # Size 1 -> 4\n        self.layer_1 = nn.Sequential(\n            SpectralNormLayer(nn.ConvTranspose2d(noise_size, feature_size * 8, kernel_size = 4)),\n            nn.BatchNorm2d(feature_size * 8),\n            nn.ReLU()\n        )\n        \n        # Size 4 -> 8\n        self.layer_2 = nn.Sequential(\n            SpectralNormLayer(nn.ConvTranspose2d(feature_size * 8, feature_size * 4, kernel_size = 4, stride = 2, padding = 1)),\n            nn.BatchNorm2d(feature_size * 4),\n            nn.ReLU()\n        )\n        \n        # Size 8 -> 16\n        self.layer_3 = nn.Sequential(\n            SpectralNormLayer(nn.ConvTranspose2d(feature_size * 4, feature_size * 2, kernel_size = 4, stride = 2, padding = 1)),\n            nn.BatchNorm2d(feature_size * 2),\n            nn.ReLU()\n        )\n        \n        # Attention Layer 1\n        self.attention_1 = SelfAttentionLayer(feature_size * 2)\n        \n        # Size 16 -> 32\n        self.layer_4 = nn.Sequential(\n            SpectralNormLayer(nn.ConvTranspose2d(feature_size * 2, feature_size * 2, kernel_size = 4, stride = 2, padding = 1)),\n            nn.BatchNorm2d(feature_size * 2),\n            nn.ReLU()\n        )\n        \n        # Attention Layer 2\n        self.attention_2 = SelfAttentionLayer(feature_size * 2)\n        \n        # Output layer 32 -> 64\n        self.output_layer = nn.Sequential(\n            nn.ConvTranspose2d(feature_size * 2, 3, kernel_size = 4, stride = 2, padding = 1),\n            nn.Tanh()\n        )\n        \n        self.optimizer = optim.Adam(self.parameters(), lr = 0.0002, betas = (0.5, 0.999))\n        \n    def forward(self, noise, labels):\n        \n        label_embed = self.label_embedding(labels)\n        linear_out = self.linear_layer(label_embed)\n        \n        x = torch.mul(linear_out, noise)\n        x = x.view(x.shape[0], -1, 1, 1)\n        x = self.layer_1(x)\n        x = self.layer_2(x)\n        x = self.layer_3(x)\n\n        if self.attn:\n            x, _ = self.attention_1(x)\n        \n        x = self.layer_4(x)\n        \n        if self.attn:\n            x, _ = self.attention_2(x)\n            \n        outputs = self.output_layer(x)\n        return outputs","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class Discriminator(nn.Module):\n    def __init__(self, embedding_dim, num_classes, num_channels, image_size, attn = True, feature_size = 64):\n        super().__init__()\n        self.embedding_dim = embedding_dim\n        self.attn = attn\n        \n        assert image_size == 64, \"Cannot create model for images other than size 64x64!\"\n        self.image_size = image_size\n        \n        # Size 64 -> 32\n        self.layer_1 = nn.Sequential(\n            SpectralNormLayer(nn.Conv2d(num_channels, feature_size, 4, 2, 1)),\n            nn.LeakyReLU(0.1)\n        )\n        \n        # Size 32 -> 16\n        self.layer_2 = nn.Sequential(\n            SpectralNormLayer(nn.Conv2d(feature_size, feature_size * 2, 4, 2, 1)),\n            nn.LeakyReLU(0.1)\n        )\n        \n        # Size 16 -> 8\n        self.layer_3 = nn.Sequential(\n            SpectralNormLayer(nn.Conv2d(feature_size * 2, feature_size * 2, 4, 2, 1)),\n            nn.LeakyReLU(0.1)\n        )\n        \n        # Attention Layer 1\n        self.attention_1 = SelfAttentionLayer(feature_size * 2)\n        \n        # Size 8 -> 4\n        self.layer_4 = nn.Sequential(\n            SpectralNormLayer(nn.Conv2d(feature_size * 2, feature_size * 4, 4, 2, 1)),\n            nn.LeakyReLU(0.1)\n        )\n        \n        # Attention Layer 2\n        self.attention_2 = SelfAttentionLayer(feature_size * 4)\n        \n        # Embedding Layer\n        self.label_embedding = nn.Embedding(num_classes, embedding_dim)\n        \n        # Linear Layer\n        self.image_label_layer = nn.Sequential(\n            nn.Conv2d(embedding_dim + feature_size * 4, feature_size * 8, 1, 1, 0, bias = False),\n            nn.BatchNorm2d(feature_size * 8),\n            nn.LeakyReLU(0.2)\n        )\n        \n        # Output Layer\n        self.output_layer = nn.Sequential(\n            nn.Conv2d(feature_size * 8, 1, 4, 1, 0, bias = False),\n            nn.Sigmoid()\n        )\n        \n        self.optimizer = optim.Adam(self.parameters(), lr = 0.0002, betas = (0.5, 0.999))\n        \n    def forward(self, inputs, labels):\n        x = self.layer_1(inputs)\n        x = self.layer_2(x)\n        x = self.layer_3(x)\n        \n        if self.attn:\n            x, _ = self.attention_1(x)\n            \n        x = self.layer_4(x)\n        \n        if self.attn:\n            x, _ = self.attention_2(x)\n            \n        label_embed = self.label_embedding(labels)\n        label_embed = label_embed.unsqueeze(2).unsqueeze(2).repeat(1, 1, 4, 4)\n        \n        x = torch.cat([x, label_embed], dim = 1)\n        x = self.image_label_layer(x)\n        outputs = self.output_layer(x)\n        return outputs.view(-1, 1)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"generator = Generator(train_config.noise_size, train_config.embedding_dim, n_classes, train_config.num_channels, train_config.image_size)\ngenerator.apply(initialize_weights)\ngenerator.to(train_config.device)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"discriminator = Discriminator(train_config.noise_size, n_classes, train_config.num_channels, train_config.image_size)\ndiscriminator.apply(initialize_weights)\ndiscriminator.to(train_config.device)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Train Model","metadata":{}},{"cell_type":"code","source":"adversarial_loss = nn.BCELoss().to(train_config.device)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"fixed_noise = torch.randn(size=(n_classes, train_config.noise_size)).to(train_config.device)\nfixed_labels = torch.arange(0, n_classes, dtype=torch.long).to(train_config.device)\n\ndef plot_output(epoch):\n  plt.clf()\n  with torch.no_grad():\n\n    generator.eval()\n    test_images = generator(fixed_noise, fixed_labels)\n    generator.train()\n    grid = torchvision.utils.make_grid(test_images.cpu(), normalize=True)\n    show_grid(grid)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"pbar = tqdm()\n\ndevice = train_config.device\n\nfor epoch in range(train_config.epochs):\n    print(f\"Epoch: {epoch + 1} / {train_config.epochs}\")\n    pbar.reset(total=len(train_loader))\n    \n    # Setting up losses\n    discriminator_losses = []\n    generator_losses = []\n    \n    for i, data in enumerate(train_loader):\n        \n        # Bring to device\n        real_images = data[\"images\"].to(device)\n        real_labels = data[\"labels\"].to(device)\n        \n        # Get batch size\n        current_batch_size = real_images.size()[0]\n        \n        # For real vs fake\n        real_valid = torch.ones(current_batch_size, 1).to(device)\n        fake_valid = torch.zeros(current_batch_size, 1).to(device)\n        \n        # Train Generator\n        generator.zero_grad()\n        input_noise = torch.randn(size=(current_batch_size, train_config.noise_size)).to(device)\n        fake_images = generator(input_noise, real_labels)\n        disc_fake_valid = discriminator(fake_images, real_labels)\n        \n        generator_loss = adversarial_loss(disc_fake_valid, real_valid)\n        generator_loss.backward()\n        generator.optimizer.step()\n        generator_losses.append(generator_loss)\n        \n        # Train Discriminator\n        discriminator.zero_grad()\n        \n        ## Calculate real loss\n        disc_real_valid = discriminator(real_images, real_labels)\n        disc_real_loss = adversarial_loss(disc_real_valid, real_valid)\n        \n        ## Calculate wrong loss\n        wrong_labels = torch.randint(0, n_classes, (current_batch_size, )).to(device)\n        disc_wrong_valid = discriminator(real_images, wrong_labels)\n        disc_wrong_loss = adversarial_loss(disc_wrong_valid, fake_valid)\n        \n        ## Calculate fake loss\n        disc_fake_valid = discriminator(fake_images.detach(), real_labels)\n        disc_fake_loss = adversarial_loss(disc_fake_valid, fake_valid)\n        \n        ## Calculating total loss\n        discriminator_loss = disc_real_loss + disc_wrong_loss + disc_fake_loss\n        if discriminator_loss > 0.5:\n            discriminator_loss.backward()\n            discriminator.optimizer.step()\n        else:\n            discriminator_loss = discriminator_loss.detach()\n        discriminator_losses.append(discriminator_loss)\n        \n        pbar.update()\n        \n    print(f\"Discriminator Loss: {torch.mean(torch.FloatTensor(discriminator_losses)):.3f}\")\n    print(f\"Generator Loss: {torch.mean(torch.FloatTensor(generator_losses)):.3f}\")\n    \n    if (epoch + 1) % train_config.save_epoch == 0:\n        plot_output(epoch + 1)\n        \n        \npbar.refresh()","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"with torch.no_grad():\n\n    generator.eval()\n    save_images = generator(fixed_noise, fixed_labels)\n    generator.train()\n    grid = torchvision.utils.make_grid(save_images.cpu(), normalize=True)\n    torchvision.utils.save_image(grid, '/kaggle/working/output_images.png')","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}]}