{ "cells": [ { "cell_type": "code", "execution_count": null, "id": "6698279f", "metadata": {}, "outputs": [], "source": [ "#pip install torch" ] }, { "cell_type": "code", "execution_count": 1, "id": "201b312f", "metadata": {}, "outputs": [], "source": [ "import torch\n", "import torch.nn as nn\n", "import torch.nn.functional as F\n", "from torchvision import datasets, transforms\n", "from torch.utils.data import DataLoader, random_split\n", "from tqdm import tqdm" ] }, { "cell_type": "code", "execution_count": 2, "id": "822023ca", "metadata": {}, "outputs": [], "source": [ "def window_partition(x, win):\n", " B, H, W, C = x.shape\n", " x = x.view(B,H//win,win,W//win,win,C)\n", " x = x.permute(0,1,3,2,4,5)\n", " x = x.reshape(-1,win,win,C)\n", " return x\n", "\n", "def window_reverse(windows, win, H, W):\n", " B = windows.shape[0]//(H//win * W//win)\n", " x = windows.view(B,H//win,W//win, win,win,-1)\n", " x = x.permute(0,1,3,2,4,5)\n", " x=x.reshape(B,H,W,-1)\n", " return x" ] }, { "cell_type": "code", "execution_count": 3, "id": "0b7847fa", "metadata": {}, "outputs": [], "source": [ "class WindowAttention(nn.Module):\n", " def __init__(self, dim, num_heads, win):\n", " super().__init__()\n", " \n", " self.dim = dim\n", " self.num_heads = num_heads\n", " self.head_dim = self.dim//self.num_heads\n", " self.scale = self.head_dim**-0.5\n", " self.win = win\n", " \n", " self.q = nn.Linear(dim,dim)\n", " self.k = nn.Linear(dim,dim)\n", " self.v = nn.Linear(dim,dim)\n", " \n", " self.proj = nn.Linear(dim,dim)\n", " \n", " coords = torch.stack(torch.meshgrid(torch.arange(win), torch.arange(win), indexing = \"ij\"))\n", " coords_flat = coords.flatten(1)\n", " \n", " rel = coords_flat[:,:,None] - coords_flat[:,None,:]\n", " rel = rel.permute(1,2,0)\n", " \n", " rel[:,:,0] =rel[:,:,0] + (win-1)\n", " rel[:,:,1] =rel[:,:,1] + (win-1)\n", " rel[:,:,0] =rel[:,:,0] * (2*win-1)\n", " index = rel.sum(-1)\n", " \n", " self.register_buffer(\"pos_index\", index)\n", " self.rel_bias = nn.Parameter(torch.zeros((2*win-1) * (2*win-1), num_heads))\n", " \n", " \n", " def forward(self,x,mask=None):\n", " B_,N,C =x.shape\n", " q = self.q(x)\n", " k = self.k(x)\n", " v = self.v(x)\n", " \n", " q = q.view(B_, N, self.num_heads, C // self.num_heads).transpose(1, 2)\n", " k = k.view(B_, N, self.num_heads, C // self.num_heads).transpose(1, 2)\n", " v = v.view(B_, N, self.num_heads, C // self.num_heads).transpose(1, 2)\n", " \n", " q = q * self.scale\n", " attn = q @ k.transpose(-2,-1)\n", " \n", " rb = self.rel_bias[self.pos_index.view(-1)]\n", " rb = rb.view(N,N, self.num_heads).permute(2,0,1)\n", " attn = attn + rb.unsqueeze(0)\n", " \n", " \n", " if mask is not None:\n", " nw = mask.shape[0]\n", " attn = attn.view(B_//nw,nw,self.num_heads,N,N) + mask.unsqueeze(1).unsqueeze(0)\n", " attn = attn.view(-1,self.num_heads,N,N)\n", "\n", " attn = attn.softmax(dim=-1)\n", " out = (attn @ v).transpose(1,2).reshape(B_,N,C)\n", " out = self.proj(out)\n", " return out" ] }, { "cell_type": "code", "execution_count": 4, "id": "d34e7b3e", "metadata": {}, "outputs": [], "source": [ "class SwinBlock(nn.Module):\n", " def __init__(self, dim, res, heads, win, shift):\n", " super().__init__()\n", " \n", " self.dim = dim\n", " self.res = res\n", " self.win = win\n", " self.shift = shift\n", "\n", " self.norm1 = nn.LayerNorm(dim)\n", " self.attn = WindowAttention(dim, heads, win)\n", " self.norm2 = nn.LayerNorm(dim)\n", " self.mlp = nn.Sequential(\n", " nn.Linear(dim, 4 * dim),\n", " nn.GELU(),\n", " nn.Linear(4 * dim, dim),\n", " )\n", "\n", " H, W = res\n", " if shift > 0:\n", " self.mask = self.create_mask(H, W, win, shift)\n", " else:\n", " self.mask = None\n", " \n", " def create_mask(self, H, W, win, shift):\n", " img_mask = torch.zeros((1, H, W, 1))\n", " count = 0\n", "\n", " for h in (slice(0,-win), slice(-win,-shift), slice(-shift, None)):\n", " for w in (slice(0,-win), slice(-win, -shift), slice(-shift, None)):\n", " img_mask[:,h,w,:] = count\n", " count += 1\n", "\n", "\n", " mask = window_partition(img_mask, win)\n", " mask = mask.view(-1, win*win)\n", " mask = mask.unsqueeze(1) - mask.unsqueeze(2)\n", " mask = mask.masked_fill(mask!=0, -10000.0)\n", " return mask\n", " \n", " def forward(self, x):\n", " B, L, C = x.shape\n", " H, W = self.res\n", " residual = x\n", " x = self.norm1(x)\n", " x = x.view(B,H,W,C)\n", "\n", " if self.shift > 0:\n", " x = torch.roll(x, shifts=(-self.shift, -self.shift), dims = (1,2))\n", " \n", " win_x = window_partition(x, self.win).view(-1, self.win*self.win, C)\n", " \n", " attn_out = self.attn(win_x, mask = self.mask.to(x.device) if self.mask is not None else None)\n", " \n", " x = window_reverse(attn_out, self.win, H, W)\n", "\n", " if self.shift > 0:\n", " x = torch.roll(x, shifts=(self.shift, self.shift), dims = (1,2))\n", "\n", " x = residual + x.view(B, L, C)\n", " residual2 = x\n", " x = self.norm2(x)\n", " x = self.mlp(x)\n", " return residual2 + x" ] }, { "cell_type": "code", "execution_count": 5, "id": "9bb962d8", "metadata": {}, "outputs": [], "source": [ "class PatchMerging(nn.Module):\n", " def __init__(self, dim):\n", " super().__init__()\n", " self.dim = dim\n", "\n", " self.reduction = nn.Linear(4*dim, 2*dim, bias = False)\n", " self.norm = nn.LayerNorm(4*dim)\n", "\n", " def forward(self, x, H, W):\n", " B, L, C = x.shape\n", " x = x.view(B, H, W, C)\n", "\n", " x0 = x[:,0::2,0::2,:]\n", " x1 = x[:,1::2,0::2,:]\n", " x2 = x[:,0::2,1::2,:]\n", " x3 = x[:,1::2,1::2,:]\n", "\n", " x0 = x0.reshape(B, -1, C)\n", " x1 = x1.reshape(B, -1, C)\n", " x2 = x2.reshape(B, -1, C)\n", " x3 = x3.reshape(B, -1, C)\n", "\n", " x = torch.cat([x0,x1,x2,x3], -1)\n", "\n", " x = self.norm(x)\n", " x= self.reduction(x)\n", " return x, H//2, W//2" ] }, { "cell_type": "code", "execution_count": 6, "id": "e9fc4137", "metadata": {}, "outputs": [], "source": [ "class SwinTransformer(nn.Module):\n", " def __init__(self, embed_dim = 96, win = 7 , num_classes = 45):\n", " super().__init__()\n", " self.patch_embed = nn.Conv2d(3, embed_dim, kernel_size= 4, stride = 4)\n", " self.patch_norm = nn.LayerNorm(embed_dim) \n", " \n", " \n", " dims = [embed_dim, embed_dim*2, embed_dim*4, embed_dim*8]\n", " heads = [3, 6, 12, 24]\n", " shift = win//2\n", "\n", " #stage_1\n", "\n", " self.stage1_blocks = nn.Sequential(\n", " SwinBlock(dims[0], (56,56), heads = heads[0], win = win, shift = 0),\n", " SwinBlock(dims[0], (56,56), heads = heads[0], win = win, shift = shift)\n", " )\n", "\n", " self.patch_merge1 = PatchMerging(dims[0])\n", "\n", " #stage_2\n", " self.stage2_blocks = nn.Sequential(\n", " SwinBlock(dims[1], (28,28), heads = heads[1], win = win, shift = 0),\n", " SwinBlock(dims[1], (28,28), heads = heads[1], win = win, shift = shift)\n", " )\n", " \n", " self.patch_merge2 = PatchMerging(dims[1])\n", " \n", " #stage_3\n", " self.stage3_blocks = nn.Sequential(*[\n", " SwinBlock(dims[2], (14,14), heads = heads[2], win = win, shift = 0 if i%2 == 0 else shift)\n", " for i in range(6)\n", " ])\n", " \n", " self.patch_merge3 = PatchMerging(dims[2])\n", " \n", " #stage_4\n", " self.stage4_blocks = nn.Sequential(\n", " SwinBlock(dims[3], (7,7), heads = heads[3], win = win, shift = 0),\n", " SwinBlock(dims[3], (7,7), heads = heads[3], win = win, shift = 0)\n", " )\n", " \n", " self.norm =nn.LayerNorm(dims[3])\n", " self.fc = nn.Linear(dims[3], num_classes)\n", "\n", " \n", " def forward(self, x):\n", " x = self.patch_embed(x) # (B, 96, 56, 56)\n", " B, C, H, W = x.shape\n", " x = self.patch_norm(x.flatten(2).transpose(1,2)) # (B, 3136, 96)\n", "\n", " x = self.stage1_blocks(x)\n", " x, H, W = self.patch_merge1(x, H, W) # (B, 784, 192)\n", "\n", " x = self.stage2_blocks(x)\n", " x, H, W = self.patch_merge2(x, H, W) # (B, 196, 384)\n", "\n", " x = self.stage3_blocks(x)\n", " x, H, W = self.patch_merge3(x, H, W) # (B, 49, 768)\n", "\n", " x = self.stage4_blocks(x)\n", " x = self.norm(x).mean(dim=1)\n", " return self.fc(x)" ] }, { "cell_type": "code", "execution_count": 33, "id": "a515cd9d", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "Classes : ['airplane', 'airport', 'baseball_diamond', 'basketball_court', 'beach', 'bridge', 'chaparral', 'church', 'circular_farmland', 'cloud', 'commercial_area', 'dense_residential', 'desert', 'forest', 'freeway', 'golf_course', 'ground_track_field', 'harbor', 'industrial_area', 'intersection', 'island', 'lake', 'meadow', 'medium_residential', 'mobile_home_park', 'mountain', 'overpass', 'palace', 'parking_lot', 'railway', 'railway_station', 'rectangular_farmland', 'river', 'roundabout', 'runway', 'sea_ice', 'ship', 'snowberg', 'sparse_residential', 'stadium', 'storage_tank', 'tennis_court', 'terrace', 'thermal_power_station', 'wetland']\n", "Train : 25200 | Test: 6300\n" ] } ], "source": [ "torch.cuda.empty_cache()\n", "import gc; gc.collect()\n", "\n", "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n", "\n", "transform_train = transforms.Compose([\n", " transforms.Resize((224,224)),\n", " transforms.ToTensor(),\n", " transforms.Normalize([0.485, 0.456, 0.406],\n", " [0.229, 0.224, 0.225]),\n", "])\n", "\n", "transform_test = transforms.Compose([\n", " transforms.Resize((224,224)),\n", " transforms.ToTensor(),\n", " transforms.Normalize([0.485, 0.456, 0.406],\n", " [0.229, 0.224, 0.225]),\n", "])\n", "\n", "full_dataset = datasets.ImageFolder(r\"D:\\NWPU\\NWPU-RESISC45\\NWPU-RESISC45\", transform=transform_train)\n", "\n", "train_size = int(0.8 * len(full_dataset))\n", "test_size = len(full_dataset) - train_size\n", "\n", "train_dataset, test_dataset = random_split(full_dataset, [train_size, test_size], generator=torch.Generator().manual_seed(42))\n", "test_dataset.dataset.transform = transform_test\n", "\n", "train_loader = DataLoader(train_dataset, batch_size = 32,shuffle=True, pin_memory=False)\n", "test_loader = DataLoader(test_dataset, batch_size = 64,shuffle=False, pin_memory=False)\n", "\n", "print(f\"Classes : {full_dataset.classes}\")\n", "print(f\"Train : {len(train_dataset)} | Test: {len(test_dataset)}\")\n", "\n", "def lr_lambda(epoch): \n", " warmup = 3\n", " if epoch < warmup:\n", " return (epoch + 1) / warmup\n", " return 0.5 * (1 + math.cos(math.pi * (epoch - warmup) / (EPOCHS - warmup)))\n", "\n", "model = SwinTransformer().to(device)\n", "optimizer = torch.optim.AdamW(model.parameters(), lr = 3e-4, weight_decay = 0.05)\n", "scaler = GradScaler()\n", "loss_fn = nn.CrossEntropyLoss(label_smoothing = 0.1)\n", "scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda)" ] }, { "cell_type": "code", "execution_count": 34, "id": "e62997ce", "metadata": {}, "outputs": [ { "name": "stderr", "output_type": "stream", "text": [ "Epoch 1/20: 100%|████████████████████████████████████████████████████████████████████| 788/788 [03:01<00:00, 4.34it/s]\n" ] }, { "name": "stdout", "output_type": "stream", "text": [ "Epoch 01 | Loss: 2.7651 | Acc: 45.56% | LR: 2.00e-04\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ "Epoch 2/20: 100%|████████████████████████████████████████████████████████████████████| 788/788 [02:46<00:00, 4.73it/s]\n" ] }, { "name": "stdout", "output_type": "stream", "text": [ "Epoch 02 | Loss: 2.1457 | Acc: 53.30% | LR: 3.00e-04\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ "Epoch 3/20: 100%|████████████████████████████████████████████████████████████████████| 788/788 [02:45<00:00, 4.76it/s]\n" ] }, { "name": "stdout", "output_type": "stream", "text": [ "Epoch 03 | Loss: 1.9506 | Acc: 58.30% | LR: 3.00e-04\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ "Epoch 4/20: 100%|████████████████████████████████████████████████████████████████████| 788/788 [02:45<00:00, 4.76it/s]\n" ] }, { "name": "stdout", "output_type": "stream", "text": [ "Epoch 04 | Loss: 1.7673 | Acc: 65.73% | LR: 2.97e-04\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ "Epoch 5/20: 100%|████████████████████████████████████████████████████████████████████| 788/788 [02:45<00:00, 4.75it/s]\n" ] }, { "name": "stdout", "output_type": "stream", "text": [ "Epoch 05 | Loss: 1.6530 | Acc: 66.76% | LR: 2.90e-04\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ "Epoch 6/20: 100%|████████████████████████████████████████████████████████████████████| 788/788 [02:45<00:00, 4.75it/s]\n" ] }, { "name": "stdout", "output_type": "stream", "text": [ "Epoch 06 | Loss: 1.5569 | Acc: 69.49% | LR: 2.78e-04\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ "Epoch 7/20: 100%|████████████████████████████████████████████████████████████████████| 788/788 [02:45<00:00, 4.76it/s]\n" ] }, { "name": "stdout", "output_type": "stream", "text": [ "Epoch 07 | Loss: 1.4877 | Acc: 70.49% | LR: 2.61e-04\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ "Epoch 8/20: 100%|████████████████████████████████████████████████████████████████████| 788/788 [02:45<00:00, 4.75it/s]\n" ] }, { "name": "stdout", "output_type": "stream", "text": [ "Epoch 08 | Loss: 1.4043 | Acc: 72.35% | LR: 2.40e-04\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ "Epoch 9/20: 100%|████████████████████████████████████████████████████████████████████| 788/788 [02:46<00:00, 4.74it/s]\n" ] }, { "name": "stdout", "output_type": "stream", "text": [ "Epoch 09 | Loss: 1.3383 | Acc: 72.68% | LR: 2.17e-04\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ "Epoch 10/20: 100%|███████████████████████████████████████████████████████████████████| 788/788 [02:46<00:00, 4.74it/s]\n" ] }, { "name": "stdout", "output_type": "stream", "text": [ "Epoch 10 | Loss: 1.2748 | Acc: 75.05% | LR: 1.91e-04\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ "Epoch 11/20: 100%|███████████████████████████████████████████████████████████████████| 788/788 [02:46<00:00, 4.75it/s]\n" ] }, { "name": "stdout", "output_type": "stream", "text": [ "Epoch 11 | Loss: 1.2021 | Acc: 75.14% | LR: 1.64e-04\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ "Epoch 12/20: 100%|███████████████████████████████████████████████████████████████████| 788/788 [02:46<00:00, 4.75it/s]\n" ] }, { "name": "stdout", "output_type": "stream", "text": [ "Epoch 12 | Loss: 1.1198 | Acc: 76.59% | LR: 1.36e-04\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ "Epoch 13/20: 100%|███████████████████████████████████████████████████████████████████| 788/788 [02:46<00:00, 4.74it/s]\n" ] }, { "name": "stdout", "output_type": "stream", "text": [ "Epoch 13 | Loss: 1.0493 | Acc: 78.35% | LR: 1.09e-04\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ "Epoch 14/20: 100%|███████████████████████████████████████████████████████████████████| 788/788 [02:42<00:00, 4.84it/s]\n" ] }, { "name": "stdout", "output_type": "stream", "text": [ "Epoch 14 | Loss: 0.9696 | Acc: 78.70% | LR: 8.31e-05\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ "Epoch 15/20: 100%|███████████████████████████████████████████████████████████████████| 788/788 [02:32<00:00, 5.16it/s]\n" ] }, { "name": "stdout", "output_type": "stream", "text": [ "Epoch 15 | Loss: 0.9006 | Acc: 80.17% | LR: 5.96e-05\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ "Epoch 16/20: 100%|███████████████████████████████████████████████████████████████████| 788/788 [02:32<00:00, 5.16it/s]\n" ] }, { "name": "stdout", "output_type": "stream", "text": [ "Epoch 16 | Loss: 0.8427 | Acc: 80.57% | LR: 3.91e-05\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ "Epoch 17/20: 100%|███████████████████████████████████████████████████████████████████| 788/788 [02:33<00:00, 5.14it/s]\n" ] }, { "name": "stdout", "output_type": "stream", "text": [ "Epoch 17 | Loss: 0.7990 | Acc: 80.98% | LR: 2.25e-05\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ "Epoch 18/20: 100%|███████████████████████████████████████████████████████████████████| 788/788 [02:37<00:00, 4.99it/s]\n" ] }, { "name": "stdout", "output_type": "stream", "text": [ "Epoch 18 | Loss: 0.7704 | Acc: 81.30% | LR: 1.01e-05\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ "Epoch 19/20: 100%|███████████████████████████████████████████████████████████████████| 788/788 [02:32<00:00, 5.15it/s]\n" ] }, { "name": "stdout", "output_type": "stream", "text": [ "Epoch 19 | Loss: 0.7537 | Acc: 81.19% | LR: 2.55e-06\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ "Epoch 20/20: 100%|███████████████████████████████████████████████████████████████████| 788/788 [02:45<00:00, 4.77it/s]\n" ] }, { "name": "stdout", "output_type": "stream", "text": [ "Epoch 20 | Loss: 0.7461 | Acc: 81.14% | LR: 0.00e+00\n" ] } ], "source": [ "import math\n", "from torch.amp import GradScaler, autocast\n", "\n", "EPOCHS = 20\n", "\n", "def evaluate():\n", " model.eval()\n", " correct = total = 0\n", " with torch.no_grad():\n", " for img, label in test_loader:\n", " img, label = img.to(device), label.to(device)\n", " with autocast('cuda'):\n", " out = model(img)\n", " correct += out.argmax(1).eq(label).sum().item()\n", " total += label.size(0)\n", " return correct / total * 100\n", "\n", "\n", "for epoch in range(EPOCHS):\n", " model.train()\n", " running_loss = 0.0\n", "\n", " for img, label in tqdm(train_loader, desc=f\"Epoch {epoch+1}/{EPOCHS}\"):\n", " img, label = img.to(device), label.to(device)\n", "\n", " with autocast('cuda'):\n", " out = model(img)\n", " loss = loss_fn(out, label)\n", "\n", " optimizer.zero_grad()\n", " scaler.scale(loss).backward()\n", " scaler.unscale_(optimizer)\n", " torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)\n", " scaler.step(optimizer)\n", " scaler.update()\n", " running_loss += loss.item()\n", "\n", " scheduler.step()\n", " acc = evaluate()\n", " avg_loss = running_loss / len(train_loader)\n", " current_lr = optimizer.param_groups[0]['lr']\n", " print(f\"Epoch {epoch+1:02d} | Loss: {avg_loss:.4f} | Acc: {acc:.2f}% | LR: {current_lr:.2e}\")" ] }, { "cell_type": "code", "execution_count": 36, "id": "8470d242", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "Saved to: C:\\Users\\kompe\\swin_resisc45.pth\n" ] } ], "source": [ "import os\n", "\n", "torch.save({\n", " 'model_state_dict': model.state_dict(),\n", " 'classes' : full_dataset.classes,\n", " 'num_classes' : 45,\n", "}, \"swin_resisc45.pth\")\n", "\n", "print(\"Saved to:\", os.path.abspath(\"swin_resisc45.pth\"))" ] }, { "cell_type": "code", "execution_count": 37, "id": "2f0fe44f", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "Checkpoint updated\n", "dict_keys(['model_state_dict', 'classes', 'num_classes', 'embed_dim'])\n" ] } ], "source": [ "# load existing checkpoint\n", "checkpoint = torch.load(r\"C:\\Users\\kompe\\swin_resisc45.pth\", map_location='cpu')\n", "\n", "# add missing keys with correct values\n", "checkpoint['embed_dim'] = 96\n", "checkpoint['num_classes'] = 45\n", "\n", "# resave\n", "torch.save(checkpoint, r\"D:\\NWPU\\HF\\swin-resisc45.pth\")\n", "print(\"Checkpoint updated\")\n", "print(checkpoint.keys())" ] }, { "cell_type": "code", "execution_count": null, "id": "4efe3ddf", "metadata": {}, "outputs": [], "source": [] } ], "metadata": { "kernelspec": { "display_name": "Python 3 (ipykernel)", "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.11.4" } }, "nbformat": 4, "nbformat_minor": 5 }