{"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":14122737,"sourceType":"datasetVersion","datasetId":8997717},{"sourceId":4537,"sourceType":"modelInstanceVersion","isSourceIdPinned":true,"modelInstanceId":3329,"modelId":986}],"dockerImageVersionId":31193,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"import os\nimport numpy as np\nimport pandas as pd\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom torch.utils.data import Dataset, DataLoader\nimport torchvision.transforms as transforms\nfrom PIL import Image\nfrom sklearn.model_selection import KFold\nfrom sklearn.metrics import r2_score\nfrom tqdm import tqdm\nimport sys\nfrom typing import Dict, List, Tuple, Optional\n\n# Bypass transformers torch version check\nimport transformers.utils.import_utils\ntransformers.utils.import_utils.check_torch_load_is_safe = lambda: None\n\n# Kaggle model path (default: /kaggle/input/dinov2-giant/)\n# If not available, will download from Hugging Face\nmodel_path = '/kaggle/input/dinov2/pytorch/giant/1'\nfrom transformers import Dinov2Model\n\n# Configuration Class (Kaggle Paths)\nclass CFG:\n    seed = 42\n    batch_size = 32\n    epochs = 100\n    patience = 15\n    lr = 1e-3\n    rho = 0.05\n    weight_decay = 1e-3\n    n_folds = 5\n    device = 'cuda' if torch.cuda.is_available() else 'cpu'\n    root_dir = \"/kaggle/input/csiro-biomass\"  # Kaggle data root\n    train_csv = \"/kaggle/input/csiro-biomass/train.csv\"  # Kaggle train CSV\n    targets = [\"Dry_Clover_g\", \"Dry_Dead_g\", \"Dry_Green_g\", \"Dry_Total_g\", \"GDM_g\"]\n    img_size = 224\n    MODEL_DIR = './model'  # Kaggle working dir for saving models\n\ndef seed_everything(seed: int) -> None:\n    torch.manual_seed(seed)\n    torch.cuda.manual_seed(seed)\n    torch.cuda.manual_seed_all(seed)\n    np.random.seed(seed)\n    os.environ['PYTHONHASHSEED'] = str(seed)\n    torch.backends.cudnn.deterministic = True\n    torch.backends.cudnn.benchmark = False\n\nseed_everything(CFG.seed)\n\n# Import SAM optimizer\nsys.path.append('/kaggle/input/samsam/')  # 将 SAM 所在目录添加到系统路径\nfrom sam import SAM  # 从 sam.py 中导入 SAM 类\n# ------------------------------\n# Utility Classes (No Duplication)\n# ------------------------------\nclass ImageFeatureExtractor:\n    def __init__(self, config: CFG, model_path: str = None):\n        self.config = config\n        self.mean = [0.485, 0.456, 0.406]\n        self.std = [0.229, 0.224, 0.225]\n        self.transform = self._build_transform()\n        self.model = self._load_dinov2(model_path or 'facebook/dinov2-giant')\n        self.embed_dim = self.model.config.hidden_size\n\n    def _build_transform(self) -> transforms.Compose:\n        return transforms.Compose([\n            transforms.ToTensor(),\n            transforms.Resize((self.config.img_size, self.config.img_size)),\n            transforms.Normalize(self.mean, self.std)\n        ])\n\n    def _load_dinov2(self, model_name_or_path: str) -> Dinov2Model:\n        print(\"Loading DINOv2 Model...\")\n        try:\n            model = Dinov2Model.from_pretrained(model_name_or_path)\n        except OSError:\n            print(f\"Failed to load from {model_name_or_path}, trying to download...\")\n            model = Dinov2Model.from_pretrained('facebook/dinov2-giant')\n        \n        model.head = nn.Identity()\n        model.to(self.config.device)\n        model.eval()\n        return model\n\n    def extract_features(self, df: pd.DataFrame) -> Dict[str, torch.Tensor]:\n        class TempImageDataset(Dataset):\n            def __init__(self, df, root_dir, transform):\n                self.df = df\n                self.root_dir = root_dir\n                self.transform = transform\n            def __len__(self): return len(self.df)\n            def __getitem__(self, idx):\n                row = self.df.iloc[idx]\n                img_path = os.path.join(self.root_dir, row['image_path'])\n                try:\n                    image = Image.open(img_path).convert('RGB')\n                except:\n                    image = Image.new('RGB', (self.config.img_size, self.config.img_size))\n                return self.transform(image), row['sample_id']\n\n        batch_size = 8 if self.config.img_size <= 500 else 4  # Adjust for Kaggle GPU\n        loader = DataLoader(\n            TempImageDataset(df, self.config.root_dir, self.transform),\n            batch_size=batch_size, shuffle=False, num_workers=2  # Kaggle num_workers=2 to avoid OOM\n        )\n\n        feature_cache = {}\n        with torch.no_grad():\n            for imgs, ids in tqdm(loader, desc=\"Extracting Features\"):\n                imgs = imgs.to(self.config.device, non_blocking=True)\n                outputs = self.model(imgs)\n                dense_features = outputs.last_hidden_state[:, 1:, :].cpu()\n                for i, sample_id in enumerate(ids):\n                    feature_cache[sample_id] = dense_features[i].clone()\n        \n        del self.model\n        torch.cuda.empty_cache()\n        print(\"Feature extraction complete.\")\n        return feature_cache\n\nclass CSVProcessor:\n    @staticmethod\n    def process_csv(df: pd.DataFrame, required_targets: List[str]) -> pd.DataFrame:\n        if all(col in df.columns for col in required_targets):\n            return df\n        \n        print(\"Detected Long-Format CSV. Converting to Wide-Format...\")\n        try:\n            df['target_type'] = df['sample_id'].apply(lambda x: x.split('__')[-1])\n            df['base_id'] = df['sample_id'].apply(lambda x: x.split('__')[0])\n        except IndexError:\n            sys.exit(\"Error parsing sample_id.\")\n        \n        exclude_cols = ['sample_id', 'image_path', 'Sampling_Date', 'State', 'Species', 'target_type', 'base_id']\n        numeric_cols = [c for c in df.columns if c not in exclude_cols and pd.api.types.is_numeric_dtype(df[c])]\n        value_col = 'weight' if 'weight' in df.columns else numeric_cols[-1]\n        \n        df_wide = df.pivot_table(\n            index=['image_path', 'base_id'], \n            columns='target_type', \n            values=value_col, \n            aggfunc='first'\n        ).reset_index()\n        df_wide.rename(columns={'base_id': 'sample_id'}, inplace=True)\n        \n        for t in required_targets:\n            df_wide[t] = df_wide[t].fillna(0.0) if t in df_wide.columns else 0.0\n        \n        print(f\"Conversion successful. Data shape: {df_wide.shape}\")\n        return df_wide\n\nclass CachedFeatureDataset(Dataset):\n    def __init__(self, df: pd.DataFrame, features_dict: Dict[str, torch.Tensor], target_col: str, input_size: int):\n        self.df = df\n        self.features_dict = features_dict\n        self.target_col = target_col\n        self.input_size = input_size\n\n    def __len__(self) -> int:\n        return len(self.df)\n\n    def __getitem__(self, idx: int) -> Tuple[torch.Tensor, torch.Tensor]:\n        row = self.df.iloc[idx]\n        sample_id = row['sample_id']\n        target = row[self.target_col]\n        feat = self.features_dict.get(sample_id, torch.zeros((1, self.input_size)))\n        return feat, torch.tensor(target, dtype=torch.float32)\n\n# ------------------------------\n# Model Architecture\n# ------------------------------\nclass MultiHeadSelfAttention(nn.Module):\n    def __init__(self, dim: int, num_heads: int = 8, dropout: float = 0.1):\n        super().__init__()\n        assert dim % num_heads == 0, \"dim must be divisible by num_heads\"\n        \n        self.dim = dim\n        self.num_heads = num_heads\n        self.head_dim = dim // num_heads\n        self.scale = self.head_dim ** -0.5\n        self.dropout = nn.Dropout(dropout)\n        \n        self.qkv = nn.Linear(dim, dim * 3)\n        self.proj = nn.Linear(dim, dim)\n\n    def forward(self, x: torch.Tensor) -> torch.Tensor:\n        B, N, C = x.shape\n        \n        qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, self.head_dim).permute(2, 0, 3, 1, 4)\n        q, k, v = qkv.unbind(0)\n        \n        attn = (q @ k.transpose(-2, -1)) * self.scale\n        attn = attn.softmax(dim=-1)\n        attn = self.dropout(attn)\n        \n        x = (attn @ v).transpose(1, 2).reshape(B, N, C)\n        x = self.proj(x)\n        x = self.dropout(x)\n        \n        return x\n\nclass GatedResidualBlock(nn.Module):\n    def __init__(self, dim: int, expansion: int = 2, dropout: float = 0.2):\n        super().__init__()\n        self.norm1 = nn.LayerNorm(dim)\n        self.linear1 = nn.Linear(dim, dim * expansion)\n        self.act = nn.GELU()\n        self.dropout1 = nn.Dropout(dropout)\n        \n        self.norm2 = nn.LayerNorm(dim * expansion)\n        self.linear2 = nn.Linear(dim * expansion, dim)\n        self.dropout2 = nn.Dropout(dropout)\n        \n        self.gate = nn.Sequential(\n            nn.Linear(dim, dim),\n            nn.Sigmoid()\n        )\n\n    def forward(self, x: torch.Tensor) -> torch.Tensor:\n        residual = x\n        x = self.norm1(x)\n        x = self.linear1(x)\n        x = self.act(x)\n        x = self.dropout1(x)\n        \n        x = self.norm2(x)\n        x = self.linear2(x)\n        x = self.dropout2(x)\n        \n        gate_weight = self.gate(residual)\n        x = gate_weight * x + (1 - gate_weight) * residual\n        return x\n\nclass FeatureAggregation(nn.Module):\n    def __init__(self, dim: int):\n        super().__init__()\n        self.attn_pool = nn.Sequential(\n            nn.Linear(dim, dim),\n            nn.Tanh(),\n            nn.Linear(dim, 1)\n        )\n        self.conv_pool = nn.Conv1d(dim, dim, kernel_size=3, padding=1)\n        self.norm = nn.LayerNorm(dim * 3)\n        self.proj = nn.Linear(dim * 3, dim)\n\n    def forward(self, x: torch.Tensor) -> torch.Tensor:\n        B, N, C = x.shape\n        \n        attn_scores = self.attn_pool(x)\n        attn_weights = F.softmax(attn_scores, dim=1)\n        attn_feat = torch.sum(attn_weights * x, dim=1)\n        \n        max_feat = torch.max(x, dim=1)[0]\n        \n        conv_feat = self.conv_pool(x.transpose(1, 2)).transpose(1, 2)\n        conv_feat = torch.mean(conv_feat, dim=1)\n        \n        agg_feat = torch.cat([attn_feat, max_feat, conv_feat], dim=1)\n        agg_feat = self.norm(agg_feat)\n        agg_feat = self.proj(agg_feat)\n        \n        return agg_feat\n\nclass EnhancedBiomassMLP(nn.Module):\n    def __init__(\n        self, \n        input_dim: int = 1536, \n        hidden_dim: int = 768,\n        num_heads: int = 8,\n        num_blocks: int = 3,\n        dropout: float = 0.2\n    ):\n        super().__init__()\n        \n        self.input_proj = nn.Sequential(\n            nn.Linear(input_dim, hidden_dim),\n            nn.LayerNorm(hidden_dim),\n            nn.GELU(),\n            nn.Dropout(dropout)\n        )\n        self.noise_scale = nn.Parameter(torch.tensor(0.01))\n        \n        self.self_attn = MultiHeadSelfAttention(hidden_dim, num_heads, dropout)\n        \n        self.res_blocks = nn.Sequential(\n            *[GatedResidualBlock(hidden_dim, expansion=2, dropout=dropout) \n              for _ in range(num_blocks)]\n        )\n        \n        self.aggregation = FeatureAggregation(hidden_dim)\n        \n        self.output_head = nn.Sequential(\n            nn.Linear(hidden_dim, hidden_dim // 2),\n            nn.LayerNorm(hidden_dim // 2),\n            nn.GELU(),\n            nn.Dropout(dropout * 0.5),\n            nn.Linear(hidden_dim // 2, 1)\n        )\n\n    def forward(self, x: torch.Tensor) -> torch.Tensor:\n        B, N, C = x.shape\n        \n        if self.training:\n            noise = torch.randn_like(x) * self.noise_scale\n            x = x + noise\n        \n        x = self.input_proj(x)\n        x_attn = self.self_attn(x)\n        x = x + x_attn\n        x = self.res_blocks(x)\n        feat = self.aggregation(x)\n        out = self.output_head(feat)\n        \n        return out.squeeze(-1)\n\n# ------------------------------\n# Trainer Class (No Duplication)\n# ------------------------------\nclass Trainer:\n    def __init__(self, config: CFG, feature_cache: Dict[str, torch.Tensor]):\n        self.config = config\n        self.feature_cache = feature_cache\n        self.device = config.device\n        self.kf = KFold(n_splits=config.n_folds, shuffle=True, random_state=config.seed)\n        os.makedirs(config.MODEL_DIR, exist_ok=True)\n\n    def get_data_loaders(\n        self, df: pd.DataFrame, train_idx: np.ndarray, val_idx: np.ndarray, target_col: str, input_dim: int\n    ) -> Tuple[DataLoader, DataLoader]:\n        train_df = df.iloc[train_idx]\n        val_df = df.iloc[val_idx]\n        \n        train_dataset = CachedFeatureDataset(train_df, self.feature_cache, target_col, input_dim)\n        val_dataset = CachedFeatureDataset(val_df, self.feature_cache, target_col, input_dim)\n        \n        train_loader = DataLoader(\n            train_dataset, batch_size=self.config.batch_size, shuffle=True, pin_memory=True, num_workers=2\n        )\n        val_loader = DataLoader(\n            val_dataset, batch_size=self.config.batch_size, shuffle=False, pin_memory=True, num_workers=2\n        )\n        \n        return train_loader, val_loader\n\n    def train_fold(\n        self, \n        model: nn.Module, \n        train_loader: DataLoader, \n        val_loader: DataLoader, \n        criterion: nn.Module, \n        optimizer: torch.optim.Optimizer,\n        scheduler: torch.optim.lr_scheduler._LRScheduler,\n        save_path: str,\n        fold: int,\n        target_name: str\n    ) -> float:\n        best_r2 = -float('inf')\n        best_model_path = \"\"\n        patience_counter = 0\n        \n        for epoch in range(self.config.epochs):\n            # Training\n            model.train()\n            train_loss = 0.0\n            for features, targets in train_loader:\n                features = features.to(self.device, non_blocking=True)\n                targets = targets.to(self.device, non_blocking=True)\n                \n                outputs = model(features)\n                loss = criterion(outputs, targets)\n                loss.backward()\n                optimizer.first_step(zero_grad=True)\n                \n                criterion(model(features), targets).backward()\n                optimizer.second_step(zero_grad=True)\n                \n                train_loss += loss.item()\n            \n            # Validation\n            model.eval()\n            val_loss = 0.0\n            all_preds, all_targets = [], []\n            \n            with torch.no_grad():\n                for features, targets in val_loader:\n                    features = features.to(self.device, non_blocking=True)\n                    targets = targets.to(self.device, non_blocking=True)\n                    \n                    outputs = model(features)\n                    loss = criterion(outputs, targets)\n                    val_loss += loss.item()\n                    \n                    all_preds.extend(outputs.cpu().numpy())\n                    all_targets.extend(targets.cpu().numpy())\n            \n            # Metrics\n            avg_train_loss = train_loss / len(train_loader)\n            avg_val_loss = val_loss / len(val_loader)\n            current_r2 = r2_score(all_targets, all_preds)\n            scheduler.step()\n            \n            # Save best model\n            save_msg = \"\"\n            if current_r2 > best_r2:\n                best_r2 = current_r2\n                if best_model_path and os.path.exists(best_model_path):\n                    try:\n                        os.remove(best_model_path)\n                    except:\n                        pass\n                best_model_path = os.path.join(save_path, f\"fold{fold}_R2_{best_r2:.4f}.pt\")\n                torch.save(model.state_dict(), best_model_path)\n                save_msg = \"Saved (New Best)\"\n                patience_counter = 0\n            else:\n                patience_counter += 1\n            \n            # Log\n            print(\n                f\"  Epoch {epoch+1:02d} | Train Loss: {avg_train_loss:.4f} | \"\n                f\"Val Loss: {avg_val_loss:.4f} | Val R2: {current_r2:.4f} | \"\n                f\"Best R2: {best_r2:.4f} | Patience: {patience_counter} {save_msg}\"\n            )\n            \n            # Early stopping\n            if patience_counter >= self.config.patience:\n                print(f\"  Early stopping at epoch {epoch+1}\")\n                break\n        \n        print(f\"  Fold {fold + 1} Best R2: {best_r2:.4f}\")\n        return best_r2\n\n    def train_target(self, df: pd.DataFrame, target_name: str, input_dim: int) -> float:\n        print(f\"\\n{'='*65}\")\n        print(f\"START Training Target: {target_name}\")\n        print(f\"{'='*65}\")\n        \n        save_dir = os.path.join(self.config.MODEL_DIR, target_name)\n        os.makedirs(save_dir, exist_ok=True)\n        \n        fold_r2s = []\n        for fold, (train_idx, val_idx) in enumerate(self.kf.split(df)):\n            print(f\"\\n[Target: {target_name}] >>> Fold {fold + 1}/{self.config.n_folds}\")\n            \n            # Data loaders\n            train_loader, val_loader = self.get_data_loaders(\n                df, train_idx, val_idx, target_name, input_dim\n            )\n            \n            # Model init\n            model = EnhancedBiomassMLP(\n                input_dim=input_dim,\n                hidden_dim=768,\n                num_heads=8,\n                num_blocks=3,\n                dropout=0.2\n            ).to(self.device)\n            \n            # Optimizer & Scheduler\n            base_optimizer = torch.optim.AdamW\n            optimizer = SAM(\n                model.parameters(),\n                base_optimizer,\n                lr=self.config.lr,\n                rho=self.config.rho,\n                weight_decay=self.config.weight_decay\n            )\n            criterion = nn.SmoothL1Loss(beta=0.1)\n            scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(\n                optimizer.base_optimizer, T_max=self.config.epochs, eta_min=1e-6\n            )\n            \n            # Train fold\n            fold_best_r2 = self.train_fold(\n                model, train_loader, val_loader, criterion, optimizer, scheduler,\n                save_dir, fold, target_name\n            )\n            fold_r2s.append(fold_best_r2)\n        \n        avg_r2 = np.mean(fold_r2s)\n        print(f\"Target [{target_name}] Average 5-Fold R2: {avg_r2:.4f}\")\n        return avg_r2\n\n# ------------------------------\n# Result Summary Utility\n# ------------------------------\ndef print_result_summary(all_targets_avg_r2: Dict[str, float]) -> None:\n    print(\"\\n\" + \"=\"*50)\n    print(\"FINAL RESULTS SUMMARY\")\n    print(\"=\"*50)\n    \n    weights = {\n        \"Dry_Clover_g\": 0.1, \n        \"Dry_Dead_g\": 0.1, \n        \"Dry_Green_g\": 0.1, \n        \"Dry_Total_g\": 0.5, \n        \"GDM_g\": 0.2\n    }\n    final_weighted_score = sum(\n        all_targets_avg_r2.get(t, 0.0) * weights.get(t, 0.0) \n        for t in CFG.targets\n    )\n    \n    # Print table\n    print(f\"{'Target':<15} | {'Weight':<6} | {'Avg R2':<10}\")\n    print(\"-\" * 38)\n    for t_name in CFG.targets:\n        r2_val = all_targets_avg_r2.get(t_name, 0.0)\n        weight = weights.get(t_name, 0.0)\n        print(f\"{t_name:<15} | {weight:<6.1f} | {r2_val:.4f}\")\n    \n    print(\"-\" * 38)\n    print(f\"Final Weighted Score: {final_weighted_score:.4f}\")\n    print(\"=\"*50)\n\n# ------------------------------\n# Main Workflow (No Duplication)\n# ------------------------------\ndef main():\n    # 1. Load and process data\n    df = pd.read_csv(CFG.train_csv)\n    df = CSVProcessor.process_csv(df, CFG.targets)\n    \n    # 2. Extract features\n    feature_extractor = ImageFeatureExtractor(CFG, model_path)\n    feature_cache = feature_extractor.extract_features(df)\n    input_dim = feature_extractor.embed_dim\n    \n    # 3. Train all targets\n    trainer = Trainer(CFG, feature_cache)\n    all_targets_avg_r2 = {\n        target: trainer.train_target(df, target, input_dim)\n        for target in CFG.targets\n    }\n    \n    # 4. Print summary\n    print_result_summary(all_targets_avg_r2)\n\nif __name__ == \"__main__\":\n    main()","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2025-12-12T07:55:00.467909Z","iopub.execute_input":"2025-12-12T07:55:00.4685Z"}},"outputs":[],"execution_count":null}]}