{"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":"gpu","dataSources":[{"sourceId":112509,"databundleVersionId":14254895,"sourceType":"competition"},{"sourceId":4535,"sourceType":"modelInstanceVersion","isSourceIdPinned":true,"modelInstanceId":3327,"modelId":986}],"dockerImageVersionId":31240,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# Intro\n\nTraining notebook for CSIRO Image2Biomass competition.\n\nFor the **EDA** refer to https://www.kaggle.com/code/simonedegasperis/starter-eda-baseline-brainstorming","metadata":{}},{"cell_type":"markdown","source":"# Imports","metadata":{}},{"cell_type":"code","source":"import os\nimport sys\n\nos.environ[\"USE_TF\"] = \"0\"\nos.environ[\"PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION\"] = \"python\"","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-09T10:43:14.779683Z","iopub.execute_input":"2026-01-09T10:43:14.780473Z","iopub.status.idle":"2026-01-09T10:43:14.784548Z","shell.execute_reply.started":"2026-01-09T10:43:14.780436Z","shell.execute_reply":"2026-01-09T10:43:14.783758Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import os\nimport pandas as pd\nimport random\nimport numpy as np\nimport torch\nimport gc\n\nimport math\nfrom PIL import Image\nfrom collections import Counter\nfrom sklearn.model_selection import KFold, StratifiedKFold, StratifiedGroupKFold\nimport matplotlib.pyplot as plt\nimport albumentations as A\nimport torch.optim as optim\n\nimport cv2\nfrom albumentations.pytorch import ToTensorV2\n\nimport os\nimport pandas as pd\nfrom PIL import Image\n\nimport torch\nimport torch.nn as nn\nfrom torch.utils.data import Dataset, DataLoader\nfrom torchvision import transforms, models\nfrom transformers import AutoModel","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-09T10:43:14.785821Z","iopub.execute_input":"2026-01-09T10:43:14.786409Z","iopub.status.idle":"2026-01-09T10:43:14.800356Z","shell.execute_reply.started":"2026-01-09T10:43:14.786383Z","shell.execute_reply":"2026-01-09T10:43:14.799604Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Config","metadata":{}},{"cell_type":"code","source":"base_path = '/kaggle/input/csiro-biomass/'","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-09T10:43:14.801228Z","iopub.execute_input":"2026-01-09T10:43:14.801483Z","iopub.status.idle":"2026-01-09T10:43:14.813150Z","shell.execute_reply.started":"2026-01-09T10:43:14.801460Z","shell.execute_reply":"2026-01-09T10:43:14.812466Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class CONFIG():\n    train_path = os.path.join(base_path, 'train')\n    test_path = os.path.join(base_path, 'test')\n    test_csv_path = os.path.join(base_path, 'test.csv')\n    train_csv_path = os.path.join(base_path, 'train.csv')\n    data_dir = os.path.join(base_path,  \"\")\n    nFolds = 5\n    seed = 42\n    pretrained = False\n    pretrained_weights_path = os.path.join(base_path, \"\")\n    best_model_dir = \"\"\n    device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n    epochs = 20\n    batch_size = 8\n    lr = 1e-3\n    eta_min = 1e-5\n    weight_decay = 1e-6\n    \n    # image\n    img_size_h = 448\n    img_size_w = 448\n    in_chans = 3\n    \n    # target columns\n    target_cols = [\n        \"Dry_Clover_g\",\n        \"Dry_Dead_g\",\n        \"Dry_Green_g\",\n        \"Dry_Total_g\",\n        \"GDM_g\"\n    ]\n    n_targets = 3\n    # we can predict multiple configs\n    targets_configs = [\"Dry_Clover_g\", \"Dry_Dead_g\", \"Dry_Green_g\"] # [\"Dry_Green_g\", \"Dry_Total_g\", \"GDM_g\"]\n    weights = np.array([0.1, 0.1, 0.1, 0.5, 0.2])\n    mapping = {\"Dry_Clover_g\": 0, \"Dry_Dead_g\": 1, \"Dry_Green_g\": 2, \"Dry_Total_g\": 3, \"GDM_g\": 4}\n    train_backbone = False","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-09T10:43:14.814538Z","iopub.execute_input":"2026-01-09T10:43:14.814759Z","iopub.status.idle":"2026-01-09T10:43:14.825715Z","shell.execute_reply.started":"2026-01-09T10:43:14.814744Z","shell.execute_reply":"2026-01-09T10:43:14.825095Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Functions","metadata":{}},{"cell_type":"code","source":"class BiomassTwoStreamDataset(Dataset):\n    \"\"\"\n    Dataset that splits each image into two square crops (left & right)\n    Returns:\n        - x_left: (3, 224, 224)\n        - x_right: (3, 224, 224)\n        - y: (3,) targets if training/validation, else None\n    \"\"\"\n    def __init__(self, df, transforms=None, is_test=False, input_res=224, targets=[\"Dry_Green_g\", \"Dry_Total_g\", \"GDM_g\"]):\n        self.df = df.reset_index(drop=True)\n        self.image_paths = df[\"image_path\"].values\n        self.transforms = transforms\n        self.is_test = is_test\n        self.input_res = input_res\n\n        if not is_test:\n            self.targets_3 = df[targets].values.astype(\"float32\")\n\n    def __len__(self):\n        return len(self.df)\n\n    def to_tensor(self, img):\n        \"\"\"\n        Convert image to tensor and normalize to ImageNet\n        \"\"\"\n        img = torch.from_numpy(img).permute(2, 0, 1).float() #/ 255.0\n        # mean = torch.tensor(IMAGENET_MEAN).view(3,1,1)\n        # std  = torch.tensor(IMAGENET_STD).view(3,1,1)\n        # img = (img - mean) / std\n        return img\n\n    def split_into_two_squares(self, img):\n        \"\"\"\n        Split rectangular image into two square crops (left & right)\n        \"\"\"\n        H, W, _ = img.shape\n        if W < H:\n            raise ValueError(f\"Expected W >= H, got image shape {img.shape}\")\n        \n        left  = img[:, :H, :]\n        right = img[:, W-H:, :]\n        return left, right\n\n    def __getitem__(self, idx):\n        # Load image\n        image_path = self.image_paths[idx]\n        full_path = os.path.join(config.data_dir, image_path)\n        img = cv2.imread(full_path)\n        if img is None:\n            raise FileNotFoundError(f\"Image not found at {full_path}\")\n        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)\n\n        # Split into two squares\n        x_left, x_right = self.split_into_two_squares(img)\n\n        # Resize to input resolution\n        # x_left  = cv2.resize(x_left, (self.input_res, self.input_res), interpolation=cv2.INTER_LINEAR)\n        # x_right = cv2.resize(x_right, (self.input_res, self.input_res), interpolation=cv2.INTER_LINEAR)\n\n        # Apply transforms if provided\n        if self.transforms:\n            x_left  = self.transforms(image=x_left)[\"image\"]\n            x_right = self.transforms(image=x_right)[\"image\"]\n        else:\n            x_left  = self.to_tensor(x_left)\n            x_right = self.to_tensor(x_right)\n\n        if self.is_test:\n            return x_left, x_right\n        else:\n            y = torch.tensor(self.targets_3[idx], dtype=torch.float32)\n            return x_left, x_right, y\n\nclass DINOv2TwoStreamRegressor(nn.Module):\n    def __init__(self, n_targets=3, model_path=\"/kaggle/input/dinov2/pytorch/large/1\",\n                 freeze_backbone=True, hidden_dim=512, dropout=0.1):\n        super().__init__()\n\n        self.backbone = AutoModel.from_pretrained(model_path)\n        embed_dim = self.backbone.config.hidden_size\n\n        pooled_dim = embed_dim * 3  # CLS + avg + max\n\n        self.regressor = nn.Sequential(\n            nn.Linear(pooled_dim * 2, hidden_dim),\n            nn.GELU(),\n            nn.Dropout(dropout),\n            nn.Linear(hidden_dim, n_targets)\n        )\n\n        if freeze_backbone:\n            for p in self.backbone.parameters():\n                p.requires_grad = False\n\n    def _pool_tokens(self, outputs):\n        last_hidden = outputs.last_hidden_state\n        cls = last_hidden[:, 0]\n        patches = last_hidden[:, 1:]\n\n        avg_pool = patches.mean(dim=1)\n        max_pool = patches.max(dim=1)[0]\n\n        return torch.cat([cls, avg_pool, max_pool], dim=1)\n\n    def forward(self, x_left, x_right):\n        out_l = self.backbone(pixel_values=x_left)\n        out_r = self.backbone(pixel_values=x_right)\n\n        emb_l = self._pool_tokens(out_l)\n        emb_r = self._pool_tokens(out_r)\n\n        combined = torch.cat([emb_l, emb_r], dim=1)\n        return self.regressor(combined)\n\n# -------------------------------------------------------------\n#  Training & Validation Loops\n# -------------------------------------------------------------\ndef train_one_epoch(model, loader, optimizer, criterion, device):\n    model.train()\n    running_loss = 0.0\n\n    for img_left, img_right, labels in loader:\n        img_left, img_right, labels = img_left.to(device), img_right.to(device), labels.to(device)\n\n        optimizer.zero_grad()\n        preds = model(img_left, img_right)\n        loss = criterion(preds, labels)\n\n        loss.backward()\n        optimizer.step()\n        running_loss += loss.item() * img_left.size(0)\n\n    return running_loss / len(loader.dataset)\n\ndef val_fn(model, dataloader, criterion, device, targets):\n    \"\"\"\n    Validation / evaluation for 1 epoch.\n    Assumes the dataset returns only 3 independent targets per sample.\n    \"\"\"\n    model.eval()\n    total_loss = 0\n    all_targets_5_np = []\n    all_preds_5_np = []\n\n    with torch.no_grad():\n        for img_left, img_right, targets_3 in dataloader:\n            img_left, img_right, targets_3 = img_left.to(device), img_right.to(device), targets_3.to(device)  # shape (B, 3)\n\n            # Model forward\n            outputs_3 = model(img_left, img_right)  # (B, 3)\n\n            # Compute loss directly on independent targets\n            loss = criterion(outputs_3, targets_3)\n            total_loss += loss.item() * img_left.size(0)  # scale by batch size\n\n            # Expand predictions to 5 for metrics\n            preds_5 = expand_predictions_torch(outputs_3, targets)\n            all_preds_5_np.append(preds_5.cpu().numpy())\n\n            # Expand targets to 5 as well\n            targets_5 = expand_predictions_torch(targets_3, targets)\n            all_targets_5_np.append(targets_5.cpu().numpy())\n\n    # Average loss over dataset\n    val_loss = total_loss / len(dataloader.dataset)\n\n    # Concatenate all predictions / targets\n    y_true_5 = np.concatenate(all_targets_5_np, axis=0)\n    y_pred_5 = np.concatenate(all_preds_5_np, axis=0)\n\n    # Compute weighted R² metric on full 5 targets\n    weighted_r2, _ = weighted_r2_score(y_true_5, y_pred_5)\n\n    return val_loss, weighted_r2\n\n\ndef weighted_r2_score(y_true: np.ndarray, y_pred: np.ndarray):\n    \"\"\"\n    Metric\n    y_true, y_pred: shape (N, 5)\n    \"\"\"\n    weights = config.weights\n    r2_scores = []\n    \n    for i in range(y_true.shape[1]):\n        y_t = y_true[:, i]\n        y_p = y_pred[:, i]\n        ss_res = np.sum((y_t - y_p) ** 2)\n        ss_tot = np.sum((y_t - np.mean(y_t)) ** 2)\n        r2 = 1 - ss_res / ss_tot if ss_tot > 0 else 0.0\n        r2_scores.append(r2)\n        \n    r2_scores = np.array(r2_scores)\n    weighted_r2 = np.sum(r2_scores * weights) / np.sum(weights)\n    return weighted_r2, r2_scores\n\n# ======== 3. Preprocessing ========\ndef get_processed_data():\n    \"\"\"\n    'long' -> 'wide'\n    \"\"\"\n    train_df = pd.read_csv(config.train_csv_path)\n    \n    # unique id\n    train_df[\"image_id\"] = train_df[\"image_path\"].apply(lambda x: x.split('/')[-1].split('.')[0])\n    \n    # pivot\n    train_pivot = train_df.pivot(\n        index=\"image_id\", \n        columns=\"target_name\", \n        values=\"target\"\n    ).reset_index()\n    \n    meta_df = train_df.drop_duplicates(subset=\"image_id\").drop(\n        columns=[\"sample_id\", \"target_name\", \"target\"]\n    )\n    \n    train_processed_df = meta_df.merge(train_pivot, on=\"image_id\", how=\"left\")\n    \n    # CV fold\n    #kf = KFold(n_splits=config.nFolds, shuffle=True, random_state=config.seed)\n    #train_processed_df[\"fold\"] = -1\n    #for fold, (train_idx, val_idx) in enumerate(kf.split(train_processed_df)):\n    #    train_processed_df.loc[val_idx, \"fold\"] = fold\n\n    skf = StratifiedKFold(n_splits=config.nFolds, shuffle=True, random_state=config.seed)\n    train_processed_df[\"fold\"] = -1\n    # Add the stratification column (State)\n    for fold, (train_idx, val_idx) in enumerate(skf.split(train_processed_df, train_processed_df[\"State\"])):\n        train_processed_df.loc[val_idx, \"fold\"] = fold\n\n    # StratifiedKFold\n    #sgkf = StratifiedGroupKFold(n_splits=config.nFolds, shuffle=True, random_state=config.seed)\n    #train_processed_df[\"fold\"] = -1\n    #groups = train_processed_df[\"Sampling_Date\"]\n    #y = train_processed_df[\"State\"]\n    #for fold, (train_idx, val_idx) in enumerate(sgkf.split(train_processed_df, y, groups=groups)):\n    #    train_processed_df.loc[val_idx, \"fold\"] = fold\n        \n    return train_processed_df\n\ndef get_transforms(is_train):\n    if is_train:\n        return A.Compose([\n            A.RandomResizedCrop((config.img_size_h, config.img_size_w), scale=(0.8, 1.0)),\n            A.HorizontalFlip(p=0.5),\n            A.Rotate(limit=15, p=0.5),\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        ])\n    else:\n        return A.Compose([\n            A.Resize(config.img_size_h, config.img_size_w),\n            A.Normalize(mean=[0.485,0.456,0.406], std=[0.229,0.224,0.225]),\n            ToTensorV2(),\n        ])\n    \ndef get_two_stream_transforms(is_train):\n    if is_train:\n        return A.Compose([\n            A.HorizontalFlip(p=0.5),\n            A.Rotate(limit=15, p=0.5),\n            A.ColorJitter(brightness=0.1, contrast=0.1, saturation=0.1, hue=0.05, p=0.5),\n            A.Resize(224, 224),  # ensure DINOv2 input size\n            A.Normalize(mean=[0.485,0.456,0.406], std=[0.229,0.224,0.225]),\n            ToTensorV2(),\n        ])\n    else:\n        return A.Compose([\n            A.Resize(224, 224),\n            A.Normalize(mean=[0.485,0.456,0.406], std=[0.229,0.224,0.225]),\n            ToTensorV2(),\n        ])\n\ndef expand_predictions_torch(preds_3, targets):\n    \"\"\"\n    [torch] Expand the 3 NN predictions (N, 3) to 5 predictions (N, 5).\n    \"\"\"\n    if targets == [\"Dry_Green_g\", \"Dry_Total_g\", \"GDM_g\"]:\n        # clip\n        P_Green = torch.clamp(preds_3[:, 0], min=0)\n        P_Total = torch.clamp(preds_3[:, 1], min=0)\n        P_GDM = torch.clamp(preds_3[:, 2], min=0)\n        \n        # Compute derived targets based on constraints.\n        P_Clover = torch.clamp(P_GDM - P_Green, min=0)\n        P_Dead = torch.clamp(P_Total - P_GDM, min=0)\n    \n    elif targets == [\"Dry_Clover_g\", \"Dry_Dead_g\", \"Dry_Green_g\"]:\n        P_Clover = torch.clamp(preds_3[:, 0], min=0)\n        P_Dead = torch.clamp(preds_3[:, 1], min=0)\n        P_Green = torch.clamp(preds_3[:, 2], min=0)\n\n        # Compute derived targets based on constraints.\n        P_GDM = torch.clamp(P_Green + P_Clover, min=0)\n        P_Total = torch.clamp(P_GDM + P_Dead, min=0)\n    elif targets == [\"Dry_Clover_g\", \"Dry_Dead_g\", \"GDM_g\"]:\n        P_Clover = torch.clamp(preds_3[:, 0], min=0)\n        P_Dead =torch.clamp(preds_3[:, 1], min=0)\n        P_GDM = torch.clamp(preds_3[:, 2], min=0)\n\n        # Compute derived targets based on constraints.\n        P_Green = torch.clamp(P_GDM - P_Clover, min=0)\n        P_Total = torch.clamp(P_GDM + P_Dead, min=0)\n\n    \n    preds_5 = torch.stack(\n        [\n            P_Clover, # Index 0\n            P_Dead,   # Index 1\n            P_Green,  # Index 2\n            P_Total,  # Index 3\n            P_GDM     # Index 4\n        ],\n        dim=1\n    )\n    return preds_5\n\n\ndef expand_predictions_np(preds_3, targets):\n    \"\"\"\n    [Numpy] Expand the three NN predictions (N, 3) to five predictions (N, 5).\n    \"\"\"\n\n    if targets == [\"Dry_Green_g\", \"Dry_Total_g\", \"GDM_g\"]:\n        # clip\n        P_Green = np.clip(preds_3[:, 0], a_min=0, a_max=None)\n        P_Total = np.clip(preds_3[:, 1], a_min=0, a_max=None)\n        P_GDM = np.clip(preds_3[:, 2], a_min=0, a_max=None)\n        \n        # Compute derived targets based on constraints.\n        P_Clover = np.clip(P_GDM - P_Green, a_min=0, a_max=None)\n        P_Dead = np.clip(P_Total - P_GDM, a_min=0, a_max=None)\n    \n    elif targets == [\"Dry_Clover_g\", \"Dry_Dead_g\", \"Dry_Green_g\"]:\n        P_Clover = np.clip(preds_3[:, 0], a_min=0, a_max=None)\n        P_Dead = np.clip(preds_3[:, 1], a_min=0, a_max=None)\n        P_Green = np.clip(preds_3[:, 2], a_min=0, a_max=None)\n\n        # Compute derived targets based on constraints.\n        P_GDM = np.clip(P_Green + P_Clover, a_min=0, a_max=None)\n        P_Total = np.clip(P_GDM + P_Dead, a_min=0, a_max=None)\n\n    elif targets == [\"Dry_Clover_g\", \"Dry_Dead_g\", \"GDM_g\"]:\n        P_Clover = np.clip(preds_3[:, 0], a_min=0, a_max=None)\n        P_Dead = np.clip(preds_3[:, 1], a_min=0, a_max=None)\n        P_GDM = np.clip(preds_3[:, 2], a_min=0, a_max=None)\n\n        # Compute derived targets based on constraints.\n        P_Green = np.clip(P_GDM - P_Clover, a_min=0, a_max=None)\n        P_Total = np.clip(P_GDM + P_Dead, a_min=0, a_max=None)\n    \n    preds_5 = np.stack(\n        [P_Clover, P_Dead, P_Green, P_Total, P_GDM],\n        axis=1\n    )\n    return preds_5\n        \n# ======== 8. Inference ========\ndef run_inference(targets):\n    print(\"\\n======== Starting Inference ========\")\n    \n    # data\n    test_df = pd.read_csv(config.test_csv_path)\n    test_unique_df = test_df.drop_duplicates(subset=\"image_path\").reset_index(drop=True)\n    \n    test_dataset = BiomassTwoStreamDataset(\n        test_unique_df, \n        transforms=get_two_stream_transforms(is_train=False), \n        is_test=True,\n        targets=targets\n    )\n    test_loader = DataLoader(\n        test_dataset, batch_size=config.batch_size * 2, shuffle=False,\n        num_workers=os.cpu_count(), pin_memory=True\n    )\n    \n    # Prediction using 5-fold models\n    all_fold_preds = []\n    for fold in range(config.nFolds):\n        print(f\"Predicting with Fold {fold+1}...\")\n        model_path = os.path.join(config.best_model_dir, f\"model_fold_{fold}.pth\")\n\n        model = DINOv2TwoStreamRegressor(n_targets=3, freeze_backbone=True).to(config.device)\n\n        try:\n            model.load_state_dict(torch.load(model_path))\n        except FileNotFoundError:\n            print(f\"Warning: Model file {model_path} not found. Skipping fold {fold+1}.\")\n            continue\n            \n        model.eval()\n        \n        fold_preds = []\n        with torch.no_grad():\n            for img_left, img_right in test_loader:\n                img_left, img_right = img_left.to(config.device), img_right.to(config.device)\n                outputs = model(img_left, img_right)\n                fold_preds.append(outputs.cpu().numpy())\n        \n        all_fold_preds.append(np.concatenate(fold_preds, axis=0))\n    \n    if not all_fold_preds:\n        print(\"Error: No models were loaded. Cannot perform inference.\")\n        return\n\n    # average\n    # (nFolds, n_test_images, n_targets) -> (n_test_images, n_targets)\n    avg_preds_3 = np.mean(all_fold_preds, axis=0) # (n_test_images, 3)\n    avg_preds_5 = expand_predictions_np(avg_preds_3, targets) # (n_test_images, 5)\n    \n    # submission\n    \n    preds_df = pd.DataFrame(avg_preds_5, columns=config.target_cols)\n    test_unique_df = pd.concat([test_unique_df, preds_df], axis=1)\n    \n    test_pred_long_df = test_unique_df.melt(\n        id_vars=[\"image_path\"], \n        value_vars=config.target_cols,\n        var_name=\"target_name\",\n        value_name=\"target\"\n    )\n\n    submission_df = test_df[[\"sample_id\", \"image_path\", \"target_name\"]].merge(\n        test_pred_long_df,\n        on=[\"image_path\", \"target_name\"],\n        how=\"left\"\n    )\n    \n    final_submission = submission_df[[\"sample_id\", \"target\"]].copy()\n\n    # sanitize\n    final_submission[\"target\"] = final_submission[\"target\"].fillna(0)\n    final_submission[\"target\"] = final_submission[\"target\"].replace([np.inf, -np.inf], 0)\n    final_submission[\"target\"] = final_submission[\"target\"].clip(lower=0)\n    \n    # submission\n    final_submission.to_csv(\"submission.csv\", index=False)\n    print(\"Inference complete. submission.csv saved.\")\n    print(final_submission.head())\n\nclass EarlyStopping:\n    def __init__(self, patience=5, min_delta=1e-4):\n        self.patience = patience\n        self.min_delta = min_delta\n        self.counter = 0\n        self.best_score = None\n\n    def step(self, score):\n        if self.best_score is None or score < self.best_score - self.min_delta:\n            self.best_score = score\n            self.counter = 0\n            return True\n        else:\n            self.counter += 1\n            return False\n\n    def should_stop(self):\n        return self.counter >= self.patience","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-09T10:43:14.879608Z","iopub.execute_input":"2026-01-09T10:43:14.880177Z","iopub.status.idle":"2026-01-09T10:43:14.918147Z","shell.execute_reply.started":"2026-01-09T10:43:14.880148Z","shell.execute_reply":"2026-01-09T10:43:14.917497Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Training","metadata":{}},{"cell_type":"code","source":"config = CONFIG()\ntrain_df = get_processed_data()\n\noof_predictions = np.zeros((len(train_df), len(config.target_cols)))\n\nfor fold in range(config.nFolds):\n    print(f\"\\n======== FOLD {fold+1}/{config.nFolds} ========\")\n    # targets = [\"Dry_Green_g\", \"Dry_Total_g\", \"GDM_g\"]\n    targets = config.targets_configs # [\"Dry_Clover_g\", \"Dry_Dead_g\", \"Dry_Green_g\"]\n\n    # -----------------------------\n    # Split fold \n    # -----------------------------\n    train_fold_df = train_df[train_df.fold != fold].reset_index(drop=True)\n    val_fold_df_with_idx = train_df[train_df.fold == fold]\n    val_indices = val_fold_df_with_idx.index\n    val_fold_df = val_fold_df_with_idx.reset_index(drop=True)\n\n    # -----------------------------\n    # Datasets & loaders\n    # -----------------------------\n    train_dataset = BiomassTwoStreamDataset(\n        train_fold_df,\n        transforms=get_transforms(is_train=True),\n        targets=targets\n    )\n    val_dataset = BiomassTwoStreamDataset(\n        val_fold_df,\n        transforms=get_transforms(is_train=False),\n        targets=targets\n    )\n\n    train_loader = DataLoader(\n        train_dataset,\n        batch_size=config.batch_size,\n        shuffle=True,\n        num_workers=os.cpu_count(),\n        pin_memory=True,\n    )\n\n    val_loader = DataLoader(\n        val_dataset,\n        batch_size=config.batch_size * 2,\n        shuffle=False,\n        num_workers=os.cpu_count(),\n        pin_memory=True,\n    )\n\n    # -----------------------------\n    # Model\n    # -----------------------------\n    model = DINOv2TwoStreamRegressor(\n        n_targets=3,\n        freeze_backbone=True,\n    ).to(config.device)\n\n    criterion = nn.SmoothL1Loss(beta=0.5)\n    best_val = -np.inf\n    model_path = f\"model_fold_{fold}.pth\"\n\n    # ======================================================\n    # STAGE 1 — Train head only\n    # ======================================================\n    print(\"Training head only...\")\n\n    early_stop = EarlyStopping(patience=7)\n\n    optimizer = torch.optim.AdamW(\n        model.regressor.parameters(),\n        lr=config.lr,\n        weight_decay=1e-3,\n    )\n\n    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(\n        optimizer,\n        T_max=config.epochs,\n        eta_min=config.eta_min,\n    )\n\n    for epoch in range(config.epochs):\n        train_loss = train_one_epoch(\n            model,\n            train_loader,\n            optimizer,\n            criterion,\n            config.device,\n        )\n\n        val_loss, weighted_r2 = val_fn(\n            model,\n            val_loader,\n            criterion,\n            config.device,\n            targets\n        )\n\n        print(\n            f\"[HEAD] Epoch {epoch+1}/{config.epochs} \"\n            f\"- Train: {train_loss:.4f} | Val: {val_loss:.4f} |  Val R2: {weighted_r2:.4f}\"\n        )\n\n        scheduler.step()\n\n        # ---- Save best model ----\n        if weighted_r2 > best_val:\n            best_val = weighted_r2\n            torch.save(model.state_dict(), model_path)\n\n        # ---- Early stopping logic ----\n        early_stop.step(val_loss)\n        if early_stop.should_stop():\n            print(\"Early stopping (head)\")\n            break\n\n    if config.train_backbone:\n        # ======================================================\n        # STAGE 2 — LayerNorm-only fine-tuning\n        # ======================================================\n        print(\"Fine-tuning LayerNorm only...\")\n    \n        # Freeze entire backbone\n        for param in model.backbone.parameters():\n            param.requires_grad = False\n    \n        # Unfreeze only LayerNorms\n        for name, param in model.backbone.named_parameters():\n            if \"norm\" in name.lower():\n                param.requires_grad = True\n    \n        early_stop = EarlyStopping(patience=4)\n    \n        optimizer = torch.optim.AdamW(\n            [\n                {\n                    \"params\": model.regressor.parameters(),\n                    \"lr\": config.lr,\n                    \"weight_decay\": 1e-3,\n                },\n                {\n                    \"params\": filter(\n                        lambda p: p.requires_grad,\n                        model.backbone.parameters(),\n                    ),\n                    \"lr\": config.lr * 0.01,\n                    \"weight_decay\": 1e-5,\n                },\n            ]\n        )\n    \n        scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(\n            optimizer,\n            mode=\"min\",\n            factor=0.3,\n            patience=2,\n        )\n    \n        for epoch in range(config.epochs // 2):\n            train_loss = train_one_epoch(\n                model,\n                train_loader,\n                optimizer,\n                criterion,\n                config.device,\n            )\n    \n            val_loss, weighted_r2 = val_fn(\n                model,\n                val_loader,\n                criterion,\n                config.device,\n                targets\n            )\n    \n            print(\n                f\"[FT] Epoch {epoch+1}/{config.epochs//2} \"\n                f\"- Train: {train_loss:.4f} | Val: {val_loss:.4f} | Val R2: {weighted_r2:.4f} \"\n            )\n    \n            scheduler.step(val_loss)\n    \n            # ---- Save best model ----\n            if weighted_r2 > best_val:\n                best_val = weighted_r2\n                torch.save(model.state_dict(), model_path)\n            # ---- Early stopping logic ----\n            early_stop.step(val_loss)\n            if early_stop.should_stop():\n                print(\"Early stopping (head)\")\n                break\n\n    # -----------------------------\n    # OOF predictions\n    # -----------------------------\n    print(f\"Loading best model from {model_path}\")\n\n    model.load_state_dict(torch.load(model_path))\n    model.eval()\n\n    fold_preds_3 = []\n\n    with torch.no_grad():\n        for img_left, img_right, _ in val_loader:\n            img_left, img_right = img_left.to(config.device), img_right.to(config.device)\n            preds = model(img_left, img_right)\n            fold_preds_3.append(preds.cpu().numpy())\n\n    fold_preds_3 = np.concatenate(fold_preds_3, axis=0)\n\n    # Expand 3 → 5\n    fold_preds_5 = expand_predictions_np(fold_preds_3, targets)\n    oof_predictions[val_indices] = fold_preds_5\n\n    # -----------------------------\n    # Cleanup\n    # -----------------------------\n    del model, train_dataset, val_dataset, train_loader, val_loader\n    gc.collect()\n    torch.cuda.empty_cache()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-09T10:43:14.919318Z","iopub.execute_input":"2026-01-09T10:43:14.919550Z","execution_failed":"2026-01-09T10:50:34.254Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# OOF evaluation","metadata":{}},{"cell_type":"code","source":"# -----------------------------\n# Final OOF score\n# -----------------------------\noof_score, oof_scores_by_target = weighted_r2_score(train_df[config.target_cols].values, oof_predictions)\nprint(f\"\\n======== CV Finished ========\")\nprint(f\"Overall OOF Weighted R2 Score: {oof_score:.4f}\")\nfor i, col in enumerate(config.target_cols):\n    print(f\"  - {col}: {oof_scores_by_target[i]:.4f}\")\n\nprint(\"----------------------------------------------------------------------------------------------\")\nprint(\"--------------------------------------------END-----------------------------------------------\")\nprint(\"----------------------------------------------------------------------------------------------\")\n\nrun_inference(targets)","metadata":{"trusted":true,"execution":{"execution_failed":"2026-01-09T10:50:34.254Z"}},"outputs":[],"execution_count":null}]}