oom fix

config/config.yaml

@@ -59,12 +59,14 @@ model:
     anomaly_smooth_sigma: 1.5
 
 train:
-    batch_size_mvtec: 20
-    batch_size_visa: 12
+    # Paper uses 20/12. On ~14–16GB use 2–4 + AMP; raise if you have headroom.
+    batch_size_mvtec: 4
+    batch_size_visa: 2
+    use_amp: true
     optimizer: "adam"
     lr: 0.0001
     epochs: 100
-    num_workers: 4
+    num_workers: 2
     device: "cuda"
     seed: 42
     unified: true # every batch samples across ALL classes

notebooks/exploratory.ipynb

@@ -199,7 +199,7 @@
       "metadata": {},
       "outputs": [],
       "source": [
-        "# تست Reconstruction Sub-network به تنهایی\n",
+        "# Component smoke tests stay on CPU so they do not eat VRAM before the real model.\n",
         "recon = ReconstructionNet(norm_type=cfg[\"model\"][\"reconstruction_norm\"])\n",
         "x = torch.randn(2, 3, 256, 256)\n",
         "rec_img, rec_jnd, feats_inp, feats_rec, feats_jnd = recon(x)\n",
@@ -210,7 +210,6 @@
         "print(\"feats_rec:\", [f.shape for f in feats_rec])\n",
         "print(\"feats_jnd:\", [f.shape for f in feats_jnd])\n",
         "\n",
-        "# تست DiffNeck: مسیرهای (Rec,Inp), (JND,Inp), (Rec,JND,Inp)\n",
         "diffneck = DiffNeck(\n",
         "    channels_list=recon.level_channels,\n",
         "    out_ch=cfg[\"model\"][\"base_channels\"],\n",
@@ -221,7 +220,6 @@
         "assert desc.shape[0] == 2 and desc.shape[1] == cfg[\"model\"][\"base_channels\"]\n",
         "print(\"DiffNeck descriptor:\", desc.shape)\n",
         "\n",
-        "# تست LocalizationNet به تنهایی\n",
         "loc = LocalizationNet(\n",
         "    in_ch=3 + cfg[\"model\"][\"base_channels\"],\n",
         "    base_ch=cfg[\"model\"][\"base_channels\"],\n",
@@ -229,16 +227,16 @@
         "    dilation_rates=tuple(cfg[\"model\"][\"dilation_rates\"]),\n",
         "    norm_type=cfg[\"model\"][\"localization_norm\"],\n",
         ")\n",
-        "# pad descriptor to full res if needed\n",
-        "if desc.shape[-2:] != x.shape[-2:]:\n",
-        "    import torch.nn.functional as F\n",
-        "    desc_up = F.interpolate(desc, size=x.shape[-2:], mode=\"bilinear\", align_corners=False)\n",
-        "else:\n",
-        "    desc_up = desc\n",
-        "loc_in = torch.cat([x, desc_up], dim=1)\n",
-        "seg = loc(loc_in)\n",
+        "import torch.nn.functional as F\n",
+        "desc_up = desc if desc.shape[-2:] == x.shape[-2:] else F.interpolate(\n",
+        "    desc, size=x.shape[-2:], mode=\"bilinear\", align_corners=False\n",
+        ")\n",
+        "seg = loc(torch.cat([x, desc_up], dim=1))\n",
         "assert seg.shape == (2, 4, 256, 256)\n",
-        "print(\"Localization seg_logits:\", seg.shape)"
+        "print(\"Localization seg_logits:\", seg.shape)\n",
+        "\n",
+        "# Free unit-test nets / activations before putting OmniAL on GPU\n",
+        "del recon, diffneck, loc, x, rec_img, rec_jnd, feats_inp, feats_rec, feats_jnd, desc, desc_up, seg\n"
       ]
     },
     {
@@ -248,12 +246,22 @@
       "metadata": {},
       "outputs": [],
       "source": [
-        "# تست کامل مدل OmniAL (forward pass end-to-end)\n",
+        "# Full OmniAL forward (small batch). Clear leftover CUDA junk first.\n",
+        "if torch.cuda.is_available():\n",
+        "    torch.cuda.empty_cache()\n",
+        "\n",
         "model, optimizer, criterion, device = build_model_and_optimizer(cfg)\n",
-        "x = torch.randn(2, 3, *size).to(device)\n",
-        "out = model(x)\n",
+        "n_params = sum(p.numel() for p in model.parameters()) / 1e6\n",
+        "print(f\"params: {n_params:.1f}M | device={device}\")\n",
+        "\n",
+        "with torch.no_grad():\n",
+        "    x = torch.randn(2, 3, *size, device=device)\n",
+        "    out = model(x)\n",
         "for k, v in out.items():\n",
-        "    print(k, v.shape)"
+        "    print(k, v.shape)\n",
+        "del x, out\n",
+        "if torch.cuda.is_available():\n",
+        "    torch.cuda.empty_cache()\n"
       ]
     },
     {
@@ -271,17 +279,24 @@
       "metadata": {},
       "outputs": [],
       "source": [
+        "# Smoke-test loss on 2 samples only — NOT the full train batch (paper BS=20).\n",
+        "# Full BS=20 + autograd + this UNet can exceed ~14GB VRAM.\n",
+        "if torch.cuda.is_available():\n",
+        "    torch.cuda.empty_cache()\n",
+        "\n",
         "batch = next(iter(train_loader))\n",
-        "anomaly_img = batch[\"anomaly_img\"].to(device)\n",
+        "n = 2  # keep tiny for notebook testing\n",
+        "anomaly_img = batch[\"anomaly_img\"][:n].to(device)\n",
         "targets = {\n",
-        "    \"image\": batch[\"image\"].to(device),\n",
-        "    \"jnd\": batch[\"jnd\"].to(device),\n",
-        "    \"severity\": batch[\"severity\"].to(device),\n",
+        "    \"image\": batch[\"image\"][:n].to(device),\n",
+        "    \"jnd\": batch[\"jnd\"][:n].to(device),\n",
+        "    \"severity\": batch[\"severity\"][:n].to(device),\n",
         "}\n",
         "\n",
-        "outputs = model(anomaly_img)\n",
-        "losses = criterion(outputs, targets)\n",
-        "losses"
+        "with torch.no_grad():  # forward-only is enough to verify the loss pipeline\n",
+        "    outputs = model(anomaly_img)\n",
+        "    losses = criterion(outputs, targets)\n",
+        "losses\n"
       ]
     },
     {
@@ -290,7 +305,8 @@
       "metadata": {},
       "source": [
         "## 6. آموزش مدل (Training Loop)\n",
-        "برای تست سریع، ابتدا `epochs` را در `config.yaml` کم بگذار (مثلا 2) تا فقط pipeline را چک کنی."
+        "برای تست سریع فقط ۲ epoch. **حتماً DataLoader را دوباره بساز** (batch قبلی احتمالاً هنوز ۲۰ بوده).\n",
+        "اگر OOM دیدی: kernel را Restart کن، بعد `batch_size_mvtec` را روی ۲ بگذار.\n"
       ]
     },
     {
@@ -300,10 +316,22 @@
       "metadata": {},
       "outputs": [],
       "source": [
-        "# برای تست سریع فقط چند epoch:\n",
+        "# Quick pipeline check (2 epochs). Rebuild loader — changing cfg alone does NOT resize an existing DataLoader.\n",
+        "# Restart kernel first if VRAM is already nearly full from earlier cells.\n",
         "cfg[\"train\"][\"epochs\"] = 2\n",
+        "cfg[\"train\"][\"batch_size_mvtec\"] = 2   # safe on ~14GB with this architecture\n",
+        "cfg[\"train\"][\"use_amp\"] = True\n",
+        "cfg[\"train\"][\"num_workers\"] = 2\n",
+        "\n",
+        "if torch.cuda.is_available():\n",
+        "    torch.cuda.empty_cache()\n",
+        "    print(\"VRAM allocated before train:\", round(torch.cuda.memory_allocated()/1e9, 2), \"GB\")\n",
+        "\n",
+        "train_loader, train_dataset = build_dataloader(cfg, dataset_name=\"mvtec\", shuffle=True)\n",
+        "print(\"train batch_size:\", train_loader.batch_size, \"| samples:\", len(train_dataset))\n",
+        "\n",
         "history = train_model(model, train_loader, optimizer, criterion, device, cfg)\n",
-        "plot_loss_curves(history)"
+        "plot_loss_curves(history)\n"
       ]
     },
     {
@@ -388,4 +416,4 @@
   },
   "nbformat": 4,
   "nbformat_minor": 5
-}
+}

scripts/train.py

@@ -2,6 +2,7 @@ import os
 import sys
 import torch
 from torch.optim import Adam
+from torch.cuda.amp import GradScaler, autocast
 
 sys.path.append(os.path.join(os.path.dirname(__file__), ".."))
 from models.model import OmniAL
@@ -18,35 +19,44 @@ def build_model_and_optimizer(cfg: dict, device: str = None):
     return model, optimizer, criterion, device
 
 
-def run_one_epoch(model, dataloader, optimizer, criterion, device, train: bool = True):
+def run_one_epoch(model, dataloader, optimizer, criterion, device, train: bool = True,
+                  use_amp: bool = False, scaler: GradScaler = None):
     model.train() if train else model.eval()
     epoch_losses = {"total": 0.0, "l2_img": 0.0, "ssim_img": 0.0,
                      "l2_jnd": 0.0, "ssim_jnd": 0.0, "focal": 0.0}
     n_batches = 0
+    use_amp = bool(use_amp and device != "cpu" and torch.cuda.is_available())
+    if scaler is None:
+        scaler = GradScaler(enabled=use_amp)
 
     context = torch.enable_grad() if train else torch.no_grad()
     with context:
         for batch in dataloader:
-            anomaly_img = batch["anomaly_img"].to(device)
+            anomaly_img = batch["anomaly_img"].to(device, non_blocking=True)
             targets = {
-                "image": batch["image"].to(device),
-                "jnd": batch["jnd"].to(device),
-                "severity": batch["severity"].to(device),
+                "image": batch["image"].to(device, non_blocking=True),
+                "jnd": batch["jnd"].to(device, non_blocking=True),
+                "severity": batch["severity"].to(device, non_blocking=True),
             }
 
-            outputs = model(anomaly_img)
-            losses = criterion(outputs, targets)
+            with autocast(enabled=use_amp):
+                outputs = model(anomaly_img)
+                losses = criterion(outputs, targets)
 
             if train:
-                optimizer.zero_grad()
-                losses["total"].backward()
-                optimizer.step()
+                optimizer.zero_grad(set_to_none=True)
+                scaler.scale(losses["total"]).backward()
+                scaler.step(optimizer)
+                scaler.update()
 
             for k in epoch_losses:
                 v = losses[k]
                 epoch_losses[k] += v.item() if torch.is_tensor(v) else v
             n_batches += 1
 
+            # drop refs so peak memory can reclaim between steps
+            del anomaly_img, targets, outputs, losses
+
     for k in epoch_losses:
         epoch_losses[k] /= max(n_batches, 1)
     return epoch_losses
@@ -56,19 +66,30 @@ def train_model(model, dataloader, optimizer, criterion, device, cfg: dict, val_
     epochs = cfg["train"]["epochs"]
     ckpt_dir = cfg["logging"]["ckpt_dir"]
     save_every = cfg["logging"]["save_every"]
+    use_amp = bool(cfg["train"].get("use_amp", True))
     os.makedirs(ckpt_dir, exist_ok=True)
 
+    scaler = GradScaler(enabled=use_amp and device != "cpu" and torch.cuda.is_available())
     history = {"train_total": []}
     if val_loader is not None:
         history["val_total"] = []
 
+    print(f"Training: epochs={epochs}, amp={use_amp}, "
+          f"batch≈{getattr(dataloader, 'batch_size', '?')}, device={device}")
+
     for epoch in range(1, epochs + 1):
-        train_losses = run_one_epoch(model, dataloader, optimizer, criterion, device, train=True)
+        train_losses = run_one_epoch(
+            model, dataloader, optimizer, criterion, device,
+            train=True, use_amp=use_amp, scaler=scaler,
+        )
         history["train_total"].append(train_losses["total"])
         msg = f"[Epoch {epoch}/{epochs}] train_total={train_losses['total']:.4f}"
 
         if val_loader is not None:
-            val_losses = run_one_epoch(model, dataloader, optimizer, criterion, device, train=False)
+            val_losses = run_one_epoch(
+                model, val_loader, optimizer, criterion, device,
+                train=False, use_amp=use_amp, scaler=scaler,
+            )
             history["val_total"].append(val_losses["total"])
             msg += f" | val_total={val_losses['total']:.4f}"
 
@@ -79,4 +100,7 @@ def train_model(model, dataloader, optimizer, criterion, device, cfg: dict, val_
             torch.save(model.state_dict(), ckpt_path)
             print(f"  -> checkpoint saved: {ckpt_path}")
 
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+
     return history