{"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":13772481,"sourceType":"datasetVersion","datasetId":8758997},{"sourceId":4535,"sourceType":"modelInstanceVersion","isSourceIdPinned":true,"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":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-19T16:22:32.413101Z","iopub.execute_input":"2025-11-19T16:22:32.414047Z","iopub.status.idle":"2025-11-19T16:22:32.419454Z","shell.execute_reply.started":"2025-11-19T16:22:32.414016Z","shell.execute_reply":"2025-11-19T16:22:32.418608Z"}},"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":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-19T16:22:32.421035Z","iopub.execute_input":"2025-11-19T16:22:32.42127Z","iopub.status.idle":"2025-11-19T16:22:32.453769Z","shell.execute_reply.started":"2025-11-19T16:22:32.421252Z","shell.execute_reply":"2025-11-19T16:22:32.453037Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class BiomassConfig:\n    \"\"\"配置参数类\"\"\"\n\n    def __init__(self):\n        # 基础设置\n        self.seed = 2025\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.test_image_dir = '/kaggle/input/csiro-biomass'\n\n        # DINOv2模型路径\n        self.dinov2_path = '/kaggle/input/dinov2/pytorch/large/1/'\n\n        # 训练参数\n        self.batch_size = 16  # 减小batch size以适应ViT\n        self.num_epochs = 40\n        self.learning_rate = 1e-4  # 更小的学习率\n        self.fc_lr = 1e-2  # 全连接层学习率\n        self.weight_decay = 1e-4\n        self.num_workers = 4\n\n        # 图像处理 - 适应ViT输入尺寸\n        self.hidden_dim = 1024\n        self.image_size = 224  # ViT标准输入尺寸\n        # 使用ImageNet标准化参数，因为DINOv2是在ImageNet上预训练的\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 = True\n\n        # 混合精度训练\n        self.use_amp = True\n\n        # TTA (测试时增强)\n        self.use_tta = False\n        self.tta_scales = [1.0, 0.9, 1.1]  # 多尺度\n        self.tta_flip = True  # 水平翻转\n        self.tta_vflip = True # 垂直翻转\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":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-19T16:22:32.454646Z","iopub.execute_input":"2025-11-19T16:22:32.45489Z","iopub.status.idle":"2025-11-19T16:22:32.475053Z","shell.execute_reply.started":"2025-11-19T16:22:32.45487Z","shell.execute_reply":"2025-11-19T16:22:32.474386Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class BiomassDataset(Dataset):\n    \"\"\"生物量数据集类\"\"\"\n\n    def __init__(self, image_folder: str, y_data: torch.Tensor = None,\n                 transform: transforms.Compose = None, is_train: bool = True,\n                 df_imgs: pd.DataFrame = None, indices: list = None):\n        self.image_folder = image_folder\n        self.transform = transform\n        self.is_train = is_train\n        self.df_imgs = df_imgs\n        self.indices = indices  # 新增：指定使用的索引\n\n        if is_train:\n            # 收集所有图像路径\n            self.image_paths = []\n            for root, dirs, files in os.walk(image_folder):\n                for file in files:\n                    if file.lower().endswith('.jpg'):\n                        self.image_paths.append(os.path.join(root, file))\n            self.image_paths.sort()  # 确保顺序一致\n            \n            # 如果指定了索引，只使用这些索引对应的图像和标签\n            if indices is not None:\n                self.image_paths = [self.image_paths[i] for i in indices]\n                self.y_data = y_data[indices] if y_data is not None else None\n            else:\n                self.y_data = y_data\n        else:\n            self.image_paths = df_imgs['image_path'].values.tolist() if df_imgs is not None else []\n\n    def __len__(self):\n        return len(self.image_paths)\n\n    def __getitem__(self, idx):\n        image_path = self.image_paths[idx]\n        image = Image.open(image_path).convert('RGB')\n\n        if self.transform:\n            image = self.transform(image)\n\n        if self.is_train and self.y_data is not None:\n            label = self.y_data[idx]\n            return image, label\n        else:\n            return image, image_path\n\nclass TestDataset(Dataset):\n    \"\"\"测试数据集类\"\"\"\n\n    def __init__(self, df_imgs: pd.DataFrame, image_root: str, transform: transforms.Compose):\n        self.df = df_imgs.reset_index(drop=True)\n        self.root = Path(image_root)\n        self.transform = transform\n\n    def __len__(self):\n        return len(self.df)\n\n    def __getitem__(self, idx):\n        row = self.df.iloc[idx]\n        img_path = self.root / row['image_path']\n        with Image.open(img_path) as im:\n            im = im.convert('RGB')\n        x = self.transform(im)\n        return x, row['image_path']","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-19T16:22:32.475816Z","iopub.execute_input":"2025-11-19T16:22:32.476053Z","iopub.status.idle":"2025-11-19T16:22:32.501463Z","shell.execute_reply.started":"2025-11-19T16:22:32.476033Z","shell.execute_reply":"2025-11-19T16:22:32.500816Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class BiomassModel(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.autocast_device_type = 'cuda' if self.device.type == 'cuda' else 'cpu'\n\n        # 加载DINOv2模型\n        self.backbone = AutoModel.from_pretrained(config.dinov2_path)\n        \n        # 更复杂的全连接层\n        self.fc = self._build_complex_fc_layer()\n        \n        # 冻结backbone参数（更全面的冻结策略）\n        self.freeze_backbone()\n\n    def _build_complex_fc_layer(self):\n        \"\"\"构建更复杂的全连接层\"\"\"\n        hidden_dim1 = 512\n        hidden_dim2 = 256\n        hidden_dim3 = 128\n        dropout_rate = 0.4\n        \n        fc_layers = nn.Sequential(\n            # 第一层\n            nn.Linear(self.config.hidden_dim, hidden_dim1),\n            nn.BatchNorm1d(hidden_dim1),\n            nn.GELU(),\n            nn.Dropout(dropout_rate),\n            \n            # 第二层\n            nn.Linear(hidden_dim1, hidden_dim2),\n            nn.BatchNorm1d(hidden_dim2),\n            nn.GELU(),\n            nn.Dropout(dropout_rate),\n            \n            # 第三层\n            nn.Linear(hidden_dim2, hidden_dim3),\n            nn.BatchNorm1d(hidden_dim3),\n            nn.GELU(),\n            nn.Dropout(dropout_rate * 0.5),  # 最后一层dropout较小\n            \n            # 输出层\n            nn.Linear(hidden_dim3, self.config.num_targets)\n        )\n        \n        # 初始化权重\n        self._init_fc_weights(fc_layers)\n        return fc_layers\n\n    def _init_fc_weights(self, model):\n        \"\"\"初始化全连接层权重\"\"\"\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 freeze_backbone(self, freeze_strategy=\"full\"):\n        \"\"\"\n        冻结backbone参数\n        \n        Args:\n            freeze_strategy: \n                \"full\" - 完全冻结backbone\n                \"partial\" - 冻结大部分层，微调最后几层\n                \"none\" - 不冻结，全部微调\n        \"\"\"\n        if freeze_strategy == \"full\":\n            # 完全冻结backbone\n            for name, param in self.backbone.named_parameters():\n                param.requires_grad = False\n            print(\"Backbone fully frozen\")\n            \n        elif freeze_strategy == \"partial\":\n            # 冻结策略：只微调最后2层，其他全部冻结\n            frozen_layers = 10  # ViT base有12层，冻结前10层\n            \n            if hasattr(self.backbone, 'encoder'):\n                encoder = self.backbone.encoder\n                if hasattr(encoder, 'layer'):\n                    layers = encoder.layer\n                    for i, layer in enumerate(layers):\n                        if i < frozen_layers:\n                            # 冻结前10层\n                            for param in layer.parameters():\n                                param.requires_grad = False\n                        else:\n                            # 最后2层可训练\n                            for param in layer.parameters():\n                                param.requires_grad = True\n                    \n                    print(f\"Backbone partially frozen: {frozen_layers}/12 layers frozen\")\n            \n            # 冻结patch embedding和position embedding\n            for name, param in self.backbone.named_parameters():\n                if 'embeddings' in name:\n                    param.requires_grad = False\n            \n        elif freeze_strategy == \"none\":\n            # 全部可训练\n            for param in self.backbone.parameters():\n                param.requires_grad = True\n            print(\"Backbone fully trainable\")\n        \n        # 打印可训练参数统计\n        self._print_trainable_parameters()\n\n    def _print_trainable_parameters(self):\n        \"\"\"打印可训练参数统计\"\"\"\n        total_params = sum(p.numel() for p in self.parameters())\n        trainable_params = sum(p.numel() for p in self.parameters() if p.requires_grad)\n        frozen_params = total_params - trainable_params\n        \n        print(f\"Total parameters: {total_params:,}\")\n        print(f\"Trainable parameters: {trainable_params:,}\")\n        print(f\"Frozen parameters: {frozen_params:,}\")\n        print(f\"Trainable ratio: {trainable_params/total_params:.2%}\")\n\n    def unfreeze_backbone_layers(self, num_layers_to_unfreeze=2):\n        \"\"\"\n        逐步解冻backbone层（用于训练后期）\n        \n        Args:\n            num_layers_to_unfreeze: 要解冻的层数（从最后开始）\n        \"\"\"\n        if hasattr(self.backbone, 'encoder'):\n            encoder = self.backbone.encoder\n            if hasattr(encoder, 'layer'):\n                layers = encoder.layer\n                total_layers = len(layers)\n                \n                # 解冻最后几层\n                start_layer = total_layers - num_layers_to_unfreeze\n                for i in range(max(0, start_layer), total_layers):\n                    for param in layers[i].parameters():\n                        param.requires_grad = True\n                \n                print(f\"Unfrozen last {num_layers_to_unfreeze} backbone layers\")\n\n\n    def forward(self, x):\n        with autocast(device_type=self.autocast_device_type, enabled=self.config.use_amp):\n            # 使用DINOv2 backbone提取特征\n            outputs = self.backbone(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            # 通过复杂的全连接层\n            output = self.fc(features)\n            \n            # 可选：对输出进行约束（生物量不能为负）\n            output = F.relu(output)  # 或者使用 softplus 获得更平滑的正值\n            # output = F.softplus(output)\n            \n        return output","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-19T16:22:32.503441Z","iopub.execute_input":"2025-11-19T16:22:32.503922Z","iopub.status.idle":"2025-11-19T16:22:32.527335Z","shell.execute_reply.started":"2025-11-19T16:22:32.503898Z","shell.execute_reply":"2025-11-19T16:22:32.526681Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class WeightedMSELoss(nn.Module):\n    def __init__(self, weights=[1.0, 1.0, 1.0, 5.0, 2.0]):\n        super().__init__()\n        self.weights = torch.tensor(weights)\n\n    def forward(self, input, target):\n        # 计算每个维度的MSE\n        mse_per_dim = (input - target) ** 2\n\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        # 求平均\n        return weighted_mse.mean()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-19T16:22:32.528022Z","iopub.execute_input":"2025-11-19T16:22:32.528226Z","iopub.status.idle":"2025-11-19T16:22:32.559008Z","shell.execute_reply.started":"2025-11-19T16:22:32.528209Z","shell.execute_reply":"2025-11-19T16:22:32.558239Z"}},"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\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    def prepare_data(self) -> Tuple[torch.Tensor, DataLoader]:\n        \"\"\"准备数据\"\"\"\n        # 读取训练数据\n        data = pd.read_csv(self.config.train_csv)\n        data1 = pd.read_excel(self.config.train1_xlsx)\n        y_data = torch.reshape(\n            torch.tensor(data.iloc[:, -1]),\n            (int(len(data) / 5), 5)\n        )\n\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        # 数据变换 - 使用适合ViT的变换\n        transform = transforms.Compose([\n            transforms.Resize((self.config.image_size, self.config.image_size)),\n            transforms.ToTensor(),\n            transforms.Normalize(\n                mean=self.config.normalize_mean,\n                std=self.config.normalize_std\n            )\n        ])\n\n        # 创建数据集和数据加载器\n        train_dataset = BiomassDataset(\n            self.config.train_image_dir, y_data, transform, is_train=True\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        return y_data, y, groups, train_loader\n\n    def create_model_and_optimizer(self) -> Tuple[nn.Module, optim.Optimizer]:\n        \"\"\"创建模型和优化器\"\"\"\n        model = BiomassModel(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            if 'fc' in name:\n                fc_params.append(param)\n            else:\n                backbone_params.append(param)\n\n        optimizer = optim.AdamW([\n            {'params': backbone_params, 'lr': self.config.learning_rate, 'weight_decay': self.config.weight_decay},\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        for i in range(5):\n            r2 = r2_score(targets[:, i], predictions[:, i])\n            r2_scores.append(r2)\n\n        # 计算加权总分\n        r2_total = sum([r * w for r, w in zip(r2_scores, self.config.target_weights)])\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 in train_loader:\n            # 1. 数据移动到设备并确保数据类型正确\n            batch_x = batch_x.to(self.device)\n            batch_y = batch_y.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)\n                loss = criterion(output, batch_y.float())\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 in train_loader:\n            batch_x = batch_x.to(self.device)\n            output = model(batch_x)\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 = np.mean((all_predictions - all_targets) ** 2)  # MSE\n\n        return loss, r2_scores, r2_total\n\n    def cross_validation(self, y_data: torch.Tensor, y: torch.Tensor, groups: np.array) -> 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        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            transform = transforms.Compose([\n                transforms.Resize((self.config.image_size, self.config.image_size)),\n                transforms.ToTensor(),\n                transforms.Normalize(\n                    mean=self.config.normalize_mean,\n                    std=self.config.normalize_std\n                )\n            ])\n\n            # 创建数据加载器\n            train_dataset = BiomassDataset(\n                self.config.train_image_dir, \n                y_data,  # 传入完整的y_data\n                transform, \n                is_train=True,\n                indices=train_idx.tolist()  # 指定使用哪些索引\n            )\n            val_dataset = BiomassDataset(\n                self.config.train_image_dir, \n                y_data,  # 传入完整的y_data\n                transform, \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'dinov2_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()\n\n        fold_history = {\n            'train_loss': [], 'val_loss': [], 'r2_scores': [],\n            'r2_total': [], 'epoch_time': [], 'best_val_loss': 1e3,\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        return fold_history, model, optimizer\n\n    def train(self) -> Dict:\n        \"\"\"训练主函数\"\"\"\n        y_data, y, group, train_loader = self.prepare_data()\n\n        if self.config.use_cross_validation:\n            return self.cross_validation(y_data, y, group)\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 in dataloader:\n                if isinstance(batch, list):\n                    batch_x = batch[0].to(self.device)\n                else:\n                    batch_x = batch.to(self.device)\n\n                # 多尺度 + 翻转 TTA\n                tta_predictions = []\n\n                for scale in self.config.tta_scales:\n                    # 缩放\n                    if scale != 1.0:\n                        new_size = int(self.config.image_size * scale)\n                        batch_scaled = F.interpolate(\n                            batch_x, size=(new_size, new_size),\n                            mode='bilinear', align_corners=False\n                        )\n                        # 中心裁剪\n                        batch_scaled = transforms.CenterCrop(\n                            self.config.image_size\n                        )(batch_scaled)\n                    else:\n                        batch_scaled = batch_x\n\n                    # 原始预测\n                    pred = model(batch_scaled)\n                    tta_predictions.append(pred)\n\n                    # 水平翻转\n                    if self.config.tta_flip:\n                        batch_flipped = torch.flip(batch_scaled, dims=[3])\n                        pred_flipped = model(batch_flipped)\n                        tta_predictions.append(pred_flipped)\n\n                    # 垂直翻转\n                    if self.config.tta_vflip:  # 新增垂直翻转配置\n                        batch_vflipped = torch.flip(batch_scaled, dims=[2])  # 垂直翻转\n                        pred_vflipped = model(batch_vflipped)\n                        tta_predictions.append(pred_vflipped)\n\n                # 平均所有TTA预测\n                final_pred = torch.stack(tta_predictions).mean(dim=0)\n                all_predictions.append(final_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\n        for model in self.models:\n            # 预测\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                        if isinstance(batch, list):\n                            batch_x = batch[0].to(self.device)\n                        else:\n                            batch_x = batch.to(self.device)\n                        output = model(batch_x)\n                        predictions.append(output.cpu().numpy())\n                predictions = np.concatenate(predictions, axis=0)\n\n            all_predictions.append(predictions)\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[1].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[1].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[1].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[1].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":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-19T16:22:32.560006Z","iopub.execute_input":"2025-11-19T16:22:32.560264Z","iopub.status.idle":"2025-11-19T16:22:32.606375Z","shell.execute_reply.started":"2025-11-19T16:22:32.560236Z","shell.execute_reply":"2025-11-19T16:22:32.605675Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# 创建配置\nconfig = BiomassConfig()\n\n# 创建训练器\ntrainer = BiomassTrainer(config)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-19T16:22:32.607064Z","iopub.execute_input":"2025-11-19T16:22:32.607244Z","iopub.status.idle":"2025-11-19T16:22:32.630386Z","shell.execute_reply.started":"2025-11-19T16:22:32.60723Z","shell.execute_reply":"2025-11-19T16:22:32.629673Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# 训练模型\nhistory = trainer.train()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-19T16:22:32.631149Z","iopub.execute_input":"2025-11-19T16:22:32.63139Z","iopub.status.idle":"2025-11-19T16:23:57.056014Z","shell.execute_reply.started":"2025-11-19T16:22:32.631374Z","shell.execute_reply":"2025-11-19T16:23:57.055063Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# 绘制训练历史并保存\ntrainer.plot_training_history(history)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-19T16:23:57.057294Z","iopub.execute_input":"2025-11-19T16:23:57.057904Z","iopub.status.idle":"2025-11-19T16:23:57.797439Z","shell.execute_reply.started":"2025-11-19T16:23:57.057864Z","shell.execute_reply":"2025-11-19T16:23:57.796797Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# 准备测试数据\ndf_test = pd.read_csv(config.test_csv)\ndf_imgs = df_test[['image_path']].drop_duplicates().reset_index(drop=True)\n\ntransform = transforms.Compose([\n    transforms.Resize((config.image_size, config.image_size)),\n    transforms.ToTensor(),\n    transforms.Normalize(mean=config.normalize_mean, std=config.normalize_std)\n])\n\ntest_dataset = TestDataset(df_imgs, config.test_image_dir, transform)\ntest_loader = DataLoader(test_dataset, batch_size=config.batch_size,\n                         shuffle=False, num_workers=config.num_workers)\n\n# 预测\npredictions = trainer.predict(test_loader)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-19T16:23:57.798275Z","iopub.execute_input":"2025-11-19T16:23:57.798609Z","iopub.status.idle":"2025-11-19T16:23:59.56878Z","shell.execute_reply.started":"2025-11-19T16:23:57.798554Z","shell.execute_reply":"2025-11-19T16:23:59.567877Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# 准备提交文件\nflattened_predictions = predictions.flatten().tolist()\ndf_test['target'] = flattened_predictions\ndf_test.drop(columns=['image_path', 'target_name'], inplace=True)\ndf_test.to_csv('submission.csv', index=False)\n\nprint(\"DINOv2 training and prediction completed!\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-19T16:23:59.570146Z","iopub.execute_input":"2025-11-19T16:23:59.570584Z","iopub.status.idle":"2025-11-19T16:23:59.588587Z","shell.execute_reply.started":"2025-11-19T16:23:59.570549Z","shell.execute_reply":"2025-11-19T16:23:59.587567Z"}},"outputs":[],"execution_count":null}]}