{"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":"# Get Data","metadata":{}},{"cell_type":"code","source":"!mkdir /kaggle/working/generative-dog-images\n!unzip /kaggle/input/generative-dog-images/all-dogs.zip -d /kaggle/working/generative-dog-images > /dev/null 2>&1\n!unzip /kaggle/input/generative-dog-images/Annotation.zip -d /kaggle/working/generative-dog-images > /dev/null 2>&1","metadata":{"execution":{"iopub.status.busy":"2021-09-05T13:24:35.427632Z","iopub.execute_input":"2021-09-05T13:24:35.427978Z","iopub.status.idle":"2021-09-05T13:24:52.082049Z","shell.execute_reply.started":"2021-09-05T13:24:35.427904Z","shell.execute_reply":"2021-09-05T13:24:52.08101Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import os\nimport torch\nimport torch.nn as nn\nfrom torchvision import transforms\nfrom torchvision.utils import make_grid\nimport matplotlib.pyplot as plt\nfrom PIL import Image\nimport xml.etree.ElementTree as ET ","metadata":{"execution":{"iopub.status.busy":"2021-09-05T13:24:52.084928Z","iopub.execute_input":"2021-09-05T13:24:52.085273Z","iopub.status.idle":"2021-09-05T13:24:53.830224Z","shell.execute_reply.started":"2021-09-05T13:24:52.085238Z","shell.execute_reply":"2021-09-05T13:24:53.829355Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"all_dogs_dir = '/kaggle/working/generative-dog-images/all-dogs'\nannotation_dir = '/kaggle/working/generative-dog-images/Annotation'","metadata":{"execution":{"iopub.status.busy":"2021-09-05T13:24:53.832007Z","iopub.execute_input":"2021-09-05T13:24:53.832336Z","iopub.status.idle":"2021-09-05T13:24:53.837795Z","shell.execute_reply.started":"2021-09-05T13:24:53.832302Z","shell.execute_reply":"2021-09-05T13:24:53.835874Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def get_bndbox(filename, square=False):\n    tree = ET.parse(filename)\n    root = tree.getroot()\n    box = root.find('object').find('bndbox')\n    xmin = int(box.find('xmin').text)\n    ymin = int(box.find('ymin').text)\n    xmax = int(box.find('xmax').text)\n    ymax = int(box.find('ymax').text)\n    \n    if square:\n        center_x, center_y = (xmin + xmax)//2, (ymin+ymax)//2\n        max_w = max(xmax-xmin, ymax-ymin)\n        xmin = center_x - max_w//2\n        xmax = xmin + max_w\n        ymin = center_y - max_w//2\n        ymax = ymin + max_w\n        \n    return xmin, ymin, xmax, ymax  ","metadata":{"execution":{"iopub.status.busy":"2021-09-05T13:24:53.839402Z","iopub.execute_input":"2021-09-05T13:24:53.839756Z","iopub.status.idle":"2021-09-05T13:24:53.852219Z","shell.execute_reply.started":"2021-09-05T13:24:53.839723Z","shell.execute_reply":"2021-09-05T13:24:53.85136Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"dogs_dir = '/kaggle/working/dogs'\nif not os.path.exists(dogs_dir):\n    os.makedirs(dogs_dir)\n\nfor breed in os.listdir(annotation_dir):\n    for dog in os.listdir(os.path.join(annotation_dir, breed)):\n        bndbox = get_bndbox(os.path.join(annotation_dir, breed, dog), square=True)\n        jpg_name = os.path.join(all_dogs_dir, dog+'.jpg')\n        if os.path.exists(jpg_name):\n            img = Image.open(jpg_name).crop(bndbox)\n            img.save(os.path.join(dogs_dir, dog+'.jpg'))\nprint('number of dogs:', len(os.listdir(dogs_dir)))","metadata":{"execution":{"iopub.status.busy":"2021-09-05T13:24:53.853349Z","iopub.execute_input":"2021-09-05T13:24:53.853938Z","iopub.status.idle":"2021-09-05T13:27:20.077804Z","shell.execute_reply.started":"2021-09-05T13:24:53.853835Z","shell.execute_reply":"2021-09-05T13:27:20.076716Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Dataset","metadata":{}},{"cell_type":"code","source":"from torch.utils.data import Dataset, DataLoader\n\nclass DogDataset(Dataset):\n    def __init__(self, data_dir, transforms=None):\n        self.files = [os.path.join(data_dir, file) for file in os.listdir(data_dir)]\n        self.transforms = transforms\n        \n    def __len__(self):\n        return len(self.files)\n    \n    def __getitem__(self, index):\n        img = Image.open(self.files[index])\n        if self.transforms is not None:\n            img = self.transforms(img)\n        return img","metadata":{"execution":{"iopub.status.busy":"2021-09-05T13:27:20.079145Z","iopub.execute_input":"2021-09-05T13:27:20.079641Z","iopub.status.idle":"2021-09-05T13:27:20.087189Z","shell.execute_reply.started":"2021-09-05T13:27:20.079604Z","shell.execute_reply":"2021-09-05T13:27:20.086244Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"img_size = (64, 64)\nimage_transforms = transforms.Compose([\n    transforms.ToTensor(),\n    transforms.Resize(img_size),\n    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))\n])\n\nbatch_size = 128\ntrainloader = DataLoader(\n    DogDataset(data_dir=dogs_dir, transforms=image_transforms),\n    batch_size = batch_size,\n    shuffle = True,\n    num_workers = 3,\n)","metadata":{"execution":{"iopub.status.busy":"2021-09-05T13:27:20.08845Z","iopub.execute_input":"2021-09-05T13:27:20.088832Z","iopub.status.idle":"2021-09-05T13:27:20.145305Z","shell.execute_reply.started":"2021-09-05T13:27:20.088797Z","shell.execute_reply":"2021-09-05T13:27:20.144578Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Model","metadata":{}},{"cell_type":"code","source":"class Generator(nn.Module):\n    def __init__(self, z_channels, out_channels=3):\n        super(Generator, self).__init__()\n        \n        convs = []\n        channels = [z_channels, 1024, 512, 256, 128, 64]\n        for i in range(1, len(channels)):\n            convs.append(nn.ConvTranspose2d(channels[i-1], channels[i], 2, stride=2, bias=False))\n            convs.append(nn.BatchNorm2d(channels[i]))\n            convs.append(nn.ReLU(inplace=True))\n        convs.append(nn.ConvTranspose2d(channels[-1], out_channels, 2, stride=2, bias=False))\n        convs.append(nn.Tanh())\n        \n        self.convs = nn.Sequential(*convs)\n        \n    def forward(self, x):\n        return self.convs(x)","metadata":{"execution":{"iopub.status.busy":"2021-09-05T13:27:20.147504Z","iopub.execute_input":"2021-09-05T13:27:20.147858Z","iopub.status.idle":"2021-09-05T13:27:20.155416Z","shell.execute_reply.started":"2021-09-05T13:27:20.147825Z","shell.execute_reply":"2021-09-05T13:27:20.154241Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class Discriminator(nn.Module):\n    def __init__(self):\n        super(Discriminator, self).__init__()\n        channels = [3, 64, 128, 256, 512] # size: 64->32->16->8->4\n        convs = []\n        for i in range(1, len(channels)):\n            convs.append(nn.Conv2d(channels[i-1], channels[i], 3, padding=1, stride=2, bias=False))\n            #convs.append(nn.InstanceNorm2d(channels[i], affine=True))\n            #convs.append(nn.LayerNorm([channels[i], 64//2**i, 64//2**i]))\n            convs.append(nn.LeakyReLU(0.2, inplace=True))\n        \n        convs.append(nn.Conv2d(channels[-1], 1, 4, bias=False))\n        # convs.append(nn.Sigmoid())\n        \n        self.convs = nn.Sequential(*convs)\n    \n    def forward(self, x):\n        x = self.convs(x)\n        return x.view(-1)","metadata":{"execution":{"iopub.status.busy":"2021-09-05T13:27:20.157444Z","iopub.execute_input":"2021-09-05T13:27:20.157874Z","iopub.status.idle":"2021-09-05T13:27:20.166744Z","shell.execute_reply.started":"2021-09-05T13:27:20.157816Z","shell.execute_reply":"2021-09-05T13:27:20.165725Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def weights_init(m):\n    classname = m.__class__.__name__\n    if classname.find('Conv') != -1:\n        m.weight.data.normal_(0.0, 0.02)\n    elif classname.find('BatchNorm') != -1:\n        m.weight.data.normal_(1.0, 0.02)\n        m.bias.data.fill_(0)","metadata":{"execution":{"iopub.status.busy":"2021-09-05T13:27:20.168109Z","iopub.execute_input":"2021-09-05T13:27:20.168436Z","iopub.status.idle":"2021-09-05T13:27:20.179453Z","shell.execute_reply.started":"2021-09-05T13:27:20.168403Z","shell.execute_reply":"2021-09-05T13:27:20.178665Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"z_channels = 100\nG = Generator(z_channels, 3)\nG.apply(weights_init)\nD = Discriminator()\nD.apply(weights_init)\n\ncuda = torch.cuda.is_available()\nif cuda:\n    print('Use GPU')\n    G = G.cuda()\n    D = D.cuda()\nelse:\n    print('No GPU')","metadata":{"execution":{"iopub.status.busy":"2021-09-05T13:27:20.180571Z","iopub.execute_input":"2021-09-05T13:27:20.181048Z","iopub.status.idle":"2021-09-05T13:27:24.564277Z","shell.execute_reply.started":"2021-09-05T13:27:20.181013Z","shell.execute_reply":"2021-09-05T13:27:24.563318Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Loss and Optimizer","metadata":{}},{"cell_type":"code","source":"lr = 1e-4\noptimizer_G = torch.optim.Adam(G.parameters(), lr=lr, betas=(0.5, 0.9))\noptimizer_D = torch.optim.Adam(D.parameters(), lr=lr, betas=(0.5, 0.9))","metadata":{"execution":{"iopub.status.busy":"2021-09-05T13:27:24.567527Z","iopub.execute_input":"2021-09-05T13:27:24.567817Z","iopub.status.idle":"2021-09-05T13:27:24.574639Z","shell.execute_reply.started":"2021-09-05T13:27:24.56779Z","shell.execute_reply":"2021-09-05T13:27:24.573619Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def gradient_penalty(netD, real_img, fake_img, LAMDA=10, cuda=True):\n    batch_size = real_img.size(0)\n    alpha = torch.rand(batch_size, 1)\n    alpha= alpha.expand(batch_size, real_img.nelement()//batch_size).reshape(real_img.shape)\n    if cuda:\n        alpha = alpha.cuda()\n    x = (alpha * real_img + (1-alpha) * fake_img).requires_grad_(True)\n    if cuda:\n        x = x.cuda()\n    out = netD(x)\n    \n    grad_outputs = torch.ones(out.shape)\n    if cuda:\n        grad_outputs = grad_outputs.cuda()\n        \n    gradients = torch.autograd.grad(outputs=out, inputs=x, grad_outputs=grad_outputs, create_graph=True, only_inputs=True)[0]\n    gradients = gradients.reshape(batch_size, -1)\n    \n    return LAMDA * ((gradients.norm(2, dim=1)-1)**2).mean()","metadata":{"execution":{"iopub.status.busy":"2021-09-05T13:27:24.57771Z","iopub.execute_input":"2021-09-05T13:27:24.578011Z","iopub.status.idle":"2021-09-05T13:27:24.587588Z","shell.execute_reply.started":"2021-09-05T13:27:24.577979Z","shell.execute_reply":"2021-09-05T13:27:24.586858Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Train","metadata":{}},{"cell_type":"code","source":"noise_mean = 0\nnoise_std = 0.1\nepoches = 100\nG_loss_list = []\nD_loss_list = []\nfor epoch in range(epoches):\n    for i, img in enumerate(trainloader):\n        z = torch.normal(noise_mean, noise_std, size=(img.size(0), z_channels, 1, 1))\n        if cuda:\n            img, z = img.cuda(), z.cuda()\n        \n        # train D\n        D.zero_grad()\n        loss_real = -D(img).mean()\n        fake_img = G(z).detach()\n        loss_fake = D(fake_img).mean()\n        gp = gradient_penalty(D, img.detach(), fake_img, LAMDA=10, cuda=cuda)\n        loss_D = loss_real + loss_fake + gp\n        loss_D.backward()\n        optimizer_D.step()\n        \n        if (i+1) % 5 == 0:\n            # train G\n            z = torch.normal(0, 1, size=(img.size(0), z_channels, 1, 1))\n            if cuda:\n                z = z.cuda()\n            G.zero_grad()\n            loss_G = -D(G(z)).mean()\n            loss_G.backward()\n            optimizer_G.step()\n            \n            D_loss_list.append(loss_D.item())\n            G_loss_list.append(loss_G.item())\n            \n    print(f'[Epoch {epoch+1}/{epoches}] [G loss: {loss_G.item()}] [D loss: {loss_D.item()} | loss_real: {loss_real.item()} loss_fake: {loss_fake.item()}]')\n    ","metadata":{"_kg_hide-output":true,"execution":{"iopub.status.busy":"2021-09-05T13:27:24.588994Z","iopub.execute_input":"2021-09-05T13:27:24.589339Z","iopub.status.idle":"2021-09-05T13:28:28.183061Z","shell.execute_reply.started":"2021-09-05T13:27:24.589305Z","shell.execute_reply":"2021-09-05T13:28:28.181513Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Plot","metadata":{}},{"cell_type":"code","source":"plt.plot(G_loss_list, label='G')\nplt.plot(D_loss_list, label='D')\nplt.legend()\nplt.show()","metadata":{"execution":{"iopub.status.busy":"2021-09-05T13:28:28.184592Z","iopub.execute_input":"2021-09-05T13:28:28.18496Z","iopub.status.idle":"2021-09-05T13:28:28.336461Z","shell.execute_reply.started":"2021-09-05T13:28:28.184921Z","shell.execute_reply":"2021-09-05T13:28:28.335523Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"noise = torch.normal(noise_mean, noise_std, size=(64, z_channels, 1, 1))\nif cuda:\n    noise = noise.cuda()\n\nwith torch.no_grad():\n    fig = plt.figure(figsize=(10,10))\n    plt.axis(\"off\")\n    plt.imshow(make_grid(G(noise), padding=2, normalize=True).cpu().permute(1,2,0))\n    plt.show()","metadata":{"execution":{"iopub.status.busy":"2021-09-05T13:28:28.337722Z","iopub.execute_input":"2021-09-05T13:28:28.338091Z","iopub.status.idle":"2021-09-05T13:28:28.594106Z","shell.execute_reply.started":"2021-09-05T13:28:28.338054Z","shell.execute_reply":"2021-09-05T13:28:28.593222Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}]}