Add train code with dataset.sh

Pix2Pix.ipynb

@@ -0,0 +1,375 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 10,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import torch\n",
+    "from torch import nn\n",
+    "from tqdm.auto import tqdm\n",
+    "import torchvision\n",
+    "from torchvision import transforms\n",
+    "from torchvision.utils import make_grid\n",
+    "from torch.utils.data import DataLoader\n",
+    "import matplotlib.pyplot as plt\n",
+    "torch.manual_seed(0)\n",
+    "\n",
+    "def show_tensor_images(image_tensor, num_images=25, size=(1, 28, 28)):\n",
+    "    '''\n",
+    "    Function for visualizing images: Given a tensor of images, number of images, and\n",
+    "    size per image, plots and prints the images in an uniform grid.\n",
+    "    '''\n",
+    "    image_shifted = image_tensor\n",
+    "    image_unflat = image_shifted.detach().cpu().view(-1, *size)\n",
+    "    image_grid = make_grid(image_unflat[:num_images], nrow=5)\n",
+    "    plt.imshow(image_grid.permute(1, 2, 0).squeeze())\n",
+    "    plt.show()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def crop(image, new_shape):\n",
+    "    middle_height = image.shape[2] // 2\n",
+    "    middle_width = image.shape[3] // 2\n",
+    "    starting_height = middle_height - new_shape[2] // 2\n",
+    "    final_height = starting_height + new_shape[2]\n",
+    "    starting_width = middle_width - new_shape[3] // 2\n",
+    "    final_width = starting_width + new_shape[3]\n",
+    "    cropped_image = image[:, :, starting_height:final_height, starting_width:final_width]    \n",
+    "    return cropped_image"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "class Encoder(nn.Module):\n",
+    "    def  __init__(self,input_channels, use_bn=True, use_dropout = False):\n",
+    "        super(Encoder, self).__init__()\n",
+    "        self.conv1 = nn.Conv2d(input_channels, 2 * input_channels, kernel_size=3)\n",
+    "        self.conv2 = nn.Conv2d(input_channels * 2, 2 * input_channels, kernel_size = 3)\n",
+    "        self.activation = nn.ReLU()\n",
+    "        self.pooling = nn.MaxPool2d(kernel_size=2, stride=2)\n",
+    "        if use_bn:\n",
+    "            self.bn = nn.BatchNorm2d(input_channels * 2)\n",
+    "        self.use_bn = use_bn\n",
+    "        if use_dropout:\n",
+    "            self.dropout = nn.Dropout()\n",
+    "        self.use_dp = use_dropout\n",
+    "\n",
+    "    def forward(self, x):\n",
+    "        x = self.conv1(x)\n",
+    "        if self.use_bn:\n",
+    "            x = self.bn(x)\n",
+    "        if self.use_dp:\n",
+    "            x = self.dropout(x)\n",
+    "        x = self.activation(x)\n",
+    "        x = self.conv2(x)\n",
+    "        if self.use_bn:\n",
+    "            x = self.bn(x)\n",
+    "        if self.use_dp:\n",
+    "            x = self.dropout(x)\n",
+    "        x = self.activation(x)\n",
+    "        x = self.pooling(x)\n",
+    "        return x"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "class Decoder(nn.Module):\n",
+    "    def __init__(self, input_channels, use_bn=True, use_dropout = False):\n",
+    "        super(Decoder, self).__init__()\n",
+    "        self.upsample = nn.Upsample(scale_factor = 2, mode='bilinear', align_corners=True)\n",
+    "        self.conv1 = nn.Conv2d(input_channels, input_channels // 2, kernel_size = 2)\n",
+    "        self.conv2 = nn.Conv2d(input_channels , input_channels //2 , kernel_size = 3,padding=1)\n",
+    "        self.conv3 = nn.Conv2d(input_channels //2, input_channels // 2, kernel_size = 2,padding=1)\n",
+    "        self.activation = nn.ReLU()\n",
+    "        if use_bn:\n",
+    "            self.bn = nn.BatchNorm2d(input_channels // 2)\n",
+    "        self.use_bn = use_bn\n",
+    "        if use_dropout:\n",
+    "            self.dropout = nn.Dropout()\n",
+    "        self.use_dp = use_dropout\n",
+    "        \n",
+    "\n",
+    "    def forward(self, x, skip_con_x):\n",
+    "        x = self.upsample(x)\n",
+    "        x = self.conv1(x)\n",
+    "        skip_con_x = crop(skip_con_x, x.shape)\n",
+    "        x = torch.cat([x, skip_con_x], axis = 1)\n",
+    "        x = self.conv2(x)\n",
+    "        if self.use_bn:\n",
+    "            x = self.bn(x)\n",
+    "        if self.use_dp:\n",
+    "            x = self.dropout(x)\n",
+    "        x = self.activation(x)\n",
+    "        x = self.conv3(x)\n",
+    "        if self.use_bn:\n",
+    "            x = self.bn(x)\n",
+    "        if self.use_dp:\n",
+    "            x = self.dropout(x)\n",
+    "        x = self.activation(x)\n",
+    "        return x\n",
+    "        "
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "class FeatureMapBlock(nn.Module):\n",
+    "    def __init__(self, input_channels, output_channels):\n",
+    "        super(FeatureMapBlock, self).__init__()\n",
+    "        self.conv = nn.Conv2d(input_channels, output_channels, kernel_size=1)\n",
+    "    def forward(self, x):\n",
+    "        x = self.conv(x)\n",
+    "        return x"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "class UNet(nn.Module):\n",
+    "    def __init__(self, input_channels, output_channels, hidden_channels = 64):\n",
+    "        super(UNet, self).__init__()\n",
+    "        self.upfeature = FeatureMapBlock(input_channels, hidden_channels)\n",
+    "        self.encoder1 = Encoder(hidden_channels, use_dropout=True)\n",
+    "        self.encoder2 = Encoder(hidden_channels * 2,use_dropout=True)\n",
+    "        self.encoder3 = Encoder(hidden_channels * 4, use_dropout=True)\n",
+    "        self.encoder4 = Encoder(hidden_channels * 8)\n",
+    "        self.encoder5 = Encoder(hidden_channels * 16)\n",
+    "        self.encoder6 = Encoder(hidden_channels * 32)\n",
+    "        self.decoder1 = Decoder(hidden_channels * 64)\n",
+    "        self.decoder2 = Decoder(hidden_channels * 32)\n",
+    "        self.decoder3 = Decoder(hidden_channels * 16)\n",
+    "        self.decoder4 = Decoder(hidden_channels * 8)\n",
+    "        self.decoder5 = Decoder(hidden_channels * 4 )\n",
+    "        self.decoder6 = Decoder(hidden_channels * 2)\n",
+    "        self.downfeature = FeatureMapBlock(hidden_channels , output_channels)\n",
+    "        self.sigmoid = nn.Sigmoid()\n",
+    "\n",
+    "\n",
+    "    def forward(self, x):\n",
+    "        x0 = self.upfeature(x)\n",
+    "        x1 = self.encoder1(x0)\n",
+    "        x2 = self.encoder2(x1)\n",
+    "        x3 = self.encoder3(x2)\n",
+    "        x4 = self.encoder4(x3)\n",
+    "        x5 = self.encoder5(x4)\n",
+    "        x6 = self.encoder6(x5)\n",
+    "        x7 = self.decoder1(x6, x5)\n",
+    "        x8 = self.decoder2(x7, x4)\n",
+    "        x9 = self.decoder3(x8, x3)\n",
+    "        x10 = self.decoder4(x9, x2)\n",
+    "        x11 = self.decoder5(x10, x1)\n",
+    "        x12 = self.decoder6(x11, x0)\n",
+    "        xn = self.downfeature(x12)\n",
+    "\n",
+    "        return self.sigmoid(xn)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "class Discriminator(nn.Module):\n",
+    "    def __init__(self, input_channels, hidden_channels=8) -> None:\n",
+    "        super(Discriminator, self).__init__()\n",
+    "        self.upfeature = FeatureMapBlock(input_channels,hidden_channels)\n",
+    "        self.encoder1 = Encoder(hidden_channels, use_bn=False)\n",
+    "        self.encoder2 = Encoder(hidden_channels*2)\n",
+    "        self.encoder3 = Encoder(hidden_channels*4)\n",
+    "        self.encoder4 = Encoder(hidden_channels*8)\n",
+    "        self.final = nn.Conv2d(hidden_channels * 16,1,kernel_size=1)\n",
+    "\n",
+    "    def forward(self, x,y):\n",
+    "        x = torch.cat([x,y],axis=1)\n",
+    "        x0 = self.upfeature(x)\n",
+    "        x1 = self.encoder1(x0)\n",
+    "        x2 = self.encoder2(x1)\n",
+    "        x3 = self.encoder3(x2)\n",
+    "        x4 = self.encoder4(x3)\n",
+    "        xn = self.final(x4)\n",
+    "\n",
+    "        return xn\n",
+    "\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 9,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import torch.nn.functional as F\n",
+    "# New parameters\n",
+    "adv_criterion = nn.BCEWithLogitsLoss() \n",
+    "recon_criterion = nn.L1Loss() \n",
+    "lambda_recon = 200\n",
+    "\n",
+    "n_epochs = 20\n",
+    "input_dim = 3\n",
+    "real_dim = 3\n",
+    "display_step = 200\n",
+    "batch_size = 4\n",
+    "lr = 0.0002\n",
+    "target_shape = 256\n",
+    "device = 'cuda'"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "transform = transforms.Compose([\n",
+    "    transforms.ToTensor(),\n",
+    "])\n",
+    "\n",
+    "dataset = torchvision.datasets.ImageFolder(\"/datsets/maps\", transform=transform)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 12,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "gen = UNet(input_dim, real_dim).to(device)\n",
+    "gen_opt = torch.optim.Adam(gen.parameters(),lr=lr)\n",
+    "disc = Discriminator(input_dim+real_dim).to(device)\n",
+    "disc_opt = torch.optim.Adam(disc.parameters(),lr=lr)\n",
+    "\n",
+    "def weights_init(m):\n",
+    "    if isinstance(m,nn.Conv2d) or isinstance(m,nn.ConvTranspose2d):\n",
+    "        torch.nn.init.normal_(m.weight,0.0,0.02)\n",
+    "    if isinstance(m,nn.BatchNorm2d):\n",
+    "        torch.nn.init.normal_(m.weight,0.0,0.02)\n",
+    "        torch.nn.init.constant_(m.bias,0.0)\n",
+    "\n",
+    "gen =gen.apply(weights_init)\n",
+    "disc = disc.apply(weights_init)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 13,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def get_gen_loss(gen, disc, real, condition, dev_criterion, recon_criterion, lambda_recon):\n",
+    "    gen_pred = gen(condition)\n",
+    "    disc_pred = disc(gen_pred, real)\n",
+    "    gen_adv_loss = dev_criterion(disc_pred, torch.ones_like(disc_pred))\n",
+    "    gen_recon_loss = recon_criterion(real,gen_pred)\n",
+    "    gen_loss = gen_adv_loss + lambda_recon * gen_recon_loss\n",
+    "\n",
+    "    return gen_loss"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def train():\n",
+    "    mean_gen_loss = 0\n",
+    "    mean_disc_loss = 0\n",
+    "    dataloader = DataLoader(dataset,batch_size=batch_size,shuffle=True)\n",
+    "\n",
+    "    cur_step = 0\n",
+    "    for epoch in range(n_epochs):\n",
+    "        for image,_ in tqdm(dataloader):\n",
+    "            image_width = image.shape[3]\n",
+    "            condition = image[:, :, :, :image_width // 2]\n",
+    "            condition = nn.functional.interpolate(condition, size=target_shape)\n",
+    "            real = image[:, :, :, image_width // 2:]\n",
+    "            real = nn.functional.interpolate(real, size=target_shape)\n",
+    "            cur_batch_size = len(real)\n",
+    "            condition = condition.to(device)\n",
+    "            real = real.to(device)\n",
+    "\n",
+    "            disc_opt.zero_grad()\n",
+    "            with torch.no_grad():\n",
+    "                fake = gen(condition)\n",
+    "            disc_fake_pred = disc(fake.detach(),condition)\n",
+    "            disc_fake_loss = adv_criterion(disc_fake_pred,torch.zeros_like(disc_fake_pred))\n",
+    "            disc_real_pred = disc(real,condition)\n",
+    "            disc_real_loss = adv_criterion(disc_real_pred,torch.ones_like(disc_real_pred))\n",
+    "            disc_loss = (disc_fake_loss + disc_real_loss) / 2\n",
+    "            disc_loss.backward(retain_graph= True)\n",
+    "            disc_opt.step()\n",
+    "\n",
+    "            gen.zero_grad()\n",
+    "            gen_loss = get_gen_loss(gen,disc,real,condition,adv_criterion,recon_criterion,lambda_recon)\n",
+    "            gen_loss.backward()\n",
+    "            gen_opt.step()\n",
+    "\n",
+    "            mean_gen_loss += gen_loss.item() / display_step\n",
+    "            mean_disc_loss += disc_loss.item() / display_step\n",
+    "\n",
+    "            if cur_step % display_step == 0:\n",
+    "                print(f\"Epoch: {epoch} | Step: {cur_step} | Gen-Loss: {mean_gen_loss} | Disc-loss: {mean_disc_loss}\")\n",
+    "                show_tensor_images(condition, size=(input_dim, target_shape, target_shape))\n",
+    "                show_tensor_images(real, size=(real_dim, target_shape, target_shape))\n",
+    "                show_tensor_images(fake, size=(real_dim, target_shape, target_shape))\n",
+    "                mean_gen_loss = 0\n",
+    "                mean_disc_loss = 0\n",
+    "\n",
+    "                torch.save({\n",
+    "                    \"gen\":gen,\n",
+    "                    \"disc\":disc,\n",
+    "                    \"gen_opt\": gen_opt,\n",
+    "                    \"disc_opt\": disc_opt\n",
+    "                },f\"checkpoints/Pix2Pix_Epoch{epoch}.pth\")\n",
+    "            cur_step += 1\n",
+    "\n",
+    "train()\n",
+    "\n",
+    "            "
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.12.3"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

UNet.py

@@ -0,0 +1,141 @@
+import torch.nn as nn
+from utils import *
+
+
+
+class Encoder(nn.Module):
+    def  __init__(self,input_channels, use_bn=True, use_dropout = False):
+        super(Encoder, self).__init__()
+        self.conv1 = nn.Conv2d(input_channels, 2 * input_channels, kernel_size=3)
+        self.conv2 = nn.Conv2d(input_channels * 2, 2 * input_channels, kernel_size = 3)
+        self.activation = nn.ReLU()
+        self.pooling = nn.MaxPool2d(kernel_size=2, stride=2)
+        if use_bn:
+            self.bn = nn.BatchNorm2d(input_channels * 2)
+        self.use_bn = use_bn
+        if use_dropout:
+            self.dropout = nn.Dropout()
+        self.use_dp = use_dropout
+
+    def forward(self, x):
+        x = self.conv1(x)
+        if self.use_bn:
+            x = self.bn(x)
+        if self.use_dp:
+            x = self.dropout(x)
+        x = self.activation(x)
+        x = self.conv2(x)
+        if self.use_bn:
+            x = self.bn(x)
+        if self.use_dp:
+            x = self.dropout(x)
+        x = self.activation(x)
+        x = self.pooling(x)
+        return x
+    
+
+class Decoder(nn.Module):
+    def __init__(self, input_channels, use_bn=True, use_dropout = False):
+        super(Decoder, self).__init__()
+        self.upsample = nn.Upsample(scale_factor = 2, mode='bilinear', align_corners=True)
+        self.conv1 = nn.Conv2d(input_channels, input_channels // 2, kernel_size = 2)
+        self.conv2 = nn.Conv2d(input_channels , input_channels //2 , kernel_size = 3,padding=1)
+        self.conv3 = nn.Conv2d(input_channels //2, input_channels // 2, kernel_size = 2,padding=1)
+        self.activation = nn.ReLU()
+        if use_bn:
+            self.bn = nn.BatchNorm2d(input_channels // 2)
+        self.use_bn = use_bn
+        if use_dropout:
+            self.dropout = nn.Dropout()
+        self.use_dp = use_dropout
+        
+
+    def forward(self, x, skip_con_x):
+        x = self.upsample(x)
+        x = self.conv1(x)
+        skip_con_x = crop(skip_con_x, x.shape)
+        x = torch.cat([x, skip_con_x], axis = 1)
+        x = self.conv2(x)
+        if self.use_bn:
+            x = self.bn(x)
+        if self.use_dp:
+            x = self.dropout(x)
+        x = self.activation(x)
+        x = self.conv3(x)
+        if self.use_bn:
+            x = self.bn(x)
+        if self.use_dp:
+            x = self.dropout(x)
+        x = self.activation(x)
+        return x
+        
+
+class FeatureMapBlock(nn.Module):
+    def __init__(self, input_channels, output_channels):
+        super(FeatureMapBlock, self).__init__()
+        self.conv = nn.Conv2d(input_channels, output_channels, kernel_size=1)
+    def forward(self, x):
+        x = self.conv(x)
+        return x
+    
+
+class UNet(nn.Module):
+    def __init__(self, input_channels, output_channels, hidden_channels = 64):
+        super(UNet, self).__init__()
+        self.upfeature = FeatureMapBlock(input_channels, hidden_channels)
+        self.encoder1 = Encoder(hidden_channels, use_dropout=True)
+        self.encoder2 = Encoder(hidden_channels * 2,use_dropout=True)
+        self.encoder3 = Encoder(hidden_channels * 4, use_dropout=True)
+        self.encoder4 = Encoder(hidden_channels * 8)
+        self.encoder5 = Encoder(hidden_channels * 16)
+        self.encoder6 = Encoder(hidden_channels * 32)
+        self.decoder1 = Decoder(hidden_channels * 64)
+        self.decoder2 = Decoder(hidden_channels * 32)
+        self.decoder3 = Decoder(hidden_channels * 16)
+        self.decoder4 = Decoder(hidden_channels * 8)
+        self.decoder5 = Decoder(hidden_channels * 4 )
+        self.decoder6 = Decoder(hidden_channels * 2)
+        self.downfeature = FeatureMapBlock(hidden_channels , output_channels)
+        self.sigmoid = nn.Sigmoid()
+
+
+    def forward(self, x):
+        x0 = self.upfeature(x)
+        x1 = self.encoder1(x0)
+        x2 = self.encoder2(x1)
+        x3 = self.encoder3(x2)
+        x4 = self.encoder4(x3)
+        x5 = self.encoder5(x4)
+        x6 = self.encoder6(x5)
+        x7 = self.decoder1(x6, x5)
+        x8 = self.decoder2(x7, x4)
+        x9 = self.decoder3(x8, x3)
+        x10 = self.decoder4(x9, x2)
+        x11 = self.decoder5(x10, x1)
+        x12 = self.decoder6(x11, x0)
+        xn = self.downfeature(x12)
+
+        return self.sigmoid(xn)
+    
+
+class Discriminator(nn.Module):
+    def __init__(self, input_channels, hidden_channels=8) -> None:
+        super(Discriminator, self).__init__()
+        self.upfeature = FeatureMapBlock(input_channels,hidden_channels)
+        self.encoder1 = Encoder(hidden_channels, use_bn=False)
+        self.encoder2 = Encoder(hidden_channels*2)
+        self.encoder3 = Encoder(hidden_channels*4)
+        self.encoder4 = Encoder(hidden_channels*8)
+        self.final = nn.Conv2d(hidden_channels * 16,1,kernel_size=1)
+
+    def forward(self, x,y):
+        x = torch.cat([x,y],axis=1)
+        x0 = self.upfeature(x)
+        x1 = self.encoder1(x0)
+        x2 = self.encoder2(x1)
+        x3 = self.encoder3(x2)
+        x4 = self.encoder4(x3)
+        xn = self.final(x4)
+
+        return xn
+

train.py

@@ -0,0 +1,95 @@
+import torch.nn.functional as F
+import torch.nn as nn
+from torchvision import transforms
+import torchvision
+from torch.utils.data import DataLoader
+from UNet import *
+from tqdm.auto import tqdm
+
+adv_criterion = nn.BCEWithLogitsLoss() 
+recon_criterion = nn.L1Loss() 
+lambda_recon = 200
+
+n_epochs = 50
+input_dim = 3
+real_dim = 3
+display_step = 200
+batch_size = 4
+lr = 0.0002
+target_shape = 256
+device = 'cuda'
+
+
+transform = transforms.Compose([
+    transforms.ToTensor(),
+])
+
+dataset = torchvision.datasets.ImageFolder("/datasets/maps", transform=transform)
+
+
+gen = UNet(input_dim, real_dim).to(device)
+gen_opt = torch.optim.Adam(gen.parameters(),lr=lr)
+disc = Discriminator(input_dim+real_dim).to(device)
+disc_opt = torch.optim.Adam(disc.parameters(),lr=lr)
+
+def weights_init(m):
+    if isinstance(m,nn.Conv2d) or isinstance(m,nn.ConvTranspose2d):
+        torch.nn.init.normal_(m.weight,0.0,0.02)
+    if isinstance(m,nn.BatchNorm2d):
+        torch.nn.init.normal_(m.weight,0.0,0.02)
+        torch.nn.init.constant_(m.bias,0.0)
+
+gen =gen.apply(weights_init)
+disc = disc.apply(weights_init)
+
+
+mean_gen_loss = 0
+mean_disc_loss = 0
+dataloader = DataLoader(dataset,batch_size=batch_size,shuffle=True)
+
+cur_step = 0
+for epoch in range(n_epochs):
+    for image,_ in tqdm(dataloader):
+        image_width = image.shape[3]
+        condition = image[:, :, :, :image_width // 2]
+        condition = nn.functional.interpolate(condition, size=target_shape)
+        real = image[:, :, :, image_width // 2:]
+        real = nn.functional.interpolate(real, size=target_shape)
+        cur_batch_size = len(real)
+        condition = condition.to(device)
+        real = real.to(device)
+
+        disc_opt.zero_grad()
+        with torch.no_grad():
+            fake = gen(condition)
+        disc_fake_pred = disc(fake.detach(),condition)
+        disc_fake_loss = adv_criterion(disc_fake_pred,torch.zeros_like(disc_fake_pred))
+        disc_real_pred = disc(real,condition)
+        disc_real_loss = adv_criterion(disc_real_pred,torch.ones_like(disc_real_pred))
+        disc_loss = (disc_fake_loss + disc_real_loss) / 2
+        disc_loss.backward(retain_graph= True)
+        disc_opt.step()
+
+        gen.zero_grad()
+        gen_loss = get_gen_loss(gen,disc,real,condition,adv_criterion,recon_criterion,lambda_recon)
+        gen_loss.backward()
+        gen_opt.step()
+
+        mean_gen_loss += gen_loss.item() / display_step
+        mean_disc_loss += disc_loss.item() / display_step
+
+        if cur_step % display_step == 0:
+            print(f"Epoch: {epoch} | Step: {cur_step} | Gen-Loss: {mean_gen_loss} | Disc-loss: {mean_disc_loss}")
+            show_tensor_images(condition,epoch,cur_step,"condition", size=(input_dim, target_shape, target_shape))
+            show_tensor_images(real, epoch,cur_step,"real",size=(real_dim, target_shape, target_shape))
+            show_tensor_images(fake, epoch,cur_step,"generated",size=(real_dim, target_shape, target_shape))
+            mean_gen_loss = 0
+            mean_disc_loss = 0
+
+            torch.save({
+                "gen":gen,
+                "disc":disc,
+                "gen_opt": gen_opt,
+                "disc_opt": disc_opt
+            },f"checkpoints/Pix2Pix_Epoch{epoch}.pth")
+        cur_step += 1

utils.py

@@ -0,0 +1,34 @@
+import os
+import matplotlib.pyplot as plt
+from torchvision.utils import make_grid
+import torch
+
+def show_tensor_images(image_tensor, epoch,step,otype,num_images=25, size=(1, 28, 28)):
+    image_shifted = image_tensor
+    image_unflat = image_shifted.detach().cpu().view(-1, *size)
+    image_grid = make_grid(image_unflat[:num_images], nrow=5)
+    if not os.path.exists(f"/outputs/Epoch{epoch}"):
+        os.makedirs(f"/outputs/Epoch{epoch}")
+    plt.savefig(os.path.join(f"/outputs/Epoch{epoch}_step_{step}_{otype}"))
+    plt.close()
+
+
+def crop(image, new_shape):
+    middle_height = image.shape[2] // 2
+    middle_width = image.shape[3] // 2
+    starting_height = middle_height - new_shape[2] // 2
+    final_height = starting_height + new_shape[2]
+    starting_width = middle_width - new_shape[3] // 2
+    final_width = starting_width + new_shape[3]
+    cropped_image = image[:, :, starting_height:final_height, starting_width:final_width]    
+    return cropped_image
+
+
+def get_gen_loss(gen, disc, real, condition, dev_criterion, recon_criterion, lambda_recon):
+    gen_pred = gen(condition)
+    disc_pred = disc(gen_pred, real)
+    gen_adv_loss = dev_criterion(disc_pred, torch.ones_like(disc_pred))
+    gen_recon_loss = recon_criterion(real,gen_pred)
+    gen_loss = gen_adv_loss + lambda_recon * gen_recon_loss
+
+    return gen_loss