Optimize notebook for CPU training: MobileNetV2 encoder, 15 epochs, batch 4

train_change_detection_model.ipynb

@@ -6,13 +6,12 @@
       "source": [
         "# Satellite Change Detection — Siamese U-Net Training\n",
         "\n",
-        "This notebook trains a **Siamese U-Net** on the **LEVIR-CD** dataset for pixel-level\n",
+        "This notebook trains a **Siamese U-Net** on the **LEVIR-CD+** dataset for pixel-level\n",
         "satellite image change detection. The exported model plugs directly into the\n",
         "AI Change Detection web app.\n",
         "\n",
-        "**Run in Google Colab** with a GPU runtime (Runtime → Change runtime type → T4 GPU).\n",
-        "\n",
-        "Training takes ~2-3 hours on a free T4."
+        "**Optimized for CPU** — uses a lightweight MobileNetV2 encoder and 15 epochs.\n",
+        "Training takes ~3-4 hours on a Colab CPU runtime."
       ]
     },
     {
@@ -197,10 +196,10 @@
         "val_ds = LEVIRCDDataset(DATA_ROOT, \"val\", val_transform)\n",
         "test_ds = LEVIRCDDataset(DATA_ROOT, \"test\", val_transform)\n",
         "\n",
-        "BATCH = 8\n",
-        "train_dl = DataLoader(train_ds, batch_size=BATCH, shuffle=True, num_workers=2, pin_memory=True)\n",
-        "val_dl = DataLoader(val_ds, batch_size=BATCH, shuffle=False, num_workers=2, pin_memory=True)\n",
-        "test_dl = DataLoader(test_ds, batch_size=BATCH, shuffle=False, num_workers=2, pin_memory=True)\n",
+        "BATCH = 4  # smaller batch for CPU\n",
+        "train_dl = DataLoader(train_ds, batch_size=BATCH, shuffle=True, num_workers=0, pin_memory=False)\n",
+        "val_dl = DataLoader(val_ds, batch_size=BATCH, shuffle=False, num_workers=0, pin_memory=False)\n",
+        "test_dl = DataLoader(test_ds, batch_size=BATCH, shuffle=False, num_workers=0, pin_memory=False)\n",
         "\n",
         "print(f\"Train: {len(train_ds)}, Val: {len(val_ds)}, Test: {len(test_ds)}\")"
       ],
@@ -214,7 +213,7 @@
         "## 4. Siamese U-Net Model\n",
         "\n",
         "Architecture:\n",
-        "- **Shared encoder** (ResNet34, ImageNet pretrained) processes both images\n",
+        "- **Shared encoder** (MobileNetV2, ImageNet pretrained) — lightweight and fast on CPU\n",
         "- Feature maps from both branches are **concatenated** at each decoder level\n",
         "- Standard U-Net decoder produces a binary change mask"
       ]
@@ -228,6 +227,9 @@
         "import segmentation_models_pytorch as smp\n",
         "\n",
         "\n",
+        "ENCODER_NAME = \"mobilenet_v2\"  # lightweight encoder for CPU training\n",
+        "\n",
+        "\n",
         "class SiameseUNet(nn.Module):\n",
         "    \"\"\"\n",
         "    Siamese U-Net for change detection.\n",
@@ -235,9 +237,8 @@
         "    concatenated features are decoded into a binary change mask.\n",
         "    \"\"\"\n",
         "\n",
-        "    def __init__(self, encoder_name=\"resnet34\", pretrained=True):\n",
+        "    def __init__(self, encoder_name=ENCODER_NAME, pretrained=True):\n",
         "        super().__init__()\n",
-        "        # Build a standard U-Net to reuse its encoder and decoder pieces\n",
         "        aux = smp.Unet(\n",
         "            encoder_name=encoder_name,\n",
         "            encoder_weights=\"imagenet\" if pretrained else None,\n",
@@ -246,8 +247,7 @@
         "        )\n",
         "        self.encoder = aux.encoder\n",
         "\n",
-        "        # The decoder expects concatenated features (2x channels at each level)\n",
-        "        encoder_channels = self.encoder.out_channels  # e.g. (3,64,64,128,256,512)\n",
+        "        encoder_channels = self.encoder.out_channels\n",
         "        doubled = tuple(c * 2 for c in encoder_channels)\n",
         "\n",
         "        self.decoder = smp.decoders.unet.decoder.UnetDecoder(\n",
@@ -261,23 +261,21 @@
         "        self.head = nn.Conv2d(16, 1, kernel_size=1)\n",
         "\n",
         "    def forward(self, img_a, img_b):\n",
-        "        # Shared encoder for both temporal images\n",
         "        feats_a = self.encoder(img_a)\n",
         "        feats_b = self.encoder(img_b)\n",
-        "\n",
-        "        # Concatenate features at every level\n",
         "        feats_cat = [torch.cat([fa, fb], dim=1) for fa, fb in zip(feats_a, feats_b)]\n",
-        "\n",
         "        decoded = self.decoder(*feats_cat)\n",
         "        logits = self.head(decoded)\n",
         "        return logits\n",
         "\n",
         "\n",
         "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
-        "model = SiameseUNet(encoder_name=\"resnet34\", pretrained=True).to(device)\n",
+        "model = SiameseUNet(encoder_name=ENCODER_NAME, pretrained=True).to(device)\n",
         "\n",
         "total_params = sum(p.numel() for p in model.parameters()) / 1e6\n",
-        "print(f\"Model on {device}, {total_params:.1f}M parameters\")"
+        "print(f\"Model on {device}, {total_params:.1f}M parameters\")\n",
+        "if device.type == \"cpu\":\n",
+        "    print(\"Running on CPU — training will take ~3-4 hours for 15 epochs\")"
       ],
       "execution_count": null,
       "outputs": []
@@ -337,10 +335,11 @@
       "cell_type": "code",
       "metadata": {},
       "source": [
+        "import time\n",
         "from tqdm.auto import tqdm\n",
         "\n",
-        "NUM_EPOCHS = 50\n",
-        "LR = 1e-4\n",
+        "NUM_EPOCHS = 15  # fewer epochs for CPU training\n",
+        "LR = 3e-4       # slightly higher LR to converge faster\n",
         "\n",
         "criterion = BCEDiceLoss(bce_weight=0.5)\n",
         "optimizer = torch.optim.Adam(model.parameters(), lr=LR, weight_decay=1e-4)\n",
@@ -348,8 +347,11 @@
         "\n",
         "best_f1 = 0.0\n",
         "history = {\"train_loss\": [], \"val_loss\": [], \"val_f1\": [], \"val_iou\": []}\n",
+        "train_start = time.time()\n",
         "\n",
         "for epoch in range(1, NUM_EPOCHS + 1):\n",
+        "    epoch_start = time.time()\n",
+        "\n",
         "    # --- Train ---\n",
         "    model.train()\n",
         "    running_loss = 0.0\n",
@@ -395,13 +397,18 @@
         "    history[\"val_f1\"].append(metrics[\"f1\"])\n",
         "    history[\"val_iou\"].append(metrics[\"iou\"])\n",
         "\n",
+        "    elapsed_min = (time.time() - epoch_start) / 60\n",
+        "    total_min = (time.time() - train_start) / 60\n",
+        "    eta_min = elapsed_min * (NUM_EPOCHS - epoch)\n",
+        "\n",
         "    print(\n",
         "        f\"Epoch {epoch:02d} | \"\n",
         "        f\"train_loss={train_loss:.4f} | \"\n",
         "        f\"val_loss={val_loss:.4f} | \"\n",
         "        f\"F1={metrics['f1']:.4f} | \"\n",
         "        f\"IoU={metrics['iou']:.4f} | \"\n",
-        "        f\"P={metrics['precision']:.4f} R={metrics['recall']:.4f}\"\n",
+        "        f\"P={metrics['precision']:.4f} R={metrics['recall']:.4f} | \"\n",
+        "        f\"{elapsed_min:.1f}min (ETA: {eta_min:.0f}min)\"\n",
         "    )\n",
         "\n",
         "    if metrics[\"f1\"] > best_f1:\n",
@@ -409,7 +416,8 @@
         "        torch.save(model.state_dict(), \"best_siamese_unet.pth\")\n",
         "        print(f\"  >> Saved best model (F1={best_f1:.4f})\")\n",
         "\n",
-        "print(f\"\\nTraining complete. Best val F1: {best_f1:.4f}\")"
+        "total_time = (time.time() - train_start) / 60\n",
+        "print(f\"\\nTraining complete in {total_time:.1f} minutes. Best val F1: {best_f1:.4f}\")"
       ],
       "execution_count": null,
       "outputs": []