{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.10.14","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"none","dataSources":[{"sourceId":21755,"databundleVersionId":1475600,"sourceType":"competition"},{"sourceId":1217826,"sourceType":"datasetVersion","datasetId":298806},{"sourceId":3956508,"sourceType":"datasetVersion","datasetId":850761}],"dockerImageVersionId":30787,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"import torch\nimport torch.nn as nn\nfrom PIL import Image\nimport os\n#!pip install config\nimport config\nimport numpy as np\nimport torch\nimport albumentations as A\nfrom albumentations.pytorch import ToTensorV2","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# **Configuration**","metadata":{}},{"cell_type":"code","source":"def save_checkpoint(model, optimizer, filename=\"my_checkpoint.pth.tar\"):\n    print(\"=> Saving checkpoint\")\n    checkpoint = {\n        \"state_dict\": model.state_dict(),\n        \"optimizer\": optimizer.state_dict(),\n    }\n    torch.save(checkpoint, filename)\n\n\ndef load_checkpoint(checkpoint_file, model, optimizer, lr):\n    print(\"=> Loading checkpoint\")\n    checkpoint = torch.load(checkpoint_file, map_location=config.DEVICE)\n    model.load_state_dict(checkpoint[\"state_dict\"])\n    optimizer.load_state_dict(checkpoint[\"optimizer\"])\n\n    # If we don't do this then it will just have learning rate of old checkpoint\n    # and it will lead to many hours of debugging \\:\n    for param_group in optimizer.param_groups:\n        param_group[\"lr\"] = lr\n\n\ndef seed_everything(seed=42):\n    os.environ[\"PYTHONHASHSEED\"] = str(seed)\n    random.seed(seed)\n    np.random.seed(seed)\n    torch.manual_seed(seed)\n    torch.cuda.manual_seed(seed)\n    torch.cuda.manual_seed_all(seed)\n    torch.backends.cudnn.deterministic = True\n    torch.backends.cudnn.benchmark = False\n\nDEVICE = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n#TRAIN_DIR = \"data/train\"\n#VAL_DIR = \"data/val\"\nBATCH_SIZE = 1\nLEARNING_RATE = 1e-5\nLAMBDA_IDENTITY = 0.0\nLAMBDA_CYCLE = 10\nNUM_WORKERS = 4\nNUM_EPOCHS = 100\nLOAD_MODEL = False\nSAVE_MODEL = True\nCHECKPOINT_GEN_H = \"genh.pth.tar\"\nCHECKPOINT_GEN_Z = \"genz.pth.tar\"\nCHECKPOINT_CRITIC_H = \"critich.pth.tar\"\nCHECKPOINT_CRITIC_Z = \"criticz.pth.tar\"\n\ntransforms = A.Compose(\n    [\n        A.Resize(width=256, height=256),\n        A.HorizontalFlip(p=0.5),\n        A.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5], max_pixel_value=255),\n        ToTensorV2(),\n    ],\n    additional_targets={\"image0\": \"image\"},\n)","metadata":{"execution":{"iopub.status.busy":"2024-11-05T21:01:21.292453Z","iopub.execute_input":"2024-11-05T21:01:21.293147Z","iopub.status.idle":"2024-11-05T21:01:21.309443Z","shell.execute_reply.started":"2024-11-05T21:01:21.293099Z","shell.execute_reply":"2024-11-05T21:01:21.307871Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# **Generator**","metadata":{}},{"cell_type":"code","source":"class ConvBlock(nn.Module):\n    def __init__(self,in_channels,out_channels,down=True,use_act=True,**kwargs):\n        super().__init__\n        self.conv=nn.Sequential(\n            Conv2d(in_channels,\n            out_channels, padding_mode='reflect', **kwargs ) if down else nn.Conv2dTranspose2d(in_channels,out_channels,**kwargs),\n            nn.InstanceNorm2d(out_channels),\n            nn.ReLU(inplace=True) if use_act else nn.Identity()\n        )\n        def forward(self,x):\n            return self.conv(x)\nclass ResidualBlock(nn.Module):\n    def __init__(self,channels):\n        self.Block=nn.Sequential(\n            ConvBlock(channels,channels,kernel_size=3,padding=1),\n            ConvBlock(channels,channels,use_act=False,kernel_size=3,padding=1)   \n        )\n        def forward(self,x):\n            return x+conv(x)\nclass Generator(nn.Module):\n    def __init__(self,img_channels,num_features=64,num_residual=9):\n        super().__init__()\n        self.initail=nn.Sequential(\n            nn.Conv2d(img_channels,num_features,kernel_size=7,stride=1,padding=3,padding_mode='reflect'),\n            nn.ReLU(inplace=True)\n        )\n        self.down_blocks=nn.ModuleList([\n            ConvBlock(num_features,num_features*2,kernel_size=3,padding=1,stride=2),\n            ConvBlock(num_features*2,num_features*4,kernel_size=3,padding=1,stride=2)\n        ])\n        self.residual_blocks=nn.Sequential(\n            *[ResidualBlock(num_features*4 for _ in range (num_residual))]\n        )\n        self.up_blocks=nn.ModuleList([\n            ConvBlock(num_features*4,num_features*2,down=False,kernel_size=3,padding=1,stride=2),\n            ConvBlock(num_features*2,num_features*1,down=False,kernel_size=3,padding=1,stride=2)\n        ])\n        self.last_layer=nn.Conv2d(num_features,img_channels,kernel_size=7,stride=1,padding=3,padding_mode='reflect')\n        \n        def forward(self,x):\n            x=self.initail(x)\n            for layer in self.down_blocks:\n                x= layer(x)\n            x=self.residual_blocks(x)\n            for layer in self.up_blocks:\n                x= layer(x)\n            return torch.tanh(last_layer(x))","metadata":{"execution":{"iopub.status.busy":"2024-11-05T20:42:02.262186Z","iopub.execute_input":"2024-11-05T20:42:02.263497Z","iopub.status.idle":"2024-11-05T20:42:02.285610Z","shell.execute_reply.started":"2024-11-05T20:42:02.263445Z","shell.execute_reply":"2024-11-05T20:42:02.284179Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# **Discriminator**","metadata":{}},{"cell_type":"code","source":"\nclass Block(nn.Module):\n    def __init__(self , in_channels, out_channels, stride):\n        super().__init__()\n        self.conv=nn.Sequential(\n            nn.Conv2d(in_channels,out_channels,4,stride,1,bias=True,padding_mode='reflect'),\n            nn.InstanceNorm2d(out_channels),\n            nn.LeakyReLU(0.2)\n        )\n        def forward(self,x):\n            return self.conv(x)\nclass Discriminator(nn.Module):\n    def __init__(self,in_channels=3,features=[64,128,256,512]):\n        super().__init__()\n        self.initial= nn.Sequential(\n        nn.Conv2d(\n            in_channels,\n            features[0],\n            kernerl_size=4,\n            stride=2,\n            padding=1,\n            padding_mode='reflect'\n        ),   \n            nn.LeakyReLU(0.2)\n        )\n        \n        layers=[]\n        in_channels=feature[0]\n        for feature in features[1:]:\n            layers.append(Block(in_channels,feature,stride=1 if feature==features[-1] else 2))\n            in_channels=feature\n        layers.append(nn.Conv2d(in_channels,1,kernel_size=4,stride=1,padding=1,padding_mode='reflect'))\n        self.Model=(nn.Sequential(*layers))\n    \n    def forward(self,x):\n        x=self.initial(x)\n        return torch.sigmoid(self.Model(x))","metadata":{"execution":{"iopub.status.busy":"2024-11-05T20:42:05.446124Z","iopub.execute_input":"2024-11-05T20:42:05.446595Z","iopub.status.idle":"2024-11-05T20:42:05.459468Z","shell.execute_reply.started":"2024-11-05T20:42:05.446547Z","shell.execute_reply":"2024-11-05T20:42:05.458383Z"},"trusted":true},"execution_count":null,"outputs":[]}]}