{"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":12786005,"sourceType":"datasetVersion","datasetId":8083628},{"sourceId":12779055,"sourceType":"datasetVersion","datasetId":8079015},{"sourceId":12786116,"sourceType":"datasetVersion","datasetId":8083695}],"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# 2. 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_model_path=None):\n        super().__init__()\n        # 如果没有提供本地路径，则使用timm的默认ImageNet预训练权重\n        use_timm_pretrained = pretrained_model_path is None\n        self.model = timm.create_model(model_name, pretrained=use_timm_pretrained)\n        \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.3),\n            nn.Linear(in_features, 512),\n            nn.Mish(),\n            nn.Dropout(p=0.5),\n            nn.Linear(512, 1)\n        )\n        \n        # 如果提供了本地路径，则加载它\n        if pretrained_model_path and os.path.exists(pretrained_model_path):\n            print(f\"加载自定义本地预训练权重: {pretrained_model_path}\")\n            # 使用 strict=False 增加加载的灵活性\n            self.model.load_state_dict(torch.load(pretrained_model_path), strict=False)\n        elif pretrained_model_path:\n            print(f\"警告: 预训练模型路径未找到: {pretrained_model_path}\")\n\n    def forward(self, x):\n        return self.model(x)\n\n# 这个Dataset类与您的图片目录结构完全匹配\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, str(img_id) + '.jpeg')\n        img_path_png = os.path.join(self.img_dir, folder_name, str(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            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:\n            # 如果读取文件出错，同样递归加载下一个\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:\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# 3. 训练函数\n# ==============================================================================\ndef pretrain_on_2015(model, train_loader, device, num_epochs, lr, save_path):\n    print(\"\\n--- 开始在2015数据集上进行预训练 ---\")\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    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        scheduler.step()\n        \n        print(f\"Epoch {epoch + 1}/{num_epochs} | 训练损失: {train_loss:.4f} | 学习率: {optimizer.param_groups[0]['lr']:.6f}\")\n\n    torch.save(model.state_dict(), save_path)\n    print(f\"--- 预训练完成，模型已保存至: {save_path} ---\")\n\n# ==============================================================================\n# 4. 主执行逻辑\n# ==============================================================================\nif __name__ == \"__main__\":\n    # --- 步骤 1: 定义路径和参数 ---\n    print(\"--- 步骤 1: 初始化路径和参数 ---\")\n    MODEL_NAME = 'efficientnet_b6'\n    BATCH_SIZE = 4\n    NUM_WORKERS = 2\n    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n    print(f\"使用的设备: {device}\")\n\n    # <--- 修改: 定义您下载的b6模型权重路径\n    BASE_MODEL_PATH = '/kaggle/input/efficientnet-b6/pytorch_model_effi_b6.bin'\n\n    # --- 2015 预训练阶段参数 ---\n    PRETRAIN_IMAGE_SIZE = 256\n    PRETRAIN_EPOCHS = 20\n    PRETRAIN_LR = 3e-4\n    # 确保这个路径是包含 No_DR, Mild 等子文件夹的父目录\n    PRETRAIN_DATA_DIR = '/kaggle/input/new-data-preprocessed-images-256/new_data_preprocessed_images_256' \n    PRETRAIN_LABEL_FILE = '/kaggle/input/new-data-labels/trainLabels.csv'\n    PRETRAINED_2015_MODEL_PATH = os.path.join(os.getcwd(), f'pretrained_2015_{MODEL_NAME}_{PRETRAIN_IMAGE_SIZE}.pth')\n    \n    # --- 阶段一: 在2015数据集上进行预训练 ---\n    if not os.path.exists(PRETRAINED_2015_MODEL_PATH):\n        # 1. 加载2015标签文件\n        # <--- 修正: 直接使用 'image' 和 'level' 列，不再重命名\n        pretrain_df = pd.read_csv(PRETRAIN_LABEL_FILE)\n        \n        # 2. 准备数据集和加载器\n        # 这个Dataset会自动处理子文件夹和文件格式\n        pretrain_dataset = CustomImageDataset(\n            PRETRAIN_DATA_DIR, \n            pretrain_df, \n            transform=get_transforms('train', image_size=PRETRAIN_IMAGE_SIZE)\n        )\n        pretrain_loader = DataLoader(pretrain_dataset, batch_size=BATCH_SIZE, shuffle=True, num_workers=NUM_WORKERS, pin_memory=True)\n        \n        # 3. 初始化模型 (加载您指定的本地.bin文件)\n        pretrain_model = EfficientNetModel(MODEL_NAME, pretrained_model_path=BASE_MODEL_PATH).to(device)\n        \n        # 4. 执行预训练\n        pretrain_on_2015(pretrain_model, pretrain_loader, device, PRETRAIN_EPOCHS, PRETRAIN_LR, PRETRAINED_2015_MODEL_PATH)\n    else:\n        print(f\"--- 找到已存在的2015预训练模型，跳过预训练: {PRETRAINED_2015_MODEL_PATH} ---\")\n\n    print(\"\\n--- 脚本执行结束 ---\")\n    print(\"已成功执行预训练流程。如果您想继续进行微调，请取消后续代码的注释。\")","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}]}