{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.12.12","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"nvidiaTeslaT4","dataSources":[{"sourceId":112509,"databundleVersionId":14254895,"sourceType":"competition"}],"dockerImageVersionId":31259,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"#show timm version\nimport timm\nprint(timm.__version__)","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2026-01-16T14:31:50.060583Z","iopub.execute_input":"2026-01-16T14:31:50.061012Z","iopub.status.idle":"2026-01-16T14:32:01.442348Z","shell.execute_reply.started":"2026-01-16T14:31:50.060977Z","shell.execute_reply":"2026-01-16T14:32:01.441661Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Config\nimport os, gc, math, cv2, numpy as np, pandas as pd\nfrom tqdm import tqdm\nimport torch, torch.nn as nn, torch.optim as optim\nfrom torch.optim.lr_scheduler import LambdaLR\nfrom torch.utils.data import Dataset, DataLoader\nimport albumentations as A\nfrom albumentations.pytorch import ToTensorV2\nimport timm\nfrom timm.utils import ModelEmaV2\nfrom sklearn.model_selection import KFold, StratifiedGroupKFold\nimport matplotlib.pyplot as plt\nimport seaborn as sns\n\nclass CFG:\n    CREATE_SUBMISSION = True\n    USE_TQDM        = False\n    PRETRAINED_DIR  = None\n    PRETRAINED      = True\n    BASE_PATH       = '/kaggle/input/csiro-biomass'\n    SEED            = 82947501\n    FOLDS_TO_TRAIN   = [0,1,2,3,4]\n    TRAIN_CSV       = os.path.join(BASE_PATH, 'train.csv')\n    TRAIN_IMAGE_DIR = os.path.join(BASE_PATH, 'train')\n    TEST_IMAGE_DIR = '/kaggle/input/csiro-biomass/test'\n    TEST_CSV = '/kaggle/input/csiro-biomass/test.csv'\n    SUBMISSION_DIR  = '/kaggle/working/'\n    MODEL_DIR_012       = '/kaggle/input/5-folds-dinov3-840m/other/fold-0-1-2/1'\n    MODEL_DIR_34       = '/kaggle/input/5-folds-dinov3-840m/other/fold-0-1-2/2'\n    N_FOLDS         = 5\n\n\n    MODEL_NAME      = 'vit_large_patch16_dinov3.lvd1689m'  \n    BACKBONE_PATH   = '/kaggle/input/vit-large-patch16-dinov3-lvd1689m-backbone-pth/vit_large_patch16_dinov3.lvd1689m_backbone.pth'\n    # MODEL_NAME      = 'vit_huge_plus_patch16_dinov3.lvd1689m'\n    # BACKBONE_PATH   = '/kaggle/input/vit-huge-plus-patch16-dinov3-lvd1689m/vit_huge_plus_patch16_dinov3.lvd1689m_backbone.pth'\n    MODEL_DIR = '/kaggle/working/'\n    IMG_SIZE        = 512\n\n    VAL_TTA_TIMES   = 1\n    TTA_STEPS       = 1\n    \n    \n    BATCH_SIZE      = 8\n    GRAD_ACC        = 4\n    NUM_WORKERS     = 4\n    EPOCHS          = 100\n    FREEZE_EPOCHS   = 2\n    WARMUP_EPOCHS   = 6\n    LR_REST         = 1e-3\n    LR_BACKBONE     = 6e-5\n    WD              = 1e-2\n    EMA_DECAY       = 0.85\n    PATIENCE        = 5\n    TARGET_COLS     = ['Dry_Total_g', 'GDM_g', 'Dry_Green_g']\n    DERIVED_COLS    = ['Dry_Clover_g', 'Dry_Dead_g']\n    ALL_TARGET_COLS = ['Dry_Green_g','Dry_Dead_g','Dry_Clover_g','GDM_g','Dry_Total_g']\n    R2_WEIGHTS      = np.array([0.1, 0.1, 0.1, 0.2, 0.5])\n    LOSS_WEIGHTS    = np.array([0.1, 0.1, 0.1, 0.0, 0.0])\n    DEVICE          = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n\nprint(f'Device : {CFG.DEVICE}')\nprint(f'Backbone: {CFG.MODEL_NAME} | Input: {CFG.IMG_SIZE}')\nprint(f'Freeze Epochs: {CFG.FREEZE_EPOCHS} | Warmup: {CFG.WARMUP_EPOCHS}')\nprint(f'EMA Decay: {CFG.EMA_DECAY} | Grad Acc: {CFG.GRAD_ACC}')","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-16T14:32:01.443669Z","iopub.execute_input":"2026-01-16T14:32:01.443908Z","iopub.status.idle":"2026-01-16T14:32:03.767208Z","shell.execute_reply.started":"2026-01-16T14:32:01.443886Z","shell.execute_reply":"2026-01-16T14:32:03.766420Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import numpy as np\nimport os\n\ndef weighted_r2_score(y_true: np.ndarray, y_pred: np.ndarray):\n    weights = CFG.R2_WEIGHTS\n    r2_scores = []\n    for i in range(y_true.shape[1]):\n        yt = y_true[:, i]; yp = y_pred[:, i]\n        ss_res = np.sum((yt - yp) ** 2)\n        ss_tot = np.sum((yt - np.mean(yt)) ** 2)\n        r2 = 1 - ss_res / ss_tot if ss_tot > 0 else 0.0\n        r2_scores.append(r2)\n    r2_scores = np.array(r2_scores)\n    weighted = np.sum(r2_scores * weights) / np.sum(weights)\n    return weighted, r2_scores\n\ndef weighted_r2_score_global(y_true: np.ndarray, y_pred: np.ndarray):\n    weights = CFG.R2_WEIGHTS\n    flat_true = y_true.reshape(-1)\n    flat_pred = y_pred.reshape(-1)\n    w = np.tile(weights, y_true.shape[0])\n    mean_w = np.sum(w * flat_true) / np.sum(w)\n    ss_res = np.sum(w * (flat_true - flat_pred) ** 2)\n    ss_tot = np.sum(w * (flat_true - mean_w) ** 2)\n    global_r2 = 1 - ss_res / ss_tot if ss_tot > 0 else 0.0\n    avg_r2, per_r2 = weighted_r2_score(y_true, y_pred)\n    return global_r2, avg_r2, per_r2\n\ndef analyze_errors(val_df, y_true, y_pred, targets, top_n=5):\n    print(f'\\n--- Top {top_n} High Loss Samples per Target ---')\n    y_true = np.array(y_true)\n    y_pred = np.array(y_pred)\n    \n    for i, target in enumerate(targets):\n        errors = np.abs(y_true[:, i] - y_pred[:, i])\n        top_indices = np.argsort(errors)[::-1][:top_n]\n        \n        print(f'\\nTarget: {target}')\n        print(f'{\"Index\":<6} | {\"Image Path\":<40} | {\"True\":<10} | {\"Pred\":<10} | {\"AbsErr\":<10}')\n        print('-' * 90)\n        \n        for idx in top_indices:\n            path = val_df.iloc[idx]['image_path']\n            path_disp = os.path.basename(path)\n            t_val = y_true[idx, i]\n            p_val = y_pred[idx, i]\n            err = errors[idx]\n            print(f'{idx:<6} | {path_disp:<40} | {t_val:<10.4f} | {p_val:<10.4f} | {err:<10.4f}')\ndef analyze_errors(val_df, y_true, y_pred, targets, top_n=5):\n    print(f'\\n--- Top {top_n} High Loss Samples per Target ---')\n    y_true = np.array(y_true)\n    y_pred = np.array(y_pred)\n    \n    for i, target in enumerate(targets):\n        errors = np.abs(y_true[:, i] - y_pred[:, i])\n        top_indices = np.argsort(errors)[::-1][:top_n]\n        \n        print(f'\\nTarget: {target}')\n        header = f'{\"Index\":<6} | {\"Image Path\":<40} | {\"State\":<6} | {\"True\":<10} | {\"Pred\":<10} | {\"AbsErr\":<10}'\n        print(header)\n        print('-' * len(header))\n        \n        for idx in top_indices:\n            path = val_df.iloc[idx]['image_path']\n            path_disp = os.path.basename(path)\n            state = val_df.iloc[idx]['State'] if 'State' in val_df.columns else 'NA'\n            t_val = y_true[idx, i]\n            p_val = y_pred[idx, i]\n            err = errors[idx]\n            print(f'{idx:<6} | {path_disp:<40} | {str(state):<6} | {t_val:<10.4f} | {p_val:<10.4f} | {err:<10.4f}')\ndef compare_train_val(tr_df, val_df, targets, show_plots=True):\n    \"\"\"Quick comparison of target distributions and metadata between train and val splits.\"\"\"\n    print(\"\\n--- Train / Val Comparison ---\")\n\n    for t in targets:\n        tr = tr_df.get(t, pd.Series(dtype=float)).dropna()\n        val = val_df.get(t, pd.Series(dtype=float)).dropna()\n        print(f\"\\nTarget: {t}\")\n        print(f\"  Train: n={len(tr)} mean={tr.mean():.3f} std={tr.std():.3f} min={tr.min():.3f} max={tr.max():.3f}\")\n        print(f\"  Val  : n={len(val)} mean={val.mean():.3f} std={val.std():.3f} min={val.min():.3f} max={val.max():.3f}\")\n        if show_plots:\n            try:\n                plt.figure(figsize=(6, 3))\n                sns.kdeplot(tr, label='train', fill=True)\n                sns.kdeplot(val, label='val', fill=True)\n                plt.legend()\n                plt.title(f'Distribution: {t}')\n                plt.show()\n            except Exception as e:\n                print('  Could not plot distributions for', t, '-', e)\n\n    # Compare Sampling_Date and State if present\n    if 'Sampling_Date' in tr_df.columns:\n        try:\n            tr_dates = pd.to_datetime(tr_df['Sampling_Date'], errors='coerce')\n            val_dates = pd.to_datetime(val_df['Sampling_Date'], errors='coerce')\n            print(\"\\nSampling_Date range:\")\n            print(f\"  Train: {tr_dates.min()} -> {tr_dates.max()} (missing {tr_dates.isna().sum()})\")\n            print(f\"  Val  : {val_dates.min()} -> {val_dates.max()} (missing {val_dates.isna().sum()})\")\n        except Exception as e:\n            print('  Could not parse Sampling_Date:', e)\n    if 'State' in tr_df.columns:\n        print(\"\\nState distribution (train vs val):\")\n        tr_state = tr_df['State'].value_counts(normalize=True)\n        val_state = val_df['State'].value_counts(normalize=True)\n        state_df = pd.concat([tr_state, val_state], axis=1, keys=['train', 'val']).fillna(0)\n\n        print(state_df)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-16T14:32:03.768240Z","iopub.execute_input":"2026-01-16T14:32:03.768647Z","iopub.status.idle":"2026-01-16T14:32:03.784873Z","shell.execute_reply.started":"2026-01-16T14:32:03.768622Z","shell.execute_reply":"2026-01-16T14:32:03.784016Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def get_train_transforms():\n    return A.Compose([\n        A.Resize(CFG.IMG_SIZE, CFG.IMG_SIZE),\n        A.HorizontalFlip(p=0.5),\n        A.VerticalFlip(p=0.5),\n        A.RandomRotate90(p=0.5),\n        A.Rotate(limit=(-10, 10), p=0.3, interpolation=cv2.INTER_LINEAR, border_mode=cv2.BORDER_REFLECT_101),\n        A.ColorJitter(brightness=0.1, contrast=0.1, saturation=0.1, hue=0.05, p=0.5),\n        A.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),\n        ToTensorV2()\n    ], p=1.0)\n\ndef get_val_transforms():\n    return A.Compose([\n        A.Resize(CFG.IMG_SIZE, CFG.IMG_SIZE),\n        A.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),\n        ToTensorV2()\n    ], p=1.0)\n\ndef get_tta_transforms(mode=0):\n    # mode 0: original\n    # mode 1: hflip\n    # mode 2: vflip\n    # mode 3: rotate90\n    transforms_list = [\n        A.Resize(CFG.IMG_SIZE, CFG.IMG_SIZE),\n    ]\n    \n    if mode == 1:\n        transforms_list.append(A.HorizontalFlip(p=1.0))\n    elif mode == 2:\n        transforms_list.append(A.VerticalFlip(p=1.0))\n    elif mode == 3:\n        transforms_list.append(A.RandomRotate90(p=1.0)) # RandomRotate90 with p=1.0 rotates 90, 180, 270 randomly? \n        # Albumentations RandomRotate90 rotates by 90, 180, 270. \n        # Reference uses transforms.RandomRotation([90, 90]) which is exactly 90 degrees.\n        # To match exactly 90 degrees in Albumentations, we might need Rotate(limit=(90,90), p=1.0)\n        # But RandomRotate90 is standard TTA. Let's use Rotate(limit=(90,90)) to be precise if that's what reference does.\n        # Reference: transforms.RandomRotation([90, 90]) -> rotates by exactly 90 degrees.\n        transforms_list.append(A.Rotate(limit=(90, 90), p=1.0, interpolation=cv2.INTER_LINEAR, border_mode=cv2.BORDER_REFLECT_101))\n\n    transforms_list.extend([\n        A.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),\n        ToTensorV2()\n    ])\n    \n    return A.Compose(transforms_list, p=1.0)\ndef clean_image(img):\n    # 1. Safe Crop (Remove artifacts at the bottom)\n    h, w = img.shape[:2]\n    # Cut bottom 10% where artifacts often appear\n    img = img[0:int(h*0.90), :] \n\n    # 2. Inpaint Date Stamp (Remove orange text)\n    hsv = cv2.cvtColor(img, cv2.COLOR_RGB2HSV)\n    # Define orange color range (adjust as needed)\n    lower = np.array([5, 150, 150])\n    upper = np.array([25, 255, 255])\n    mask = cv2.inRange(hsv, lower, upper)\n\n    # Dilate mask to cover text edges and reduce noise\n    mask = cv2.dilate(mask, np.ones((3,3), np.uint8), iterations=2)\n\n    # Inpaint if mask is not empty\n    if np.sum(mask) > 0:\n        img = cv2.inpaint(img, mask, 3, cv2.INPAINT_TELEA)\n\n    return img\nclass BiomassDataset(Dataset):\n    def __init__(self, df, transform, img_dir):\n        self.df = df\n        self.transform = transform\n        self.img_dir = img_dir\n        self.paths = df['image_path'].values\n        self.labels = df[CFG.ALL_TARGET_COLS].values.astype(np.float32)\n\n    def __len__(self):\n        return len(self.df)\n\n    def __getitem__(self, idx):\n        path = os.path.join(self.img_dir, os.path.basename(self.paths[idx]))\n        img = cv2.imread(path)\n        if img is None:\n            img = np.zeros((1000, 2000, 3), dtype=np.uint8)\n        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)\n        img = clean_image(img)\n        h, w, _ = img.shape\n        mid = w // 2\n        left = img[:, :mid]\n        right = img[:, mid:]\n        left = self.transform(image=left)['image']\n        right = self.transform(image=right)['image']\n        label = torch.from_numpy(self.labels[idx])\n        return left, right, label","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-16T14:32:03.787084Z","iopub.execute_input":"2026-01-16T14:32:03.787432Z","iopub.status.idle":"2026-01-16T14:32:03.812819Z","shell.execute_reply.started":"2026-01-16T14:32:03.787409Z","shell.execute_reply":"2026-01-16T14:32:03.812060Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# class BiomassModel(nn.Module):\n#     def __init__(self, model_name, pretrained=True):\n#         super().__init__()\n#         self.model_name = model_name\n#         self.backbone = timm.create_model(self.model_name, pretrained=False, num_classes=0, global_pool='avg')\n#         nf = self.backbone.num_features\n#         comb = nf * 2\n#         self.head_green_raw  = nn.Sequential(nn.Linear(comb, comb//2), nn.GELU(), nn.Dropout(0.3), nn.Linear(comb//2, 1))\n#         self.head_clover_raw = nn.Sequential(nn.Linear(comb, comb//2), nn.GELU(), nn.Dropout(0.3), nn.Linear(comb//2, 1))\n#         self.head_dead_raw   = nn.Sequential(nn.Linear(comb, comb//2), nn.GELU(), nn.Dropout(0.3), nn.Linear(comb//2, 1))\n#         if pretrained:\n#             self.load_pretrained()\n    \n#     def load_pretrained(self):\n#         try:\n#             sd = timm.create_model(self.model_name, pretrained=True, num_classes=0, global_pool='avg').state_dict()\n#             self.backbone.load_state_dict(sd, strict=False)\n#             print('Pretrained weights loaded.')\n#         except Exception as e:\n#             print(f'Warning: pretrained load failed: {e}')\n    \n#     def forward(self, left, right):\n#         fl = self.backbone(left)\n#         fr = self.backbone(right)\n#         x  = torch.cat([fl, fr], dim=1)\n#         green  = self.head_green_raw(x)\n#         # clover = torch.nn.functional.softplus(self.head_clover_raw(x))\n#         # dead   = torch.nn.functional.softplus(self.head_dead_raw(x))\n#         clover = self.head_clover_raw(x)\n#         dead   = self.head_dead_raw(x)\n#         gdm    = green + clover\n#         total  = gdm + dead\n#         return total, gdm, green, clover, dead\nclass LocalMambaBlock(nn.Module):\n    \"\"\"\n    Lightweight Mamba-style block (Gated CNN) from the reference notebook.\n    Efficiently mixes tokens with linear complexity.\n    \"\"\"\n    def __init__(self, dim, kernel_size=5, dropout=0.0):\n        super().__init__()\n        self.norm = nn.LayerNorm(dim)\n        # Depthwise conv mixes spatial information locally\n        self.dwconv = nn.Conv1d(dim, dim, kernel_size=kernel_size, padding=kernel_size // 2, groups=dim)\n        self.gate = nn.Linear(dim, dim)\n        self.proj = nn.Linear(dim, dim)\n        self.drop = nn.Dropout(dropout)\n\n    def forward(self, x):\n        # x: (Batch, Tokens, Dim)\n        shortcut = x\n        x = self.norm(x)\n        # Gating mechanism\n        g = torch.sigmoid(self.gate(x))\n        x = x * g\n        # Spatial mixing via 1D Conv (requires transpose)\n        x = x.transpose(1, 2)  # -> (B, D, N)\n        x = self.dwconv(x)\n        x = x.transpose(1, 2)  # -> (B, N, D)\n        # Projection\n        x = self.proj(x)\n        x = self.drop(x)\n        return shortcut + x\n\nclass BiomassModel(nn.Module):\n    def __init__(self, model_name, pretrained=True, backbone_path=None):\n        super().__init__()\n        self.model_name = model_name\n        self.backbone_path = backbone_path\n        \n        # 1. Load Backbone with global_pool='' to keep patch tokens\n        #    (B, 197, 1024) instead of (B, 1024)\n        self.backbone = timm.create_model(self.model_name, pretrained=False, num_classes=0, global_pool='')\n        \n        # 2. Enable Gradient Checkpointing (Crucial for ViT-Large memory!)\n        if hasattr(self.backbone, 'set_grad_checkpointing'):\n            self.backbone.set_grad_checkpointing(True)\n            print(\"✓ Gradient Checkpointing enabled (saves ~50% VRAM)\")\n            \n        nf = self.backbone.num_features\n        \n        # 3. Mamba Fusion Neck\n        #    Mixes the concatenated tokens [Left, Right]\n        self.fusion = nn.Sequential(\n            LocalMambaBlock(nf, kernel_size=5, dropout=0.1),\n            LocalMambaBlock(nf, kernel_size=5, dropout=0.1)\n        )\n        \n        # 4. Pooling & Heads\n        self.pool = nn.AdaptiveAvgPool1d(1)\n        \n        # Heads (using the same logic as before, but on fused features)\n        self.head_green_raw  = nn.Sequential(\n            nn.Linear(nf, nf//2), nn.GELU(), nn.Dropout(0.2), \n            nn.Linear(nf//2, 1), nn.Softplus()\n        )\n        self.head_clover_raw = nn.Sequential(\n            nn.Linear(nf, nf//2), nn.GELU(), nn.Dropout(0.2), \n            nn.Linear(nf//2, 1), nn.Softplus()\n        )\n        self.head_dead_raw   = nn.Sequential(\n            nn.Linear(nf, nf//2), nn.GELU(), nn.Dropout(0.2), \n            nn.Linear(nf//2, 1), nn.Softplus()\n        )\n        \n        if pretrained:\n            self.load_pretrained()\n    \n    def load_pretrained(self):\n        try:\n            # Load weights normally\n            if self.backbone_path and os.path.exists(self.backbone_path):\n                print(f\"Loading backbone weights from local file: {self.backbone_path}\")\n                sd = torch.load(self.backbone_path, map_location='cpu')\n                # Handle common checkpoint wrappers (e.g. if saved with 'model' key)\n                if 'model' in sd: sd = sd['model']\n                elif 'state_dict' in sd: sd = sd['state_dict']\n            else:\n                # Original behavior: Download from internet\n                print(\"Downloading backbone weights...\")\n                sd = timm.create_model(self.model_name, pretrained=True, num_classes=0, global_pool='').state_dict()\n            \n            # Interpolate pos_embed if needed (for 256x256 vs 224x224)\n            if 'pos_embed' in sd and hasattr(self.backbone, 'pos_embed'):\n                pe_ck = sd['pos_embed']\n                pe_m  = self.backbone.pos_embed\n                if pe_ck.shape != pe_m.shape:\n                    print(f\"Interpolating pos_embed: {pe_ck.shape} -> {pe_m.shape}\")\n                    # (Simple interpolation logic here or rely on timm's load if strict=False handles it well enough)\n                    # For robust interpolation, use the snippet provided in previous turn\n            \n            self.backbone.load_state_dict(sd, strict=False)\n            print('Pretrained weights loaded.')\n        except Exception as e:\n            print(f'Warning: pretrained load failed: {e}')\n    \n    def forward(self, left, right):\n        # 1. Extract Tokens (B, N, D)\n        #    Note: ViT usually returns [CLS, Patch1, Patch2...]\n        #    We remove CLS token for spatial mixing, or keep it. Let's keep it.\n        x_l = self.backbone(left)\n        x_r = self.backbone(right)\n        \n        # 2. Concatenate Left and Right tokens along sequence dimension\n        #    (B, N, D) + (B, N, D) -> (B, 2N, D)\n        x_cat = torch.cat([x_l, x_r], dim=1)\n        \n        # 3. Apply Mamba Fusion\n        #    This allows tokens from Left image to interact with tokens from Right image\n        x_fused = self.fusion(x_cat)\n        \n        # 4. Global Pooling\n        #    (B, 2N, D) -> (B, D, 2N) -> (B, D, 1) -> (B, D)\n        x_pool = self.pool(x_fused.transpose(1, 2)).flatten(1)\n        \n        # 5. Prediction Heads\n        green  = self.head_green_raw(x_pool)\n        clover = self.head_clover_raw(x_pool)\n        dead   = self.head_dead_raw(x_pool)\n        \n        # Summation logic\n        gdm    = green + clover\n        total  = gdm + dead\n        \n        return total, gdm, green, clover, dead\ndef biomass_loss(outputs, labels, w=None):\n    total, gdm, green, clover, dead = outputs\n    mse = nn.MSELoss()\n    huber = nn.SmoothL1Loss(beta=5.0) # Huber loss for robust regression (beta=5.0 as recommended)\n    \n    l_green  = huber(green.squeeze(),  labels[:,0])\n    l_dead   = huber(dead.squeeze(), labels[:,1]) # Use Huber loss for Dead\n    l_clover = huber(clover.squeeze(), labels[:,2])\n    l_gdm    = huber(gdm.squeeze(),    labels[:,3])\n    l_total  = huber(total.squeeze(),  labels[:,4])\n\n    # Stack per-target losses in the SAME order as CFG.ALL_TARGET_COLS\n    losses = torch.stack([l_green, l_dead, l_clover, l_gdm, l_total])\n    # losses = torch.stack([l_green, l_dead, l_clover])\n\n    # Use provided weights, or default to CFG.R2_WEIGHTS\n    if w is None:\n        return losses.mean()\n    w = torch.as_tensor(w, device=losses.device, dtype=losses.dtype)\n    w = w / w.sum()\n    return (losses * w).sum()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-16T14:32:03.813594Z","iopub.execute_input":"2026-01-16T14:32:03.813798Z","iopub.status.idle":"2026-01-16T14:32:03.835751Z","shell.execute_reply.started":"2026-01-16T14:32:03.813776Z","shell.execute_reply":"2026-01-16T14:32:03.835162Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"from contextlib import nullcontext\nscaler = torch.cuda.amp.GradScaler(enabled=torch.cuda.is_available())\n\n@torch.no_grad()\ndef valid_epoch(eval_model, loader, device):\n    eval_model.eval()\n    running = 0.0\n    preds_total, preds_gdm, preds_green, preds_clover, preds_dead, all_labels = [], [], [], [], [], []\n    amp_ctx = (lambda: torch.amp.autocast(device_type='cuda')) if torch.cuda.is_available() else (lambda: nullcontext())\n    \n    for l, r, lab in loader:\n        l, r, lab = l.to(device, non_blocking=True), r.to(device, non_blocking=True), lab.to(device, non_blocking=True)\n        with amp_ctx():\n            p_total, p_gdm, p_green, p_clover, p_dead = eval_model(l, r)\n            loss = biomass_loss((p_total, p_gdm, p_green, p_clover, p_dead), lab, w=CFG.LOSS_WEIGHTS)\n        running += loss.item() * l.size(0)\n        preds_total.extend(p_total.cpu().numpy().ravel())\n        preds_gdm.extend(p_gdm.cpu().numpy().ravel())\n        preds_green.extend(p_green.cpu().numpy().ravel())\n        preds_clover.extend(p_clover.cpu().numpy().ravel())\n        preds_dead.extend(p_dead.cpu().numpy().ravel())\n        all_labels.extend(lab.cpu().numpy())\n    \n    pred_total  = np.array(preds_total)\n    pred_gdm    = np.array(preds_gdm)\n    pred_green  = np.array(preds_green)\n    pred_clover = np.array(preds_clover)\n    pred_dead   = np.array(preds_dead)\n    true_labels = np.stack(all_labels)\n    \n    pred_all = np.stack([pred_green, pred_dead, pred_clover, pred_gdm, pred_total], axis=1)\n    global_r2, avg_r2, per_r2 = weighted_r2_score_global(true_labels, pred_all)\n    return running / len(loader.dataset), global_r2, avg_r2, per_r2, pred_all, true_labels\n\n@torch.no_grad()\ndef valid_epoch_tta(eval_model, loaders, device):\n    eval_model.eval()\n    amp_ctx = (lambda: torch.amp.autocast(device_type='cuda')) if torch.cuda.is_available() else (lambda: nullcontext())\n    \n    # We need to aggregate predictions from all loaders\n    # Assuming all loaders have same order and size (which they should if shuffle=False)\n    \n    all_preds_accum = None\n    all_labels = None\n    total_loss = 0.0\n    \n    for loader_idx, loader in enumerate(loaders):\n        preds_total, preds_gdm, preds_green, preds_clover, preds_dead = [], [], [], [], []\n        current_labels = []\n        running_loss = 0.0\n        \n        for l, r, lab in loader:\n            l, r, lab = l.to(device, non_blocking=True), r.to(device, non_blocking=True), lab.to(device, non_blocking=True)\n            with amp_ctx():\n                p_total, p_gdm, p_green, p_clover, p_dead = eval_model(l, r)\n                loss = biomass_loss((p_total, p_gdm, p_green, p_clover, p_dead), lab, w=CFG.LOSS_WEIGHTS)\n            \n            running_loss += loss.item() * l.size(0)\n            \n            preds_total.extend(p_total.cpu().numpy().ravel())\n            preds_gdm.extend(p_gdm.cpu().numpy().ravel())\n            preds_green.extend(p_green.cpu().numpy().ravel())\n            preds_clover.extend(p_clover.cpu().numpy().ravel())\n            preds_dead.extend(p_dead.cpu().numpy().ravel())\n            \n            if loader_idx == 0:\n                current_labels.extend(lab.cpu().numpy())\n        \n        total_loss += (running_loss / len(loader.dataset))\n        \n        # Stack predictions for this loader: (N, 5)\n        # Order: Green, Dead, Clover, GDM, Total (matching CFG.ALL_TARGET_COLS order roughly, but let's be precise)\n        # CFG.ALL_TARGET_COLS = ['Dry_Green_g','Dry_Dead_g','Dry_Clover_g','GDM_g','Dry_Total_g']\n        # preds lists are just raw outputs.\n        # Let's stack them in the order expected by weighted_r2_score_global which expects:\n        # y_true, y_pred where columns match.\n        # The model returns: total, gdm, green, clover, dead\n        # We need to stack them to match true_labels which comes from CFG.ALL_TARGET_COLS\n        # CFG.ALL_TARGET_COLS is ['Dry_Green_g','Dry_Dead_g','Dry_Clover_g','GDM_g','Dry_Total_g']\n        \n        pred_stack = np.stack([\n            np.array(preds_green),\n            np.array(preds_dead),\n            np.array(preds_clover),\n            np.array(preds_gdm),\n            np.array(preds_total)\n        ], axis=1)\n        \n        if all_preds_accum is None:\n            all_preds_accum = pred_stack\n            all_labels = np.stack(current_labels)\n        else:\n            all_preds_accum += pred_stack\n            \n    # Average predictions\n    avg_preds = all_preds_accum / len(loaders)\n    avg_loss = total_loss / len(loaders)\n    \n    global_r2, avg_r2, per_r2 = weighted_r2_score_global(all_labels, avg_preds)\n    return avg_loss, global_r2, avg_r2, per_r2, avg_preds, all_labels\n\ndef set_backbone_requires_grad(model: BiomassModel, requires_grad: bool):\n    for p in model.backbone.parameters():\n        p.requires_grad = requires_grad\n\ndef build_optimizer(model: BiomassModel):\n    head_params = (list(model.head_green_raw.parameters()) +\n                   list(model.head_clover_raw.parameters()) +\n                   list(model.head_dead_raw.parameters()))\n    backbone_params = list(model.backbone.parameters())\n    return optim.AdamW([\n        {'params': backbone_params, 'lr': CFG.LR_BACKBONE, 'weight_decay': CFG.WD},\n        {'params': head_params,     'lr': CFG.LR_HEAD,     'weight_decay': CFG.WD},\n    ])\ndef build_optimizer(model: BiomassModel):\n    # 1. Get backbone parameter IDs for exclusion\n    backbone_ids = {id(p) for p in model.backbone.parameters()}\n    \n    # 2. Separate params into backbone vs. everything else (heads, fusion, etc.)\n    backbone_params = []\n    rest_params = []\n    \n    for p in model.parameters():\n        if p.requires_grad:\n            if id(p) in backbone_ids:\n                backbone_params.append(p)\n            else:\n                rest_params.append(p)\n    \n    return optim.AdamW([\n        {'params': backbone_params, 'lr': CFG.LR_BACKBONE, 'weight_decay': CFG.WD},\n        {'params': rest_params,     'lr': CFG.LR_REST,     'weight_decay': CFG.WD},\n])\ndef build_scheduler(optimizer):\n    def lr_lambda(epoch):\n        e = max(0, epoch - 1)\n        if e < CFG.WARMUP_EPOCHS:\n            return float(e + 1) / float(max(1, CFG.WARMUP_EPOCHS))\n        progress = (e - CFG.WARMUP_EPOCHS) / float(max(1, CFG.EPOCHS - CFG.WARMUP_EPOCHS))\n        return 0.5 * (1.0 + math.cos(math.pi * progress))\n    return LambdaLR(optimizer, lr_lambda)\n\ndef train_epoch(model, loader, opt, scheduler, device, ema: ModelEmaV2 | None = None):\n    model.train()\n    running = 0.0\n    opt.zero_grad()\n    amp_ctx = (lambda: torch.amp.autocast(device_type='cuda')) if torch.cuda.is_available() else (lambda: nullcontext())\n    itera = tqdm(loader, desc='train', leave=False) if CFG.USE_TQDM else loader\n    for i, (l, r, lab) in enumerate(itera):\n        l, r, lab = l.to(device, non_blocking=True), r.to(device, non_blocking=True), lab.to(device, non_blocking=True)\n        with amp_ctx():\n            total, gdm, green, clover, dead = model(l, r)\n            loss = biomass_loss((total, gdm, green, clover, dead), lab, w=CFG.LOSS_WEIGHTS) / CFG.GRAD_ACC\n        scaler.scale(loss).backward()\n        running += loss.item() * l.size(0) * CFG.GRAD_ACC\n        \n        if (i + 1) % CFG.GRAD_ACC == 0 or (i + 1) == len(loader):\n            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)\n            scaler.step(opt)\n            scaler.update()\n            if ema is not None:\n                ema.update(model.module if hasattr(model, 'module') else model)\n            opt.zero_grad()\n    scheduler.step()\n    return running / len(loader.dataset)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-16T14:32:03.836463Z","iopub.execute_input":"2026-01-16T14:32:03.836669Z","iopub.status.idle":"2026-01-16T14:32:03.860436Z","shell.execute_reply.started":"2026-01-16T14:32:03.836649Z","shell.execute_reply":"2026-01-16T14:32:03.859675Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"from tqdm.auto import tqdm\nimport os\n\n# --- DODANO: Inicjalizacja pliku logów ---\nlog_file = os.path.join(CFG.MODEL_DIR, \"training_log.txt\")\nwith open(log_file, \"w\") as f:\n    f.write(\"Training Start\\n\" + \"=\"*30 + \"\\n\")\n\nprint('Loading data...')\ndf_long = pd.read_csv(CFG.TRAIN_CSV)\ndf_wide = df_long.pivot(index='image_path', columns='target_name', values='target').reset_index()\nassert df_wide['image_path'].is_unique, 'Leakage risk: duplicate image_path rows'\n\n# Merge metadata (Sampling_Date, State) for stratification\nif 'Sampling_Date' in df_long.columns and 'State' in df_long.columns:\n    print('Merging metadata for stratification...')\n    meta_df = df_long[['image_path', 'Sampling_Date', 'State']].drop_duplicates()\n    df_wide = df_wide.merge(meta_df, on='image_path', how='left')\n\n# Keep necessary columns\ndf_wide = df_wide[['image_path', 'Sampling_Date', 'State'] + CFG.ALL_TARGET_COLS]\nprint(f'{len(df_wide)} training images')\n\n# Use StratifiedGroupKFold\nsgkf = StratifiedGroupKFold(n_splits=CFG.N_FOLDS, shuffle=True, random_state=CFG.SEED)\noof_true, oof_pred, fold_summary = [], [], []\n\n# Split based on groups (Sampling_Date) and stratification target (State)\ngroups = df_wide['Sampling_Date']\ny_stratify = df_wide['State']\n\nfor fold, (tr_idx, val_idx) in enumerate(sgkf.split(df_wide, y_stratify, groups=groups)):\n    if fold not in CFG.FOLDS_TO_TRAIN:\n        print(f'Skipping fold {fold} as per configuration.')\n        continue\n    print('\\n' + '='*70)\n    print(f'FOLD {fold+1}/{CFG.N_FOLDS} | {len(tr_idx)} train / {len(val_idx)} val')\n    print('='*70)\n    \n    # Log fold start\n    with open(log_file, \"a\") as f:\n        f.write(f\"\\nFOLD {fold+1} | Train: {len(tr_idx)} Val: {len(val_idx)}\\n\")\n        \n    torch.cuda.empty_cache(); gc.collect()\n    \n    tr_df  = df_wide.iloc[tr_idx].reset_index(drop=True)\n    val_df = df_wide.iloc[val_idx].reset_index(drop=True)\n\n    try:\n        compare_train_val(tr_df, val_df, CFG.ALL_TARGET_COLS, show_plots=True)\n    except Exception as e:\n        print('Warning: compare_train_val failed:', e)\n\n    tr_set = BiomassDataset(tr_df,  get_train_transforms(), CFG.TRAIN_IMAGE_DIR)\n    \n    val_loaders = []\n    for mode in range(CFG.VAL_TTA_TIMES):\n        val_set_tta = BiomassDataset(val_df, get_tta_transforms(mode), CFG.TRAIN_IMAGE_DIR)\n        val_loader_tta = DataLoader(val_set_tta, batch_size=CFG.BATCH_SIZE, shuffle=False, num_workers=CFG.NUM_WORKERS, pin_memory=True)\n        val_loaders.append(val_loader_tta)\n\n    tr_loader  = DataLoader(tr_set, batch_size=CFG.BATCH_SIZE, shuffle=True, num_workers=CFG.NUM_WORKERS, pin_memory=True, drop_last=True)\n    \n    print('Building model...')\n    backbone_path = getattr(CFG, 'BACKBONE_PATH', None)\n    base_model = BiomassModel(CFG.MODEL_NAME, pretrained=CFG.PRETRAINED, backbone_path=backbone_path).to(CFG.DEVICE)\n    \n    if getattr(CFG, 'PRETRAINED_DIR', None) and os.path.isdir(CFG.PRETRAINED_DIR):\n        pretrained_path = os.path.join(CFG.PRETRAINED_DIR, f'best_model_fold{fold}.pth')\n        if os.path.exists(pretrained_path):\n            try:\n                state = torch.load(pretrained_path, map_location='cpu')\n                if isinstance(state, dict) and ('model_state_dict' in state or 'state_dict' in state):\n                    key = 'model_state_dict' if 'model_state_dict' in state else 'state_dict'\n                    sd = state[key]\n                else:\n                    sd = state\n                base_model.load_state_dict(sd, strict=False)\n                base_model.to(CFG.DEVICE)\n                print(f'  ✓ Loaded pretrained weights for fold {fold} from {pretrained_path}')\n            except Exception as e:\n                print(f'  ✗ Failed to load pretrained fold {fold}: {e}')\n    \n    model = nn.DataParallel(base_model)\n    set_backbone_requires_grad(base_model, False)\n    optimizer = build_optimizer(base_model)\n    scheduler = build_scheduler(optimizer)\n    ema = ModelEmaV2(base_model, decay=CFG.EMA_DECAY)\n    \n    best_global_r2 = -np.inf\n    patience = 0\n    best_fold_preds = None; best_fold_true = None\n    best_avg_r2 = -np.inf\n    \n    # --- DODANO: Bufor dla SWA ---\n    last_states = []\n    \n    save_path = os.path.join(CFG.MODEL_DIR, f'best_model_fold{fold}.pth')\n    \n    epoch_pbar = tqdm(range(1, CFG.EPOCHS + 1), desc=f'Fold {fold} Training')\n    for epoch in epoch_pbar:\n        if epoch == CFG.FREEZE_EPOCHS + 1:\n            patience = 0\n            set_backbone_requires_grad(base_model, True)\n            print(f'Epoch {epoch}: backbone unfrozen')\n        \n        tr_loss = train_epoch(model, tr_loader, optimizer, scheduler, CFG.DEVICE, ema)\n        eval_model = ema.module if ema is not None else (model.module if hasattr(model, 'module') else model)\n\n        lrs = [pg['lr'] for pg in optimizer.param_groups]\n        lr_str = f\"LR_BB: {lrs[0]:.2e} | LR_Head: {lrs[-1]:.2e}\"\n        \n        val_loss, global_r2, avg_r2, per_r2, preds_fold, true_fold = valid_epoch_tta(eval_model, val_loaders, CFG.DEVICE)\n        \n        per_r2_str = ' | '.join([f'{CFG.ALL_TARGET_COLS[i][:5]}: {r2:.3f}' for i, r2 in enumerate(per_r2)])\n        \n        epoch_pbar.set_postfix({'VLoss': f'{val_loss:.3f}', 'GlobR2': f'{global_r2:.3f}'})\n        \n        log_msg = f'Fold {fold} | Epoch {epoch:02d} | {lr_str} | TLoss {tr_loss:.5f} | VLoss {val_loss:.5f} | GlobalR² {global_r2:.4f}'\n        print(log_msg + (\" [BEST]\" if global_r2 > best_global_r2 else \"\"))\n        print(f'  → {per_r2_str}')\n        \n        with open(log_file, \"a\") as f:\n            f.write(log_msg + \"\\n\")\n            f.write(f\"  → {per_r2_str}\\n\")\n        \n        # --- DODANO: Zapisywanie stanu do SWA na koniec każdej epoki ---\n        current_state = {k: v.cpu().clone() for k, v in eval_model.state_dict().items()}\n        last_states.append(current_state)\n        if len(last_states) > 3:\n            last_states.pop(0)\n        \n        if global_r2 > best_global_r2:\n            best_global_r2 = global_r2\n            best_avg_r2 = avg_r2\n            best_state = {k: v.cpu().clone() for k, v in eval_model.state_dict().items()}\n            torch.save(best_state, save_path)\n            print(f'  → SAVED EMA weights to {save_path} (GlobalR²: {best_global_r2:.4f})')\n            del best_state \n            \n            patience = 0\n            best_fold_preds = preds_fold; best_fold_true = true_fold\n        else:\n            patience += 1\n            if patience >= CFG.PATIENCE:\n                print(f'  → EARLY STOP (no improvement in {CFG.PATIENCE} epochs)')\n                with open(log_file, \"a\") as f: f.write(f\"Early stop at epoch {epoch}\\n\")\n                break\n                \n        del preds_fold, true_fold\n        torch.cuda.empty_cache()\n        gc.collect()\n    \n    # --- DODANO: Obliczanie i zapisywanie SWA po pętli epok ---\n    if len(last_states) > 1:\n        print(f'\\nApplying Tail Averaging (SWA) over last {len(last_states)} epochs...')\n        averaged_state = last_states[0]\n        for i in range(1, len(last_states)):\n            for key in averaged_state.keys():\n                averaged_state[key] += last_states[i][key]\n        \n        for key in averaged_state.keys():\n            averaged_state[key] = torch.div(averaged_state[key], len(last_states))\n        \n        swa_save_path = os.path.join(CFG.MODEL_DIR, f'swa_model_fold{fold}.pth')\n        torch.save(averaged_state, swa_save_path)\n        print(f'  ✓ SWA model saved to {swa_save_path}')\n        del averaged_state, last_states\n    \n    if best_fold_preds is not None:\n        oof_true.append(best_fold_true); oof_pred.append(best_fold_preds)\n        fold_summary.append({'fold': fold, 'global_r2': best_global_r2,'avg_r2':best_avg_r2})\n    \n    del model, base_model, tr_loader, val_loaders, optimizer, scheduler, ema\n    if 'eval_model' in locals(): del eval_model\n    torch.cuda.empty_cache(); gc.collect()\n\nif oof_true:\n    oof_true_arr = np.concatenate(oof_true, axis=0)\n    oof_pred_arr = np.concatenate(oof_pred, axis=0)\n    oof_global_r2, oof_avg_r2, oof_per_r2 = weighted_r2_score_global(oof_true_arr, oof_pred_arr)\n\n    final_msg = f'\\nOOF Global Weighted R²: {oof_global_r2:.4f} | OOF Avg Target R²: {oof_avg_r2:.4f}'\n    print('\\nTraining complete! Models saved in:', CFG.MODEL_DIR)\n    print(final_msg)\n    \n    with open(log_file, \"a\") as f:\n        f.write(\"\\n\" + \"=\"*30 + \"\\nFINAL SUMMARY\\n\")\n        f.write(final_msg + \"\\n\")\n        for fs in fold_summary:\n            f.write(f\"Fold {fs['fold']}: Global R2 {fs['global_r2']:.4f}\\n\")\nelse:\n    print('No OOF predictions collected.')","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-16T19:14:03.579369Z","iopub.status.idle":"2026-01-16T19:14:03.579640Z","shell.execute_reply.started":"2026-01-16T19:14:03.579519Z","shell.execute_reply":"2026-01-16T19:14:03.579536Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}