{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.11.13","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"nvidiaTeslaT4","dataSources":[{"sourceId":99552,"databundleVersionId":13441085,"sourceType":"competition"},{"sourceId":12937280,"sourceType":"datasetVersion","datasetId":8119423}],"dockerImageVersionId":31089,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"# ====================================================\n# RSNA INTRACRANIAL ANEURYSM - STAGE 1 TRAINING (v2)\n# Uses Stage-0 prebuilt v2 cache (volumes, masks, pseudo_masks, brainmasks, manifest)\n# Two-phase training with per-sample segmentation weights and rich progress logs:\n#   Phase 1: real masks weighted (real_seg_weight), synthetic seg weight = 0.0\n#   Phase 2: real same, synthetic seg weight = small (default 0.075)\n# Saves: stage1_phase1_best.pth, stage1_phase2_best.pth, stage1_segmentation_best.pth\n# ====================================================\n\nimport os\nimport math\nimport time\nimport random\nimport numpy as np\nimport pandas as pd\nimport cv2\nfrom scipy import ndimage\nfrom typing import Optional, Tuple\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom torch.utils.data import Dataset, DataLoader\nfrom sklearn.model_selection import StratifiedKFold\nfrom sklearn.metrics import roc_auc_score\nfrom tqdm import tqdm","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-08-26T11:11:26.586994Z","iopub.execute_input":"2025-08-26T11:11:26.587252Z","iopub.status.idle":"2025-08-26T11:11:33.990205Z","shell.execute_reply.started":"2025-08-26T11:11:26.587232Z","shell.execute_reply":"2025-08-26T11:11:33.989649Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ====================================================\n# Config\n# ====================================================\nclass Config:\n    # --- Paths ---\n    PREBUILT_ROOT = \"/kaggle/input/rsna2025-v2-intracranial-aneurysm-detection-nb153/stage1_AneurysmNet_prebuilt_v2\"\n    MANIFEST_PATH = os.path.join(PREBUILT_ROOT, \"meta/manifest.csv\")\n    VOLUMES_DIR   = os.path.join(PREBUILT_ROOT, \"volumes\")\n    MASKS_DIR     = os.path.join(PREBUILT_ROOT, \"masks\")          # real\n    PSEUDO_DIR    = os.path.join(PREBUILT_ROOT, \"pseudo_masks\")   # synthetic\n    BRAINMASKS_DIR= os.path.join(PREBUILT_ROOT, \"brainmasks\")\n    MANIFEST_EXTRA_FIELDS = True  # stage-0 adds brainmask_relpath, brain_voxel_fraction\n\n    # --- Data ---\n    TARGET_SIZE = (48, 112, 112)  # (D,H,W)\n    USE_BRAINMASKS = True\n    BRAINMASK_KEY = 'm'\n    BRAINMASK_MIN_FRAC = 0.02  # if below, skip masking\n\n    # --- Training ---\n    DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n    MIXED_PRECISION = True\n    STAGE1_BATCH_SIZE = 12  # bump to 12/16 if GPU mem allows\n    NUM_WORKERS = 4        # bump to 2 original setting /4/8 unsupported on kaggle to reduce CPU bottlenecks\n    PREFETCH_FACTOR = 4\n    PERSISTENT_WORKERS = True\n    STAGE1_LR = 2e-4\n    WEIGHT_DECAY = 1e-4\n    EPOCHS_PHASE1 = 15\n    EPOCHS_PHASE2 = 10\n    EARLY_STOP_PATIENCE = 5\n    GRAD_ACCUM_STEPS = 8\n    # Validation throughput\n    VAL_BATCH_MULT = 1  # keep validation batch moderate to avoid I/O stalls\n    VAL_NUM_WORKERS = 4  # allow more workers for validation to feed GPUs\n\n    # --- Segmentation weights ---\n    REAL_SEG_DEFAULT_W = 0.9      # strengthen real supervision by default\n    PHASE1_SYNTH_SEG_W = 0.0\n    PHASE2_SYNTH_SEG_W = 0.05     # slightly lower synthetic weight in fine-tune\n    FOCAL_LOSS_WEIGHT = 0.2\n\n    # --- Loss variants ---\n    USE_TVERSKY = True\n    TV_ALPHA = 0.3\n    TV_BETA  = 0.7\n\n    # --- Augmentation ---\n    ZOOM_AUG_ENABLED = True        # enable positive-centric zoom aug\n    ZOOM_AUG_POS_FRAC = 0.65       # probability to apply on positives with mask\n    ZOOM_JITTER_VOX = 3            # jitter center by ±voxels\n\n\n    # --- Splits ---\n    FOLDS = 1   # set >1 later if you want CV here\n    SEED = 42\n\n# ====================================================\n# Utils\n# ====================================================\n\ndef set_seed(seed: int = 42):\n    random.seed(seed)\n    np.random.seed(seed)\n    torch.manual_seed(seed)\n    torch.cuda.manual_seed_all(seed)\n    torch.backends.cudnn.deterministic = True\n    torch.backends.cudnn.benchmark = True  # speeds up fixed-size convs\n\n\ndef load_manifest_df() -> pd.DataFrame:\n    df = pd.read_csv(Config.MANIFEST_PATH)\n    # Required columns: series_id, label, vol_relpath\n    for col in [\"series_id\", \"label\", \"vol_relpath\"]:\n        if col not in df.columns:\n            raise RuntimeError(f\"Manifest missing required column: {col}\")\n    return df\n\n\ndef gpu_mem_str():\n    if not torch.cuda.is_available():\n        return \"cpu\"\n    try:\n        a = torch.cuda.memory_allocated() / (1024**3)\n        r = torch.cuda.memory_reserved() / (1024**3)\n        return f\"{a:.2f}G/{r:.2f}G\"\n    except Exception:\n        return \"gpu\"","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-08-26T11:11:43.37067Z","iopub.execute_input":"2025-08-26T11:11:43.370949Z","iopub.status.idle":"2025-08-26T11:11:43.379109Z","shell.execute_reply.started":"2025-08-26T11:11:43.370928Z","shell.execute_reply":"2025-08-26T11:11:43.378548Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ====================================================\n# Dataset\n# ====================================================\nclass PrebuiltDataset(Dataset):\n    def __init__(self, df: pd.DataFrame, phase_synth_w: float):\n        self.df = df.reset_index(drop=True)\n        self.phase_synth_w = float(phase_synth_w)\n        # Oversample real-mask rows 3x to increase real supervision frequency\n        try:\n            has_real = self.df['mask_relpath'].fillna('').str.startswith('masks/')\n            real_df = self.df[has_real]\n            if len(real_df) > 0:\n                self.df = pd.concat([self.df, real_df, real_df, real_df], ignore_index=True)\n        except Exception:\n            pass\n\n    def __len__(self):\n        return len(self.df)\n\n    def _load_volume(self, sid: str) -> np.ndarray:\n        path = os.path.join(Config.VOLUMES_DIR, f\"{sid}.npy\")\n        v = np.load(path).astype(np.float32)  # (D,H,W), 0..1\n        v = np.nan_to_num(v, nan=0.0, posinf=1.0, neginf=0.0)\n        return v\n\n    def _load_brainmask(self, sid: str, frac: Optional[float], shp: Tuple[int,int,int]):\n        if not Config.USE_BRAINMASKS:\n            return None\n        if frac is not None and float(frac) < Config.BRAINMASK_MIN_FRAC:\n            return None\n        p = os.path.join(Config.BRAINMASKS_DIR, f\"{sid}_brainmask.npz\")\n        if not os.path.exists(p):\n            return None\n        try:\n            bm = np.load(p)[Config.BRAINMASK_KEY].astype(np.float32)\n            bm = np.nan_to_num(bm, nan=0.0, posinf=1.0, neginf=0.0)\n            if bm.shape != shp or bm.sum() <= 0:\n                return None\n            # guarantee a small safety margin by eroding high-frequency holes\n            try:\n                ker = np.ones((3,3), np.uint8)\n                for z in range(bm.shape[0]):\n                    bm[z] = cv2.morphologyEx(bm[z].astype(np.uint8), cv2.MORPH_CLOSE, ker)\n            except Exception:\n                pass\n            return bm\n        except Exception:\n            return None\n\n    def _load_mask(self, sid: str, mask_rel: str, is_synth: int, label: int) -> Tuple[np.ndarray, bool]:\n        # Returns (mask[D,H,W] float32 in {0,1}, is_synthetic: bool)\n        if isinstance(mask_rel, str) and len(mask_rel) > 0:\n            if mask_rel.startswith('masks/'):\n                p = os.path.join(Config.PREBUILT_ROOT, mask_rel)\n                if os.path.exists(p):\n                    m = np.load(p).astype(np.float32)\n                    m = np.nan_to_num(m, nan=0.0, posinf=1.0, neginf=0.0)\n                    return (m > 0).astype(np.float32), False\n            elif mask_rel.startswith('pseudo_masks/'):\n                p = os.path.join(Config.PREBUILT_ROOT, mask_rel)\n                if os.path.exists(p):\n                    m = np.load(p).astype(np.float32)\n                    m = np.nan_to_num(m, nan=0.0, posinf=1.0, neginf=0.0)\n                    return (m > 0).astype(np.float32), True\n        # Fallbacks\n        D,H,W = Config.TARGET_SIZE\n        if int(label) == 1:\n            return np.zeros((D,H,W), dtype=np.float32), True\n        else:\n            return np.zeros((D,H,W), dtype=np.float32), False\n\n    def _largest_component_bbox(self, mask: np.ndarray):\n        if mask is None or mask.max() <= 0:\n            return None\n        labeled, num = ndimage.label(mask > 0)\n        if num == 0:\n            return None\n        best_ct = 0\n        best_cid = 0\n        for cid in range(1, num+1):\n            ct = int((labeled == cid).sum())\n            if ct > best_ct:\n                best_ct = ct\n                best_cid = cid\n        comp = (labeled == best_cid)\n        idx = np.argwhere(comp)\n        if idx.size == 0:\n            return None\n        z0,y0,x0 = idx.min(axis=0)\n        z1,y1,x1 = idx.max(axis=0)\n        return int(z0), int(z1), int(y0), int(y1), int(x0), int(x1)\n\n    def _resize_volume_mask(self, vol: np.ndarray, msk: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:\n        td, th, tw = Config.TARGET_SIZE\n        D,H,W = vol.shape\n        # depth index selection\n        if D != td:\n            idx = np.linspace(0, max(D-1,0), num=td).astype(int) if D>0 else np.zeros(td, dtype=int)\n            vol = vol[idx]\n            msk = msk[idx]\n        # per-slice resize\n        if (H, W) != (th, tw):\n            outv = np.empty((td, th, tw), dtype=np.float32)\n            outm = np.empty((td, th, tw), dtype=np.float32)\n            for i in range(td):\n                outv[i] = cv2.resize(vol[i].astype(np.float32), (tw, th))\n                outm[i] = cv2.resize(msk[i].astype(np.float32), (tw, th), interpolation=cv2.INTER_NEAREST)\n            vol, msk = outv, (outm > 0.5).astype(np.float32)\n        return vol.astype(np.float32), msk.astype(np.float32)\n\n    def __getitem__(self, idx):\n        r = self.df.iloc[idx]\n        sid = str(r['series_id'])\n        label = int(r['label'])\n        mask_rel = r.get('mask_relpath', '') if isinstance(r.get('mask_relpath', ''), str) else ''\n        is_synth_col = int(r.get('is_synthetic', 0))\n        real_seg_weight = r.get('real_seg_weight', np.nan)\n        brain_frac = r.get('brain_voxel_fraction', np.nan)\n\n        vol = self._load_volume(sid)  # (D,H,W)\n        bm = self._load_brainmask(sid, brain_frac if pd.notna(brain_frac) else None, vol.shape)\n        if bm is not None:\n            vol = vol * bm  # gate\n\n        mask, is_synth = self._load_mask(sid, mask_rel, is_synth_col, label)\n\n        # Positive-centric zoom augmentation\n        if Config.ZOOM_AUG_ENABLED and (mask.max() > 0) and (random.random() < Config.ZOOM_AUG_POS_FRAC):\n            bbox = self._largest_component_bbox(mask)\n            if bbox is not None:\n                z0,z1,y0,y1,x0,x1 = bbox\n                # expand bbox with jitter\n                j = int(Config.ZOOM_JITTER_VOX)\n                zc = max(0, min(vol.shape[0]-1, (z0+z1)//2 + random.randint(-j, j)))\n                yc = max(0, min(vol.shape[1]-1, (y0+y1)//2 + random.randint(-j*2, j*2)))\n                xc = max(0, min(vol.shape[2]-1, (x0+x1)//2 + random.randint(-j*2, j*2)))\n                # choose cube edge roughly covering bbox\n                dz = max(4, z1 - z0 + 6)\n                dy = max(16, y1 - y0 + 24)\n                dx = max(16, x1 - x0 + 24)\n                edge_z = min(vol.shape[0], dz)\n                edge_y = min(vol.shape[1], dy)\n                edge_x = min(vol.shape[2], dx)\n                z1a = max(0, zc - edge_z//2); z2a = min(vol.shape[0], z1a + edge_z)\n                y1a = max(0, yc - edge_y//2); y2a = min(vol.shape[1], y1a + edge_y)\n                x1a = max(0, xc - edge_x//2); x2a = min(vol.shape[2], x1a + edge_x)\n                v_crop = vol[z1a:z2a, y1a:y2a, x1a:x2a]\n                m_crop = mask[z1a:z2a, y1a:y2a, x1a:x2a]\n                # resize back to target size\n                vol, mask = self._resize_volume_mask(v_crop, m_crop)\n\n        # per-sample seg weight\n        if is_synth:\n            seg_w = self.phase_synth_w\n        else:\n            if pd.notna(real_seg_weight):\n                try:\n                    rsw = float(real_seg_weight)\n                except Exception:\n                    rsw = Config.REAL_SEG_DEFAULT_W\n            else:\n                rsw = Config.REAL_SEG_DEFAULT_W\n            seg_w = float(np.clip(rsw, 0.2, 1.0))\n\n        # to tensors\n        vol_t = torch.from_numpy(vol).unsqueeze(0)         # [1,D,H,W]\n        mask_t = torch.from_numpy((mask > 0).astype(np.float32)).unsqueeze(0)\n        label_t = torch.tensor([float(label)], dtype=torch.float32)\n        segw_t  = torch.tensor([float(seg_w)], dtype=torch.float32)\n\n        return {\n            'series_id': sid,\n            'volume': vol_t,\n            'mask': mask_t,\n            'label': label_t,\n            'seg_weight': segw_t,\n            'is_synthetic_mask': torch.tensor([1.0 if is_synth else 0.0], dtype=torch.float32),\n        }\n\n# ====================================================\n# Simple 3D UNet + classifier head\n# ====================================================\nclass ConvBlock3D(nn.Module):\n    def __init__(self, in_ch, out_ch):\n        super().__init__()\n        self.net = nn.Sequential(\n            nn.Conv3d(in_ch, out_ch, 3, padding=1), nn.GroupNorm(num_groups=8, num_channels=out_ch), nn.ReLU(inplace=True),\n            nn.Conv3d(out_ch, out_ch, 3, padding=1), nn.GroupNorm(num_groups=8, num_channels=out_ch), nn.ReLU(inplace=True)\n        )\n    def forward(self, x):\n        return self.net(x)\n\nclass UNet3D(nn.Module):\n    def __init__(self, in_ch=1, base=24):\n        super().__init__()\n        b = base\n        self.enc1 = ConvBlock3D(in_ch, b)\n        self.pool1 = nn.MaxPool3d(2)\n        self.enc2 = ConvBlock3D(b, b*2)\n        self.pool2 = nn.MaxPool3d(2)\n        self.enc3 = ConvBlock3D(b*2, b*4)\n        self.pool3 = nn.MaxPool3d((2,2,2))\n        self.bott = ConvBlock3D(b*4, b*8)\n        self.up3 = nn.ConvTranspose3d(b*8, b*4, 2, stride=2)\n        self.dec3 = ConvBlock3D(b*8, b*4)\n        self.up2 = nn.ConvTranspose3d(b*4, b*2, 2, stride=2)\n        self.dec2 = ConvBlock3D(b*4, b*2)\n        self.up1 = nn.ConvTranspose3d(b*2, b, 2, stride=2)\n        self.dec1 = ConvBlock3D(b*2, b)\n        self.seg_head = nn.Conv3d(b, 1, 1)\n        # classification head from bottleneck features\n        self.cls_pool = nn.AdaptiveAvgPool3d(1)\n        self.cls_head = nn.Linear(b*8, 1)\n\n    def forward(self, x):  # x: [B,1,D,H,W]\n        e1 = self.enc1(x)\n        e2 = self.enc2(self.pool1(e1))\n        e3 = self.enc3(self.pool2(e2))\n        b  = self.bott(self.pool3(e3))\n        # decoder\n        d3 = self.up3(b)\n        d3 = torch.cat([d3, e3], dim=1)\n        d3 = self.dec3(d3)\n        d2 = self.up2(d3)\n        d2 = torch.cat([d2, e2], dim=1)\n        d2 = self.dec2(d2)\n        d1 = self.up1(d2)\n        d1 = torch.cat([d1, e1], dim=1)\n        d1 = self.dec1(d1)\n        seg = self.seg_head(d1)  # [B,1,D,H,W]\n        # classifier from bottleneck\n        cls = self.cls_head(self.cls_pool(b).flatten(1))  # [B,1]\n        return seg, cls\n\n# ====================================================\n# Losses\n# ====================================================\nclass DiceLoss(nn.Module):\n    def __init__(self, eps=1e-6):\n        super().__init__()\n        self.eps = eps\n    def forward(self, logits, targets, reduction='mean'):\n        probs = torch.sigmoid(logits)\n        num = 2 * (probs * targets).sum(dim=(2,3,4)) + self.eps\n        den = (probs.pow(2) + targets.pow(2)).sum(dim=(2,3,4)) + self.eps\n        dice = 1 - (num / den)  # per-sample\n        if reduction == 'none':\n            return dice\n        return dice.mean()\n\nclass FocalLoss(nn.Module):\n    def __init__(self, alpha=0.25, gamma=2.0, eps=1e-6):\n        super().__init__()\n        self.alpha = alpha\n        self.gamma = gamma\n        self.eps = eps\n    def forward(self, logits, targets, reduction='mean'):\n        probs = torch.sigmoid(logits).clamp(self.eps, 1-self.eps)\n        ce = -(targets*torch.log(probs) + (1-targets)*torch.log(1-probs))\n        pt = torch.where(targets==1, probs, 1-probs)\n        loss = self.alpha * (1-pt).pow(self.gamma) * ce\n        loss = loss.mean(dim=(2,3,4))  # per-sample\n        if reduction == 'none':\n            return loss\n        return loss.mean()\n\nclass EnhancedCombinedLoss(nn.Module):\n    def __init__(self):\n        super().__init__()\n        self.dice_loss = DiceLoss()\n        self.focal_loss = FocalLoss(alpha=0.25, gamma=2.0)\n        # foreground-weighted BCE: compute pos_weight dynamically per-batch\n        self.bce_vox = nn.BCEWithLogitsLoss(reduction='none')\n        self.bce_cls = nn.BCEWithLogitsLoss()\n    def forward(self, seg_logits, cls_logits, seg_targets, cls_targets, seg_weights: torch.Tensor):\n        # clamp seg logits to avoid AMP overflow\n        seg_logits = torch.nan_to_num(seg_logits, nan=0.0, posinf=20.0, neginf=-20.0)\n        B = seg_logits.shape[0]\n        dice_ps = self.dice_loss(seg_logits, seg_targets, reduction='none')\n        focal_ps= self.focal_loss(seg_logits, seg_targets, reduction='none')\n        # compute foreground weighting\n        with torch.no_grad():\n            fg = seg_targets.sum(dim=(2,3,4)).clamp(min=1.0)\n            tot = torch.tensor(seg_targets[0,0].numel(), device=seg_targets.device, dtype=seg_targets.dtype)\n            bg = (tot - fg).clamp(min=1.0)\n            pos_w = (bg / fg).view(-1, 1, 1, 1, 1)\n        bce_elem = F.binary_cross_entropy_with_logits(seg_logits, seg_targets, weight=pos_w.expand_as(seg_targets), reduction='none')\n        bce_ps  = bce_elem.view(B, -1).mean(dim=1)\n        dice_ps = torch.nan_to_num(dice_ps, nan=0.0)\n        focal_ps= torch.nan_to_num(focal_ps, nan=0.0)\n        bce_ps  = torch.nan_to_num(bce_ps,  nan=0.0)\n        seg_ps  = 0.5*dice_ps + 0.3*bce_ps + Config.FOCAL_LOSS_WEIGHT*focal_ps\n        seg_ps  = torch.nan_to_num(seg_ps, nan=0.0)\n        seg_w   = seg_weights.view(-1)\n        if (seg_w == 0).all():\n            seg_loss = seg_ps.new_tensor(0.0)\n        else:\n            seg_loss = (seg_ps * seg_w).mean()\n        # classification loss\n        cls_logits = torch.nan_to_num(cls_logits, nan=0.0, posinf=20.0, neginf=-20.0)\n        cls_loss= self.bce_cls(cls_logits.view(-1), cls_targets.view(-1))\n        total   = seg_loss + cls_loss\n        return total, seg_loss.detach(), cls_loss.detach()\n\n# ====================================================\n# Train / Validate with progress bars\n# ====================================================\n\ndef train_epoch(model, loader, optimizer, criterion, scaler, epoch=None, phase_name=\"P1\"):\n    model.train()\n    t_loss = t_seg = t_cls = 0.0\n    n = 0\n    pbar = tqdm(loader, desc=f\"Train {phase_name}{'' if epoch is None else f' [ep {epoch}]'}\", leave=False, mininterval=0.1)\n    optimizer.zero_grad(set_to_none=True)\n    for iter_idx, batch in enumerate(pbar):\n        vol   = batch['volume'].to(Config.DEVICE, non_blocking=True)\n        try:\n            vol = vol.to(memory_format=torch.channels_last_3d)\n        except Exception:\n            pass\n        mask  = batch['mask'].to(Config.DEVICE, non_blocking=True)\n        label = batch['label'].to(Config.DEVICE, non_blocking=True)\n        segw  = batch['seg_weight'].to(Config.DEVICE, non_blocking=True)\n        with torch.amp.autocast('cuda', enabled=Config.MIXED_PRECISION):\n            seg_logits, cls_logits = model(vol)\n            loss, seg_loss, cls_loss = criterion(seg_logits, cls_logits, mask, label, segw)\n        try:\n            scaler.scale(loss / Config.GRAD_ACCUM_STEPS).backward()\n            if ((iter_idx + 1) % Config.GRAD_ACCUM_STEPS) == 0:\n                torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)\n                scaler.step(optimizer)\n                scaler.update()\n                optimizer.zero_grad(set_to_none=True)\n        except Exception:\n            optimizer.zero_grad(set_to_none=True)\n            continue\n        bs = vol.size(0)\n        t_loss += loss.item()*bs; t_seg += seg_loss.item()*bs; t_cls += cls_loss.item()*bs; n += bs\n        pbar.set_postfix(loss=f\"{t_loss/max(n,1):.4f}\", seg=f\"{t_seg/max(n,1):.4f}\", cls=f\"{t_cls/max(n,1):.4f}\", lr=f\"{optimizer.param_groups[0]['lr']:.2e}\")\n    # finalize leftover accumulation\n    if (len(loader) % Config.GRAD_ACCUM_STEPS) != 0:\n        try:\n            torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)\n            scaler.step(optimizer)\n            scaler.update()\n        except Exception:\n            pass\n        finally:\n            optimizer.zero_grad(set_to_none=True)\n    return t_loss/n, t_seg/n, t_cls/n\n\n@torch.no_grad()\ndef validate_epoch(model, loader, criterion, epoch=None, phase_name=\"P1\"):\n    model.eval()\n    t_loss = t_seg = t_cls = 0.0\n    n = 0\n    all_probs = []\n    all_labels= []\n    pbar = tqdm(loader, desc=f\"Valid {phase_name}{'' if epoch is None else f' [ep {epoch}]'}\", leave=False, mininterval=0.1)\n    for batch in pbar:\n        vol   = batch['volume'].to(Config.DEVICE, non_blocking=True)\n        try:\n            vol = vol.to(memory_format=torch.channels_last_3d)\n        except Exception:\n            pass\n        mask  = batch['mask'].to(Config.DEVICE, non_blocking=True)\n        label = batch['label'].to(Config.DEVICE, non_blocking=True)\n        segw  = batch['seg_weight'].to(Config.DEVICE, non_blocking=True)\n        with torch.amp.autocast('cuda', enabled=Config.MIXED_PRECISION):\n            seg_logits, cls_logits = model(vol)\n            loss, seg_loss, cls_loss = criterion(seg_logits, cls_logits, mask, label, segw)\n        bs = vol.size(0)\n        t_loss += loss.item()*bs; t_seg += seg_loss.item()*bs; t_cls += cls_loss.item()*bs; n += bs\n        pbar.set_postfix(loss=f\"{t_loss/max(n,1):.4f}\", seg=f\"{t_seg/max(n,1):.4f}\", cls=f\"{t_cls/max(n,1):.4f}\")\n        all_probs.append(torch.sigmoid(cls_logits).detach().cpu().view(-1).numpy())\n        all_labels.append(label.detach().cpu().view(-1).numpy())\n    all_probs = np.concatenate(all_probs) if len(all_probs)>0 else np.array([])\n    all_labels = np.concatenate(all_labels) if len(all_labels)>0 else np.array([])\n    auc = np.nan\n    try:\n        if len(all_probs)>0 and len(np.unique(all_labels)) > 1:\n            auc = float(roc_auc_score(all_labels, all_probs))\n    except Exception:\n        pass\n    return t_loss/n, t_seg/n, t_cls/n, auc\n\n# ====================================================\n# Main\n# ====================================================\n\ndef run_training():\n    set_seed(Config.SEED)\n    # TF32 for better throughput\n    try:\n        torch.backends.cuda.matmul.allow_tf32 = True\n        torch.backends.cudnn.allow_tf32 = True\n        torch.set_float32_matmul_precision('high')\n    except Exception:\n        pass\n    df = load_manifest_df()\n\n    # Build a single stratified split (can expand to CV later)\n    y = df['label'].astype(int).values\n    skf = StratifiedKFold(n_splits=5, shuffle=True, random_state=Config.SEED)\n    train_idx, val_idx = next(skf.split(np.zeros_like(y), y))\n    train_df = df.iloc[train_idx].reset_index(drop=True)\n    val_df   = df.iloc[val_idx].reset_index(drop=True)\n\n    # Phase 1 datasets/loads\n    ds_train_p1 = PrebuiltDataset(train_df, phase_synth_w=Config.PHASE1_SYNTH_SEG_W)\n    ds_val_p1   = PrebuiltDataset(val_df,   phase_synth_w=Config.PHASE1_SYNTH_SEG_W)\n    dl_train = DataLoader(ds_train_p1, batch_size=Config.STAGE1_BATCH_SIZE, shuffle=True,\n                          num_workers=Config.NUM_WORKERS, pin_memory=True,\n                          prefetch_factor=Config.PREFETCH_FACTOR,\n                          persistent_workers=Config.PERSISTENT_WORKERS)\n    dl_val   = DataLoader(ds_val_p1,   batch_size=Config.STAGE1_BATCH_SIZE * Config.VAL_BATCH_MULT, shuffle=False,\n                          num_workers=Config.VAL_NUM_WORKERS, pin_memory=True,\n                          prefetch_factor=Config.PREFETCH_FACTOR,\n                          persistent_workers=Config.PERSISTENT_WORKERS)\n\n    model = UNet3D(in_ch=1, base=24).to(Config.DEVICE)\n    try:\n        model = model.to(memory_format=torch.channels_last_3d)\n    except Exception:\n        pass\n    # Multi-GPU (if available): enable DP for train and val\n    dp_enabled = False\n    if torch.cuda.is_available() and torch.cuda.device_count() > 1:\n        dp_enabled = True\n        print(f\"Using {torch.cuda.device_count()} GPUs via DataParallel\")\n        model = nn.DataParallel(model)\n\n    optimizer = torch.optim.AdamW(model.parameters(), lr=Config.STAGE1_LR, weight_decay=Config.WEIGHT_DECAY)\n    criterion = EnhancedCombinedLoss().to(Config.DEVICE)\n    scaler = torch.amp.GradScaler('cuda', enabled=Config.MIXED_PRECISION)\n\n    best_loss = float('inf'); best_state = None; patience = 0\n    for epoch in range(1, Config.EPOCHS_PHASE1+1):\n        print(f\"\\n[Phase 1] Epoch {epoch}/{Config.EPOCHS_PHASE1}\")\n        tr_loss, tr_seg, tr_cls = train_epoch(model, dl_train, optimizer, criterion, scaler, epoch=epoch, phase_name='P1')\n        # Always validate every epoch\n        va_loss, va_seg, va_cls, va_auc = validate_epoch(model, dl_val, criterion, epoch=epoch, phase_name='P1')\n        print(f\"Train Loss: {tr_loss:.4f} | Seg: {tr_seg:.4f} | Cls: {tr_cls:.4f} | GPU {gpu_mem_str()}\")\n        print(f\" Val  Loss: {va_loss:.4f} | Seg: {va_seg:.4f} | Cls: {va_cls:.4f} | AUC: {va_auc if not np.isnan(va_auc) else 'NA'} | GPU {gpu_mem_str()}\")\n        if va_loss < best_loss - 1e-5:\n            best_loss = va_loss; best_state = {k:v.detach().cpu() for k,v in (model.module.state_dict() if isinstance(model, nn.DataParallel) else model.state_dict()).items()}; patience = 0\n            torch.save(best_state, 'stage1_phase1_best.pth')\n            print(\"💾 Saved Phase 1 best checkpoint\")\n        else:\n            patience += 1\n            if patience >= Config.EARLY_STOP_PATIENCE:\n                print(\"Early stopping Phase 1\")\n                break\n\n    if best_state is not None:\n        if isinstance(model, nn.DataParallel):\n            model.module.load_state_dict(best_state)\n        else:\n            model.load_state_dict(best_state)\n\n    # Phase 2: small synthetic weight\n    print(\"\\n====== PHASE 2: enabling small synthetic seg supervision ======\")\n    ds_train_p2 = PrebuiltDataset(train_df, phase_synth_w=Config.PHASE2_SYNTH_SEG_W)\n    ds_val_p2   = PrebuiltDataset(val_df,   phase_synth_w=Config.PHASE2_SYNTH_SEG_W)\n    dl_train2 = DataLoader(ds_train_p2, batch_size=Config.STAGE1_BATCH_SIZE, shuffle=True,\n                           num_workers=Config.NUM_WORKERS, pin_memory=True,\n                           prefetch_factor=Config.PREFETCH_FACTOR,\n                           persistent_workers=Config.PERSISTENT_WORKERS)\n    dl_val2   = DataLoader(ds_val_p2,   batch_size=Config.STAGE1_BATCH_SIZE * Config.VAL_BATCH_MULT, shuffle=False,\n                           num_workers=Config.VAL_NUM_WORKERS, pin_memory=True,\n                           prefetch_factor=Config.PREFETCH_FACTOR,\n                           persistent_workers=Config.PERSISTENT_WORKERS)\n\n    # optional: lower LR a bit for fine-tune\n    for g in optimizer.param_groups:\n        g['lr'] = Config.STAGE1_LR * 0.5\n\n    best2 = float('inf'); best2_state = None; patience = 0\n    for epoch in range(1, Config.EPOCHS_PHASE2+1):\n        print(f\"\\n[Phase 2] Epoch {epoch}/{Config.EPOCHS_PHASE2}\")\n        tr_loss, tr_seg, tr_cls = train_epoch(model, dl_train2, optimizer, criterion, scaler, epoch=epoch, phase_name='P2')\n        va_loss, va_seg, va_cls, va_auc = validate_epoch(model, dl_val2, criterion, epoch=epoch, phase_name='P2')\n        print(f\"Train Loss: {tr_loss:.4f} | Seg: {tr_seg:.4f} | Cls: {tr_cls:.4f} | GPU {gpu_mem_str()}\")\n        print(f\" Val  Loss: {va_loss:.4f} | Seg: {va_seg:.4f} | Cls: {va_cls:.4f} | AUC: {va_auc if not np.isnan(va_auc) else 'NA'} | GPU {gpu_mem_str()}\")\n        if va_loss < best2 - 1e-5:\n            best2 = va_loss; best2_state = {k:v.detach().cpu() for k,v in (model.module.state_dict() if isinstance(model, nn.DataParallel) else model.state_dict()).items()}; patience = 0\n            torch.save(best2_state, 'stage1_phase2_best.pth')\n            print(\"💾 Saved Phase 2 best checkpoint\")\n        else:\n            patience += 1\n            if patience >= Config.EARLY_STOP_PATIENCE:\n                print(\"Early stopping Phase 2\")\n                break\n\n    final_state = best2_state or best_state or (model.module.state_dict() if isinstance(model, nn.DataParallel) else model.state_dict())\n    torch.save(final_state, 'stage1_segmentation_best.pth')\n    print(\"\\n✅ Stage 1 complete. Saved: stage1_segmentation_best.pth\")\n\n\nif __name__ == '__main__':\n    run_training()\n","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2025-08-26T11:11:46.894597Z","iopub.execute_input":"2025-08-26T11:11:46.895158Z","iopub.status.idle":"2025-08-26T18:42:36.014619Z","shell.execute_reply.started":"2025-08-26T11:11:46.895137Z","shell.execute_reply":"2025-08-26T18:42:36.013925Z"}},"outputs":[],"execution_count":null}]}