{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.10.13","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"none","dataSources":[{"sourceId":15062,"databundleVersionId":545987,"sourceType":"competition"}],"dockerImageVersionId":30665,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# Import Package","metadata":{}},{"cell_type":"code","source":"import torch, torchvision, os, PIL, pdb\nfrom torch import nn\nfrom torch.utils.data import Dataset\nfrom torch.utils.data import DataLoader\nfrom torchvision import transforms\nfrom torchvision.utils import make_grid\nfrom tqdm.auto import tqdm\nfrom PIL import Image\nimport matplotlib.pyplot as plt\nimport numpy as np\nfrom tqdm import tqdm","metadata":{"execution":{"iopub.status.busy":"2024-03-06T02:48:50.481950Z","iopub.execute_input":"2024-03-06T02:48:50.482397Z","iopub.status.idle":"2024-03-06T02:49:02.797730Z","shell.execute_reply.started":"2024-03-06T02:48:50.482366Z","shell.execute_reply":"2024-03-06T02:49:02.796804Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# show function","metadata":{}},{"cell_type":"code","source":"def show(tensor, num=25, name=''):\n    data = tensor.detach().cpu()\n    grid = make_grid(data[:num], nrow=5).permute(1,2,0)  \n        \n    plt.imshow(grid.numpy().clip(0,1))\n    plt.show()","metadata":{"execution":{"iopub.status.busy":"2024-03-06T02:49:02.799411Z","iopub.execute_input":"2024-03-06T02:49:02.800028Z","iopub.status.idle":"2024-03-06T02:49:02.805367Z","shell.execute_reply.started":"2024-03-06T02:49:02.800000Z","shell.execute_reply":"2024-03-06T02:49:02.804413Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Configuration","metadata":{}},{"cell_type":"code","source":"class CFG:\n    # hyperpaameters and general parameters\n    n_epochs = 10000\n    batch_size = 128\n    lr = 1e-4\n    z_dim = 200\n    device = 'cuda'\n    \n    cur_step = 0\n    crit_cycles = 5 # The number of iterations of the critic\n    gen_losses = []\n    crit_losses = []\n    show_step = 35\n    save_step = 35 # for checkpoint","metadata":{"execution":{"iopub.status.busy":"2024-03-06T03:03:40.753071Z","iopub.execute_input":"2024-03-06T03:03:40.753469Z","iopub.status.idle":"2024-03-06T03:03:40.759622Z","shell.execute_reply.started":"2024-03-06T03:03:40.753439Z","shell.execute_reply":"2024-03-06T03:03:40.758573Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Unzip data","metadata":{}},{"cell_type":"code","source":"!mkdir -p /kaggle/working/all-dogs","metadata":{"execution":{"iopub.status.busy":"2024-03-06T02:49:02.820410Z","iopub.execute_input":"2024-03-06T02:49:02.820692Z","iopub.status.idle":"2024-03-06T02:49:03.857692Z","shell.execute_reply.started":"2024-03-06T02:49:02.820669Z","shell.execute_reply":"2024-03-06T02:49:03.856319Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"!unzip /kaggle/input/generative-dog-images/all-dogs.zip -d /kaggle/working/all-dogs","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"len(os.listdir('/kaggle/working/all-dogs/all-dogs'))","metadata":{"execution":{"iopub.status.busy":"2024-03-06T02:49:21.333080Z","iopub.execute_input":"2024-03-06T02:49:21.333426Z","iopub.status.idle":"2024-03-06T02:49:21.356458Z","shell.execute_reply.started":"2024-03-06T02:49:21.333397Z","shell.execute_reply":"2024-03-06T02:49:21.355407Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Generate Model","metadata":{}},{"cell_type":"code","source":"class Generator(nn.Module):\n    def __init__(self, z_dim=64, d_dim=16):\n        super(Generator, self).__init__()\n        self.z_dim = z_dim\n        self.gen = nn.Sequential(\n            nn.ConvTranspose2d(z_dim, d_dim*32, 4, 1, 0), \n            nn.BatchNorm2d(d_dim*32),\n            nn.ReLU(True),\n\n            nn.ConvTranspose2d(d_dim*32, d_dim*16, 4, 2, 1), \n            nn.BatchNorm2d(d_dim*16),\n            nn.ReLU(),\n\n            nn.ConvTranspose2d(d_dim*16, d_dim*8, 4, 2, 1),\n            nn.BatchNorm2d(d_dim*8),\n            nn.ReLU(),\n\n            nn.ConvTranspose2d(d_dim*8, d_dim*4, 4, 2, 1), \n            nn.BatchNorm2d(d_dim*4),\n            nn.ReLU(),\n\n            nn.ConvTranspose2d(d_dim*4, d_dim*2, 4, 2, 1), \n            nn.BatchNorm2d(d_dim*2),\n            nn.ReLU(),\n\n            nn.ConvTranspose2d(d_dim*2, 3, 4, 2, 1), \n            nn.Tanh()\n            \n        )\n\n    def forward(self, noise):\n        x = noise.view(len(noise), self.z_dim, 1, 1) \n        return self.gen(x)\n\ndef gen_noise(num, z_dim, device='cuda'):\n    return torch.randn(num, z_dim, device=device)","metadata":{"execution":{"iopub.status.busy":"2024-03-06T02:49:21.357507Z","iopub.execute_input":"2024-03-06T02:49:21.357787Z","iopub.status.idle":"2024-03-06T02:49:21.369137Z","shell.execute_reply.started":"2024-03-06T02:49:21.357763Z","shell.execute_reply":"2024-03-06T02:49:21.368197Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Critic Model","metadata":{}},{"cell_type":"code","source":"class Critic(nn.Module):\n    def __init__(self, d_dim=16):\n        super(Critic, self).__init__()\n\n        self.crit = nn.Sequential(\n            nn.Conv2d(3, d_dim, 4, 2, 1), \n            nn.InstanceNorm2d(d_dim), \n            nn.LeakyReLU(0.2),\n\n            nn.Conv2d(d_dim, d_dim*2, 4, 2, 1), \n            nn.InstanceNorm2d(d_dim*2), \n            nn.LeakyReLU(0.2),\n\n            nn.Conv2d(d_dim*2, d_dim*4, 4, 2, 1), \n            nn.InstanceNorm2d(d_dim*4), \n            nn.LeakyReLU(0.2),\n\n            nn.Conv2d(d_dim*4, d_dim*8, 4, 2, 1), \n            nn.InstanceNorm2d(d_dim*8), \n            nn.LeakyReLU(0.2),\n\n            nn.Conv2d(d_dim*8, d_dim*16, 4, 2, 1), \n            nn.InstanceNorm2d(d_dim*16), \n            nn.LeakyReLU(0.2),\n\n            nn.Conv2d(d_dim*16, 1, 4, 1, 0), \n\n        )\n\n    def forward(self, image):\n        \n        crit_pred = self.crit(image) \n        return crit_pred.view(len(crit_pred), -1) ","metadata":{"execution":{"iopub.status.busy":"2024-03-06T02:49:21.370391Z","iopub.execute_input":"2024-03-06T02:49:21.371484Z","iopub.status.idle":"2024-03-06T02:49:21.389445Z","shell.execute_reply.started":"2024-03-06T02:49:21.371448Z","shell.execute_reply":"2024-03-06T02:49:21.388566Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class Dataset(Dataset):\n    \n    def __init__(self, path, size=128, lim=10000): \n        self.sizes=[size, size]\n        items, labels=[], []\n\n        for data in os.listdir(path)[:lim]:\n        # path = '/kaggle/working/all-dogs/all-dogs'\n        # data = 'n02110185_11131.jpg'\n            item=os.path.join(path, data)\n            items.append(item)\n            labels.append(data)\n\n        self.items=items\n        self.labels=labels\n\n    \n    def __len__(self):\n        return len(self.items)\n\n\n    \n    def __getitem__(self, idx):\n        data = PIL.Image.open(self.items[idx]).convert('RGB') \n        ## reize image\n        data = np.asarray(torchvision.transforms.Resize(self.sizes)(data)) \n        ## 轉換image to torch  form\n        data = np.transpose(data, (2,0,1)).astype(np.float32, copy=False) \n        ## transform numpy to tensor\n        data = torch.from_numpy(data).div(255) # from 0 to 1\n\n        return data, self.labels[idx]\n               \n\n## Dataset\ndata_path = '/kaggle/working/all-dogs/all-dogs'\nds = (data_path, size=128, lim=25000)\n\n## Dataloader\ndataloader = DataLoader(ds, batch_size=CFG.batch_size, shuffle=True)\n\n## Models\ngen = Generator(CFG.z_dim).to(CFG.device)\ncrit = Critic().to(CFG.device)\n\n## Optimizer\n# gen_opt = torch.optim.Adam(gen.parameters(), lr=CFG.lr, betas=(0.5,0.9))\ngen_opt = torch.optim.RMSprop(gen.parameters(), lr=CFG.lr)\n# crit_opt = torch.optim.Adam(crit.parameters(), lr=CFG.lr, betas=(0.5,0.9))\ncrit_opt = torch.optim.RMSprop(crit.parameters(), lr=CFG.lr)\n\n\nx, y = next(iter(dataloader))\nshow(x)","metadata":{"execution":{"iopub.status.busy":"2024-03-06T03:03:52.117640Z","iopub.execute_input":"2024-03-06T03:03:52.118322Z","iopub.status.idle":"2024-03-06T03:03:53.078923Z","shell.execute_reply.started":"2024-03-06T03:03:52.118286Z","shell.execute_reply":"2024-03-06T03:03:53.077797Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# gradient penalty calculation","metadata":{}},{"cell_type":"code","source":"def get_gp(real, fake, crit, alpha, gamma=10):\n    mix_images = real * alpha + fake * (1-alpha) \n    mix_scores = crit(mix_images)\n\n    gradient = torch.autograd.grad(\n        inputs = mix_images,\n        outputs = mix_scores,\n        grad_outputs=torch.ones_like(mix_scores),\n\n        retain_graph=True,\n        create_graph=True,\n    )[0] \n\n    gradient = gradient.view(len(gradient), -1) \n    gradient_norm = gradient.norm(2, dim=1)\n    gp = gamma * ((gradient_norm-1)**2).mean()\n\n    return gp","metadata":{"execution":{"iopub.status.busy":"2024-03-06T03:05:35.436020Z","iopub.execute_input":"2024-03-06T03:05:35.436404Z","iopub.status.idle":"2024-03-06T03:05:35.443273Z","shell.execute_reply.started":"2024-03-06T03:05:35.436375Z","shell.execute_reply":"2024-03-06T03:05:35.442338Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Save and load checkpoints","metadata":{}},{"cell_type":"code","source":"root_path = '/kaggle/working/checkpoint'\ndef save_checkpoint(name):\n    torch.save({\n        'epoch': epoch,\n        'model_state_dict': gen.state_dict(),\n        'optimizer_state_dict': gen_opt.state_dict()\n    }, f'{root_path}G-{name}.pkl')\n\n    torch.save({\n        'epoch': epoch,\n        'model_state_dict': crit.state_dict(),\n        'optimizer_state_dict': crit_opt.state_dict()\n    }, f'{root_path}C-{name}.pkl')\n\n    print('Saved checkpoint')\n\ndef load_checkpoint(name):\n    checktpoint = torch.load(f'{root_path}G-{name}.pkl')\n    gen.load_state_dict(checktpoint['model_state_dict'])\n    gen_opt.load_state_dict(checktpoint['optimizer_state_dict'])\n    \n    checktpoint = torch.load(f'{root_path}C-{name}.pkl')\n    crit.load_state_dict(checktpoint['model_state_dict'])\n    crit_opt.load_state_dict(checktpoint['optimizer_state_dict'])\n\n    print('Loaded checkpoint')","metadata":{"execution":{"iopub.status.busy":"2024-03-06T02:51:16.806828Z","iopub.execute_input":"2024-03-06T02:51:16.807220Z","iopub.status.idle":"2024-03-06T02:51:16.815661Z","shell.execute_reply.started":"2024-03-06T02:51:16.807191Z","shell.execute_reply":"2024-03-06T02:51:16.814529Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Training loop","metadata":{}},{"cell_type":"code","source":"## Training loop\nfor epoch in range(CFG.n_epochs):\n    for real, _ in tqdm(dataloader):\n        cur_bs = len(real) \n        real=real.to(CFG.device)\n\n        ### Critic\n        mean_crit_loss = 0\n        for _ in range(CFG.crit_cycles):\n            crit_opt.zero_grad()\n\n            noise=gen_noise(cur_bs, CFG.z_dim)\n            fake=gen(noise)\n            crit_fake_pred = crit(fake.detach())\n            crit_real_pred = crit(real)\n\n            alpha = torch.rand(len(real), 1, 1, 1, device=CFG.device, requires_grad=True) \n            gp = get_gp(real, fake.detach(), crit, alpha)\n\n            crit_loss = crit_fake_pred.mean() - crit_real_pred.mean() + gp\n\n            mean_crit_loss += crit_loss.item() / CFG.crit_cycles\n\n            crit_loss.backward(retain_graph=True)\n\n        CFG.crit_losses += [mean_crit_loss]\n\n        ### Generate\n        gen_opt.zero_grad()\n        noise = gen_noise(cur_bs, CFG.z_dim)\n        fake = gen(noise)\n        crit_fake_pred = crit(fake)\n\n        gen_loss = -crit_fake_pred.mean()\n        gen_opt.step()\n\n        # gen_losses.append([gen_loss.item()])crit_cycles\n        CFG.gen_losses += [gen_loss.item()]\n\n        ### Stats\n        \n        if CFG.cur_step % CFG.save_step == 0 and CFG.cur_step > 0:\n            print('Saving Checkpoint: ', CFG.cur_step, CFG.save_step)\n            save_checkpoint('latest')\n        \n        if (CFG.cur_step % CFG.show_step == 0 and CFG.cur_step > 0):\n            show(fake, name='fake')\n            show(real, name='real')\n\n            gen_mean = sum(CFG.gen_losses[-CFG.show_step:]) / CFG.show_step\n            crit_mean = sum(CFG.crit_losses[-CFG.show_step:]) / CFG.show_step\n            print(f'Epoch: {epoch}; Step: {CFG.cur_step}; Generataor loss: {gen_mean}, crit loss: {crit_mean}')\n\n            plt.plot(\n                range(len(CFG.gen_losses)),\n                torch.Tensor(CFG.gen_losses).numpy(),\n                label=\"Generator loss\"\n            )\n\n            plt.plot(\n                range(len(CFG.gen_losses)),\n                torch.Tensor(CFG.crit_losses).numpy(),\n                label=\"Critic loss\"\n            )\n            plt.ylim(-1000, 1000)\n            plt.legend()\n            plt.show()\n        \n        CFG.cur_step += 1","metadata":{"execution":{"iopub.status.busy":"2024-03-06T03:05:38.579133Z","iopub.execute_input":"2024-03-06T03:05:38.579532Z"},"collapsed":true,"jupyter":{"outputs_hidden":true},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}]}