{"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":14774,"databundleVersionId":875431,"sourceType":"competition"},{"sourceId":12722239,"sourceType":"datasetVersion","datasetId":8041124},{"sourceId":12725723,"sourceType":"datasetVersion","datasetId":8043467},{"sourceId":12749005,"sourceType":"datasetVersion","datasetId":8059130},{"sourceId":12779032,"sourceType":"datasetVersion","datasetId":8078999},{"sourceId":12779055,"sourceType":"datasetVersion","datasetId":8079015},{"sourceId":12779099,"sourceType":"datasetVersion","datasetId":8079047},{"sourceId":12780694,"sourceType":"datasetVersion","datasetId":8080175},{"sourceId":256313063,"sourceType":"kernelVersion"}],"dockerImageVersionId":31090,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"import cv2\nimport numpy as np\nimport os\nimport pandas as pd\nimport random\nimport timm\nimport torch\nimport torch.nn as nn\nimport torch.optim as optim\nfrom albumentations import Compose, Normalize, HorizontalFlip, Rotate, ColorJitter, Resize\nfrom albumentations.pytorch import ToTensorV2\nfrom sklearn.metrics import cohen_kappa_score\nfrom sklearn.model_selection import train_test_split\nfrom torch.optim.lr_scheduler import CosineAnnealingLR\nfrom torch.utils.data import DataLoader, Dataset\nfrom tqdm import tqdm\n\n# ==============================================================================\n# 0. 全局设置\n# ==============================================================================\ndef set_seed(seed=42):\n    random.seed(seed)\n    np.random.seed(seed)\n    torch.manual_seed(seed)\n    if torch.cuda.is_available():\n        torch.cuda.manual_seed_all(seed)\n        torch.backends.cudnn.deterministic = True\n        torch.backends.cudnn.benchmark = False\n\nset_seed(42)\n\n# ==============================================================================\n# 1. PyTorch 模型与数据类\n# ==============================================================================\ndef replace_batchnorm_with_groupnorm(module, num_groups=32):\n    for name, child in module.named_children():\n        if isinstance(child, nn.BatchNorm2d):\n            num_channels = child.num_features\n            if num_channels % num_groups == 0:\n                setattr(module, name, nn.GroupNorm(num_groups=num_groups, num_channels=num_channels))\n        else:\n            replace_batchnorm_with_groupnorm(child, num_groups)\n\nclass EfficientNetModel(nn.Module):\n    def __init__(self, model_name, pretrained=True):\n        super().__init__()\n        self.model = timm.create_model(model_name, pretrained=pretrained)\n        replace_batchnorm_with_groupnorm(self.model)\n        \n        in_features = self.model.classifier.in_features\n        self.model.classifier = nn.Sequential(\n            nn.Dropout(p=0.5),\n            nn.Linear(in_features, 1)\n        )\n\n    def forward(self, x):\n        return self.model(x)\n\nclass CustomImageDataset(Dataset):\n    def __init__(self, img_dir, labels_df, transform=None):\n        self.img_dir = img_dir\n        self.labels_df = labels_df.copy()\n        self.transform = transform\n        self.label_to_folder = {\n            0: 'No_DR', 1: 'Mild', 2: 'Moderate', 3: 'Severe', 4: 'Proliferate_DR'\n        }\n        self.labels_df['level'] = self.labels_df['level'].astype(int)\n        self.labels_df['folder_name'] = self.labels_df['level'].map(self.label_to_folder)\n        \n    def __len__(self):\n        return len(self.labels_df)\n\n    def __getitem__(self, idx):\n        row = self.labels_df.iloc[idx]\n        img_id = row['image']\n        label = float(row['level'])\n        folder_name = row['folder_name']\n\n        img_path_jpeg = os.path.join(self.img_dir, folder_name, img_id + '.jpeg')\n        img_path_png = os.path.join(self.img_dir, folder_name, img_id + '.png')\n\n        if os.path.exists(img_path_jpeg):\n            img_path = img_path_jpeg\n        elif os.path.exists(img_path_png):\n            img_path = img_path_png\n        else:\n            # 在数据量很大时，频繁打印会拖慢速度，可以注释掉\n            # print(f\"警告: 图片未找到 {img_path_jpeg} 或 {img_path_png}, 将跳过\")\n            return self.__getitem__((idx + 1) % len(self))\n\n        try:\n            img = cv2.imread(img_path)\n            img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)\n            if self.transform:\n                img = self.transform(image=img)[\"image\"]\n            return img, label\n        except Exception as e:\n            # print(f\"警告: 读取图片 {img_path} 时出错: {e}, 将跳过\")\n            return self.__getitem__((idx + 1) % len(self))\n\ndef get_transforms(data_type='train', image_size=512):\n    if data_type == 'train':\n        return Compose([\n            #Resize(image_size, image_size),\n            Rotate(limit=40, p=0.5),\n            ColorJitter(brightness=0.1, contrast=0.1, saturation=0.1, hue=0.1, p=0.3),\n            HorizontalFlip(p=0.5),\n            Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),\n            ToTensorV2(),\n        ])\n    else: # 'valid'\n        return Compose([\n            Resize(image_size, image_size),\n            Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),\n            ToTensorV2(),\n        ])\n\n# ==============================================================================\n# 2. 训练函数\n# ==============================================================================\ndef train(model, train_loader, valid_loader, device, num_epochs, lr, image_size):\n    criterion = nn.MSELoss()\n    optimizer = optim.AdamW(model.parameters(), lr=lr, weight_decay=1e-5)\n    scheduler = CosineAnnealingLR(optimizer, T_max=num_epochs, eta_min=1e-6)\n\n    best_valid_kappa = 0.0\n    epochs_no_improve = 0\n    patience = 5\n    save_path = os.path.join(os.getcwd(), f'best_model_kappa_{image_size}.pth')\n\n    for epoch in range(num_epochs):\n        model.train()\n        running_loss = 0.0\n        \n        for images, labels in tqdm(train_loader, desc=f\"Epoch {epoch+1}/{num_epochs} - 训练中\"):\n            images = images.to(device)\n            labels = labels.to(device).float().view(-1, 1)\n            optimizer.zero_grad()\n            outputs = model(images)\n            loss = criterion(outputs, labels)\n            loss.backward()\n            optimizer.step()\n            running_loss += loss.item()\n\n        train_loss = running_loss / len(train_loader)\n        \n        model.eval()\n        all_valid_labels = []\n        all_valid_preds = []\n        with torch.no_grad():\n            for images, labels in tqdm(valid_loader, desc=f\"Epoch {epoch+1}/{num_epochs} - 验证中\"):\n                images = images.to(device)\n                outputs = model(images).squeeze()\n                all_valid_labels.extend(labels.cpu().numpy())\n                preds_np = outputs.cpu().numpy()\n                all_valid_preds.extend(np.atleast_1d(preds_np))\n        \n        rounded_preds = np.clip(np.round(all_valid_preds).astype(int), 0, 4)\n        valid_kappa = cohen_kappa_score(all_valid_labels, rounded_preds, weights='quadratic')\n\n        print(f\"\\nEpoch {epoch + 1}/{num_epochs}\")\n        print(f\"  训练损失: {train_loss:.4f}\")\n        print(f\"  验证 Kappa (四舍五入): {valid_kappa:.4f}\")\n        print(f\"  当前学习率: {optimizer.param_groups[0]['lr']:.6f}\")\n\n        scheduler.step()\n\n        if valid_kappa > best_valid_kappa:\n            best_valid_kappa = valid_kappa\n            epochs_no_improve = 0\n            torch.save(model.state_dict(), save_path)\n            print(f\"  -> 已保存新最佳模型! 验证 Kappa: {best_valid_kappa:.4f}\")\n            print(f\"   模型权重已保存至: {save_path}\")\n        else:\n            epochs_no_improve += 1\n        \n        if epochs_no_improve >= patience:\n            print(f\"\\n在 {patience} 个 epoch 没有提升后触发早停。\")\n            break\n    \n    print(f\"\\n训练结束。最佳模型权重保存在: {save_path}\")\n\n# ==============================================================================\n# 3. 主执行逻辑 (已修改为 EfficientNet-B6)\n# ==============================================================================\nif __name__ == \"__main__\":\n    # --- 步骤 1: 定义路径和参数 ---\n    print(\"--- 步骤 1: 初始化路径和参数 ---\")\n    \n    # <--- 修改: 切换为 B6 模型 ---\n    MODEL_NAME = 'efficientnet_b6'\n    # <--- 修改: 为 B6 模型使用 512x512 图像尺寸 ---\n    IMAGE_SIZE = 512\n    # <--- 修改: 因模型和图像变大，大幅减小 Batch Size 以防显存溢出 ---\n    BATCH_SIZE = 4 \n    NUM_EPOCHS = 20\n    # <--- 修改: 为 B6 模型使用稍小的学习率 ---\n    LEARNING_RATE = 2e-4\n    NUM_WORKERS = 2\n    \n    # --- 路径定义 (已更新) ---\n    DATA_DIR = '/kaggle/input/newdataset512/new_data_preprocessed_images_512'\n    LABEL_FILE = '/kaggle/input/new-data-labels/trainLabels.csv'\n    \n    # <--- 修改: 使用 B6 的权重文件路径 ---\n    PRETRAINED_MODEL_PATH = '/kaggle/input/fork-of-efficientnetb6-newdataset/best_model_kappa_512.pth'\n    \n    model_path = os.path.join(os.getcwd(), f'best_model_kappa_{IMAGE_SIZE}.pth')\n\n    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n    print(f\"使用的设备: {device}\")\n    \n    # --- 步骤 2: 加载数据并创建数据集 ---\n    print(\"\\n--- 步骤 2: 加载数据 ---\")\n    all_labels_df = pd.read_csv(LABEL_FILE)\n    all_labels_df['image'] = all_labels_df['image'].astype(str)\n    print(f\"总共加载了 {len(all_labels_df)} 个标签。\")\n\n    train_df, valid_df = train_test_split(\n        all_labels_df, \n        test_size=0.2, \n        random_state=42, \n        stratify=all_labels_df['level']\n    )\n    \n    # 注意：get_transforms 现在需要传入 image_size\n    train_dataset = CustomImageDataset(DATA_DIR, train_df, transform=get_transforms('train', IMAGE_SIZE))\n    valid_dataset = CustomImageDataset(DATA_DIR, valid_df, transform=get_transforms('valid', IMAGE_SIZE))\n    \n    train_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True, num_workers=NUM_WORKERS, pin_memory=True)\n    valid_loader = DataLoader(valid_dataset, batch_size=BATCH_SIZE, shuffle=False, num_workers=NUM_WORKERS, pin_memory=True)\n    \n    print(\"数据加载器创建完毕。\")\n    \n    # --- 步骤 3: 执行模型训练 ---\n    print(\"\\n--- 步骤 3: 开始模型训练 ---\")\n    \n    model = EfficientNetModel(MODEL_NAME, pretrained=False).to(device)\n    \n    if PRETRAINED_MODEL_PATH and os.path.exists(PRETRAINED_MODEL_PATH):\n        print(f\"正在从本地路径加载预训练权重: {PRETRAINED_MODEL_PATH}\")\n        model.model.load_state_dict(torch.load(PRETRAINED_MODEL_PATH), strict=False)\n    elif PRETRAINED_MODEL_PATH:\n        print(f\"警告: 预训练模型路径未找到: {PRETRAINED_MODEL_PATH}\")\n    else:\n        print(\"未提供预训练模型路径，将使用随机初始化权重。\")\n\n    train(model, train_loader, valid_loader, device, NUM_EPOCHS, LEARNING_RATE, IMAGE_SIZE)\n\n    print(\"\\n--- 脚本执行结束 ---\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-08-10T15:22:44.252016Z","iopub.execute_input":"2025-08-10T15:22:44.252338Z"}},"outputs":[],"execution_count":null}]}