{"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":"gpu","dataSources":[{"sourceId":21755,"databundleVersionId":1475600,"sourceType":"competition"}],"dockerImageVersionId":30732,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"from platform import python_version\n\nprint(python_version())","metadata":{"_uuid":"6680a2a2-4814-4e7e-a46a-de33bad6101e","_cell_guid":"43bbe706-5f86-4117-8562-5578aaffde53","collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2024-08-04T22:48:02.615555Z","iopub.execute_input":"2024-08-04T22:48:02.615908Z","iopub.status.idle":"2024-08-04T22:48:02.621135Z","shell.execute_reply.started":"2024-08-04T22:48:02.615878Z","shell.execute_reply":"2024-08-04T22:48:02.620275Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import os\nimport torch\nimport torch.nn as nn \nfrom torchvision import datasets, transforms\nfrom tqdm import tqdm\nfrom datetime import datetime\n\nfrom torchvision.io import read_image","metadata":{"_uuid":"059ba5ca-ad22-44d3-8577-2af1fb6a0737","_cell_guid":"9505cb7d-c493-419d-90e6-92c7d6b410e5","collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2024-08-05T16:50:40.516340Z","iopub.execute_input":"2024-08-05T16:50:40.517181Z","iopub.status.idle":"2024-08-05T16:50:42.089762Z","shell.execute_reply.started":"2024-08-05T16:50:40.517146Z","shell.execute_reply":"2024-08-05T16:50:42.088773Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"print(os.getcwd())\nprint(os.listdir(os.getcwd()))","metadata":{"_uuid":"84a36d23-c199-4da3-8ebd-2055d54e06c1","_cell_guid":"65d8635b-c8e9-4696-b0e7-b1d6dd5b162e","collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2024-08-04T22:48:05.711400Z","iopub.execute_input":"2024-08-04T22:48:05.711784Z","iopub.status.idle":"2024-08-04T22:48:05.717156Z","shell.execute_reply.started":"2024-08-04T22:48:05.711758Z","shell.execute_reply":"2024-08-04T22:48:05.716239Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Load data","metadata":{}},{"cell_type":"code","source":"from torch.utils.data import Dataset\n\nclass CustomImageDataset(Dataset):\n    def __init__(self, img_dir: str, num_imgs=None):\n        self.img_dir = img_dir\n        self.images = []\n        \n        img_files = [x for x in os.listdir(self.img_dir) if x.endswith(\".jpg\")]\n        if num_imgs:\n            num_imgs = min(num_imgs, len(img_files))\n        else:\n            num_imgs = len(img_files)\n\n        img_files = img_files[:num_imgs] \n\n        for file in img_files:\n            full_path = os.path.join(self.img_dir, file)\n            self.images.append(full_path)\n        \n        self.num_imgs = num_imgs\n        \n    def __len__(self):\n        return self.num_imgs\n    \n    def __getitem__(self, idx):\n        img_path = self.images[idx]\n        image = read_image(img_path) / 127.5 - 1.0\n        return image","metadata":{"_uuid":"fb1eed6c-10d7-4e38-9294-02a5d52862a4","_cell_guid":"5b6c61b7-5868-466a-a254-98ecfb5adb6d","collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2024-08-05T16:50:46.861509Z","iopub.execute_input":"2024-08-05T16:50:46.862732Z","iopub.status.idle":"2024-08-05T16:50:46.870694Z","shell.execute_reply.started":"2024-08-05T16:50:46.862696Z","shell.execute_reply":"2024-08-05T16:50:46.869736Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"monet_ds = CustomImageDataset(\"/kaggle/input/gan-getting-started/monet_jpg\")\nmonet_data = torch.utils.data.DataLoader(monet_ds, shuffle=True)\n\nphoto_ds = CustomImageDataset(\"/kaggle/input/gan-getting-started/photo_jpg\")\nphoto_data = torch.utils.data.DataLoader(photo_ds, shuffle=True)","metadata":{"_uuid":"599de923-6297-450b-8902-f8a8931af336","_cell_guid":"ef4ad083-8f88-4656-a570-3b4e5552ce1a","collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2024-08-05T16:50:47.597319Z","iopub.execute_input":"2024-08-05T16:50:47.597683Z","iopub.status.idle":"2024-08-05T16:50:48.055745Z","shell.execute_reply.started":"2024-08-05T16:50:47.597654Z","shell.execute_reply":"2024-08-05T16:50:48.054991Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"print(\"# of monet images: \", monet_ds.num_imgs)\nprint(\"# of photos: \", photo_ds.num_imgs)","metadata":{"_uuid":"e9021436-4a7a-45a3-8b2e-32cbb7483a8d","_cell_guid":"f0cf4ee2-0642-4116-a968-d78e336c9f20","collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2024-08-04T22:48:06.063393Z","iopub.execute_input":"2024-08-04T22:48:06.063683Z","iopub.status.idle":"2024-08-04T22:48:06.068832Z","shell.execute_reply.started":"2024-08-04T22:48:06.063661Z","shell.execute_reply":"2024-08-04T22:48:06.067910Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"example_monet = next(iter(monet_data))\nexample_photo = next(iter(photo_data))","metadata":{"_uuid":"cc5f48c7-2b20-49c8-b380-e5595d286bf1","_cell_guid":"d884248a-d342-45fe-9edc-394f1d5515cd","collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2024-08-04T22:48:06.069911Z","iopub.execute_input":"2024-08-04T22:48:06.070233Z","iopub.status.idle":"2024-08-04T22:48:06.157354Z","shell.execute_reply.started":"2024-08-04T22:48:06.070210Z","shell.execute_reply":"2024-08-04T22:48:06.156638Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"print(\"monet size: \", example_monet.size())\nprint(\"photo size: \", example_photo.size())","metadata":{"_uuid":"ed221ec7-d226-4d11-8d3f-696228dd03e5","_cell_guid":"e9e1980f-10fe-457c-a3c5-0dc999443999","collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2024-08-04T22:48:06.158642Z","iopub.execute_input":"2024-08-04T22:48:06.159137Z","iopub.status.idle":"2024-08-04T22:48:06.163364Z","shell.execute_reply.started":"2024-08-04T22:48:06.159111Z","shell.execute_reply":"2024-08-04T22:48:06.162468Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import matplotlib.pyplot as plt\n\nplt.subplot(121)\nplt.title('Photo')\nplt.imshow(example_photo[0].permute(1, 2, 0) * 0.5 + 0.5)\n\nplt.subplot(122)\nplt.title('Monet')\nplt.imshow(example_monet[0].permute(1, 2, 0) * 0.5 + 0.5)","metadata":{"_uuid":"cf3577b9-b8f5-4b1f-9f7b-0ae0e1f65909","_cell_guid":"da2b62c6-3709-48d1-9317-4e3cb9036ef5","collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2024-08-04T22:48:06.164309Z","iopub.execute_input":"2024-08-04T22:48:06.164569Z","iopub.status.idle":"2024-08-04T22:48:06.787044Z","shell.execute_reply.started":"2024-08-04T22:48:06.164546Z","shell.execute_reply":"2024-08-04T22:48:06.786127Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Build generator","metadata":{"_uuid":"204bd2e7-0b9d-4040-b797-772067e5f207","_cell_guid":"06e3035d-eba3-4c54-abe3-78fa9f89dc23","execution":{"iopub.status.busy":"2024-06-12T01:49:06.693444Z","iopub.execute_input":"2024-06-12T01:49:06.693958Z","iopub.status.idle":"2024-06-12T01:49:06.699891Z","shell.execute_reply.started":"2024-06-12T01:49:06.693923Z","shell.execute_reply":"2024-06-12T01:49:06.698496Z"},"trusted":true}},{"cell_type":"code","source":"INPUT_CHANNELS = 3\nOUTPUT_CHANNELS = 3","metadata":{"_uuid":"36604566-c1d6-4b86-ab0f-2c1568fad23c","_cell_guid":"3bfa2d0c-8363-424d-b433-4a6dd50c3046","collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2024-08-05T16:50:52.841103Z","iopub.execute_input":"2024-08-05T16:50:52.841463Z","iopub.status.idle":"2024-08-05T16:50:52.846172Z","shell.execute_reply.started":"2024-08-05T16:50:52.841433Z","shell.execute_reply":"2024-08-05T16:50:52.844746Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def downsample(in_channels, filters, size, apply_instancenorm=True):\n    layers = []\n    \n    conv2d = nn.Conv2d(in_channels=in_channels, out_channels=filters, kernel_size=size, stride=2, padding=1, bias=False)\n    nn.init.normal_(conv2d.weight, mean=0.0, std=0.02)\n    layers.append(conv2d)\n    \n    if apply_instancenorm:\n        instance_norm = nn.InstanceNorm2d(num_features=filters)\n        layers.append(instance_norm)\n        \n    layers.append(nn.LeakyReLU())\n    \n    return nn.Sequential(*layers)","metadata":{"_uuid":"75c3ff34-289e-4a3c-86f1-2a072683fdf7","_cell_guid":"b17eeee1-c170-4f4c-90bd-978405dfeffd","collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2024-08-05T16:50:53.311098Z","iopub.execute_input":"2024-08-05T16:50:53.311487Z","iopub.status.idle":"2024-08-05T16:50:53.318156Z","shell.execute_reply.started":"2024-08-05T16:50:53.311457Z","shell.execute_reply":"2024-08-05T16:50:53.317056Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def upsample(in_channels, filters, size, apply_dropout=False):\n    layers = []\n    \n    conv2d = nn.ConvTranspose2d(in_channels=in_channels, out_channels=filters, kernel_size=size, stride=2, padding=1, bias=False)\n    nn.init.normal_(conv2d.weight, mean=0.0, std=0.02)\n    layers.append(conv2d)\n    \n    instance_norm = nn.InstanceNorm2d(num_features=filters)\n    layers.append(instance_norm)\n    \n    if apply_dropout:\n        layers.append(nn.Dropout(p=0.5))\n        \n    layers.append(nn.LeakyReLU())\n    \n    return nn.Sequential(*layers)","metadata":{"_uuid":"855d31e7-9535-49bf-b9fa-492ec74b6353","_cell_guid":"c934cd0e-68e4-4342-995a-28492ecbdffc","collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2024-08-05T16:50:53.918274Z","iopub.execute_input":"2024-08-05T16:50:53.918643Z","iopub.status.idle":"2024-08-05T16:50:53.925479Z","shell.execute_reply.started":"2024-08-05T16:50:53.918612Z","shell.execute_reply":"2024-08-05T16:50:53.924408Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class Generator(nn.Module):\n    def __init__(self):\n        super(Generator, self).__init__()\n\n        self.down_stack = nn.ModuleList(\n            [\n                downsample(in_channels=3, filters=64, size=4, apply_instancenorm=False),\n                downsample(in_channels=64, filters=128, size=4),\n                downsample(in_channels=128, filters=256, size=4),\n                downsample(in_channels=256, filters=512, size=4),\n                downsample(in_channels=512, filters=512, size=4),\n                downsample(in_channels=512, filters=512, size=4),\n                downsample(in_channels=512, filters=512, size=4),\n                downsample(\n                    in_channels=512, filters=512, size=4, apply_instancenorm=False\n                ),  # (bs, 1, 1, 512)\n            ]\n        )\n\n        # adding dimensions to the input channels because of the skip connections\n        self.up_stack = nn.ModuleList(\n            [\n                upsample(in_channels=512, filters=1024, size=4, apply_dropout=True),\n                upsample(\n                    in_channels=1024 + 512, filters=1024, size=4, apply_dropout=True\n                ),\n                upsample(\n                    in_channels=1024 + 512, filters=1024, size=4, apply_dropout=True\n                ),\n                upsample(in_channels=1024 + 512, filters=1024, size=4),\n                upsample(in_channels=1024 + 512, filters=512, size=4),\n                upsample(in_channels=512 + 256, filters=256, size=4),\n                upsample(\n                    in_channels=256 + 128, filters=128, size=4\n                ),  # (bs, 128, 128, 128)\n            ]\n        )\n\n        conv2d = nn.ConvTranspose2d(\n            in_channels=128 + 64,\n            out_channels=OUTPUT_CHANNELS,\n            kernel_size=4,\n            stride=2,\n            padding=1,\n        )\n        nn.init.normal_(conv2d.weight, mean=0.0, std=0.02)\n        self.last = nn.Sequential(conv2d, nn.Tanh())\n\n    def forward(self, x):\n        # Downsampling the model\n        skips = []\n        for down in self.down_stack:\n            x = down(x)\n            skips.append(x)\n\n        skips = reversed(skips[:-1])\n\n        for up, skip in zip(self.up_stack, skips):\n            x = up(x)\n            x = torch.cat((x, skip), dim=1)\n            # x = torch.cat((x, skip))\n\n        x = self.last(x)\n        return x","metadata":{"_uuid":"80d9f25d-3964-445b-a03c-1e322ac4410c","_cell_guid":"89d31423-bcc6-4226-92d3-5fb8bc596fde","collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2024-08-05T16:50:54.960277Z","iopub.execute_input":"2024-08-05T16:50:54.960642Z","iopub.status.idle":"2024-08-05T16:50:54.974839Z","shell.execute_reply.started":"2024-08-05T16:50:54.960613Z","shell.execute_reply":"2024-08-05T16:50:54.973937Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Build the Discriminator","metadata":{"_uuid":"4a5647bf-0616-4d16-86f7-07b0bd5289a7","_cell_guid":"b33fc0d4-ca3d-4efe-a8c7-b3ef1fea3daa","trusted":true}},{"cell_type":"code","source":"class Discriminator(nn.Module):\n    def __init__(self):\n        super(Discriminator, self).__init__()\n\n        self.down_stack = nn.ModuleList(\n            [\n                downsample(in_channels=3, filters=64, size=4, apply_instancenorm=False),\n                downsample(in_channels=64, filters=128, size=4),\n                downsample(in_channels=128, filters=256, size=4),\n            ]\n        )\n        self.zeropad = nn.ZeroPad2d(padding=1)\n        self.conv = nn.Conv2d(\n            in_channels=256, out_channels=512, kernel_size=4, bias=False, stride=1\n        )\n        nn.init.normal_(self.conv.weight, mean=0.0, std=0.02)\n\n        self.instance_norm = nn.InstanceNorm2d(num_features=512)\n        self.leaky_relu = nn.LeakyReLU()\n        self.last = nn.Conv2d(in_channels=512, out_channels=1, kernel_size=4, stride=1)\n        nn.init.normal_(self.last.weight, mean=0.0, std=0.02)\n\n    def forward(self, x):\n        for down in self.down_stack:\n            x = down(x)\n\n        x = self.zeropad(x)\n        x = self.conv(x)\n        x = self.instance_norm(x)\n        x = self.leaky_relu(x)\n        x = self.last(x)\n\n        return x","metadata":{"_uuid":"f6c60b65-10bd-4f68-8469-1c854cf6ccd3","_cell_guid":"8603d924-88ce-4389-919b-0379ed62ae78","collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2024-08-04T22:48:06.831058Z","iopub.execute_input":"2024-08-04T22:48:06.831375Z","iopub.status.idle":"2024-08-04T22:48:06.844253Z","shell.execute_reply.started":"2024-08-04T22:48:06.831346Z","shell.execute_reply":"2024-08-04T22:48:06.843453Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"device = (\n    \"cuda\"\n    if torch.cuda.is_available()\n    else \"mps\"\n    if torch.backends.mps.is_available()\n    else \"cpu\"\n)\nprint(f\"Using {device} device\")","metadata":{"_uuid":"741979e1-1553-4dc7-82ef-d9fe485ef9d3","_cell_guid":"e59504f7-7400-464f-8a83-e589a43ebfb4","collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2024-08-05T16:55:28.138772Z","iopub.execute_input":"2024-08-05T16:55:28.139138Z","iopub.status.idle":"2024-08-05T16:55:28.144590Z","shell.execute_reply.started":"2024-08-05T16:55:28.139106Z","shell.execute_reply":"2024-08-05T16:55:28.143680Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"example_photo.size()","metadata":{"execution":{"iopub.status.busy":"2024-08-04T22:48:06.880401Z","iopub.execute_input":"2024-08-04T22:48:06.880687Z","iopub.status.idle":"2024-08-04T22:48:06.889326Z","shell.execute_reply.started":"2024-08-04T22:48:06.880664Z","shell.execute_reply":"2024-08-04T22:48:06.888578Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"monet_generator = Generator()\nphoto_generator = Generator()\nmonet_discriminator = Discriminator()\nphoto_discriminator = Discriminator()","metadata":{"_uuid":"3ce3a45c-eba5-457c-b040-9f62747780ce","_cell_guid":"eb509e61-5bb1-4d51-b287-ba4eb96c54cc","collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2024-08-04T22:48:06.890526Z","iopub.execute_input":"2024-08-04T22:48:06.891046Z","iopub.status.idle":"2024-08-04T22:48:11.198451Z","shell.execute_reply.started":"2024-08-04T22:48:06.891021Z","shell.execute_reply":"2024-08-04T22:48:11.197471Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"monet_generator.to(device)\nexample_photo_cuda = example_photo.to(device)\nto_monet = monet_generator(example_photo_cuda)","metadata":{"_uuid":"0749cae9-073a-430f-b993-7c23ea312c17","_cell_guid":"fc58d822-899c-4f11-91a9-115eba6a883d","collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2024-08-04T22:48:11.199715Z","iopub.execute_input":"2024-08-04T22:48:11.200070Z","iopub.status.idle":"2024-08-04T22:48:12.047234Z","shell.execute_reply.started":"2024-08-04T22:48:11.200037Z","shell.execute_reply":"2024-08-04T22:48:12.046470Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"to_monet.size()","metadata":{"_uuid":"6e9bb85f-1b4e-44b5-b82d-04bb918f3178","_cell_guid":"fcd9b01a-4cd0-48df-bef0-2ddaadccc700","collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2024-08-04T22:48:12.048275Z","iopub.execute_input":"2024-08-04T22:48:12.048893Z","iopub.status.idle":"2024-08-04T22:48:12.054712Z","shell.execute_reply.started":"2024-08-04T22:48:12.048859Z","shell.execute_reply":"2024-08-04T22:48:12.053889Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"plt.subplot(1, 2, 1)\nplt.title(\"Original Photo\")\nplt.imshow(example_photo[0].permute(1, 2, 0) * 0.5 + 0.5)\n\nto_monet = to_monet[0].permute(1, 2, 0)\nto_monet = to_monet.cpu().detach().numpy()\nplt.subplot(1, 2, 2)\nplt.title(\"Monet-esque Photo\")\nplt.imshow(to_monet * 0.5 + 0.5)\nplt.show()","metadata":{"_uuid":"c84101a9-ae15-43b2-b6e5-d8feb642515c","_cell_guid":"88512652-89a3-48dd-8635-bc4c5869153e","collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2024-08-04T22:48:12.055842Z","iopub.execute_input":"2024-08-04T22:48:12.056105Z","iopub.status.idle":"2024-08-04T22:48:12.419920Z","shell.execute_reply.started":"2024-08-04T22:48:12.056083Z","shell.execute_reply":"2024-08-04T22:48:12.419107Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Build the CycleGAN model","metadata":{"_uuid":"dba72706-96b8-4c17-899e-f5ec960e58e5","_cell_guid":"a08bca4a-b116-4ac5-a2e5-b694032a30f0","trusted":true}},{"cell_type":"code","source":"def save_model(file_path, epochs, model, optimizer, curr_loss):\n    \"\"\"\n    Function to save the trained model to disk.\n    \"\"\"\n    print(f\"Saving model ...\")\n    torch.save(\n        {\n            \"epoch\": epochs,\n            \"model_state_dict\": model.state_dict(),\n            \"optimizer_state_dict\": optimizer.state_dict(),\n            \"loss\": curr_loss,\n        },\n        file_path,\n        _use_new_zipfile_serialization=False,\n    )","metadata":{"execution":{"iopub.status.busy":"2024-08-04T22:48:12.421170Z","iopub.execute_input":"2024-08-04T22:48:12.421567Z","iopub.status.idle":"2024-08-04T22:48:12.427651Z","shell.execute_reply.started":"2024-08-04T22:48:12.421538Z","shell.execute_reply":"2024-08-04T22:48:12.426847Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class CycleGAN(nn.Module):\n    def __init__(\n        self,\n        monet_generator,\n        photo_generator,\n        monet_discriminator,\n        photo_discriminator,\n        lambda_cycle=10,\n    ):\n        super(CycleGAN, self).__init__()\n        self.m_gen = monet_generator\n        self.p_gen = photo_generator\n        self.m_disc = monet_discriminator\n        self.p_disc = photo_discriminator\n        self.lambda_cycle = lambda_cycle\n        \n        # define optimizers\n        self.m_gen_optimizer = torch.optim.Adam(self.m_gen.parameters(), \n                                                lr=2e-4, betas=(0.5, 0.999))\n        self.p_gen_optimizer = torch.optim.Adam(self.p_gen.parameters(), \n                                                lr=2e-4, betas=(0.5, 0.999))\n        self.m_disc_optimizer = torch.optim.Adam(self.m_disc.parameters(), \n                                                 lr=2e-4, betas=(0.5, 0.999))        \n        self.p_disc_optimizer = torch.optim.Adam(self.p_disc.parameters(), \n                                                 lr=2e-4, betas=(0.5, 0.999))\n\n    def discriminator_loss(self, real, generated):\n        real_loss = nn.BCEWithLogitsLoss()(real, torch.ones_like(real))\n        generated_loss = nn.BCEWithLogitsLoss()(generated, torch.zeros_like(generated))\n        \n        total_disc_loss = real_loss + generated_loss\n        return total_disc_loss * 0.5 \n    \n    def generator_loss(self, generated):\n        return nn.BCEWithLogitsLoss()(generated, torch.ones_like(generated))\n    \n    def calc_cycle_loss(self, real_image, cycled_image, LAMBDA):\n        loss = torch.mean(torch.abs(real_image - cycled_image))\n        return LAMBDA * loss\n    \n    def identity_loss(self, real_image, same_image, LAMBDA):\n        loss = torch.mean(torch.abs(real_image - same_image))\n        return LAMBDA * 0.5 * loss\n    \n    def train_step(self, monet_ds, photo_ds):\n        running_monet_gen_loss = 0.0\n        running_photo_gen_loss = 0.0\n        running_monet_disc_loss = 0.0\n        running_photo_disc_loss = 0.0\n        \n        for (real_monet, real_photo) in zip(\n            tqdm(monet_ds, ncols=100, desc=\"Training for this epoch\"), \n            photo_ds):\n            \n            real_monet = real_monet.to(device=device)\n            real_photo = real_photo.to(device=device)\n            \n            # zero the gradients for batch\n            self.m_gen_optimizer.zero_grad()\n            self.p_gen_optimizer.zero_grad()\n            self.m_disc_optimizer.zero_grad()\n            self.p_disc_optimizer.zero_grad()\n        \n            # photo to monet back to photo\n            fake_monet = self.m_gen(real_photo)\n            cycled_photo = self.p_gen(fake_monet)\n            \n            # monet to photo back to monet\n            fake_photo = self.p_gen(real_monet)\n            cycled_monet = self.m_gen(fake_photo)\n            \n            # generating itself\n            same_monet = self.m_gen(real_monet)\n            same_photo = self.p_gen(real_photo)\n            \n            # discriminator used to check, inputing real images\n            disc_real_monet = self.m_disc(real_monet)\n            disc_real_photo = self.p_disc(real_photo)\n            \n            # discriminator used to check, inputing fake images\n            disc_fake_monet = self.m_disc(fake_monet)\n            disc_fake_photo = self.p_disc(fake_photo)\n            \n            # evaluate generator loss\n            monet_gen_loss = self.generator_loss(disc_fake_monet)\n            # print(\"monet gen loss size: \", monet_gen_loss.size())\n            photo_gen_loss = self.generator_loss(disc_fake_photo)\n            \n            # evaluate total cycle consistency loss\n            monet_cycle_loss = self.calc_cycle_loss(real_monet, cycled_monet, self.lambda_cycle)\n            # print(\"monet cycle loss: \", monet_cycle_loss)\n            # print(\"monet cycle loss size: \", monet_cycle_loss.size())\n            photo_cycle_loss = self.calc_cycle_loss(real_photo, cycled_photo, self.lambda_cycle)\n            total_cycle_loss = monet_cycle_loss + photo_cycle_loss\n            \n            # evaluate total generator loss\n            total_monet_gen_loss = monet_gen_loss + total_cycle_loss + self.identity_loss(real_monet, same_monet, self.lambda_cycle)\n            # running_monet_gen_loss += total_monet_gen_loss.item()\n            # print(\"total monet gen loss size: \", total_monet_gen_loss.size())\n            # print(\"total monet gen loss: \", total_monet_gen_loss.item())\n            total_monet_gen_loss.backward(retain_graph=True)\n            \n            total_photo_gen_loss = photo_gen_loss + total_cycle_loss + self.identity_loss(real_photo, same_photo, self.lambda_cycle)\n            # running_photo_gen_loss += total_photo_gen_loss.item()\n            total_photo_gen_loss.backward(retain_graph=True)\n            \n            # evaluate total discriminator loss\n            monet_disc_loss = self.discriminator_loss(disc_real_monet, disc_fake_monet)\n            # running_monet_disc_loss += monet_disc_loss.item()\n            monet_disc_loss.backward(retain_graph=True)\n            \n            photo_disc_loss = self.discriminator_loss(disc_real_photo, disc_fake_photo)\n            # running_photo_disc_loss += photo_disc_loss.item()\n            photo_disc_loss.backward()\n            \n            # Adjust the learning weights\n            self.m_gen_optimizer.step()\n            self.p_gen_optimizer.step()\n            self.m_disc_optimizer.step()\n            self.p_disc_optimizer.step()\n\n        return {\n            \"monet_gen_loss\": total_monet_gen_loss.item(),\n            \"photo_gen_loss\": total_photo_gen_loss.item(),\n            \"monet_disc_loss\": monet_disc_loss.item(),\n            \"photo_disc_loss\": photo_disc_loss.item(),\n        }\n            \n    def train(self, monet_ds, photo_ds, epochs=25, save_models=False):\n        if save_models:\n            best_val_loss = float(\"inf\")\n#             dt = datetime.now().strftime(\"%Y%m%d_%H%M%S\")\n#             out_folder = os.path.join(os.getcwd(), \"checkpoints\", dt)\n            out_folder = os.path.join(os.getcwd(), \"checkpoints\")\n            if not os.path.exists(out_folder):\n                os.makedirs(out_folder) \n        \n        for i in range(epochs):\n            print(\"Epoch {:d}\".format(i + 1))\n\n            loss = self.train_step(monet_ds, photo_ds)\n\n            print(\n                \"photo_gen_loss: {:.4f}; monet_gen_loss: {:.4f}; photo_disc_loss: {:.4f}; monet_disc_loss: {:.4f}\\n\".format(\n                    loss[\"photo_gen_loss\"],\n                    loss[\"monet_gen_loss\"],\n                    loss[\"photo_disc_loss\"],\n                    loss[\"monet_disc_loss\"],\n                )\n            )\n\n            if save_models and loss[\"monet_gen_loss\"] < best_val_loss:\n                best_val_loss = loss[\"monet_gen_loss\"]\n                model_path = os.path.join(out_folder, \"best_model.pth\")\n                save_model(model_path, i, self.m_gen, self.m_gen_optimizer, loss[\"monet_gen_loss\"])\n\n        if save_models:\n            model_path = os.path.join(out_folder, \"final_model.pth\")\n            save_model(model_path, epochs, self.m_gen, self.m_gen_optimizer, loss[\"monet_gen_loss\"])","metadata":{"_uuid":"6a67ccec-5f82-4b05-ae2b-32842133f13d","_cell_guid":"dacaa4dc-a2ef-4af1-9cfa-6770cfd3d4ab","collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2024-08-04T22:48:12.429175Z","iopub.execute_input":"2024-08-04T22:48:12.429709Z","iopub.status.idle":"2024-08-04T22:48:12.456899Z","shell.execute_reply.started":"2024-08-04T22:48:12.429683Z","shell.execute_reply":"2024-08-04T22:48:12.455992Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"DATA_SIZE = 300\n\nmonet_ds = CustomImageDataset(\"/kaggle/input/gan-getting-started/monet_jpg\", num_imgs=DATA_SIZE)\nmonet_data = torch.utils.data.DataLoader(monet_ds, shuffle=True)\n\nphoto_ds = CustomImageDataset(\"/kaggle/input/gan-getting-started/photo_jpg\", num_imgs=DATA_SIZE)\nphoto_data = torch.utils.data.DataLoader(photo_ds, shuffle=True)","metadata":{"_uuid":"e6fc9431-6138-4aec-aa56-2b8192f2da51","_cell_guid":"d394ae36-ea33-4067-a39e-50e9b3b79760","collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2024-08-04T22:48:12.457942Z","iopub.execute_input":"2024-08-04T22:48:12.458244Z","iopub.status.idle":"2024-08-04T22:48:12.478404Z","shell.execute_reply.started":"2024-08-04T22:48:12.458220Z","shell.execute_reply":"2024-08-04T22:48:12.477593Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"cycle_gan_model = CycleGAN(monet_generator=monet_generator,\n                           monet_discriminator=monet_discriminator,\n                           photo_generator=photo_generator,\n                           photo_discriminator=photo_discriminator).to(device)","metadata":{"_uuid":"992717db-457d-4037-8c86-fb62ee4821be","_cell_guid":"a3a6d412-322a-4152-917d-e77ebc424803","collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2024-08-04T22:48:12.479502Z","iopub.execute_input":"2024-08-04T22:48:12.479836Z","iopub.status.idle":"2024-08-04T22:48:12.626535Z","shell.execute_reply.started":"2024-08-04T22:48:12.479786Z","shell.execute_reply":"2024-08-04T22:48:12.625821Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"cycle_gan_model.train(monet_data, photo_data, epochs=25, save_models=True)","metadata":{"execution":{"iopub.status.busy":"2024-08-04T22:48:19.732981Z","iopub.execute_input":"2024-08-04T22:48:19.733615Z","iopub.status.idle":"2024-08-04T23:40:23.948120Z","shell.execute_reply.started":"2024-08-04T22:48:19.733585Z","shell.execute_reply":"2024-08-04T23:40:23.947294Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Visualize Monet-esque photos","metadata":{}},{"cell_type":"code","source":"import os\nimport torch\nfrom matplotlib import pyplot as plt\nimport numpy as np","metadata":{"execution":{"iopub.status.busy":"2024-08-05T16:57:45.069229Z","iopub.execute_input":"2024-08-05T16:57:45.069958Z","iopub.status.idle":"2024-08-05T16:57:45.074130Z","shell.execute_reply.started":"2024-08-05T16:57:45.069924Z","shell.execute_reply":"2024-08-05T16:57:45.073201Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"checkpts = os.path.join(os.getcwd(), \"checkpoints\")\nos.listdir(checkpts)","metadata":{"execution":{"iopub.status.busy":"2024-08-05T16:50:25.341848Z","iopub.execute_input":"2024-08-05T16:50:25.342658Z","iopub.status.idle":"2024-08-05T16:50:25.350387Z","shell.execute_reply.started":"2024-08-05T16:50:25.342622Z","shell.execute_reply":"2024-08-05T16:50:25.349394Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"model_path = os.path.join(os.getcwd(), \"checkpoints\", \"best_model.pth\")\nmodel_dict = torch.load(model_path)","metadata":{"execution":{"iopub.status.busy":"2024-08-05T16:50:26.287216Z","iopub.execute_input":"2024-08-05T16:50:26.287938Z","iopub.status.idle":"2024-08-05T16:50:27.531030Z","shell.execute_reply.started":"2024-08-05T16:50:26.287905Z","shell.execute_reply":"2024-08-05T16:50:27.530259Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"monet_generator = Generator() \nmonet_generator.load_state_dict(model_dict['model_state_dict'])\nmonet_generator.to(device)\nmonet_generator.eval()","metadata":{"execution":{"iopub.status.busy":"2024-08-05T16:56:05.243462Z","iopub.execute_input":"2024-08-05T16:56:05.243824Z","iopub.status.idle":"2024-08-05T16:56:07.579750Z","shell.execute_reply.started":"2024-08-05T16:56:05.243793Z","shell.execute_reply":"2024-08-05T16:56:07.578845Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"_, ax = plt.subplots(5, 2, figsize=(12, 12))\n\nfor i, img in enumerate(photo_data):\n    photo = img[0].permute(1,2,0).numpy()\n    photo = (photo * 127.5 + 127.5).astype(np.uint8)\n    \n    img = img.to(device)\n    pred = monet_generator(img)\n    pred = pred[0].permute(1,2,0).cpu().detach().numpy()\n    pred = (pred * 127.5 + 127.5).astype(np.uint8)\n    \n    ax[i, 0].imshow(photo)\n    ax[i, 1].imshow(pred)\n    ax[i, 0].set_title(\"Input Photo\")\n    ax[i, 1].set_title(\"Monet-esque\")\n    ax[i, 0].axis(\"off\")\n    ax[i, 1].axis(\"off\")\n    \n    if i == 4:\n        break\n    ","metadata":{"execution":{"iopub.status.busy":"2024-08-05T17:01:12.935017Z","iopub.execute_input":"2024-08-05T17:01:12.935400Z","iopub.status.idle":"2024-08-05T17:01:13.978648Z","shell.execute_reply.started":"2024-08-05T17:01:12.935371Z","shell.execute_reply":"2024-08-05T17:01:13.977710Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Create submission file","metadata":{}},{"cell_type":"code","source":"photo_ds = CustomImageDataset(\"/kaggle/input/gan-getting-started/photo_jpg\")\nphoto_data = torch.utils.data.DataLoader(photo_ds, shuffle=False)","metadata":{"execution":{"iopub.status.busy":"2024-08-05T17:04:37.087070Z","iopub.execute_input":"2024-08-05T17:04:37.087785Z","iopub.status.idle":"2024-08-05T17:04:37.112027Z","shell.execute_reply.started":"2024-08-05T17:04:37.087754Z","shell.execute_reply":"2024-08-05T17:04:37.111194Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import PIL\n! mkdir images","metadata":{"execution":{"iopub.status.busy":"2024-08-05T17:04:53.875543Z","iopub.execute_input":"2024-08-05T17:04:53.875896Z","iopub.status.idle":"2024-08-05T17:04:54.890148Z","shell.execute_reply.started":"2024-08-05T17:04:53.875867Z","shell.execute_reply":"2024-08-05T17:04:54.889088Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"i = 1\nfor img in tqdm(photo_data, ncols=100, desc=\"Progress\"):\n    img = img.to(device)\n    pred = monet_generator(img)\n    \n    pred = pred[0].permute(1, 2, 0).cpu().detach().numpy()\n    pred = (pred * 127.5 + 127.5).astype(np.uint8)\n    im = PIL.Image.fromarray(pred)\n    im.save(\"images/\" + str(i) + \".jpg\")\n    i += 1","metadata":{"execution":{"iopub.status.busy":"2024-08-05T17:08:33.213856Z","iopub.execute_input":"2024-08-05T17:08:33.214243Z","iopub.status.idle":"2024-08-05T17:11:04.475174Z","shell.execute_reply.started":"2024-08-05T17:08:33.214212Z","shell.execute_reply":"2024-08-05T17:11:04.474307Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import shutil\nshutil.make_archive(\"/kaggle/working/images\", \"zip\", \"/kaggle/working/images\")","metadata":{"execution":{"iopub.status.busy":"2024-08-05T17:15:33.810082Z","iopub.execute_input":"2024-08-05T17:15:33.810455Z","iopub.status.idle":"2024-08-05T17:15:37.331492Z","shell.execute_reply.started":"2024-08-05T17:15:33.810420Z","shell.execute_reply":"2024-08-05T17:15:37.330553Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}]}