File size: 35,028 Bytes
c59452d |
1 2 |
{"cells":[{"cell_type":"markdown","metadata":{},"source":["## DOWNLOADING DATASETS"]},{"cell_type":"code","execution_count":null,"metadata":{"trusted":true},"outputs":[],"source":["!wget https://download.visinf.tu-darmstadt.de/data/from_games/data/01_images.zip\n","!wget https://download.visinf.tu-darmstadt.de/data/from_games/data/01_labels.zip\n","!unzip /kaggle/working/01_images.zip\n","!unzip /kaggle/working/01_labels.zip\n","!rm /kaggle/working/01_images.zip\n","!rm /kaggle/working/01_labels.zip"]},{"cell_type":"code","execution_count":null,"metadata":{"trusted":true},"outputs":[],"source":["!wget https://download.visinf.tu-darmstadt.de/data/from_games/data/02_images.zip\n","!wget https://download.visinf.tu-darmstadt.de/data/from_games/data/02_labels.zip\n","!unzip /kaggle/working/02_images.zip\n","!unzip /kaggle/working/02_labels.zip\n","!rm /kaggle/working/02_images.zip\n","!rm /kaggle/working/02_labels.zip"]},{"cell_type":"markdown","metadata":{},"source":["## IN-MEMORY DOWNLOAD AND ZIPPING: REDUCING SECONDARY STORAGE FOOTPRINT & OVERHEAD"]},{"cell_type":"code","execution_count":null,"metadata":{"trusted":true},"outputs":[],"source":["import requests\n","import zipfile\n","import io\n","\n","# URL of the zip file\n","url = 'https://download.visinf.tu-darmstadt.de/data/from_games/data/03_labels.zip'\n","\n","# Create a streaming GET request\n","response = requests.get(url, stream=True)\n","\n","# Create an in-memory bytes buffer for the zip file\n","zip_buffer = io.BytesIO()\n","\n","# Download the file in chunks and write to the in-memory buffer\n","for chunk in response.iter_content(chunk_size=1024):\n"," if chunk:\n"," zip_buffer.write(chunk)\n","\n","# Seek to the beginning of the buffer\n","zip_buffer.seek(0)\n","\n","# Open the zip file in memory\n","with zipfile.ZipFile(zip_buffer, 'r') as zip_ref:\n"," # Extract all files directly to your desired directory\n"," zip_ref.extractall('/kaggle/working/')\n","\n","# No need to explicitly delete the zip file, it's only in memory"]},{"cell_type":"code","execution_count":null,"metadata":{"trusted":true},"outputs":[],"source":["import os\n","\n","_, _, files = next(os.walk(\"/kaggle/working/labels/\"))\n","print(len(files))"]},{"cell_type":"code","execution_count":null,"metadata":{"trusted":true},"outputs":[],"source":["import torch\n","import sys\n","print('__Python VERSION:', sys.version)\n","print('__pyTorch VERSION:', torch.__version__)\n","print('__CUDA VERSION')\n","from subprocess import call\n","# call([\"nvcc\", \"--version\"]) does not work\n","! nvcc --version\n","print('__CUDNN VERSION:', torch.backends.cudnn.version())\n","print('__Number CUDA Devices:', torch.cuda.device_count())\n","print('__Devices')\n","call([\"nvidia-smi\", \"--format=csv\", \"--query-gpu=index,name,driver_version,memory.total,memory.used,memory.free\"])\n","print('Active CUDA Device: GPU', torch.cuda.current_device())\n","print ('Available devices ', torch.cuda.device_count())\n","print ('Current cuda device ', torch.cuda.current_device())"]},{"cell_type":"markdown","metadata":{},"source":["## DEFINING STYLE TRANSFER VISION TRANSFORMER ARCHITECTURE"]},{"cell_type":"code","execution_count":143,"metadata":{"_kg_hide-output":true,"execution":{"iopub.execute_input":"2024-09-21T23:16:41.009559Z","iopub.status.busy":"2024-09-21T23:16:41.008700Z","iopub.status.idle":"2024-09-21T23:16:41.019269Z","shell.execute_reply":"2024-09-21T23:16:41.017876Z","shell.execute_reply.started":"2024-09-21T23:16:41.009514Z"},"trusted":true},"outputs":[],"source":["import torch\n","import torch.nn as nn\n","import torch.nn.functional as F\n","import math\n","\n","class PatchEmbedding(nn.Module):\n"," def __init__(self, in_channels=3, patch_size=8, emb_dim=768, img_size=256):\n"," super(PatchEmbedding, self).__init__()\n"," self.patch_size = patch_size\n"," num_patches = (img_size // patch_size) ** 2\n"," self.proj = nn.Conv2d(in_channels, emb_dim, kernel_size=patch_size, stride=patch_size)\n"," self.positional_encoding = nn.Parameter(torch.zeros(1, num_patches, emb_dim))\n","\n"," def forward(self, x):\n"," # Shape of x: (batch_size, channels, height, width)\n"," batch_size = x.shape[0]\n"," x = self.proj(x) # (batch_size, emb_dim, H/P, W/P)\n"," x = x.flatten(2) # Flatten spatial dimensions\n"," x = x.transpose(1, 2) # (batch_size, num_patches, emb_dim)\n"," x += self.positional_encoding # Add positional encoding\n"," return x"]},{"cell_type":"code","execution_count":144,"metadata":{"execution":{"iopub.execute_input":"2024-09-21T23:16:41.036391Z","iopub.status.busy":"2024-09-21T23:16:41.035657Z","iopub.status.idle":"2024-09-21T23:16:41.055122Z","shell.execute_reply":"2024-09-21T23:16:41.053886Z","shell.execute_reply.started":"2024-09-21T23:16:41.036334Z"},"trusted":true},"outputs":[],"source":["class TransformerEncoderBlock(nn.Module):\n"," def __init__(self, emb_dim=768, num_heads=8, hidden_dim=2048, dropout=0.1):\n"," super(TransformerEncoderBlock, self).__init__()\n"," self.attn = LocationBasedMultiheadAttention(emb_dim, num_heads, dropout=dropout)\n"," self.ff = nn.Sequential(\n"," nn.Linear(emb_dim, hidden_dim),\n"," nn.ReLU(),\n"," nn.Linear(hidden_dim, emb_dim),\n"," )\n"," self.norm1 = nn.LayerNorm(emb_dim)\n"," self.norm2 = nn.LayerNorm(emb_dim)\n"," self.adain = AdaIN(emb_dim)\n"," self.dropout = nn.Dropout(dropout)\n","\n"," def forward(self, x, style):\n"," attn_output = self.attn(x, x, x)\n"," x = x + self.dropout(attn_output)\n"," x = self.norm1(x)\n"," x = self.adain(x, style)\n"," intermediate = x # This is the intermediate representation for the skip connection\n"," ff_output = self.ff(x)\n"," x = x + self.dropout(ff_output)\n"," x = self.norm2(x)\n"," x = self.adain(x, style)\n"," return x, intermediate\n"," \n","class TransformerDecoderBlock(nn.Module):\n"," def __init__(self, emb_dim=768, num_heads=8, hidden_dim=2048, dropout=0.1):\n"," super(TransformerDecoderBlock, self).__init__()\n"," self.attn1 = LocationBasedMultiheadAttention(emb_dim, num_heads, dropout=dropout)\n"," self.attn2 = LocationBasedMultiheadAttention(emb_dim, num_heads, dropout=dropout)\n"," self.ff = nn.Sequential(\n"," nn.Linear(emb_dim, hidden_dim),\n"," nn.ReLU(),\n"," nn.Linear(hidden_dim, emb_dim),\n"," )\n"," self.norm1 = nn.LayerNorm(emb_dim)\n"," self.norm2 = nn.LayerNorm(emb_dim)\n"," self.norm3 = nn.LayerNorm(emb_dim)\n"," self.norm4 = nn.LayerNorm(emb_dim)\n"," self.adain1 = AdaIN(emb_dim)\n"," self.adain2 = AdaIN(emb_dim)\n"," self.dropout = nn.Dropout(dropout)\n","\n"," def forward(self, x, enc_output, skip_connection, style):\n"," attn_output1 = self.attn1(x, x, x)\n"," x = x + self.dropout(attn_output1)\n"," x = self.norm1(x)\n"," x = self.adain1(x, style)\n"," attn_output2 = self.attn2(x, enc_output, enc_output)\n"," x = x + self.dropout(attn_output2)\n"," x = self.norm2(x)\n"," x = self.adain2(x, style)\n"," ff_output = self.ff(x)\n"," x = x + self.dropout(ff_output)\n"," x = self.norm3(x)\n"," x = x + self.norm4(skip_connection)\n"," return x"]},{"cell_type":"code","execution_count":145,"metadata":{"execution":{"iopub.execute_input":"2024-09-21T23:16:41.057990Z","iopub.status.busy":"2024-09-21T23:16:41.057280Z","iopub.status.idle":"2024-09-21T23:16:41.072000Z","shell.execute_reply":"2024-09-21T23:16:41.070842Z","shell.execute_reply.started":"2024-09-21T23:16:41.057950Z"},"trusted":true},"outputs":[],"source":["class LocationBasedMultiheadAttention(nn.Module):\n"," def __init__(self, emb_dim, num_heads, dropout=0.1):\n"," super(LocationBasedMultiheadAttention, self).__init__()\n"," self.emb_dim = emb_dim\n"," self.num_heads = num_heads\n"," self.head_dim = emb_dim // num_heads\n"," \n"," self.q_proj = nn.Linear(emb_dim, emb_dim)\n"," self.k_proj = nn.Linear(emb_dim, emb_dim)\n"," self.v_proj = nn.Linear(emb_dim, emb_dim)\n"," self.out_proj = nn.Linear(emb_dim, emb_dim)\n"," \n"," self.dropout = nn.Dropout(dropout)\n"," \n"," # Learnable position encodings\n"," self.pos_enc = nn.Parameter(torch.randn(1, 1, emb_dim))\n","\n"," def forward(self, q, k, v):\n"," batch_size, seq_len, _ = q.shape\n"," \n"," # Add positional encodings\n"," q = q + self.pos_enc\n"," k = k + self.pos_enc\n"," \n"," # Separate heads\n"," q = self.q_proj(q).view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)\n"," k = self.k_proj(k).view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)\n"," v = self.v_proj(v).view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)\n"," \n"," # Compute attention scores\n"," scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.head_dim)\n"," attn_probs = F.softmax(scores, dim=-1)\n"," attn_probs = self.dropout(attn_probs)\n"," \n"," # Compute output\n"," out = torch.matmul(attn_probs, v)\n"," out = out.transpose(1, 2).contiguous().view(batch_size, seq_len, self.emb_dim)\n"," out = self.out_proj(out)\n"," \n"," return out"]},{"cell_type":"code","execution_count":146,"metadata":{"execution":{"iopub.execute_input":"2024-09-21T23:16:41.073857Z","iopub.status.busy":"2024-09-21T23:16:41.073403Z","iopub.status.idle":"2024-09-21T23:16:41.085968Z","shell.execute_reply":"2024-09-21T23:16:41.084912Z","shell.execute_reply.started":"2024-09-21T23:16:41.073817Z"},"trusted":true},"outputs":[],"source":["class AdaIN(nn.Module):\n"," def __init__(self, emb_dim):\n"," super(AdaIN, self).__init__()\n"," self.norm = nn.InstanceNorm1d(emb_dim, affine=False)\n"," self.fc = nn.Linear(emb_dim, emb_dim * 2)\n","\n"," def forward(self, x, style):\n"," style = self.fc(style).unsqueeze(1)\n"," gamma, beta = style.chunk(2, dim=-1)\n"," out = self.norm(x.transpose(1, 2)).transpose(1, 2)\n"," out = gamma * out + beta\n"," return out"]},{"cell_type":"code","execution_count":147,"metadata":{"execution":{"iopub.execute_input":"2024-09-21T23:16:41.088410Z","iopub.status.busy":"2024-09-21T23:16:41.088010Z","iopub.status.idle":"2024-09-21T23:16:41.099022Z","shell.execute_reply":"2024-09-21T23:16:41.097846Z","shell.execute_reply.started":"2024-09-21T23:16:41.088350Z"},"trusted":true},"outputs":[],"source":["class RefinementBlock(nn.Module):\n"," def __init__(self, in_channels=768, out_channels=3, kernel_size=3, stride=1, padding=1):\n"," super(RefinementBlock, self).__init__()\n"," self.conv = nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding)\n"," self.bn = nn.BatchNorm2d(out_channels)\n"," self.relu = nn.ReLU(inplace=True)\n"," \n"," def forward(self, x):\n"," x = self.conv(x)\n"," x = self.bn(x)\n"," x = self.relu(x)\n"," return x"]},{"cell_type":"code","execution_count":148,"metadata":{"execution":{"iopub.execute_input":"2024-09-21T23:16:41.101869Z","iopub.status.busy":"2024-09-21T23:16:41.101083Z","iopub.status.idle":"2024-09-21T23:16:41.107762Z","shell.execute_reply":"2024-09-21T23:16:41.106685Z","shell.execute_reply.started":"2024-09-21T23:16:41.101815Z"},"trusted":true},"outputs":[],"source":["# !pip install einops"]},{"cell_type":"code","execution_count":149,"metadata":{"execution":{"iopub.execute_input":"2024-09-21T23:16:41.117428Z","iopub.status.busy":"2024-09-21T23:16:41.116861Z","iopub.status.idle":"2024-09-21T23:16:41.132062Z","shell.execute_reply":"2024-09-21T23:16:41.131002Z","shell.execute_reply.started":"2024-09-21T23:16:41.117388Z"},"trusted":true},"outputs":[],"source":["import torch.nn.functional as F\n","from einops import rearrange\n","\n","# Swin Transformer Block\n","class SwinTransformerBlock(nn.Module):\n"," def __init__(self, dim, num_heads, window_size=7, shift_size=2):\n"," super(SwinTransformerBlock, self).__init__()\n"," self.dim = dim\n"," self.num_heads = num_heads\n"," self.window_size = window_size\n"," self.shift_size = shift_size\n","\n"," # Layer normalization\n"," self.norm1 = nn.LayerNorm(dim)\n"," \n"," # Window-based multi-head self-attention (W-MSA) or shifted W-MSA\n"," self.attn = nn.MultiheadAttention(embed_dim=dim, num_heads=num_heads, batch_first=True)\n","\n"," # Feed-forward network\n"," self.mlp = nn.Sequential(\n"," nn.Linear(dim, 4 * dim),\n"," nn.GELU(),\n"," nn.Linear(4 * dim, dim)\n"," )\n","\n"," # Another layer normalization\n"," self.norm2 = nn.LayerNorm(dim)\n","\n"," def forward(self, x, mask=None):\n"," # Input size\n"," B, H, W, C = x.shape\n","\n"," # Partition windows\n"," x = rearrange(x, 'b (h ws1) (w ws2) c -> (b h w) (ws1 ws2) c', ws1=self.window_size, ws2=self.window_size)\n","\n"," # Add cyclic shift if shift_size > 0\n"," if self.shift_size > 0:\n"," x = torch.roll(x, shifts=(-self.shift_size, -self.shift_size), dims=(1, 2))\n","\n"," # Layer norm\n"," shortcut = x\n"," x = self.norm1(x)\n"," \n"," # Multi-head self-attention\n"," x, _ = self.attn(x, x, x, attn_mask=mask)\n"," \n"," # Residual connection\n"," x = shortcut + x\n","\n"," # MLP with another residual connection\n"," shortcut = x\n"," x = self.norm2(x)\n"," x = self.mlp(x)\n"," x = shortcut + x\n","\n"," # Reverse window partitioning\n"," x = rearrange(x, '(b h w) (ws1 ws2) c -> b (h ws1) (w ws2) c', h=H // self.window_size, w=W // self.window_size, ws1=self.window_size, ws2=self.window_size)\n","\n"," return x"]},{"cell_type":"code","execution_count":150,"metadata":{"execution":{"iopub.execute_input":"2024-09-21T23:16:41.153207Z","iopub.status.busy":"2024-09-21T23:16:41.152730Z","iopub.status.idle":"2024-09-21T23:16:41.174097Z","shell.execute_reply":"2024-09-21T23:16:41.172886Z","shell.execute_reply.started":"2024-09-21T23:16:41.153164Z"},"trusted":true},"outputs":[],"source":["class ViTImage2Image(nn.Module):\n"," def __init__(self, img_size=512, patch_size=8, emb_dim=768, num_heads=12, num_layers=12, hidden_dim=3072, window_size=8):\n"," super(ViTImage2Image, self).__init__()\n"," self.img_size = img_size\n"," self.patch_size = patch_size\n"," self.emb_dim = emb_dim\n"," self.num_heads = num_heads\n"," self.num_layers = num_layers\n"," self.hidden_dim = hidden_dim\n"," self.window_size = window_size\n"," self.num_patches = (img_size // patch_size) ** 2\n"," \n"," self.patch_embed = nn.Conv2d(3, emb_dim, kernel_size=patch_size, stride=patch_size)\n"," self.pos_embed = nn.Parameter(torch.zeros(1, self.num_patches, emb_dim))\n"," \n"," self.encoder_layers = nn.ModuleList([\n"," TransformerEncoderBlock(emb_dim, num_heads, hidden_dim, dropout=0.1)\n"," for _ in range(num_layers)\n"," ])\n"," \n"," self.decoder_layers = nn.ModuleList([\n"," TransformerDecoderBlock(emb_dim, num_heads, hidden_dim, dropout=0.1)\n"," for _ in range(num_layers)\n"," ])\n"," \n"," self.swin_layers = nn.ModuleList([\n"," SwinTransformerBlock(emb_dim, num_heads, window_size=window_size)\n"," for _ in range(num_layers)\n"," ])\n"," \n"," self.norm = nn.LayerNorm(emb_dim)\n"," self.mlp_head = nn.Sequential(\n"," nn.Linear(emb_dim, hidden_dim),\n"," nn.GELU(),\n"," nn.Linear(hidden_dim, emb_dim)\n"," )\n"," self.refinement = RefinementBlock(in_channels=emb_dim, out_channels=3, kernel_size=3, stride=1, padding=1)\n"," \n"," self.style_encoder = nn.Sequential(\n"," nn.Conv2d(3, emb_dim, kernel_size=3, stride=2, padding=1),\n"," nn.ReLU(),\n"," nn.AdaptiveAvgPool2d(1),\n"," nn.Flatten(),\n"," nn.Linear(emb_dim, emb_dim)\n"," )\n","\n"," def forward(self, content, style):\n"," x = self.patch_embed(content)\n"," B, C, H, W = x.shape\n"," x = x.flatten(2).transpose(1, 2)\n"," x = x + self.pos_embed\n"," \n"," style_features = self.style_encoder(style)\n"," \n"," skip_connections = []\n"," # Transformer encoder\n"," for encoder in self.encoder_layers:\n"," x, skip = encoder(x, style_features)\n"," skip_connections.append(skip)\n"," \n"," # Transformer decoder\n"," for decoder, skip in zip(self.decoder_layers, reversed(skip_connections)):\n"," x = decoder(x, x, skip, style_features)\n"," \n"," # Swin Transformer encoding\n"," x = x.view(B, H, W, C)\n"," for layer in self.swin_layers:\n"," x = layer(x)\n"," \n"," x = x.view(B, -1, C)\n"," x = self.norm(x)\n"," x = self.mlp_head(x)\n"," x = x.transpose(1, 2).view(B, self.emb_dim, H, W)\n"," x = self.refinement(x)\n"," x = nn.functional.interpolate(x, size=(self.img_size, self.img_size), mode='bilinear', align_corners=False)\n"," return x"]},{"cell_type":"markdown","metadata":{},"source":["## DEFINE LOSS METRICS"]},{"cell_type":"code","execution_count":151,"metadata":{"execution":{"iopub.execute_input":"2024-09-21T23:16:41.183842Z","iopub.status.busy":"2024-09-21T23:16:41.183339Z","iopub.status.idle":"2024-09-21T23:16:41.191559Z","shell.execute_reply":"2024-09-21T23:16:41.190562Z","shell.execute_reply.started":"2024-09-21T23:16:41.183800Z"},"trusted":true},"outputs":[],"source":["def total_variation_loss(x):\n"," return torch.sum(torch.abs(x[:, :, :-1, :] - x[:, :, 1:, :])) + torch.sum(torch.abs(x[:, :, :, :-1] - x[:, :, :, 1:]))\n","\n","def combined_loss(output, target):\n"," l1_loss = nn.L1Loss()(output, target)\n"," tv_loss = total_variation_loss(output)\n"," return l1_loss + 0.0001 * tv_loss"]},{"cell_type":"markdown","metadata":{},"source":["## DEFINE TRAINING METRICS"]},{"cell_type":"code","execution_count":152,"metadata":{"execution":{"iopub.execute_input":"2024-09-21T23:16:41.213172Z","iopub.status.busy":"2024-09-21T23:16:41.212746Z","iopub.status.idle":"2024-09-21T23:16:41.219196Z","shell.execute_reply":"2024-09-21T23:16:41.218074Z","shell.execute_reply.started":"2024-09-21T23:16:41.213131Z"},"trusted":true},"outputs":[],"source":["def psnr(img1, img2):\n"," mse = torch.mean((img1 - img2) ** 2)\n"," if mse == 0:\n"," return float('inf')\n"," return 20 * torch.log10(1.0 / torch.sqrt(mse))"]},{"cell_type":"markdown","metadata":{},"source":["## DATALOADERS"]},{"cell_type":"code","execution_count":181,"metadata":{"execution":{"iopub.execute_input":"2024-09-21T23:24:17.558825Z","iopub.status.busy":"2024-09-21T23:24:17.557763Z","iopub.status.idle":"2024-09-21T23:24:17.588295Z","shell.execute_reply":"2024-09-21T23:24:17.587164Z","shell.execute_reply.started":"2024-09-21T23:24:17.558778Z"},"trusted":true},"outputs":[],"source":["import os\n","from PIL import Image\n","from torch.utils.data import Dataset, DataLoader\n","from torchvision import transforms\n","\n","# Custom Dataset class\n","class ImageLabelDataset(Dataset):\n"," def __init__(self, images_dir, labels_dir, transform=None):\n"," self.images_dir = images_dir\n"," self.labels_dir = labels_dir\n"," self.image_filenames = sorted(os.listdir(images_dir))\n"," self.label_filenames = sorted(os.listdir(labels_dir))\n"," self.transform = transform\n","\n"," assert len(self.image_filenames) == len(self.label_filenames), \"Mismatch in number of images and labels\"\n","\n"," def __len__(self):\n"," return len(self.image_filenames)\n","\n"," def __getitem__(self, idx):\n"," image_path = os.path.join(self.images_dir, self.image_filenames[idx])\n"," label_path = os.path.join(self.labels_dir, self.label_filenames[idx])\n"," \n"," # Open image and label using PIL\n"," image = Image.open(image_path).convert('RGB') # Ensure image is RGB\n"," label = Image.open(label_path).convert('RGB') # Ensure label is RGB (or use 'L' for grayscale)\n","\n"," if self.transform:\n"," # Apply the same transformation to both the image and the label\n"," image = self.transform(image)\n"," label = self.transform(label)\n","\n"," return image, label\n","\n","resize_dim = 512 # resize to dimensions\n","batch_size = 4 # Adjust batch size according to your GPU memory capacity\n"," \n","# Define transformations (resize, convert to tensor, normalize)\n","transform = transforms.Compose([\n"," transforms.Resize((resize_dim, resize_dim)), # Resize images\n"," transforms.ToTensor(), # Convert images to PyTorch tensors\n"," transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]) # Normalize to [-1, 1] for each channel\n","])\n","\n","# Create dataset instances for training\n","images_dir = '/kaggle/working/images/' # Replace with your image directory path\n","labels_dir = '/kaggle/working/labels/' # Replace with your label directory path\n","dataset = ImageLabelDataset(images_dir, labels_dir, transform=transform)\n","\n","# Create DataLoader with batch size control\n","dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True, num_workers=8)"]},{"cell_type":"markdown","metadata":{},"source":["## TRAINING SCRIPT"]},{"cell_type":"code","execution_count":null,"metadata":{"execution":{"iopub.execute_input":"2024-09-21T23:25:08.871609Z","iopub.status.busy":"2024-09-21T23:25:08.871150Z"},"trusted":true},"outputs":[],"source":["import torch\n","import torch.nn as nn\n","from torch.utils.data import DataLoader\n","from tqdm import tqdm\n","import os\n","\n","# Instantiate the model\n","model = ViTImage2Image(img_size=resize_dim, patch_size=16, emb_dim=768, num_heads=16, num_layers=8, hidden_dim=3072)\n","\n","# Move model to the appropriate device (GPU if available)\n","device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n","\n","# If using 2 x T4 GPU only\n","model = nn.DataParallel(model, device_ids = [0,1])\n","\n","# Move model to device\n","model = model.to(device)\n","\n","# Optimizer\n","optimizer = torch.optim.Adam(model.parameters(), lr=0.001)\n","\n","# Number of epochs\n","num_epochs = 100\n","\n","best_loss = float('inf')\n","\n","# Path to save the best model\n","save_dir = 'saved_models'\n","os.makedirs(save_dir, exist_ok=True) # Create directory if it doesn't exist\n","best_model_path = os.path.join(save_dir, 'best_model.pth')\n","\n","# Initialize best loss as a large value\n","best_loss = float('inf')\n","\n","# Training loop with tqdm\n","for epoch in range(num_epochs):\n"," model.train()\n"," running_loss = 0.0\n"," running_psnr = 0.0\n"," \n"," # Wrap dataloader with tqdm for progress bar\n"," pbar = tqdm(enumerate(dataloader), total=len(dataloader), desc=f\"Epoch {epoch+1}/{num_epochs}\")\n"," \n"," for batch_idx, (input, target) in pbar:\n"," # Move data to the same device as the model\n"," input, target = input.to(device), target.to(device)\n"," \n"," optimizer.zero_grad() # Clear the gradients from the last step\n"," output = model(input, target) # Forward pass\n","# print(f\"Input shape: {input.shape}, Output shape: {output.shape}, Target shape: {target.shape}\")\n"," loss = combined_loss(output, target) # Compute the loss\n"," \n"," loss.backward() # Backward pass (compute gradients)\n"," optimizer.step() # Update the weights\n"," \n"," running_loss += loss.item() # Accumulate the loss for this batch\n"," \n"," # Calculate metrics\n"," current_psnr = psnr(output, target).item()\n"," running_psnr += current_psnr\n"," \n"," # Update progress bar\n"," pbar.set_postfix({\n"," 'loss': loss.item(),\n"," 'psnr': current_psnr,\n"," })\n"," \n"," # Calculate the average loss and metrics for the epoch\n"," epoch_loss = running_loss / len(dataloader)\n"," avg_psnr = running_psnr / len(dataloader)\n"," \n"," print(f\"Epoch [{epoch+1}/{num_epochs}], Loss: {epoch_loss:.4f}, PSNR: {avg_psnr:.4f}\")\n"," \n"," # Save the best model based on the lowest loss\n"," if epoch_loss < best_loss:\n"," best_loss = epoch_loss\n"," torch.save(model.state_dict(), best_model_path)\n"," print(f\"New best model saved at epoch {epoch+1} with loss {best_loss:.4f}\")\n","\n","print(\"Training complete\")"]},{"cell_type":"markdown","metadata":{},"source":["## CLEAR CUDA CACHE"]},{"cell_type":"code","execution_count":193,"metadata":{"execution":{"iopub.execute_input":"2024-09-21T23:25:05.712036Z","iopub.status.busy":"2024-09-21T23:25:05.711264Z","iopub.status.idle":"2024-09-21T23:25:05.899856Z","shell.execute_reply":"2024-09-21T23:25:05.898677Z","shell.execute_reply.started":"2024-09-21T23:25:05.711993Z"},"trusted":true},"outputs":[{"data":{"text/plain":["0"]},"execution_count":193,"metadata":{},"output_type":"execute_result"}],"source":["import gc\n","torch.cuda.empty_cache()\n","gc.collect()\n","#RuntimeError: The size of tensor a (768) must match the size of tensor b (512) at non-singleton dimension 3"]},{"cell_type":"code","execution_count":180,"metadata":{"execution":{"iopub.execute_input":"2024-09-21T23:24:04.448611Z","iopub.status.busy":"2024-09-21T23:24:04.448163Z","iopub.status.idle":"2024-09-21T23:24:04.462532Z","shell.execute_reply":"2024-09-21T23:24:04.461381Z","shell.execute_reply.started":"2024-09-21T23:24:04.448567Z"},"trusted":true},"outputs":[{"name":"stdout","output_type":"stream","text":["223543305\n","DataParallel(\n"," (module): ViTImage2Image(\n"," (patch_embed): Conv2d(3, 768, kernel_size=(16, 16), stride=(16, 16))\n"," (encoder_layers): ModuleList(\n"," (0-7): 8 x TransformerEncoderBlock(\n"," (attn): LocationBasedMultiheadAttention(\n"," (q_proj): Linear(in_features=768, out_features=768, bias=True)\n"," (k_proj): Linear(in_features=768, out_features=768, bias=True)\n"," (v_proj): Linear(in_features=768, out_features=768, bias=True)\n"," (out_proj): Linear(in_features=768, out_features=768, bias=True)\n"," (dropout): Dropout(p=0.1, inplace=False)\n"," )\n"," (ff): Sequential(\n"," (0): Linear(in_features=768, out_features=3072, bias=True)\n"," (1): ReLU()\n"," (2): Linear(in_features=3072, out_features=768, bias=True)\n"," )\n"," (norm1): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n"," (norm2): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n"," (adain): AdaIN(\n"," (norm): InstanceNorm1d(768, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)\n"," (fc): Linear(in_features=768, out_features=1536, bias=True)\n"," )\n"," (dropout): Dropout(p=0.1, inplace=False)\n"," )\n"," )\n"," (decoder_layers): ModuleList(\n"," (0-7): 8 x TransformerDecoderBlock(\n"," (attn1): LocationBasedMultiheadAttention(\n"," (q_proj): Linear(in_features=768, out_features=768, bias=True)\n"," (k_proj): Linear(in_features=768, out_features=768, bias=True)\n"," (v_proj): Linear(in_features=768, out_features=768, bias=True)\n"," (out_proj): Linear(in_features=768, out_features=768, bias=True)\n"," (dropout): Dropout(p=0.1, inplace=False)\n"," )\n"," (attn2): LocationBasedMultiheadAttention(\n"," (q_proj): Linear(in_features=768, out_features=768, bias=True)\n"," (k_proj): Linear(in_features=768, out_features=768, bias=True)\n"," (v_proj): Linear(in_features=768, out_features=768, bias=True)\n"," (out_proj): Linear(in_features=768, out_features=768, bias=True)\n"," (dropout): Dropout(p=0.1, inplace=False)\n"," )\n"," (ff): Sequential(\n"," (0): Linear(in_features=768, out_features=3072, bias=True)\n"," (1): ReLU()\n"," (2): Linear(in_features=3072, out_features=768, bias=True)\n"," )\n"," (norm1): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n"," (norm2): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n"," (norm3): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n"," (norm4): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n"," (adain1): AdaIN(\n"," (norm): InstanceNorm1d(768, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)\n"," (fc): Linear(in_features=768, out_features=1536, bias=True)\n"," )\n"," (adain2): AdaIN(\n"," (norm): InstanceNorm1d(768, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)\n"," (fc): Linear(in_features=768, out_features=1536, bias=True)\n"," )\n"," (dropout): Dropout(p=0.1, inplace=False)\n"," )\n"," )\n"," (swin_layers): ModuleList(\n"," (0-7): 8 x SwinTransformerBlock(\n"," (norm1): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n"," (attn): MultiheadAttention(\n"," (out_proj): NonDynamicallyQuantizableLinear(in_features=768, out_features=768, bias=True)\n"," )\n"," (mlp): Sequential(\n"," (0): Linear(in_features=768, out_features=3072, bias=True)\n"," (1): GELU(approximate='none')\n"," (2): Linear(in_features=3072, out_features=768, bias=True)\n"," )\n"," (norm2): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n"," )\n"," )\n"," (norm): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n"," (mlp_head): Sequential(\n"," (0): Linear(in_features=768, out_features=3072, bias=True)\n"," (1): GELU(approximate='none')\n"," (2): Linear(in_features=3072, out_features=768, bias=True)\n"," )\n"," (refinement): RefinementBlock(\n"," (conv): Conv2d(768, 3, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n"," (bn): BatchNorm2d(3, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n"," (relu): ReLU(inplace=True)\n"," )\n"," (style_encoder): Sequential(\n"," (0): Conv2d(3, 768, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))\n"," (1): ReLU()\n"," (2): AdaptiveAvgPool2d(output_size=1)\n"," (3): Flatten(start_dim=1, end_dim=-1)\n"," (4): Linear(in_features=768, out_features=768, bias=True)\n"," )\n"," )\n",")\n"]}],"source":["total_params = sum(p.numel() for p in model.parameters())\n","print(total_params)\n","print(model)"]},{"cell_type":"markdown","metadata":{},"source":["## INFERENCE"]},{"cell_type":"code","execution_count":null,"metadata":{"execution":{"iopub.status.busy":"2024-09-21T23:16:44.175025Z","iopub.status.idle":"2024-09-21T23:16:44.184902Z","shell.execute_reply":"2024-09-21T23:16:44.184625Z","shell.execute_reply.started":"2024-09-21T23:16:44.184592Z"},"trusted":true},"outputs":[],"source":["import torch\n","import torchvision.transforms as transforms\n","from PIL import Image\n","import matplotlib.pyplot as plt\n","\n","def load_image(image_path, img_size=1024):\n"," image = Image.open(image_path).convert('RGB')\n"," transform = transforms.Compose([\n"," transforms.Resize((img_size, img_size)),\n"," transforms.ToTensor(),\n"," transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])\n"," ])\n"," return transform(image).unsqueeze(0)\n","\n","def perform_inference(model, input_image):\n"," model.eval()\n"," with torch.no_grad():\n"," output = model(input_image)\n"," return output\n","\n","def visualize_tensor(tensor, title):\n"," img = tensor.squeeze().cpu().permute(1, 2, 0).numpy()\n"," img = (img - img.min()) / (img.max() - img.min()) # Normalize to [0, 1]\n"," plt.imshow(img)\n"," plt.title(title)\n"," plt.axis('off')\n"," plt.show()\n","\n","def main():\n"," # Load the model (adjust parameters as needed)\n"," model = ViTImage2Image(img_size=1024, patch_size=8, emb_dim=768, num_heads=12, num_layers=12, hidden_dim=3072)\n"," model.load_state_dict(torch.load('/kaggle/working/saved_models/best_model.pth'))\n"," \n"," device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n"," model = model.to(device)\n"," \n"," # Load and preprocess input image\n"," input_image = load_image('/kaggle/working/labels/00143.png')\n"," input_image = input_image.to(device)\n"," \n"," # Perform inference\n"," output = perform_inference(model, input_image)\n"," \n"," # Debug: Print shape and statistics of the output\n"," print(\"Output shape:\", output.shape)\n"," print(\"Output min:\", output.min().item())\n"," print(\"Output max:\", output.max().item())\n"," print(\"Output mean:\", output.mean().item())\n"," print(\"Output std:\", output.std().item())\n"," \n"," # Visualize input and output\n"," visualize_tensor(input_image, \"Input Image\")\n"," visualize_tensor(output, \"Output Image (Before Processing)\")\n"," \n"," # Try different post-processing steps\n"," # 1. Sigmoid activation\n"," output_sigmoid = torch.sigmoid(output)\n"," visualize_tensor(output_sigmoid, \"Output Image (After Sigmoid)\")\n"," \n"," # 2. Tanh activation\n"," output_tanh = torch.tanh(output)\n"," visualize_tensor(output_tanh, \"Output Image (After Tanh)\")\n"," \n"," # 3. Min-Max Normalization\n"," output_normalized = (output - output.min()) / (output.max() - output.min())\n"," visualize_tensor(output_normalized, \"Output Image (After Min-Max Normalization)\")\n","\n","if __name__ == \"__main__\":\n"," main()"]}],"metadata":{"kaggle":{"accelerator":"nvidiaTeslaT4","dataSources":[],"dockerImageVersionId":30761,"isGpuEnabled":true,"isInternetEnabled":true,"language":"python","sourceType":"notebook"},"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.10.11"}},"nbformat":4,"nbformat_minor":4}
|