{"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":112509,"databundleVersionId":14254895,"sourceType":"competition"},{"sourceId":13818387,"sourceType":"datasetVersion","datasetId":8758997},{"sourceId":4535,"sourceType":"modelInstanceVersion","modelInstanceId":3327,"modelId":986}],"dockerImageVersionId":31192,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"import os\nimport copy\nimport random\nimport time\nfrom pathlib import Path\nfrom typing import Dict, List, Tuple\n\nimport numpy as np\nimport pandas as pd\nimport json\nfrom PIL import Image\nimport matplotlib.pyplot as plt\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom torch import optim\nfrom torch.utils.data import Dataset, DataLoader\nfrom torch.amp import autocast, GradScaler\nimport torchvision.transforms as transforms\nfrom torch.optim.lr_scheduler import CosineAnnealingLR\n\nfrom transformers import AutoImageProcessor, AutoModel\n\nfrom sklearn.metrics import r2_score\nfrom sklearn.model_selection import StratifiedGroupKFold\nfrom sklearn.preprocessing import LabelEncoder","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2025-11-24T05:13:35.792897Z","iopub.execute_input":"2025-11-24T05:13:35.793580Z","iopub.status.idle":"2025-11-24T05:13:35.801390Z","shell.execute_reply.started":"2025-11-24T05:13:35.793542Z","shell.execute_reply":"2025-11-24T05:13:35.800509Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def seed_everything(seed: int = 2025):\n    \"\"\"设置随机种子\"\"\"\n    random.seed(seed)\n    np.random.seed(seed)\n    torch.manual_seed(seed)\n    torch.cuda.manual_seed_all(seed)\n    os.environ['PYTHONHASHSEED'] = str(seed)","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2025-11-24T05:13:35.802651Z","iopub.execute_input":"2025-11-24T05:13:35.802872Z","iopub.status.idle":"2025-11-24T05:13:35.818520Z","shell.execute_reply.started":"2025-11-24T05:13:35.802857Z","shell.execute_reply":"2025-11-24T05:13:35.817727Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class BiomassConfig:\n    \"\"\"配置参数类\"\"\"\n\n    def __init__(self):\n        # 基础设置\n        self.seed = 91792\n        self.device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n\n        # 数据路径\n        self.train_csv = '/kaggle/input/csiro-biomass/train.csv'\n        self.train1_xlsx = '/kaggle/input/train1/train1.xlsx'\n        self.test_csv = '/kaggle/input/csiro-biomass/test.csv'\n        # self.train_image_dir = '/kaggle/input/csiro-biomass/train'\n        self.root_image_dir = '/kaggle/input/csiro-biomass'\n\n        # 模型路径\n        self.model_path = '/kaggle/input/dinov2/pytorch/large/1/'\n\n        # 训练参数\n        self.batch_size = 16  # 减小batch size以适应ViT\n        self.num_epochs = 50\n        self.fc_lr = 1e-2\n        self.weight_decay = 1e-4\n        self.num_workers = 4\n\n        # loss权重\n        self.loss_weights = [1.0, 1.0, 1.0, 5.0, 2.0]\n        self.beta = 0.2\n\n        # 图像处理 - 适应ViT输入尺寸\n        self.season_emb_dim = 2\n        self.state_emb_dim = 2\n        self.species_emb_dim = 12\n        self.cont_dim = 2\n        self.metadata_dim = self.season_emb_dim + self.state_emb_dim + self.species_emb_dim + self.cont_dim\n        self.hidden_dim = 1024\n        self.image_size = 224 \n        self.normalize_mean = [0.485, 0.456, 0.406]\n        self.normalize_std = [0.229, 0.224, 0.225]\n\n        # 模型参数\n        self.backbone = 'dinov2_vit_large'\n        self.num_targets = 5\n        \n        # 交叉验证\n        self.n_folds = 5\n        self.use_cross_validation = False\n\n        # 混合精度训练\n        self.use_amp = True\n\n        # TTA (测试时增强)\n        self.use_tta = True\n        self.tta_scales = [1.0] \n        self.tta_hflip = True  # 水平翻转\n        self.tta_vflip = False # 垂直翻转\n\n        # 输出权重\n        self.target_weights = [0.1, 0.1, 0.1, 0.5, 0.2]  # Green, Clover, Dead, Total, GDM\n\n        # 输出目录\n        self.output_dir = '.'\n        os.makedirs(self.output_dir, exist_ok=True)","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2025-11-24T05:13:35.819864Z","iopub.execute_input":"2025-11-24T05:13:35.820118Z","iopub.status.idle":"2025-11-24T05:13:35.838173Z","shell.execute_reply.started":"2025-11-24T05:13:35.820096Z","shell.execute_reply":"2025-11-24T05:13:35.837462Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class FeatureExtractor:\n    \"\"\"特征提取器\"\"\"\n    \n    def __init__(self, config: BiomassConfig, df_imgs: pd.DataFrame, image_root: str, is_train: bool = True,\n                 tta: bool = False, dataloader: DataLoader = None):\n        self.config = config\n        self.device = self.config.device\n        self.df = df_imgs\n        self.root = Path(image_root)\n        self.image_paths = [self.root/row for row in self.df]\n        self.backbone = AutoModel.from_pretrained(self.config.model_path).to(self.device)\n        self.tta = tta\n        self.dataloader = dataloader\n        \n        # 完全冻结backbone\n        for param in self.backbone.parameters():\n            param.requires_grad = False\n        self.backbone.eval()\n        \n        # 图像预处理\n        if is_train:\n            self.transform = transforms.Compose([\n                transforms.Resize((self.config.image_size, self.config.image_size)),\n                transforms.RandomHorizontalFlip(p=0.5),\n                transforms.ToTensor(),\n                transforms.Normalize(mean=self.config.normalize_mean, \n                                   std=self.config.normalize_std)\n            ])\n        else:\n            pass\n        \n    def extract_features(self):\n        \"\"\"批量提取图像特征\"\"\"\n        print(\"Feature extraction start...\")\n        features_list = []\n        \n        with torch.no_grad():\n            if self.tta:\n                for batch in self.dataloader:\n                    batch_x = batch[0].to(self.device)\n                    outputs = self.backbone(batch_x)\n                    \n                    # 使用[CLS] token的特征\n                    if hasattr(outputs, 'last_hidden_state'):\n                        features = outputs.last_hidden_state[:, 0, :]\n                    elif hasattr(outputs, 'pooler_output'):\n                        features = outputs.pooler_output\n                    else:\n                        features = outputs.last_hidden_state.mean(dim=1)\n                    \n                    features_list.append(features.cpu())\n\n            else:\n                for i in range(0, len(self.image_paths), self.config.batch_size):\n                    batch_paths = self.image_paths[i:i+self.config.batch_size]\n                    batch_images = []\n                    \n                    # 加载和预处理图像\n                    for path in batch_paths:\n                        image = Image.open(path).convert('RGB')\n                        image = self.transform(image)\n                        batch_images.append(image)\n                \n                    # 堆叠批次\n                    batch_tensor = torch.stack(batch_images).to(self.device)\n                    \n                    # 提取特征\n                    outputs = self.backbone(batch_tensor)\n                    \n                    # 使用[CLS] token的特征\n                    if hasattr(outputs, 'last_hidden_state'):\n                        features = outputs.last_hidden_state[:, 0, :]\n                    elif hasattr(outputs, 'pooler_output'):\n                        features = outputs.pooler_output\n                    else:\n                        features = outputs.last_hidden_state.mean(dim=1)\n                    \n                    features_list.append(features.cpu())\n        \n        all_features = torch.cat(features_list, dim=0)\n        print(\"Feature extraction done.\")\n        return all_features","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-24T05:13:35.839381Z","iopub.execute_input":"2025-11-24T05:13:35.839699Z","iopub.status.idle":"2025-11-24T05:13:35.859837Z","shell.execute_reply.started":"2025-11-24T05:13:35.839683Z","shell.execute_reply":"2025-11-24T05:13:35.859181Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class BiomassDataset(Dataset):\n    \"\"\"生物量数据集类\"\"\"\n\n    def __init__(self, features: torch.Tensor, y_data: torch.Tensor = None, \n                 meta_data: torch.Tensor = None, image_paths: list = None, indices: list = None):\n        \"\"\"\n        Args:\n            features: 预提取的特征张量 [num_samples, hidden_dim]\n            y_data: 硬标签\n            y_data_soft: 软标签  \n            image_paths: 图像路径列表\n            indices: 指定使用的索引\n        \"\"\"\n        self.image_paths = image_paths\n        self.indices = indices\n        \n        # 如果指定了索引，只使用这些索引对应的特征和标签\n        if self.indices is not None:\n            self.features = features[self.indices]\n            self.y_data = y_data[self.indices] if y_data is not None else None\n            self.meta_data = meta_data[self.indices] if meta_data is not None else None\n            self.image_paths = [image_paths[i] for i in self.indices] if image_paths is not None else None\n        else:\n            self.features = features\n            self.y_data = y_data\n            self.meta_data = meta_data\n\n    def __len__(self):\n        return len(self.features)\n\n    def __getitem__(self, idx):\n        feature = self.features[idx]\n        \n        if self.y_data is not None and self.meta_data is not None:\n            label = self.y_data[idx]\n            meta_data = self.meta_data[idx]\n            return feature, label, meta_data\n        elif self.y_data is not None and self.meta_data is None:\n            label = self.y_data[idx]\n            return feature, label\n        else:\n            image_path = self.image_paths[idx] if self.image_paths is not None else \"\"\n            return feature, image_path\n\nclass TestDataset(Dataset):\n    \"\"\"测试数据集类\"\"\"\n\n    def __init__(self, features: torch.Tensor, meta_data: torch.Tensor, df_imgs: pd.DataFrame = None):\n        \"\"\"\n        Args:\n            features: 预提取的特征张量 [num_samples, hidden_dim]\n            df_imgs: 数据框（可选，用于获取image_path）\n        \"\"\"\n        self.features = features\n        self.df = df_imgs.reset_index(drop=True) if df_imgs is not None else None\n        self.meta_data = meta_data\n\n    def __len__(self):\n        return len(self.features)\n\n    def __getitem__(self, idx):\n        feature = self.features[idx]\n        meta_data = self.meta_data[idx]\n        if self.df is not None:\n            row = self.df.iloc[idx]\n            return feature, meta_data, row\n        else:\n            return feature, meta_data","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2025-11-24T05:13:35.860619Z","iopub.execute_input":"2025-11-24T05:13:35.860881Z","iopub.status.idle":"2025-11-24T05:13:35.878979Z","shell.execute_reply.started":"2025-11-24T05:13:35.860860Z","shell.execute_reply":"2025-11-24T05:13:35.878150Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class BiomassDataset_tta(Dataset):\n    \"\"\"生物量数据集类_TTA使用 - 完全预处理版本\"\"\"\n\n    def __init__(self, df_imgs: pd.DataFrame, y_data: torch.tensor, \n                 meta_data: torch.tensor, image_root: str, config: BiomassConfig):\n        self.original_df = df_imgs.reset_index(drop=True)\n        self.root = Path(image_root)\n        self.original_y_data = y_data\n        self.original_meta_data = meta_data\n        self.config = config\n        self.hflip = self.config.tta_hflip\n        self.vflip = self.config.tta_vflip\n        self.transform = transforms.Compose([\n                transforms.Resize((self.config.image_size, self.config.image_size)),\n                transforms.ToTensor(),\n                transforms.Normalize(mean=self.config.normalize_mean, \n                                   std=self.config.normalize_std)\n            ])\n        \n        # 在初始化时预处理所有数据\n        self.processed_samples = self._preprocess_all_samples()\n        \n        print(f\"TTA扩展: {len(self.original_df)} -> {len(self.processed_samples)} 个样本\")\n    \n    def _preprocess_all_samples(self):\n        \"\"\"预处理所有样本，包括TTA增强\"\"\"\n        processed_samples = []\n        \n        for i in range(len(self.original_df)):\n            row = self.original_df.iloc[i]\n            img_path = self.root / row\n            \n            # 加载原始图像\n            with Image.open(img_path) as im:\n                im = im.convert('RGB')\n\n            # 原始样本\n            processed_samples.append({\n                'x': self.transform(im),\n                'y_data': self.original_y_data[i],\n                'meta_data': self.original_meta_data[i],\n                'original_idx': i,\n                'transform_type': 'original'\n            })\n            \n            # 水平翻转样本\n            if self.hflip:\n                hflip_im = im.transpose(Image.FLIP_LEFT_RIGHT)\n                processed_samples.append({\n                    'x': self.transform(hflip_im),\n                    'y_data': self.original_y_data[i],\n                    'meta_data': self.original_meta_data[i],\n                    'original_idx': i,\n                    'transform_type': 'hflip'\n                })\n            \n            # 垂直翻转样本\n            if self.vflip:\n                vflip_im = im.transpose(Image.FLIP_TOP_BOTTOM)\n                processed_samples.append({\n                    'x': self.transform(vflip_im),\n                    'y_data': self.original_y_data[i],\n                    'meta_data': self.original_meta_data[i],\n                    'original_idx': i,\n                    'transform_type': 'vflip'\n                })\n        \n        return processed_samples\n    \n    def __len__(self):\n        return len(self.processed_samples)\n    \n    def __getitem__(self, idx):\n        \"\"\"直接返回预处理好的数据\"\"\"\n        sample = self.processed_samples[idx]\n        return sample['x'], sample['y_data'], sample['meta_data']\n    \n    def get_sample_info(self, idx):\n        \"\"\"获取样本详细信息\"\"\"\n        sample = self.processed_samples[idx]\n        row = self.original_df.iloc[sample['original_idx']]\n        return {\n            \"original_index\": sample['original_idx'],\n            \"transform_type\": sample['transform_type'],\n            \"image_path\": str(self.root / row)\n        }","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-24T05:13:35.880600Z","iopub.execute_input":"2025-11-24T05:13:35.881027Z","iopub.status.idle":"2025-11-24T05:13:35.898869Z","shell.execute_reply.started":"2025-11-24T05:13:35.881004Z","shell.execute_reply":"2025-11-24T05:13:35.898141Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class FeatureAttention(nn.Module):\n    \"\"\"一个适用于一维特征的注意力模块，类似于SEBlock\"\"\"\n    def __init__(self, channel, reduction=16):\n        super().__init__()\n        self.fc = nn.Sequential(\n            nn.Linear(channel, channel // reduction, bias=False),\n            nn.GELU(),\n            nn.Linear(channel // reduction, channel, bias=False),\n            nn.Sigmoid()\n        )\n\n    def forward(self, x):\n        # 为每个样本学习一个通道权重\n        weights = self.fc(x)\n        return x * weights\n\nclass BiomassModel_teacher(nn.Module):\n    \"\"\"生物量预测模型 - 使用DINOv2 ViT base (带注意力与门控)\"\"\"\n\n    def __init__(self, config: BiomassConfig):\n        super().__init__()\n        self.config = config\n        self.device = config.device\n        self.metadata_mlp_output_dim = 128\n        self.combined_dim = self.config.hidden_dim + self.metadata_mlp_output_dim\n        \n        self.state_embedding = nn.Embedding(num_embeddings=4, embedding_dim=self.config.state_emb_dim)\n        self.species_embedding = nn.Embedding(num_embeddings=15, embedding_dim=self.config.species_emb_dim)\n        self.metadata_mlp = self._build_metadata_mlp()\n        \n        # --- 新增：注意力模块 ---\n        # 为图片特征和元数据特征分别添加注意力\n        self.img_attention = FeatureAttention(self.config.hidden_dim)\n        self.meta_attention = FeatureAttention(128)\n        \n        # --- 修改：融合层 ---\n        # 原来的 fc 层被替换为门控融合\n        self.gate = self._build_gate_layer() # 新增的门控网络\n        self.backbone = self._build_backbone() \n\n        print(f\"Model initialized with {sum(p.numel() for p in self.parameters()):,} trainable parameters\")\n\n    def season_encoder(self, seasons):\n        season_sin = torch.sin(2 * np.pi * seasons / 4)\n        season_cos = torch.cos(2 * np.pi * seasons / 4)\n        return torch.cat([season_sin, season_cos], dim=1)\n\n    def _build_metadata_mlp(self):\n        metadata_mlp = nn.Sequential(\n            nn.Linear(self.config.metadata_dim, 64),\n            nn.GELU(),\n            nn.BatchNorm1d(64),\n            nn.Dropout(0.2),\n            nn.Linear(64, self.metadata_mlp_output_dim),\n            nn.GELU(),\n            nn.BatchNorm1d(self.metadata_mlp_output_dim)\n        )\n        self._init_fc_weights(metadata_mlp)\n        return metadata_mlp\n\n    def _build_gate_layer(self):\n        \"\"\"构建门控融合网络\"\"\"\n        # 输入是拼接后的特征，输出是相同维度的权重向量\n        gate_net = nn.Sequential(\n            nn.Linear(self.combined_dim, self.combined_dim),\n            nn.Sigmoid() # 使用Sigmoid确保权重在0-1之间\n        )\n        self._init_fc_weights(gate_net)\n        return gate_net\n\n    def _build_backbone(self):\n        \"\"\"构建共享的主干网络\"\"\"\n        self.hidden_dim1 = 512\n        self.hidden_dim2 = 256\n        self.hidden_dim3 = 128\n        self.dropout_rate = 0.4\n        input_dim = self.config.hidden_dim + self.metadata_mlp_output_dim\n\n        backbone_layers = nn.Sequential(\n            nn.Linear(input_dim, self.hidden_dim1),\n            nn.BatchNorm1d(self.hidden_dim1),\n            nn.GELU(),\n            nn.Dropout(self.dropout_rate),\n            nn.Linear(self.hidden_dim1, self.hidden_dim2),\n            nn.BatchNorm1d(self.hidden_dim2),\n            nn.GELU(),\n            nn.Dropout(self.dropout_rate),\n            nn.Linear(self.hidden_dim2, self.hidden_dim3),\n            nn.BatchNorm1d(self.hidden_dim3),\n            nn.GELU(),\n            nn.Dropout(self.dropout_rate * 0.5),\n            nn.Linear(self.hidden_dim3, self.config.num_targets)\n        )\n        self._init_fc_weights(backbone_layers)\n        return backbone_layers\n        \n    def _init_fc_weights(self, model):\n        for m in model.modules():\n            if isinstance(m, nn.Linear):\n                nn.init.xavier_uniform_(m.weight)\n                if m.bias is not None:\n                    nn.init.constant_(m.bias, 0)\n            elif isinstance(m, nn.BatchNorm1d):\n                nn.init.constant_(m.weight, 1)\n                nn.init.constant_(m.bias, 0)\n\n    def forward(self, features, meta_data):\n        # --- 步骤1: 元数据处理 ---\n        season = self.season_encoder(meta_data[:, 0].unsqueeze(1))\n        state = self.state_embedding(meta_data[:, 1].long())\n        species = self.species_embedding(meta_data[:, 2].long())\n        new_metadata = torch.cat([season, state, species, meta_data[:,-2:]], dim=1)\n        raw_meta_features = self.metadata_mlp(new_metadata) # [N, 128]\n\n        # --- 步骤2: 应用特征注意力 ---\n        # 让模型分别关注图片和元数据内部最重要的信息\n        attended_img_features = self.img_attention(features) # [N, hidden_dim]\n        attended_meta_features = self.meta_attention(raw_meta_features) # [N, 128]\n\n        # --- 步骤3: 特征拼接 ---\n        concatenated_features = torch.cat([attended_img_features, attended_meta_features], dim=1)\n\n        # --- 步骤4: 门控融合 ---\n        # 学习一个动态权重，决定两种信息的重要性\n        gate_weights = self.gate(concatenated_features) # [N, hidden_dim + 128]\n        fused_features = concatenated_features * gate_weights # [N, hidden_dim + 128]\n        \n        # --- 步骤5: 预测 ---\n        output = self.backbone(fused_features)\n        output = F.relu(output)\n        return output","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2025-11-24T05:13:35.899652Z","iopub.execute_input":"2025-11-24T05:13:35.899892Z","iopub.status.idle":"2025-11-24T05:13:35.920775Z","shell.execute_reply.started":"2025-11-24T05:13:35.899872Z","shell.execute_reply":"2025-11-24T05:13:35.920021Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class WeightedMSELoss(nn.Module):\n    def __init__(self, config: BiomassConfig):\n        super().__init__()\n        self.config = config\n        self.weights = torch.tensor(self.config.loss_weights)\n\n    def forward(self, input, target):\n        mse_per_dim = (input - target) ** 2\n\n        if self.weights.device != input.device:\n            self.weights = self.weights.to(input.device)\n\n        weighted_mse = mse_per_dim * self.weights.view(1, -1)\n\n        return weighted_mse.mean()","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2025-11-24T05:13:35.921449Z","iopub.execute_input":"2025-11-24T05:13:35.921679Z","iopub.status.idle":"2025-11-24T05:13:35.938566Z","shell.execute_reply.started":"2025-11-24T05:13:35.921659Z","shell.execute_reply":"2025-11-24T05:13:35.937896Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class BiomassTrainer:\n    \"\"\"生物量训练器\"\"\"\n\n    def __init__(self, config: BiomassConfig):\n        self.config = config\n        self.device = config.device\n        self.scaler = GradScaler('cuda') if config.use_amp else None\n        self.log_path = f'./log_dinov2_{config.num_epochs}.txt'\n        self.models = []\n        self.model_r2 = []\n        self.features = None\n        self.train_loader = None\n        self.target = None\n\n        # 训练历史记录\n        self.history = {\n            'train_loss': [], 'val_loss': [], 'r2_scores': [],\n            'r2_total': [], 'epoch_time': [], 'fold': []\n        }\n\n        seed_everything(config.seed)\n\n    @staticmethod\n    def ProcessMetadata(X):\n        season_encoder = LabelEncoder()\n        state_encoder = LabelEncoder()\n        species_encoder = LabelEncoder()\n    \n        # 对离散变量进行编码\n        season_encoded_label = torch.tensor(season_encoder.fit_transform(X[:,0])).to(torch.int).unsqueeze(1)\n        state_encoded_label = torch.tensor(state_encoder.fit_transform(X[:,1])).to(torch.int).unsqueeze(1)\n        species_encoded_label = torch.tensor(species_encoder.fit_transform(X[:,2])).to(torch.int).unsqueeze(1)\n\n        # 存储连续变量\n        ndvi = torch.tensor(X[:,3]).unsqueeze(1)\n        height = torch.tensor((X[:,4]-np.min(X[:,4])/(np.max(X[:,4])-np.min(X[:,4]))).astype(np.float32)).unsqueeze(1)\n    \n        meta_data = torch.cat([season_encoded_label, state_encoded_label, species_encoded_label, ndvi, height], dim=1)\n        return meta_data\n    \n    def prepare_data(self) -> Tuple[torch.Tensor, DataLoader]:\n        \"\"\"准备数据\"\"\"\n        # 读取训练数据\n        data = pd.read_csv(self.config.train_csv)\n        df_img = data['image_path'].drop_duplicates().reset_index(drop=True)\n        y_data = torch.reshape(\n            torch.tensor(data.iloc[:, -1]),\n            (int(len(data) / 5), 5)\n        )\n        self.target = y_data\n        \n        data1 = pd.read_excel(config.train1_xlsx)\n        meta_data = self.ProcessMetadata(data1.iloc[:,2:-3].values)\n\n        # Stratified/Group\n        state_encoder = LabelEncoder()\n        date_encoder = LabelEncoder()\n    \n        y = torch.tensor(state_encoder.fit_transform(data1['State'])).to(torch.int)\n        groups = torch.tensor(date_encoder.fit_transform(data1['Sampling_Date'])).to(torch.int)\n\n        # 特征提取\n        feature_extractor = FeatureExtractor(self.config, df_img, self.config.root_image_dir)\n        self.features = feature_extractor.extract_features()\n        \n        # 创建数据集和数据加载器\n        train_dataset = BiomassDataset(\n            self.features, y_data, meta_data\n        )\n        train_loader = DataLoader(\n            train_dataset,\n            batch_size=self.config.batch_size,\n            shuffle=True,\n            num_workers=self.config.num_workers\n        )\n\n        self.train_loader = DataLoader(\n            train_dataset,\n            batch_size=self.config.batch_size,\n            shuffle=False,\n            num_workers=self.config.num_workers\n        )\n\n        return y_data, meta_data, y, groups, train_loader\n\n    def create_model_and_optimizer(self) -> Tuple[nn.Module, optim.Optimizer]:\n        \"\"\"创建模型和优化器\"\"\"\n        model = BiomassModel_teacher(self.config).to(self.device)\n\n        # 设置不同的学习率\n        # 分离backbone参数和分类头参数\n        backbone_params = []\n        fc_params = []\n        \n        for name, param in model.named_parameters():\n            fc_params.append(param)\n\n        optimizer = optim.AdamW([\n            {'params': fc_params, 'lr': self.config.fc_lr, 'weight_decay': self.config.weight_decay}\n        ])\n\n        scheduler = CosineAnnealingLR(optimizer, T_max=10)\n\n        return model, optimizer, scheduler\n\n    def calculate_metrics(self, predictions: np.ndarray, targets: np.ndarray) -> Tuple[List[float], float]:\n        \"\"\"计算评估指标\"\"\"\n        # 计算每个目标的R²，仅用于观察拟合效果\n        r2_scores = []\n        weights = self.config.target_weights\n        num_targets = self.config.num_targets\n\n        for i in range(5):\n            r2 = r2_score(targets[:, i], predictions[:, i])\n            r2_scores.append(r2)\n\n        # 计算加权总分\n        SS_res = sum([sum(weights[i]*((predictions[:, i]-targets[:, i])**2)) for i in range(num_targets)])\n        weighted_y_sum = sum([sum(weights[i]*targets[:, i]) for i in range(num_targets)]) / len(targets)\n        SS_tot = sum([sum(weights[i]*((targets[:, i]-weighted_y_sum)**2)) for i in range(num_targets)])\n        r2_total = 1 - SS_res / SS_tot\n\n        return r2_scores, r2_total\n\n    def train_epoch(self, model: nn.Module, train_loader: DataLoader,\n                    optimizer: optim.Optimizer, criterion: nn.Module) -> float:\n        \"\"\"训练一个epoch\"\"\"\n        model.train()\n        total_loss = 0\n        num_batches = 0\n        \n        # 确定是否使用scaler（仅在CUDA + AMP时使用）\n        use_scaler = self.config.use_amp and self.device.type == 'cuda'\n        \n        for batch_x, batch_y, batch_meta in train_loader:\n            batch_x = batch_x.to(self.device)\n            batch_y = batch_y.to(self.device)\n            batch_meta = batch_meta.to(self.device)\n            \n            # 2. 梯度清零\n            optimizer.zero_grad(set_to_none=True)\n            \n            # 3. 设置Autocast（自动混合精度）\n            autocast_device_type = 'cuda' if self.device.type == 'cuda' else 'cpu'\n            with autocast(device_type=autocast_device_type, enabled=self.config.use_amp):\n                output = model(batch_x, batch_meta)\n                loss_main = criterion(output, batch_y.float())\n                loss_physics1 = nn.MSELoss()(output[:, 3], output[:, 0]+output[:, 1]+output[:, 2])\n                loss_physics2 = nn.MSELoss()(output[:, 4], output[:, 0]+output[:, 2])\n                loss = loss_main + self.config.beta * (loss_physics1 + loss_physics2)\n\n            # 4. 反向传播和优化器步骤\n            if use_scaler:\n                # 4a. CUDA + AMP路径\n                self.scaler.scale(loss).backward()\n                self.scaler.unscale_(optimizer)  # 在梯度裁剪前unscale\n                \n                # 梯度裁剪\n                torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)\n                \n                self.scaler.step(optimizer)\n                self.scaler.update()\n            else:\n                # 4b. CPU / MPS / AMP=False路径\n                loss.backward()\n                \n                # 梯度裁剪\n                torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)\n                \n                optimizer.step()\n            \n            # 5. 记录损失\n            total_loss += loss.item()\n            num_batches += 1\n        \n        return total_loss / max(1, num_batches)\n\n    @torch.no_grad()\n    def validate(self, model: nn.Module, train_loader: DataLoader) -> Tuple[float, List[float], float]:\n        \"\"\"验证模型\"\"\"\n        model.eval()\n        all_predictions = []\n        all_targets = []\n\n        for batch_x, batch_y, batch_meta in train_loader:\n            batch_x = batch_x.to(self.device)\n            batch_meta = batch_meta.to(self.device)\n            output = model(batch_x, batch_meta)\n\n            all_predictions.append(output.cpu().numpy())\n            all_targets.append(batch_y.cpu().numpy())\n\n        all_predictions = np.concatenate(all_predictions, axis=0)\n        all_targets = np.concatenate(all_targets, axis=0)\n\n        r2_scores, r2_total = self.calculate_metrics(all_predictions, all_targets)\n        loss_main = ((all_predictions - all_targets) ** 2 * self.config.loss_weights).mean()\n        loss_physics1 = ((all_predictions[:,3] - all_predictions[:,0] - \n                          all_predictions[:,1] - all_predictions[:,2]) ** 2).mean()\n        loss_physics2 = ((all_predictions[:,4] - all_predictions[:,0] - all_predictions[:,2]) ** 2).mean()\n        loss = loss_main + self.config.beta * (loss_physics1 + loss_physics2)\n        return loss, r2_scores, r2_total\n\n    def cross_validation(self, y_data: torch.Tensor, meta_data: torch.Tensor, \n                         y: torch.Tensor, groups: torch.Tensor) -> Dict:\n        \"\"\"执行交叉验证\"\"\"\n        print(\"Starting Cross-Validation...\")\n\n        # 创建分组\n        n_samples = len(y_data)\n\n        sgkf = StratifiedGroupKFold(n_splits=self.config.n_folds, shuffle=True)\n        sgkf.get_n_splits(np.arange(n_samples), y)\n        fold_results = {}\n\n        for fold, (train_idx, val_idx) in enumerate(sgkf.split(np.arange(n_samples), y, groups)):\n            print(f\"\\n=== Fold {fold + 1}/{self.config.n_folds} ===\")\n\n            # 创建数据加载器\n            train_dataset = BiomassDataset(\n                self.features, \n                y_data,\n                meta_data,\n                # is_train=True,\n                indices=train_idx.tolist() \n            )\n            val_dataset = BiomassDataset(\n                self.features, \n                y_data,\n                meta_data,\n                # is_train=True,\n                indices=val_idx.tolist() \n            )\n\n            train_loader = DataLoader(\n                train_dataset, batch_size=self.config.batch_size, shuffle=True,\n                num_workers=self.config.num_workers\n            )\n            val_loader = DataLoader(\n                val_dataset, batch_size=self.config.batch_size, shuffle=False,\n                num_workers=self.config.num_workers\n            )\n\n            # 训练模型\n            fold_history, model, optimizer = self.train_single_fold(fold, train_loader, val_loader)\n            fold_results[fold+1] = fold_history\n\n            # 保存模型\n            model_path = os.path.join(\n                self.config.output_dir, f'teacher_model_fold_{fold+1}.pth'\n            )\n            torch.save({\n                'model_state_dict': model.state_dict(),\n                'optimizer_state_dict': optimizer.state_dict()\n            }, model_path)\n        \n        with open(self.log_path, 'w') as f:\n            json.dump(fold_results, f, ensure_ascii=False, indent=2)\n        return fold_results\n\n    def train_single_fold(self, fold: int, train_loader: DataLoader,\n                          val_loader: DataLoader = None) -> Dict:\n        \"\"\"训练单个折痕\"\"\"\n        model, optimizer, scheduler = self.create_model_and_optimizer()\n        criterion = WeightedMSELoss(self.config)\n\n        fold_history = {\n            'train_loss': [], 'val_loss': [], 'r2_scores': [],\n            'r2_total': [], 'epoch_time': [], 'best_val_loss': 2e3,\n            'best_epoch': 0, 'best_r2': -1\n        }\n\n        best_model_state = None\n\n        for epoch in range(1, self.config.num_epochs + 1):\n            start_time = time.time()\n            scheduler.step()\n\n            # 训练\n            train_loss = self.train_epoch(model, train_loader, optimizer, criterion)\n\n            # 验证\n            if val_loader is not None:\n                val_loss, r2_scores, r2_total = self.validate(model, val_loader)\n            else:\n                val_loss, r2_scores, r2_total = self.validate(model, train_loader)\n\n            epoch_time = time.time() - start_time\n\n            # 记录历史\n            fold_history['train_loss'].append(train_loss)\n            fold_history['val_loss'].append(val_loss)\n            fold_history['r2_scores'].append(r2_scores)\n            fold_history['r2_total'].append(r2_total)\n            fold_history['epoch_time'].append(epoch_time)\n\n            if val_loss < fold_history['best_val_loss']:\n                fold_history['best_val_loss'] = val_loss\n                fold_history['best_r2'] = r2_total\n                fold_history['best_epoch'] = epoch\n                # 保存模型状态字典（深拷贝）\n                best_model_state = copy.deepcopy(model.state_dict())\n                print(f\"🎯 Fold {fold+1} - 新的最佳模型! val loss: {val_loss:.4f}, r2: {r2_total:.4f} (Epoch {epoch})\")\n\n            print(f\"Fold {fold+1}, Epoch {epoch}/{self.config.num_epochs}: \"\n                  f\"Train Loss: {train_loss:.4f}, Val Loss: {val_loss:.4f}, \"\n                  f\"R²: {[round(r, 4) for r in r2_scores]}, \"\n                  f\"R² Total: {r2_total:.4f}, Time: {epoch_time:.2f}s\")\n\n        if best_model_state is not None:\n            model.load_state_dict(best_model_state)\n            print(f\"✅ Fold {fold+1} - 已加载最佳模型 (Epoch {fold_history['best_epoch']}, val loss: {fold_history['best_val_loss']:.4f}, r2: {fold_history['best_r2']:.4f})\")\n        else:\n            print(f\"⚠️ Fold {fold+1} - 警告: 未找到最佳模型，使用最终epoch的模型\")\n\n        self.models.append(model)\n        self.model_r2.append(fold_history['best_r2'])\n        return fold_history, model, optimizer\n\n    def train(self) -> Dict:\n        \"\"\"训练主函数\"\"\"\n        y_data, meta_data, y, groups, train_loader = self.prepare_data()\n\n        if self.config.use_cross_validation:\n            return self.cross_validation(y_data, meta_data, y, groups)\n        else:\n            return self.train_single_fold(0, train_loader)\n\n    def predict_with_tta(self, model: nn.Module, dataloader: DataLoader) -> np.ndarray:\n        \"\"\"使用TTA进行预测\"\"\"\n        model.eval()\n        all_predictions = []\n\n        with torch.no_grad():\n            for batch_x, batch_label, batch_md in dataloader:\n                batch_x = batch_x.to(self.device)\n                batch_y = batch_label.to(self.device)\n                batch_md = batch_md.to(self.device)\n                \n                tta_predictions = []\n                pred = model(batch_x, batch_md)\n                all_predictions.append(pred.cpu().numpy())\n\n        return np.concatenate(all_predictions, axis=0)\n\n    def predict(self, test_loader: DataLoader, model_paths: List[str] = None) -> np.ndarray:\n        \"\"\"预测函数\"\"\"\n        all_predictions = []\n        lowest_r2_idx = np.argmin(self.model_r2)\n\n        for model in self.models:\n            if True:\n                if self.config.use_tta:\n                    predictions = self.predict_with_tta(model, test_loader)\n                else:\n                    model.eval()\n                    predictions = []\n                    with torch.no_grad():\n                        for batch in test_loader:\n                            batch_x = batch[0].to(self.device)\n                            batch_md = batch[2].to(self.device)\n                            output = model(batch_x, batch_md)\n                            predictions.append(output.cpu().numpy())\n                    predictions = np.concatenate(predictions, axis=0)\n                all_predictions.append(predictions)\n            else:\n                pass\n\n        # 平均所有模型的预测\n        if all_predictions:\n            return np.mean(all_predictions, axis=0)\n        else:\n            raise ValueError(\"No valid models found for prediction\")\n\n    def plot_training_history(self, history: Dict, save_path: str = None):\n        \"\"\"绘制训练历史（示例，以第1折为例）\"\"\"\n        fig, axes = plt.subplots(2, 2, figsize=(15, 10))\n        \n        # 训练损失\n        axes[0, 0].plot(history.get('train_loss', []))\n        axes[0, 0].set_title('Training Loss')\n        axes[0, 0].set_xlabel('Epoch')\n        axes[0, 0].set_ylabel('Loss')\n\n        # 验证损失\n        axes[0, 1].plot(history.get('val_loss', []))\n        axes[0, 1].set_title('Validation Loss')\n        axes[0, 1].set_xlabel('Epoch')\n        axes[0, 1].set_ylabel('Loss')\n\n        # R²总分\n        axes[1, 0].plot(history.get('r2_total', []))\n        axes[1, 0].set_title('R² Total Score')\n        axes[1, 0].set_xlabel('Epoch')\n        axes[1, 0].set_ylabel('R²')\n\n        #  epoch时间\n        axes[1, 1].plot(history.get('epoch_time', []))\n        axes[1, 1].set_title('Epoch Time')\n        axes[1, 1].set_xlabel('Epoch')\n        axes[1, 1].set_ylabel('Time (s)')\n\n        plt.tight_layout()\n\n        if save_path:\n            plt.savefig(save_path)\n        plt.show()","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2025-11-24T05:13:35.939409Z","iopub.execute_input":"2025-11-24T05:13:35.939628Z","iopub.status.idle":"2025-11-24T05:13:35.999779Z","shell.execute_reply.started":"2025-11-24T05:13:35.939609Z","shell.execute_reply":"2025-11-24T05:13:35.998879Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# 创建配置\nconfig = BiomassConfig()\n\n# 创建训练器\ntrainer = BiomassTrainer(config)","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2025-11-24T05:13:36.001806Z","iopub.execute_input":"2025-11-24T05:13:36.001985Z","iopub.status.idle":"2025-11-24T05:13:36.017193Z","shell.execute_reply.started":"2025-11-24T05:13:36.001970Z","shell.execute_reply":"2025-11-24T05:13:36.016526Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# 训练模型\nhistory, model, optimizer = trainer.train()","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2025-11-24T05:13:36.017898Z","iopub.execute_input":"2025-11-24T05:13:36.018110Z","iopub.status.idle":"2025-11-24T05:14:21.155448Z","shell.execute_reply.started":"2025-11-24T05:13:36.018091Z","shell.execute_reply":"2025-11-24T05:14:21.154485Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# 绘制训练历史并保存\ntrainer.plot_training_history(history)","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2025-11-24T05:14:21.156745Z","iopub.execute_input":"2025-11-24T05:14:21.157097Z","iopub.status.idle":"2025-11-24T05:14:21.849560Z","shell.execute_reply.started":"2025-11-24T05:14:21.157069Z","shell.execute_reply":"2025-11-24T05:14:21.848702Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# 准备测试数据，此处与训练数据相同\nif config.use_tta:\n    data = pd.read_csv(config.train_csv)\n    df_img = data['image_path'].drop_duplicates().reset_index(drop=True)\n    y_data = torch.reshape(\n        torch.tensor(data.iloc[:, -1]),\n        (int(len(data) / 5), 5)\n    )    \n    data1 = pd.read_excel(config.train1_xlsx)\n    meta_data = trainer.ProcessMetadata(data1.iloc[:,2:-3].values)\n\n    tl = BiomassDataset_tta(df_img, y_data, meta_data, config.root_image_dir, config)\n    dataloader = DataLoader(tl, \n                        batch_size=config.batch_size,\n                        shuffle=False,\n                        num_workers=config.num_workers)\n    feature_extractor = FeatureExtractor(config, df_img, config.root_image_dir, False, True, dataloader)\n    features = feature_extractor.extract_features()\n    new_y = torch.stack([tl[i][1] for i in range(len(tl))])\n    new_md = torch.stack([tl[i][2] for i in range(len(tl))])\n    new_tl = BiomassDataset(features, new_y, new_md)\n    new_dataloader = DataLoader(new_tl, \n                        batch_size=config.batch_size,\n                        shuffle=False,\n                        num_workers=config.num_workers)\n    predictions = trainer.predict(new_dataloader).reshape(-1, 2, 5).mean(axis=1)\n    tg = trainer.target.cpu().numpy()\n    print(predictions.shape, tg.shape)\nelse:\n    tl = trainer.train_loader\n    tg = trainer.target.numpy()\n    predictions = trainer.predict(tl)\n\nr2_scores, r2_total = trainer.calculate_metrics(predictions, tg)\n\nprint(f\"最终结果：r2_scores:{[round(r,4) for r in r2_scores]}, r2_total:{r2_total:.4f}\")","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2025-11-24T05:16:03.143979Z","iopub.execute_input":"2025-11-24T05:16:03.144270Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# 准备提交文件\ndf_test = pd.read_csv(config.train_csv)\nflattened_predictions = predictions.flatten().tolist()\ndf_test['target_pred'] = flattened_predictions\ndf_test.to_excel('teacher_prediction.xlsx')\nprint(\"Training and prediction done!\")","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2025-11-24T05:14:22.912587Z","iopub.status.idle":"2025-11-24T05:14:22.912889Z","shell.execute_reply.started":"2025-11-24T05:14:22.912740Z","shell.execute_reply":"2025-11-24T05:14:22.912754Z"}},"outputs":[],"execution_count":null}]}