{"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":"# Creating Monet Paintings using Gans and Pytorch\nIn this notebook I will create a simple GAN to paint MONET paintints.","metadata":{}},{"cell_type":"code","source":"# Set wandb environment variable\nfrom kaggle_secrets import UserSecretsClient\nimport os\nos.environ[\"WANDB_API_KEY\"]  = UserSecretsClient().get_secret(\"WANDB_API_KEY\")","metadata":{"execution":{"iopub.status.busy":"2021-09-07T09:27:31.494329Z","iopub.execute_input":"2021-09-07T09:27:31.494687Z","iopub.status.idle":"2021-09-07T09:27:31.644712Z","shell.execute_reply.started":"2021-09-07T09:27:31.494606Z","shell.execute_reply":"2021-09-07T09:27:31.643673Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import pandas as pd\n\nimport random\nimport matplotlib.pyplot as plt\nimport numpy as np\nfrom PIL import Image\n\nimport wandb\nimport torch\nfrom torch import nn\nimport torch.optim as optim\nimport torch.utils.data\nfrom torch.utils.data import Dataset\nimport torchvision.transforms as transforms\nimport torchvision.utils as vutils\nimport tqdm","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","execution":{"iopub.status.busy":"2021-09-07T09:27:31.64767Z","iopub.execute_input":"2021-09-07T09:27:31.647916Z","iopub.status.idle":"2021-09-07T09:27:33.691976Z","shell.execute_reply.started":"2021-09-07T09:27:31.647891Z","shell.execute_reply":"2021-09-07T09:27:33.691052Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Folders Path\ntraining_pictures = '../input/gan-getting-started/monet_jpg/'\n\n# Image Size. Necessary for defining the transforms\nIMG_SIZE = 256\n\n# Batch size during training\nBATCH_SIZE = 100\n\n# Noise Dimention\nZ_DIM = 100\n\n# Number of epochs to train\nNUM_EPOCHS = 1200\n\n# Number of steps to log results\nlog_step = 20\n\ndevice = torch.device(\"cuda:0\" if (torch.cuda.is_available()) else \"cpu\")\ndevice","metadata":{"execution":{"iopub.status.busy":"2021-09-07T09:27:33.693866Z","iopub.execute_input":"2021-09-07T09:27:33.694209Z","iopub.status.idle":"2021-09-07T09:27:33.745369Z","shell.execute_reply.started":"2021-09-07T09:27:33.69418Z","shell.execute_reply":"2021-09-07T09:27:33.744358Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Loading and visualizing the Pictures","metadata":{}},{"cell_type":"code","source":"class ImageDataset(Dataset):\n    \"\"\" Custom Dataset class \"\"\"\n    def __init__(self, root: str, transform: transforms.Compose):\n        self.root = root\n        self.transform = transform\n        self.all_imgs = tuple(os.path.join(root, p) for p in os.listdir(root) if p.endswith('.jpg'))\n\n    def __len__(self) -> int:\n        return len(self.all_imgs)\n\n    def __getitem__(self, idx: int) -> torch.Tensor:\n        return self.transform(self.load_image(self.all_imgs[idx]))\n    \n    @staticmethod\n    def load_image(path: str) -> np.ndarray:\n        with Image.open(path) as p:\n            img = np.asarray(p)\n        return img\n        \n    def plot_random_images(self, n_img: int) -> None:\n        \n        # Sample random number of images\n        sampled_images = random.sample(self.all_imgs, n_img)\n        \n        images_list = list(np.transpose(torch.from_numpy(self.load_image(p)), (2, 0, 1)) for p in sampled_images)\n        image_grid = vutils.make_grid(images_list)\n        plt.figure(figsize=(20, 20))\n        plt.imshow(np.transpose(image_grid, (1, 2, 0)))\n        plt.axis('off')\n        plt.show()\n        \n    \ndataset = ImageDataset(root=training_pictures,\n                           transform=transforms.Compose([\n                               transforms.ToTensor(),\n                               transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))\n                           ]))\n\ndataset.plot_random_images(40)","metadata":{"execution":{"iopub.status.busy":"2021-09-07T09:27:33.74758Z","iopub.execute_input":"2021-09-07T09:27:33.748182Z","iopub.status.idle":"2021-09-07T09:27:35.015897Z","shell.execute_reply.started":"2021-09-07T09:27:33.748123Z","shell.execute_reply":"2021-09-07T09:27:35.014875Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Create the GANS","metadata":{}},{"cell_type":"markdown","source":"## Generator","metadata":{}},{"cell_type":"code","source":"def get_noise(n_samples):\n    \"\"\" Get Noise function\"\"\"\n    return torch.randn(n_samples, Z_DIM, 1, 1, device=device, dtype=torch.float)\n\nclass Generator(nn.Module):\n    \"\"\" Generator Class \"\"\"\n    def __init__(self, z_dim: int, hidden_dim: int):\n        super().__init__()\n        self.main = nn.Sequential(\n            self._block(z_dim, hidden_dim*4, kernel_size=9, stride=3, padding=0, output_padding=2),\n            self._block(hidden_dim*4, hidden_dim*2, kernel_size=9, stride=3, padding=0, output_padding=1),\n            self._block(hidden_dim*2, hidden_dim, kernel_size=8, stride=3, padding=0,output_padding=0),\n            nn.ConvTranspose2d(hidden_dim, 3, kernel_size=8, stride=2, padding=0, output_padding=0),\n            nn.Tanh()  \n        )\n        \n    def _block(self, in_channels: int, out_channels: int, kernel_size=8, stride=1, dilation=1, padding=0, output_padding=0) -> nn.Sequential:\n        return nn.Sequential(\n                    nn.ConvTranspose2d(in_channels, out_channels, kernel_size=kernel_size, stride=stride, dilation=dilation, padding=padding, output_padding=output_padding),\n                    nn.BatchNorm2d(out_channels),\n                    nn.ReLU(),\n                )\n\n\n    def forward(self, input):\n        return self.main(input)\n    \n# Assert output image shape\nZ_DIM=20\ngen_shape = np.transpose(Generator(Z_DIM, 5).to(device)(get_noise(1)).detach().cpu().numpy(), (2, 3, 1, 0))[:, :, :, 0].shape\nassert gen_shape == (IMG_SIZE, IMG_SIZE, 3), gen_shape","metadata":{"execution":{"iopub.status.busy":"2021-09-07T09:27:35.017291Z","iopub.execute_input":"2021-09-07T09:27:35.017642Z","iopub.status.idle":"2021-09-07T09:27:40.316896Z","shell.execute_reply.started":"2021-09-07T09:27:35.017598Z","shell.execute_reply":"2021-09-07T09:27:40.315946Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Discriminator","metadata":{}},{"cell_type":"code","source":"class Critic(nn.Module):\n    '''\n    Critic Class I took from my coursera course.\n    Values:\n        im_chan: the number of channels in the images, fitted for the dataset used, a scalar\n              (MNIST is black-and-white, so 1 channel is your default)\n        hidden_dim: the inner dimension, a scalar\n    '''\n    def __init__(self, im_chan=3, hidden_dim=20):\n        super().__init__()\n        self.crit = nn.Sequential(\n            self.make_crit_block(im_chan, hidden_dim*6, kernel_size=8, stride=3),\n            self.make_crit_block(hidden_dim*6, hidden_dim*3, kernel_size=8, stride=3),\n            self.make_crit_block(hidden_dim*3, hidden_dim, kernel_size=8, stride=3),\n            self.make_crit_block(hidden_dim, 1, stride=1),\n            self.make_crit_block(1, 1, kernel_size=9,final_layer=True),\n        )\n\n    def make_crit_block(self, input_channels, output_channels, kernel_size=4, stride=2, final_layer=False, padding=1):\n        '''\n        Function to return a sequence of operations corresponding to a critic block of DCGAN;\n        a convolution, a batchnorm (except in the final layer), and an activation (except in the final layer).\n        Parameters:\n            input_channels: how many channels the input feature representation has\n            output_channels: how many channels the output feature representation should have\n            kernel_size: the size of each convolutional filter, equivalent to (kernel_size, kernel_size)\n            stride: the stride of the convolution\n            final_layer: a boolean, true if it is the final layer and false otherwise \n                      (affects activation and batchnorm)\n        '''\n        if not final_layer:\n            return nn.Sequential(\n                nn.Conv2d(input_channels, output_channels, kernel_size=kernel_size, stride=stride, padding=padding, bias=False),\n                nn.BatchNorm2d(output_channels),\n                nn.LeakyReLU(0.2, inplace=True),\n            )\n        else:\n            return nn.Sequential(\n                nn.Conv2d(input_channels, output_channels, kernel_size=kernel_size, stride=stride, padding=padding),\n                nn.Sigmoid()\n            )\n\n    def forward(self, image):\n        '''\n        Function for completing a forward pass of the critic: Given an image tensor, \n        returns a 1-dimension tensor representing fake/real.\n        Parameters:\n            image: a flattened image tensor with dimension (im_chan)\n        '''\n        crit_pred = self.crit(image)\n        return crit_pred.view(len(crit_pred), -1)\n    \n# Assert the output\ndisc_size = Critic()(dataset[0].view(1, 3, IMG_SIZE, IMG_SIZE)).view(-1).size()\nassert disc_size == torch.Size([1]), disc_size","metadata":{"execution":{"iopub.status.busy":"2021-09-07T09:27:40.320316Z","iopub.execute_input":"2021-09-07T09:27:40.320598Z","iopub.status.idle":"2021-09-07T09:27:40.427724Z","shell.execute_reply.started":"2021-09-07T09:27:40.320571Z","shell.execute_reply":"2021-09-07T09:27:40.426883Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Start the training","metadata":{}},{"cell_type":"code","source":"# Initialize Critic and Generator\ngen = Generator(Z_DIM, 80).to(device)\ngen_opt = torch.optim.Adam(gen.parameters())\ncrit = Critic().to(device) \ncrit_opt = torch.optim.Adam(crit.parameters(), lr=0.00005)\n\n# Loss function\ncriterion = nn.BCELoss()","metadata":{"execution":{"iopub.status.busy":"2021-09-07T09:27:40.430675Z","iopub.execute_input":"2021-09-07T09:27:40.430942Z","iopub.status.idle":"2021-09-07T09:27:40.54953Z","shell.execute_reply.started":"2021-09-07T09:27:40.430915Z","shell.execute_reply":"2021-09-07T09:27:40.548626Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Create the data loader\ndataloader = torch.utils.data.DataLoader(dataset, batch_size=BATCH_SIZE, shuffle=True)","metadata":{"execution":{"iopub.status.busy":"2021-09-07T09:27:40.553175Z","iopub.execute_input":"2021-09-07T09:27:40.553486Z","iopub.status.idle":"2021-09-07T09:27:40.557751Z","shell.execute_reply.started":"2021-09-07T09:27:40.553459Z","shell.execute_reply":"2021-09-07T09:27:40.556802Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"wandb.init(project='Monet GAN', entity='djrmarques')\nwandb.watch(gen, log=\"all\", log_freq=log_step)\nwandb.watch(crit, log=\"all\", log_freq=log_step)","metadata":{"execution":{"iopub.status.busy":"2021-09-07T09:27:40.559194Z","iopub.execute_input":"2021-09-07T09:27:40.559714Z","iopub.status.idle":"2021-09-07T09:27:47.294711Z","shell.execute_reply.started":"2021-09-07T09:27:40.559677Z","shell.execute_reply":"2021-09-07T09:27:47.29376Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"I want to have the generator trained more times that the critic. For this reason I created a function that will compute the multiplier of times the generator gets trained more than the discriminator based on the current EPOCH. ","metadata":{}},{"cell_type":"code","source":"def size_generator_training(epoch):\n    max_multiple = 2\n    return 1\n    return min(int(round((epoch/NUM_EPOCHS)*max_multiple)) + 1, max_multiple)\nsize_generator_training(1)","metadata":{"execution":{"iopub.status.busy":"2021-09-07T09:34:02.364429Z","iopub.execute_input":"2021-09-07T09:34:02.364752Z","iopub.status.idle":"2021-09-07T09:34:02.372398Z","shell.execute_reply.started":"2021-09-07T09:34:02.364724Z","shell.execute_reply":"2021-09-07T09:34:02.371508Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"for epoch in tqdm.trange(NUM_EPOCHS):\n    gen_batch_loss = []\n    crit_batch_loss = []\n    \n    gen_train_mult = size_generator_training(epoch)\n    for i, data in enumerate(dataloader):\n        \n        ## Update the Critic with Real images\n        crit.zero_grad()\n        real_data = data.to(device)\n        b_size = real_data.shape[0]\n        real_output = crit(real_data).view(-1)\n        real_error = criterion(real_output, torch.ones(b_size, dtype=torch.float, device=device))\n        real_error.backward()\n        \n        ## Train with fake image\n        noise = get_noise(b_size*gen_train_mult)  # Train the generator 2 times for each time the Critic is trained\n        fake = gen(noise)\n        fake_crit_train = fake[:b_size]\n        fake_output = crit(fake_crit_train.detach()).view(-1)\n        fake_error = criterion(fake_output, torch.zeros(b_size, dtype=torch.float, device=device))\n        fake_error.backward()\n        disc_error = real_error + fake_error\n        crit_opt.step()\n        \n        \n        ## Update the Generator\n        gen.zero_grad()\n        output = crit(fake).view(-1)\n        gen_error = criterion(output, torch.ones(b_size*gen_train_mult, dtype=torch.float, device=device))\n        gen_error.backward()\n        crit.zero_grad()\n        gen_opt.step()\n        \n        gen_batch_loss.append(gen_error.mean().item())\n        crit_batch_loss.append(disc_error.mean().item())\n        \n    if epoch%log_step == 0:\n        # Log WAB\n        wandb.log({\"Critic Error\": sum(crit_batch_loss)/len(crit_batch_loss), 'Generator Error': sum(gen_batch_loss)/len(gen_batch_loss)})\n\n        im = plt.imshow(np.transpose(vutils.make_grid(gen(torch.randn(5, Z_DIM, 1, 1, device=device)).detach().cpu()), (1, 2, 0)))\n        wandb.log({\"Generated Images\": wandb.Image(im, caption=\"Generated Images\")})","metadata":{"execution":{"iopub.status.busy":"2021-09-07T09:49:10.472774Z","iopub.execute_input":"2021-09-07T09:49:10.473098Z","iopub.status.idle":"2021-09-07T09:50:23.206449Z","shell.execute_reply.started":"2021-09-07T09:49:10.473068Z","shell.execute_reply":"2021-09-07T09:50:23.205185Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Save the trained model","metadata":{}},{"cell_type":"code","source":"torch.save(gen.state_dict(), \"generator.model\")\ntorch.save(crit.state_dict(), \"critic.model\")","metadata":{"execution":{"iopub.status.busy":"2021-09-06T18:53:59.533358Z","iopub.status.idle":"2021-09-06T18:53:59.533727Z"},"trusted":true},"execution_count":null,"outputs":[]}]}