import argparse import json import math import os import random from contextlib import nullcontext from pathlib import Path from typing import Dict, List, Optional, Tuple, Any import cv2 import numpy as np import pandas as pd import torch import torch.nn as nn import torch.nn.functional as F from PIL import Image from sklearn.metrics import ( accuracy_score, balanced_accuracy_score, classification_report, f1_score, precision_score, recall_score, ) from sklearn.model_selection import GroupKFold, StratifiedKFold from torch.utils.data import DataLoader, Dataset, WeightedRandomSampler from torchvision import models, transforms from torchvision.transforms import InterpolationMode from tqdm import tqdm try: from sklearn.model_selection import StratifiedGroupKFold HAS_STRATIFIED_GROUP_KFOLD = True except Exception: StratifiedGroupKFold = None HAS_STRATIFIED_GROUP_KFOLD = False try: import timm HAS_TIMM = True except ImportError: HAS_TIMM = False print("Warning: timm is not installed. timm-based models will be skipped.") PRIMARY_METRIC = "macro_f1" DEFAULT_INPUT_SIZE = 512 # Utilities def seed_everything(seed: int = 42, deterministic: bool = False) -> None: random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) torch.cuda.manual_seed_all(seed) os.environ["PYTHONHASHSEED"] = str(seed) if deterministic: torch.backends.cudnn.benchmark = False torch.backends.cudnn.deterministic = True try: torch.use_deterministic_algorithms(True, warn_only=True) except Exception: pass else: torch.backends.cudnn.benchmark = True torch.backends.cudnn.deterministic = False def ensure_dir(path: Path) -> None: path.mkdir(parents=True, exist_ok=True) def to_jsonable(obj: Any): if isinstance(obj, dict): return {k: to_jsonable(v) for k, v in obj.items()} if isinstance(obj, list): return [to_jsonable(v) for v in obj] if isinstance(obj, tuple): return [to_jsonable(v) for v in obj] if isinstance(obj, (np.integer, np.floating)): return obj.item() return obj def name_matches_keywords(name: str, keywords: List[str]) -> bool: if not name: return False for kw in keywords: plain_kw = kw.rstrip(".") if kw in name or name == plain_kw or name.startswith(plain_kw + "."): return True return False # Device device = torch.device("cuda" if torch.cuda.is_available() else "cpu") print(f"Using device: {device}") if device.type == "cuda": print(f"GPU: {torch.cuda.get_device_name(0)}") print(f"GPU memory: {torch.cuda.get_device_properties(0).total_memory / 1024 ** 3:.1f} GB") else: print("Warning: CUDA is not available. Training will be much slower on CPU.") # Model _VIT_KEYWORDS = [ "ViT", "Swin", "Transformer", "DeiT", "MaxViT", "CoAtNet", "EfficientFormer", "FastViT", "CaFormer", ] def _is_vit_family(model_name: str) -> bool: return any(kw.lower() in model_name.lower() for kw in _VIT_KEYWORDS) def _is_timm_model(model: nn.Module) -> bool: return hasattr(model, "get_classifier") and hasattr(model, "num_features") MODEL_INPUT_SIZES: Dict[str, int] = { "inception_v3": 299, } def get_model_input_size(model_name: str) -> int: return MODEL_INPUT_SIZES.get(model_name, DEFAULT_INPUT_SIZE) # Metrics / IO def compute_metrics( y_true: List[int], y_pred: List[int], num_classes: int, class_names: List[str], ) -> Tuple[Dict, Dict]: labels = list(range(num_classes)) report = classification_report( y_true, y_pred, labels=labels, target_names=class_names, output_dict=True, zero_division=0, ) metrics = { "accuracy": 100.0 * accuracy_score(y_true, y_pred), "balanced_accuracy": 100.0 * balanced_accuracy_score(y_true, y_pred), "macro_f1": 100.0 * f1_score(y_true, y_pred, labels=labels, average="macro", zero_division=0), "macro_precision": 100.0 * precision_score(y_true, y_pred, labels=labels, average="macro", zero_division=0), "macro_recall": 100.0 * recall_score(y_true, y_pred, labels=labels, average="macro", zero_division=0), "weighted_f1": 100.0 * f1_score(y_true, y_pred, labels=labels, average="weighted", zero_division=0), } return metrics, report def save_fold_results(results: Dict, save_dir: Path, tag: str = "best") -> None: ensure_dir(save_dir) report_df = pd.DataFrame(results["classification_report"]).transpose() with open(save_dir / f"test_report_{tag}.txt", "w", encoding="utf-8") as f: f.write(f"Primary Metric ({PRIMARY_METRIC}): {results['metrics'][PRIMARY_METRIC]:.4f}\n") f.write(f"Accuracy: {results['metrics']['accuracy']:.4f}\n") f.write(f"Balanced Accuracy: {results['metrics']['balanced_accuracy']:.4f}\n") f.write(f"Macro F1: {results['metrics']['macro_f1']:.4f}\n") f.write(f"Macro Recall: {results['metrics']['macro_recall']:.4f}\n") f.write(f"Macro Precision: {results['metrics']['macro_precision']:.4f}\n\n") f.write("Classification Report:\n") f.write(report_df.to_string()) pred_df = pd.DataFrame({ "patient": results["patients"], "image_name": results["image_names"], "True": results["targets"], "Predicted": results["predictions"], "path": results["image_path"], }) for c in range(results["num_classes"]): pred_df[f"prob_class{c}"] = [row[c] for row in results["probabilities"]] pred_df.to_csv(save_dir / f"predictions_{tag}.csv", index=False) payload = { "best_epoch": results["best_epoch"], "primary_metric": PRIMARY_METRIC, "metrics": results["metrics"], "per_class": [ results["classification_report"].get( f"class{i}", {"precision": 0, "recall": 0, "f1-score": 0} ) for i in range(results["num_classes"]) ], } with open(save_dir / f"{tag}_metrics.json", "w", encoding="utf-8") as f: json.dump(to_jsonable(payload), f, indent=2, ensure_ascii=False) def save_kfold_summary( model_name: str, fold_results: List[Dict], num_classes: int, save_dir: Path, ) -> Tuple[float, float]: ensure_dir(save_dir) metric_names = [ "accuracy", "balanced_accuracy", "macro_f1", "macro_recall", "macro_precision", "weighted_f1", ] summary = {} for name in metric_names: values = [r["metrics"][name] for r in fold_results] summary[name] = { "mean": float(np.mean(values)), "std": float(np.std(values)), } lines = [ "=" * 70, f"Model: {model_name}", "5-Fold Cross-Validation Summary", f"Primary Metric: {PRIMARY_METRIC}", "=" * 70, "", ] for i, r in enumerate(fold_results, 1): lines.append( f"Fold {i}: Macro-F1={r['metrics']['macro_f1']:.2f}% | " f"BA={r['metrics']['balanced_accuracy']:.2f}% | " f"Acc={r['metrics']['accuracy']:.2f}% | " f"BestEpoch={r['best_epoch']}" ) lines.append("") for name in metric_names: lines.append(f"{name}: {summary[name]['mean']:.2f}% +/- {summary[name]['std']:.2f}%") lines.append("") lines.append("Per-class metrics (mean +/- std)") lines.append(f"{'class':<10} {'precision':>18} {'recall':>18} {'f1-score':>18}") per_class_summary = {} for c in range(num_classes): ps = [r["per_class"][c]["precision"] for r in fold_results] rs = [r["per_class"][c]["recall"] for r in fold_results] fs = [r["per_class"][c]["f1-score"] for r in fold_results] per_class_summary[c] = { "precision_mean": float(np.mean(ps)), "precision_std": float(np.std(ps)), "recall_mean": float(np.mean(rs)), "recall_std": float(np.std(rs)), "f1_mean": float(np.mean(fs)), "f1_std": float(np.std(fs)), } lines.append( f"class{c:<5} " f"{np.mean(ps):.4f}+/-{np.std(ps):.4f}" f"{np.mean(rs):>18.4f}+/-{np.std(rs):.4f}" f"{np.mean(fs):>18.4f}+/-{np.std(fs):.4f}" ) text = "\n".join(lines) print(text) with open(save_dir / "kfold_summary.txt", "w", encoding="utf-8") as f: f.write(text) with open(save_dir / "kfold_summary.json", "w", encoding="utf-8") as f: json.dump( to_jsonable({ "model": model_name, "primary_metric": PRIMARY_METRIC, "summary": summary, "per_class": per_class_summary, }), f, indent=2, ensure_ascii=False, ) all_targets, all_predictions, all_paths = [], [], [] all_patients, all_image_names = [], [] all_probabilities = [] pooled_ready = all( "targets" in r and "predictions" in r and "image_path" in r and "probabilities" in r for r in fold_results ) if pooled_ready: for r in fold_results: all_targets.extend(r["targets"]) all_predictions.extend(r["predictions"]) all_paths.extend(r["image_path"]) all_patients.extend(r["patients"]) all_image_names.extend(r["image_names"]) all_probabilities.extend(r["probabilities"]) class_names = [f"class{i}" for i in range(num_classes)] pooled_metrics, pooled_report = compute_metrics( all_targets, all_predictions, num_classes, class_names, ) with open(save_dir / "oof_report.txt", "w", encoding="utf-8") as f: f.write("Pooled out-of-fold metrics\n") f.write(f"Primary Metric ({PRIMARY_METRIC}): {pooled_metrics[PRIMARY_METRIC]:.4f}\n") for k, v in pooled_metrics.items(): f.write(f"{k}: {v:.4f}\n") f.write("\nClassification Report:\n") f.write(pd.DataFrame(pooled_report).transpose().to_string()) oof_df = pd.DataFrame({ "patient": all_patients, "image_name": all_image_names, "True": all_targets, "Predicted": all_predictions, "path": all_paths, }) for c in range(num_classes): oof_df[f"prob_class{c}"] = [row[c] for row in all_probabilities] oof_df.to_csv(save_dir / "oof_predictions.csv", index=False) return summary[PRIMARY_METRIC]["mean"], summary[PRIMARY_METRIC]["std"] # 去掉黑边裁切，在 CLAHE + 绿色增强后增加眼底区域蒙版 class BlackBorderCrop: """Crop black borders and obvious invalid background around the fundus.""" def __init__(self, threshold: int = 10, margin_ratio: float = 0.02): self.threshold = threshold self.margin_ratio = margin_ratio def __call__(self, pil_img: Image.Image) -> Image.Image: img = np.array(pil_img.convert("RGB")) gray = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY) mask = gray > self.threshold if mask.sum() < 64: return pil_img.convert("RGB") ys, xs = np.where(mask) y1, y2 = ys.min(), ys.max() x1, x2 = xs.min(), xs.max() margin_y = int((y2 - y1 + 1) * self.margin_ratio) margin_x = int((x2 - x1 + 1) * self.margin_ratio) y1 = max(0, y1 - margin_y) y2 = min(img.shape[0], y2 + margin_y + 1) x1 = max(0, x1 - margin_x) x2 = min(img.shape[1], x2 + margin_x + 1) cropped = img[y1:y2, x1:x2] return Image.fromarray(cropped) class FundusCircularCrop: def __init__(self, threshold: int = 8, radius_pad_ratio: float = 0.03): self.threshold = threshold self.radius_pad_ratio = radius_pad_ratio def __call__(self, pil_img: Image.Image) -> Image.Image: img = np.array(pil_img.convert("RGB")) gray = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY) mask = (gray > self.threshold).astype(np.uint8) * 255 kernel = np.ones((5, 5), np.uint8) mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel) mask = cv2.medianBlur(mask, 5) contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) if not contours: return Image.fromarray(img) largest = max(contours, key=cv2.contourArea) (cx, cy), radius = cv2.minEnclosingCircle(largest) if radius < 10: return Image.fromarray(img) radius = radius * (1.0 + self.radius_pad_ratio) cx, cy, radius = float(cx), float(cy), float(radius) x1 = max(0, int(cx - radius)) y1 = max(0, int(cy - radius)) x2 = min(img.shape[1], int(cx + radius)) y2 = min(img.shape[0], int(cy + radius)) cropped = img[y1:y2, x1:x2] h, w = cropped.shape[:2] if h < 2 or w < 2: return Image.fromarray(img) local_cx = cx - x1 local_cy = cy - y1 rr = max(1, min(int(radius), min(h, w) // 2)) yy, xx = np.ogrid[:h, :w] circle_mask = ((xx - local_cx) ** 2 + (yy - local_cy) ** 2) <= (rr ** 2) out = np.zeros_like(cropped) out[circle_mask] = cropped[circle_mask] return Image.fromarray(out) class ResizeToSquare: def __init__(self, size: int): self.size = size def __call__(self, pil_img: Image.Image) -> Image.Image: return pil_img.resize((self.size, self.size), resample=Image.BILINEAR) class LightCLAHE: def __init__(self, clip_limit: float = 2.0, grid: Tuple[int, int] = (8, 8)): self.clahe = cv2.createCLAHE(clipLimit=clip_limit, tileGridSize=grid) def __call__(self, pil_img: Image.Image) -> Image.Image: img = np.array(pil_img.convert("RGB")) lab = cv2.cvtColor(img, cv2.COLOR_RGB2LAB) l, a, b = cv2.split(lab) l = self.clahe.apply(l) out = cv2.cvtColor(cv2.merge([l, a, b]), cv2.COLOR_LAB2RGB) return Image.fromarray(out) class GreenChannelEnhancement: def __init__( self, clip_limit: float = 2.5, grid: Tuple[int, int] = (8, 8), blend_alpha: float = 0.30, ): self.clahe = cv2.createCLAHE(clipLimit=clip_limit, tileGridSize=grid) self.blend_alpha = blend_alpha def __call__(self, pil_img: Image.Image) -> Image.Image: img = np.array(pil_img.convert("RGB")) r, g, b = cv2.split(img) g_eq = self.clahe.apply(g) g_new = cv2.addWeighted(g, 1.0 - self.blend_alpha, g_eq, self.blend_alpha, 0.0) out = cv2.merge([r, g_new, b]) return Image.fromarray(out) class FundusEyeMask: def __init__( self, threshold: int = 8, radius_pad_ratio: float = 0.03, morph_kernel: int = 7, blur_kernel: int = 5, ): self.threshold = threshold self.radius_pad_ratio = radius_pad_ratio self.morph_kernel = morph_kernel self.blur_kernel = blur_kernel def __call__(self, pil_img: Image.Image) -> Image.Image: img = np.array(pil_img.convert("RGB")) gray = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY) # Robust threshold against dark background after CLAHE + green enhancement _, mask = cv2.threshold(gray, self.threshold, 255, cv2.THRESH_BINARY) kernel = np.ones((self.morph_kernel, self.morph_kernel), np.uint8) mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel) mask = cv2.morphologyEx(mask, cv2.MORPH_OPEN, kernel) contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) if not contours: return Image.fromarray(img) largest = max(contours, key=cv2.contourArea) (cx, cy), radius = cv2.minEnclosingCircle(largest) if radius < 10: return Image.fromarray(img) radius = radius * (1.0 + self.radius_pad_ratio) yy, xx = np.ogrid[:img.shape[0], :img.shape[1]] circle_mask = (((xx - cx) ** 2 + (yy - cy) ** 2) <= (radius ** 2)).astype(np.uint8) * 255 if self.blur_kernel and self.blur_kernel > 1: k = self.blur_kernel if self.blur_kernel % 2 == 1 else self.blur_kernel + 1 circle_mask = cv2.GaussianBlur(circle_mask, (k, k), 0) circle_mask_f = (circle_mask.astype(np.float32) / 255.0)[..., None] out = (img.astype(np.float32) * circle_mask_f).clip(0, 255).astype(np.uint8) return Image.fromarray(out) _light_clahe = LightCLAHE() _green_enhance = GreenChannelEnhancement() _eye_mask = FundusEyeMask() _transform_cache: Dict[int, Tuple[transforms.Compose, transforms.Compose]] = {} def build_transforms(input_size: int = DEFAULT_INPUT_SIZE): """ 预处理流程: - 不再使用 BlackBorderCrop - 缩放到 input_size → CLAHE → 绿色通道增强 → 眼底区域蒙版 - 蒙版仅保留眼睛区域，屏蔽眼底边缘外的无关像素 训练增强: - 水平翻转 + 垂直翻转 - 小角度随机旋转 (±15°) + 轻微平移 + 尺度扰动 (0.85~1.15) - 适度 ColorJitter - 轻微高斯模糊 """ if input_size in _transform_cache: return _transform_cache[input_size] preprocess = [ ResizeToSquare(input_size), _light_clahe, _green_enhance, _eye_mask, ] train_tf = transforms.Compose( preprocess + [ transforms.RandomHorizontalFlip(p=0.5), transforms.RandomVerticalFlip(p=0.5), transforms.RandomAffine( degrees=15, translate=(0.05, 0.05), scale=(0.85, 1.15), interpolation=InterpolationMode.BILINEAR, fill=0, ), transforms.ColorJitter( brightness=0.20, contrast=0.20, saturation=0.10, hue=0.02, ), transforms.GaussianBlur(kernel_size=3, sigma=(0.1, 0.8)), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]), ] ) val_tf = transforms.Compose( preprocess + [ transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]), ] ) _transform_cache[input_size] = (train_tf, val_tf) return train_tf, val_tf # TTA (Test-Time Augmentation) # 增：4 路 TTA — 原图 / 水平翻转 / 垂直翻转 / 双向翻转 def predict_with_tta( model: nn.Module, inputs: torch.Tensor, amp_enabled: bool = False, ) -> torch.Tensor: amp_ctx = torch.cuda.amp.autocast if amp_enabled else nullcontext aug_variants = [ inputs, # 原图 inputs.flip(-1), # 水平翻转 inputs.flip(-2), # 垂直翻转 inputs.flip(-1).flip(-2), # 双向翻转 ] probs_list = [] for aug in aug_variants: with amp_ctx(): out = model(aug) logits = _extract_logits(out) probs_list.append(torch.softmax(logits, dim=1)) return torch.stack(probs_list, dim=0).mean(dim=0) # Dataset class ImageDataset(Dataset): def __init__(self, df: pd.DataFrame, transform=None): self.df = df.reset_index(drop=True).copy() self.transform = transform self.paths = self.df["path"].astype(str).tolist() self.labels = self.df["label"].astype(int).tolist() self.patients = self.df["patient"].astype(str).tolist() if "image_name" in self.df.columns: self.image_names = self.df["image_name"].astype(str).tolist() else: self.image_names = [Path(p).name for p in self.paths] def __len__(self) -> int: return len(self.paths) def __getitem__(self, idx: int): img_path = self.paths[idx] label = self.labels[idx] try: image = Image.open(img_path).convert("RGB") except Exception as exc: raise RuntimeError(f"Failed to open image: {img_path}") from exc if self.transform is not None: image = self.transform(image) meta = { "path": img_path, "patient": self.patients[idx], "image_name": self.image_names[idx], } return image, torch.tensor(label, dtype=torch.long), meta # Data loading / grouped splitting def validate_image_paths(df: pd.DataFrame, path_col: str = "path") -> pd.DataFrame: total = len(df) mask = df[path_col].apply(os.path.isfile) missing = total - int(mask.sum()) if missing > 0: print(f"Warning: {missing}/{total} paths do not exist and will be removed.") df = df.loc[mask].reset_index(drop=True) else: print(f"All {total} image paths are valid.") return df def load_and_prepare_data(excel_path: str, group_col: str = "patient") -> pd.DataFrame: df = pd.read_excel(excel_path, engine="openpyxl") required_cols = {"path", "label", group_col} missing_cols = required_cols - set(df.columns) if missing_cols: raise KeyError(f"Missing required columns in Excel: {sorted(missing_cols)}") df = df.copy() df[group_col] = df[group_col].astype(str).str.strip() if df[group_col].isin(["", "nan", "None"]).any(): bad_rows = int(df[group_col].isin(["", "nan", "None"]).sum()) raise ValueError(f"Found {bad_rows} rows with invalid patient/group identifiers in column '{group_col}'.") df["label"] = df["label"].replace({"AROP": 5}) df["label"] = pd.to_numeric(df["label"], errors="raise").astype(int) if df["label"].min() == 1: df["label"] = df["label"] - 1 # Merge old labels 4 and 5 into class 3 -> final 4-class setup df["label"] = df["label"].replace({4: 3, 5: 3}) df = validate_image_paths(df, path_col="path") if "patient" != group_col: df["patient"] = df[group_col].astype(str) unique_labels = sorted(df["label"].unique().tolist()) print(f"Dataset size: {len(df)} images") print(f"Unique patients: {df[group_col].nunique()}") print(f"Class distribution: {dict(df['label'].value_counts().sort_index())}") print(f"Observed labels: {unique_labels}") return df def _approximate_group_stratified_splits( df: pd.DataFrame, n_folds: int, random_seed: int, group_col: str, ): group_df = ( df.groupby(group_col)["label"] .agg(lambda x: x.value_counts().index[0]) .reset_index() ) if group_df[group_col].nunique() < n_folds: raise ValueError( f"Number of unique groups ({group_df[group_col].nunique()}) is smaller than n_folds={n_folds}." ) skf = StratifiedKFold(n_splits=n_folds, shuffle=True, random_state=random_seed) splits = [] group_ids = group_df[group_col].values group_labels = group_df["label"].values for group_train_idx, group_val_idx in skf.split(group_ids, group_labels): train_groups = set(group_ids[group_train_idx]) val_groups = set(group_ids[group_val_idx]) train_idx = df.index[df[group_col].isin(train_groups)].to_numpy() val_idx = df.index[df[group_col].isin(val_groups)].to_numpy() splits.append((train_idx, val_idx)) return splits def build_fold_splits( df: pd.DataFrame, n_folds: int, random_seed: int, group_col: str = "patient", ): groups = df[group_col].astype(str).values labels = df["label"].values if len(np.unique(groups)) < n_folds: raise ValueError( f"Unique groups in '{group_col}' = {len(np.unique(groups))}, which is smaller than n_folds={n_folds}." ) if HAS_STRATIFIED_GROUP_KFOLD: print( f"Using StratifiedGroupKFold with group_col='{group_col}', n_folds={n_folds}, seed={random_seed}." ) try: splitter = StratifiedGroupKFold( n_splits=n_folds, shuffle=True, random_state=random_seed, ) splits = list(splitter.split(df, y=labels, groups=groups)) except ValueError as exc: print(f"StratifiedGroupKFold failed: {exc}") print("Falling back to approximate grouped stratification using patient-majority labels.") splits = _approximate_group_stratified_splits(df, n_folds, random_seed, group_col) else: print("StratifiedGroupKFold is unavailable. Falling back to approximate grouped stratification.") splits = _approximate_group_stratified_splits(df, n_folds, random_seed, group_col) for fold_id, (train_idx, val_idx) in enumerate(splits, 1): train_groups = set(df.iloc[train_idx][group_col].astype(str).tolist()) val_groups = set(df.iloc[val_idx][group_col].astype(str).tolist()) overlap = train_groups & val_groups if overlap: raise RuntimeError( f"Data leakage detected in fold {fold_id}: {len(overlap)} overlapping groups." ) return splits def _compute_class_weights(train_df: pd.DataFrame, num_classes: int) -> torch.Tensor: counts = train_df["label"].value_counts().sort_index() total = len(train_df) weights = [total / (num_classes * counts.get(c, 1)) for c in range(num_classes)] return torch.tensor(weights, dtype=torch.float32, device=device) def _make_weighted_sampler(train_df: pd.DataFrame) -> WeightedRandomSampler: counts = train_df["label"].value_counts().to_dict() sample_weights = train_df["label"].map(lambda x: 1.0 / counts[x]).astype(float).values sample_weights = torch.as_tensor(sample_weights, dtype=torch.double) return WeightedRandomSampler(sample_weights, num_samples=len(sample_weights), replacement=True) def create_fold_loaders( train_df: pd.DataFrame, val_df: pd.DataFrame, input_size: int = DEFAULT_INPUT_SIZE, batch_size: int = 8, num_classes: int = 4, balance_mode: str = "loss", num_workers: int = 4, ): train_tf, val_tf = build_transforms(input_size) sampler = None class_weights = None if balance_mode == "sampler":+ sampler = _make_weighted_sampler(train_df) print("Training loader uses WeightedRandomSampler for class balancing.") elif balance_mode == "loss": class_weights = _compute_class_weights(train_df, num_classes) print(f"Training loss uses class weights: {class_weights.detach().cpu().numpy().tolist()}") else: print("No imbalance correction is used.") drop_last = (batch_size > 1) and (len(train_df) % batch_size == 1) if drop_last: print( f"Training loader will drop the last singleton batch " f"(train_size={len(train_df)}, batch_size={batch_size}) to avoid BatchNorm issues." ) pin_memory = device.type == "cuda" train_loader = DataLoader( ImageDataset(train_df, train_tf), batch_size=batch_size, shuffle=(sampler is None), sampler=sampler, num_workers=num_workers, pin_memory=pin_memory, drop_last=drop_last, persistent_workers=(num_workers > 0), ) val_loader = DataLoader( ImageDataset(val_df, val_tf), batch_size=batch_size, shuffle=False, num_workers=num_workers, pin_memory=pin_memory, persistent_workers=(num_workers > 0), ) return train_loader, val_loader, class_weights # ViT positional embedding interpolation def patch_vit_for_large_input( model: nn.Module, model_name: str, input_size: int, ) -> nn.Module: if "ViT" not in model_name: return model if not (hasattr(model, "encoder") and hasattr(model.encoder, "pos_embedding")): print(f"Warning: cannot find pos_embedding for {model_name}, skip interpolation.") return model patch_size = model.patch_size expected_patches = (input_size // patch_size) ** 2 pos_embed = model.encoder.pos_embedding current_patches = pos_embed.shape[1] - 1 if current_patches == expected_patches: print(f"[ViT] pos_embedding already matches input_size={input_size}, no interpolation needed.") return model print( f"[ViT] Interpolating pos_embedding: {current_patches} -> {expected_patches} patches " f"for input_size={input_size}." ) cls_token = pos_embed[:, :1, :] patch_tokens = pos_embed[:, 1:, :] dim = patch_tokens.shape[-1] h_old = w_old = int(math.sqrt(current_patches)) h_new = w_new = int(math.sqrt(expected_patches)) patch_tokens = ( patch_tokens .reshape(1, h_old, w_old, dim) .permute(0, 3, 1, 2) .float() ) patch_tokens = F.interpolate( patch_tokens, size=(h_new, w_new), mode="bicubic", align_corners=False, ) patch_tokens = patch_tokens.permute(0, 2, 3, 1).reshape(1, expected_patches, dim) model.encoder.pos_embedding = nn.Parameter(torch.cat([cls_token, patch_tokens], dim=1)) if hasattr(model, "image_size"): model.image_size = input_size return model # Classifier replacement def _find_last_linear(module: nn.Module): if isinstance(module, nn.Linear): return module if isinstance(module, nn.Sequential): for child in reversed(list(module.children())): result = _find_last_linear(child) if result is not None: return result if hasattr(module, "head") and isinstance(module.head, (nn.Linear, nn.Sequential)): return _find_last_linear(module.head) return None def _verify_classifier(model: nn.Module, model_name: str, expected_classes: int) -> None: for attr_name in ["fc", "head", "classifier", "heads"]: if not hasattr(model, attr_name): continue layer = getattr(model, attr_name) last_linear = _find_last_linear(layer) if last_linear is not None: if last_linear.out_features != expected_classes: raise RuntimeError( f"Classifier replacement failed for {model_name}: " f"out_features={last_linear.out_features}, expected={expected_classes}" ) print(f"Verified {model_name}: classifier -> {expected_classes} classes (in={last_linear.in_features})") return print(f"Warning: failed to automatically verify classifier for {model_name}") def replace_classifier( model_name: str, model: nn.Module, num_classes: int, dropout: float = 0.3, ) -> nn.Module: if _is_timm_model(model): in_feat = model.num_features orig_classifier = model.get_classifier() print(f"[timm] {model_name}: original classifier={type(orig_classifier).__name__}, num_features={in_feat}") model.reset_classifier(num_classes) new_fc = model.get_classifier() wrapped = False if isinstance(new_fc, nn.Linear): for parent_attr, child_attr in [ ("head", "fc"), ("head", "head"), (None, "head"), (None, "classifier"), (None, "fc"), ]: try: parent = getattr(model, parent_attr) if parent_attr else model child = getattr(parent, child_attr) if child is new_fc: setattr( parent, child_attr, nn.Sequential(nn.Dropout(dropout), nn.Linear(in_feat, num_classes)), ) wrapped = True break except AttributeError: continue if not wrapped: print(f"[timm] {model_name}: reset_classifier({num_classes}) applied (no Dropout wrapper).") _verify_classifier(model, model_name, num_classes) return model n = model_name if "VGG" in n: in_feat = model.classifier[6].in_features model.classifier[6] = nn.Sequential(nn.Dropout(dropout), nn.Linear(in_feat, num_classes)) elif n == "inception_v3": aux_in = model.AuxLogits.fc.in_features model.AuxLogits.fc = nn.Sequential(nn.Dropout(dropout), nn.Linear(aux_in, num_classes)) fc_in = model.fc.in_features model.fc = nn.Sequential(nn.Dropout(dropout), nn.Linear(fc_in, num_classes)) elif "GoogLeNet" in n: in_feat = model.fc.in_features model.fc = nn.Sequential(nn.Dropout(dropout), nn.Linear(in_feat, num_classes)) if hasattr(model, "aux1") and model.aux1 is not None and hasattr(model.aux1, "fc2"): aux1_in = model.aux1.fc2.in_features model.aux1.fc2 = nn.Sequential(nn.Dropout(dropout), nn.Linear(aux1_in, num_classes)) if hasattr(model, "aux2") and model.aux2 is not None and hasattr(model.aux2, "fc2"): aux2_in = model.aux2.fc2.in_features model.aux2.fc2 = nn.Sequential(nn.Dropout(dropout), nn.Linear(aux2_in, num_classes)) elif "ResNe" in n: in_feat = model.fc.in_features model.fc = nn.Sequential(nn.Dropout(dropout), nn.Linear(in_feat, num_classes)) elif "DenseNet" in n: in_feat = model.classifier.in_features model.classifier = nn.Sequential(nn.Dropout(dropout), nn.Linear(in_feat, num_classes)) elif "MobileNet" in n: in_feat = model.classifier[-1].in_features model.classifier[-1] = nn.Sequential(nn.Dropout(dropout), nn.Linear(in_feat, num_classes)) elif "MnasNet" in n: in_feat = model.classifier[-1].in_features model.classifier[-1] = nn.Sequential(nn.Dropout(dropout), nn.Linear(in_feat, num_classes)) elif "EfficientNet" in n: in_feat = model.classifier[-1].in_features model.classifier[-1] = nn.Sequential(nn.Dropout(dropout), nn.Linear(in_feat, num_classes)) elif "ConvNeXt" in n: in_feat = model.classifier[-1].in_features model.classifier[-1] = nn.Sequential(nn.Dropout(dropout), nn.Linear(in_feat, num_classes)) elif "RegNet" in n or "ShuffleNet" in n: in_feat = model.fc.in_features model.fc = nn.Sequential(nn.Dropout(dropout), nn.Linear(in_feat, num_classes)) elif "ViT" in n: if hasattr(model, "heads") and hasattr(model.heads, "head"): in_feat = model.heads.head.in_features model.heads = nn.Sequential(nn.Dropout(dropout), nn.Linear(in_feat, num_classes)) elif hasattr(model, "head") and isinstance(model.head, nn.Linear): in_feat = model.head.in_features model.head = nn.Sequential(nn.Dropout(dropout), nn.Linear(in_feat, num_classes)) else: raise ValueError(f"Cannot find classifier head for {n}") elif "Swin" in n: if hasattr(model, "head") and isinstance(model.head, nn.Linear): in_feat = model.head.in_features model.head = nn.Sequential(nn.Dropout(dropout), nn.Linear(in_feat, num_classes)) elif hasattr(model, "heads") and hasattr(model.heads, "head"): in_feat = model.heads.head.in_features model.heads = nn.Sequential(nn.Dropout(dropout), nn.Linear(in_feat, num_classes)) else: raise ValueError(f"Cannot find classifier head for {n}") elif _is_vit_family(n): replaced = False for attr in ["heads.head", "head", "classifier"]: parts = attr.split(".") obj = model try: for p in parts: obj = getattr(obj, p) if isinstance(obj, nn.Linear): in_feat = obj.in_features parent = model for p in parts[:-1]: parent = getattr(parent, p) setattr(parent, parts[-1], nn.Sequential(nn.Dropout(dropout), nn.Linear(in_feat, num_classes))) replaced = True break except AttributeError: continue if not replaced: raise ValueError(f"Cannot find classifier head for {n}") else: replaced = False for attr_name in ["fc", "head", "classifier"]: if not hasattr(model, attr_name): continue layer = getattr(model, attr_name) if isinstance(layer, nn.Linear): in_feat = layer.in_features setattr(model, attr_name, nn.Sequential(nn.Dropout(dropout), nn.Linear(in_feat, num_classes))) replaced = True break if isinstance(layer, nn.Sequential) and len(layer) > 0 and isinstance(layer[-1], nn.Linear): in_feat = layer[-1].in_features layer[-1] = nn.Sequential(nn.Dropout(dropout), nn.Linear(in_feat, num_classes)) replaced = True break if not replaced: raise ValueError(f"Cannot automatically replace classifier for {n}") _verify_classifier(model, model_name, num_classes) return model # Optimizer groups / freezing def _get_head_keywords(model_name: str) -> List[str]: n = model_name if "VGG" in n: return ["classifier.6"] if n == "inception_v3": return ["fc.", "AuxLogits.fc"] if "GoogLeNet" in n: return ["fc.", "aux1.fc2", "aux2.fc2"] if "ResNe" in n: return ["fc."] if "DenseNet" in n: return ["classifier."] if "MobileNet" in n: return ["classifier.3", "classifier.2", "classifier."] if "MnasNet" in n: return ["classifier.1", "classifier."] if "EfficientNet" in n or "ConvNeXt" in n: return ["classifier.", "head.fc"] if "RegNet" in n or "ShuffleNet" in n: return ["fc."] if "ViT" in n or _is_vit_family(n): return ["heads.", "head.", "classifier."] return ["fc.", "classifier.", "head.", "heads."] def get_parameter_groups( model_name: str, model: nn.Module, backbone_lr: float = 3e-5, head_lr: float = 1e-3, ): head_kw = _get_head_keywords(model_name) head_p, back_p = [], [] for name, param in model.named_parameters(): if name_matches_keywords(name, head_kw): head_p.append(param) else: back_p.append(param) if not head_p: print(f"Warning: no head parameters matched for {model_name}; all params use head_lr.") return [{"params": list(model.parameters()), "lr": head_lr}] print( f"Parameter groups | backbone: {sum(p.numel() for p in back_p):,} (lr={backbone_lr}) | " f"head: {sum(p.numel() for p in head_p):,} (lr={head_lr})" ) return [{"params": back_p, "lr": backbone_lr}, {"params": head_p, "lr": head_lr}] def set_backbone_trainable(model_name: str, model: nn.Module, train_backbone: bool) -> None: head_kw = _get_head_keywords(model_name) for name, param in model.named_parameters(): is_head = name_matches_keywords(name, head_kw) param.requires_grad = train_backbone or is_head def set_frozen_backbone_bn_eval(model_name: str, model: nn.Module) -> None: head_kw = _get_head_keywords(model_name) bn_types = (nn.BatchNorm1d, nn.BatchNorm2d, nn.BatchNorm3d, nn.SyncBatchNorm) for name, module in model.named_modules(): if isinstance(module, bn_types) and not name_matches_keywords(name, head_kw): module.eval() for param in module.parameters(): param.requires_grad = False def configure_small_batch_behavior(model_name: str, model: nn.Module, batch_size: int) -> nn.Module: if batch_size >= 2: return model if model_name == "inception_v3": print("batch_size=1 detected: disabling Inception auxiliary classifier.") if hasattr(model, "aux_logits"): model.aux_logits = False if hasattr(model, "AuxLogits"): model.AuxLogits = None elif "GoogLeNet" in model_name: print("batch_size=1 detected: disabling GoogLeNet auxiliary classifiers.") if hasattr(model, "aux_logits"): model.aux_logits = False if hasattr(model, "aux1"): model.aux1 = None if hasattr(model, "aux2"): model.aux2 = None return model # Forward helpers def _extract_logits(output): if torch.is_tensor(output): return output if hasattr(output, "logits") and torch.is_tensor(output.logits): return output.logits if isinstance(output, (tuple, list)) and len(output) > 0 and torch.is_tensor(output[0]): return output[0] raise TypeError("Unable to extract logits from model output.") def _extract_aux_outputs(output): aux_outputs = [] if isinstance(output, (tuple, list)): aux_outputs.extend([o for o in output[1:] if torch.is_tensor(o)]) else: for attr in ["aux_logits", "aux_logits2", "aux_logits1"]: if hasattr(output, attr): aux = getattr(output, attr) if torch.is_tensor(aux): aux_outputs.append(aux) return aux_outputs def forward_with_loss( model: nn.Module, inputs: torch.Tensor, labels: torch.Tensor, criterion, aux_weight: float = 0.3, ): output = model(inputs) logits = _extract_logits(output) aux_outputs = _extract_aux_outputs(output) loss = criterion(logits, labels) if model.training and aux_outputs: for aux in aux_outputs: loss = loss + aux_weight * criterion(aux, labels) return logits, loss # Losses class FocalLoss(nn.Module): def __init__(self, alpha: Optional[torch.Tensor] = None, gamma: float = 2.0, reduction: str = "mean"): super().__init__() self.alpha = alpha self.gamma = gamma self.reduction = reduction def forward(self, inputs: torch.Tensor, targets: torch.Tensor) -> torch.Tensor: log_probs = F.log_softmax(inputs, dim=1) log_pt = log_probs.gather(1, targets.unsqueeze(1)).squeeze(1) pt = log_pt.exp() loss = -((1.0 - pt) ** self.gamma) * log_pt if self.alpha is not None: alpha_t = self.alpha.to(inputs.device)[targets] loss = alpha_t * loss if self.reduction == "mean": return loss.mean() if self.reduction == "sum": return loss.sum() return loss def build_criterion( loss_type: str, class_weights: Optional[torch.Tensor] = None, focal_gamma: float = 2.0, label_smoothing: float = 0.0, ): loss_type = loss_type.lower() if loss_type == "focal": return FocalLoss(alpha=class_weights, gamma=focal_gamma, reduction="mean") if loss_type == "weighted_ce": return nn.CrossEntropyLoss(weight=class_weights, label_smoothing=label_smoothing) if loss_type == "ce": return nn.CrossEntropyLoss(label_smoothing=label_smoothing) raise ValueError(f"Unsupported loss_type: {loss_type}") # Training # - epochs 90 # - freeze_backbone_epochs # - warmup_ep = freeze_backbone_epochs # - 验证循环使用 TTA def train_one_fold( model_name: str, model: nn.Module, train_loader, val_loader, epochs: int = 90, num_classes: int = 4, backbone_lr: float = 3e-5, head_lr: float = 1e-3, class_weights: Optional[torch.Tensor] = None, fold_id: int = 1, save_dir: Optional[Path] = None, freeze_backbone_epochs: int = 8, max_grad_norm: float = 1.0, primary_metric: str = PRIMARY_METRIC, loss_type: str = "weighted_ce", focal_gamma: float = 2.0, label_smoothing: float = 0.0, use_tta: bool = True, ): if save_dir is None: save_dir = Path(model_name) else: save_dir = Path(save_dir) ensure_dir(save_dir) criterion = build_criterion( loss_type=loss_type, class_weights=class_weights, focal_gamma=focal_gamma, label_smoothing=label_smoothing, ) print( f"Fold {fold_id}: Using loss_type='{loss_type}'" f"{' with class weights' if class_weights is not None else ''}." ) print( f"Fold {fold_id}: backbone_lr={backbone_lr}, head_lr={head_lr}, " f"freeze_backbone_epochs={freeze_backbone_epochs}, " f"epochs={epochs}, use_tta={use_tta}." ) param_groups = get_parameter_groups(model_name, model, backbone_lr, head_lr) optimizer = torch.optim.AdamW(param_groups, betas=(0.9, 0.999), weight_decay=5e-4) warmup_ep = freeze_backbone_epochs sched_main = torch.optim.lr_scheduler.CosineAnnealingLR( optimizer, T_max=max(1, epochs - warmup_ep), eta_min=1e-7, ) sched_warm = torch.optim.lr_scheduler.LinearLR( optimizer, start_factor=0.1, end_factor=1.0, total_iters=warmup_ep, ) scheduler = torch.optim.lr_scheduler.SequentialLR( optimizer, schedulers=[sched_warm, sched_main], milestones=[warmup_ep], ) amp_enabled = device.type == "cuda" scaler = torch.cuda.amp.GradScaler(enabled=amp_enabled) class_names = [f"class{i}" for i in range(num_classes)] best_monitor = -float("inf") best_results = None was_backbone_trainable = None start_epoch = 0 ckpt_path = save_dir / f"fold{fold_id}_checkpoint.pth" if ckpt_path.is_file(): print(f"Fold {fold_id}: found epoch-level checkpoint, attempting to resume...") try: ckpt = torch.load(ckpt_path, map_location=device, weights_only=False) model.load_state_dict(ckpt["model_state_dict"]) optimizer.load_state_dict(ckpt["optimizer_state_dict"]) scheduler.load_state_dict(ckpt["scheduler_state_dict"]) scaler.load_state_dict(ckpt["scaler_state_dict"]) start_epoch = ckpt["epoch"] + 1 best_monitor = ckpt["best_monitor"] best_results = ckpt.get("best_results", None) print( f"Fold {fold_id}: resumed from epoch {start_epoch}/{epochs} " f"(best {primary_metric}={best_monitor:.2f})." ) except Exception as exc: print(f"Fold {fold_id}: failed to load checkpoint ({exc}), training from scratch.") start_epoch = 0 best_monitor = -float("inf") best_results = None for epoch in tqdm( range(start_epoch, epochs), desc=f"Fold {fold_id}", leave=False, initial=start_epoch, total=epochs, ): train_backbone = epoch >= freeze_backbone_epochs if was_backbone_trainable is None or was_backbone_trainable != train_backbone: set_backbone_trainable(model_name, model, train_backbone=train_backbone) stage = "unfrozen" if train_backbone else "frozen" print(f"Fold {fold_id}: backbone is now {stage} (epoch {epoch + 1}).") was_backbone_trainable = train_backbone model.train() if not train_backbone: set_frozen_backbone_bn_eval(model_name, model) run_loss = 0.0 for inputs, labels, _meta in train_loader: inputs = inputs.to(device, non_blocking=(device.type == "cuda")) labels = labels.to(device, non_blocking=(device.type == "cuda")) optimizer.zero_grad(set_to_none=True) amp_ctx = torch.cuda.amp.autocast if amp_enabled else nullcontext with amp_ctx(): logits, loss = forward_with_loss(model, inputs, labels, criterion, aux_weight=0.3) scaler.scale(loss).backward() if max_grad_norm is not None and max_grad_norm > 0: scaler.unscale_(optimizer) torch.nn.utils.clip_grad_norm_(model.parameters(), max_grad_norm) scaler.step(optimizer) scaler.update() run_loss += loss.item() * inputs.size(0) scheduler.step() ep_loss = run_loss / len(train_loader.dataset) # ---- 验证阶段：可选 TTA ---- model.eval() all_t, all_p, all_paths, all_probs = [], [], [], [] all_patients, all_image_names = [], [] with torch.no_grad(): for inputs, labels, meta in val_loader: inputs = inputs.to(device, non_blocking=(device.type == "cuda")) labels = labels.to(device, non_blocking=(device.type == "cuda")) if use_tta: # V7: 使用 TTA 推断 probs = predict_with_tta(model, inputs, amp_enabled=amp_enabled) else: amp_ctx = torch.cuda.amp.autocast if amp_enabled else nullcontext with amp_ctx(): output = model(inputs) logits = _extract_logits(output) probs = torch.softmax(logits, dim=1) pred = probs.argmax(dim=1) all_t.extend(labels.cpu().numpy().tolist()) all_p.extend(pred.cpu().numpy().tolist()) all_probs.extend(probs.cpu().numpy().tolist()) all_paths.extend(list(meta["path"])) all_patients.extend(list(meta["patient"])) all_image_names.extend(list(meta["image_name"])) metrics, report = compute_metrics(all_t, all_p, num_classes, class_names) monitor = metrics[primary_metric] if (epoch + 1) % 5 == 0 or epoch == epochs - 1 or epoch == start_epoch: print( f"F{fold_id} E{epoch + 1}/{epochs} " f"Loss={ep_loss:.4f} " f"Macro-F1={metrics['macro_f1']:.2f}% " f"BA={metrics['balanced_accuracy']:.2f}% " f"Acc={metrics['accuracy']:.2f}%" f"{' [TTA]' if use_tta else ''}" ) improved = (monitor > best_monitor) or ( np.isclose(monitor, best_monitor) and best_results is not None and metrics["balanced_accuracy"] > best_results["metrics"]["balanced_accuracy"] ) if improved: best_monitor = monitor best_results = { "best_epoch": epoch + 1, "metrics": metrics, "classification_report": report, "predictions": all_p, "targets": all_t, "image_path": all_paths, "patients": all_patients, "image_names": all_image_names, "probabilities": all_probs, "num_classes": num_classes, "per_class": [ report.get(f"class{i}", {"precision": 0, "recall": 0, "f1-score": 0}) for i in range(num_classes) ], } save_fold_results(best_results, save_dir, tag=f"fold{fold_id}_best") torch.save(model.state_dict(), save_dir / f"fold{fold_id}_best.pth") torch.save({ "epoch": epoch, "model_state_dict": model.state_dict(), "optimizer_state_dict": optimizer.state_dict(), "scheduler_state_dict": scheduler.state_dict(), "scaler_state_dict": scaler.state_dict(), "best_monitor": best_monitor, "best_results": best_results, }, ckpt_path) if ckpt_path.is_file(): ckpt_path.unlink() print(f"Fold {fold_id}: removed epoch-level checkpoint (training complete).") if best_results is None: raise RuntimeError(f"Fold {fold_id}: no valid result was produced.") return best_results def build_model_registry(): reg = {} reg["DenseNet161"] = lambda: models.densenet161(weights=models.DenseNet161_Weights.DEFAULT) reg["ConvNeXt_Tiny"] = lambda: models.convnext_tiny(weights=models.ConvNeXt_Tiny_Weights.DEFAULT) reg["ViT_B_16"] = lambda: models.vit_b_16(weights=models.ViT_B_16_Weights.DEFAULT) if HAS_TIMM: reg["SwinV2_T"] = lambda: timm.create_model( "swinv2_tiny_window8_256", pretrained=True, img_size=512, ) reg["DeiT3_S"] = lambda: timm.create_model( "deit3_small_patch16_224", pretrained=True, img_size=512, ) else: print("Skipping timm models because timm is not installed.") return reg def parse_args(): parser = argparse.ArgumentParser(description="ROP benchmark training with patient-grouped 5-fold CV (v7).") boolean_action = getattr(argparse, "BooleanOptionalAction", None) parser.add_argument( "--excel_path", type=str, default="/media/fang/9fc99a7b-15d6-4e22-ab05-fe46e6058c39/felicia/Downloads/医生审核之后第一版12-17/部分公开数据集/公开数据集训练表_调整数据1.xlsx", help="Path to Excel with at least columns: patient, path, label.", ) parser.add_argument("--group_col", type=str, default="patient", help="Grouping column for leakage-free split.") parser.add_argument("--num_classes", type=int, default=4) # V7: epochs 90 parser.add_argument("--epochs", type=int, default=90) parser.add_argument("--n_folds", type=int, default=5) parser.add_argument("--batch_size", type=int, default=8) # V7: backbone_lr 提升至 3e-5 parser.add_argument("--backbone_lr", type=float, default=3e-5) # V7: head_lr 提升至 1e-3 parser.add_argument("--head_lr", type=float, default=1e-3) parser.add_argument("--random_seed", type=int, default=42) parser.add_argument("--num_workers", type=int, default=min(8, os.cpu_count() or 2)) parser.add_argument( "--balance_mode", type=str, default="loss", choices=["none", "loss", "sampler"], help="Imbalance handling. 'loss' computes class weights; 'sampler' uses WeightedRandomSampler.", ) parser.add_argument( "--loss_type", type=str, default="weighted_ce", choices=["weighted_ce", "focal", "ce"], help="weighted_ce is the recommended default.", ) parser.add_argument("--focal_gamma", type=float, default=2.0) parser.add_argument("--label_smoothing", type=float, default=0.0) # V7: freeze_backbone_epochs 提升至 8 parser.add_argument("--freeze_backbone_epochs", type=int, default=8) parser.add_argument("--max_grad_norm", type=float, default=1.0) parser.add_argument("--output_root", type=str, default="runs_rop_V7_old") if boolean_action is not None: parser.add_argument("--use_tta", action=boolean_action, default=True, help="Enable 4-way TTA (flip) during validation.") parser.add_argument("--deterministic", action=boolean_action, default=False) else: parser.add_argument("--use_tta", action="store_true", default=True, help="Enable 4-way TTA (flip) during validation.") parser.add_argument("--no_tta", dest="use_tta", action="store_false") parser.add_argument("--deterministic", action="store_true", default=False) parser.add_argument( "--models", nargs="*", default=None, help="Optional subset of model names to train (e.g. --models DenseNet161 ViT_B_16).", ) return parser.parse_args() def main(): args = parse_args() seed_everything(args.random_seed, deterministic=args.deterministic) print("\nLoading data...") df = load_and_prepare_data(args.excel_path, group_col=args.group_col) observed_num_classes = int(df["label"].nunique()) if observed_num_classes != args.num_classes: raise ValueError( f"num_classes mismatch: args.num_classes={args.num_classes}, " f"but observed labels in Excel imply {observed_num_classes} classes after remapping." ) fold_splits = build_fold_splits( df=df, n_folds=args.n_folds, random_seed=args.random_seed, group_col=args.group_col, ) model_registry = build_model_registry() if args.models: selected = {k: v for k, v in model_registry.items() if k in set(args.models)} missing = [m for m in args.models if m not in model_registry] if missing: print(f"Warning: these models were not found and will be ignored: {missing}") model_registry = selected print(f"\nTotal models to train: {len(model_registry)}") for i, name in enumerate(model_registry, 1): print(f"{i:2d}. {name}") output_root = Path(args.output_root) ensure_dir(output_root) global_results = {} for model_idx, (model_name, model_fn) in enumerate(model_registry.items(), 1): print("\n" + "=" * 70) print(f"[{model_idx}/{len(model_registry)}] Model: {model_name}") print("=" * 70) model_dir = output_root / model_name ensure_dir(model_dir) summary_path = model_dir / "kfold_summary.json" if summary_path.is_file(): try: with open(summary_path, "r", encoding="utf-8") as f: old = json.load(f) old_summary = old.get("summary", {}) if old_summary: mean_primary = old_summary[PRIMARY_METRIC]["mean"] std_primary = old_summary[PRIMARY_METRIC]["std"] print( f"[Skip] Found existing {args.n_folds}-fold summary: " f"{PRIMARY_METRIC}={mean_primary:.2f}% +/- {std_primary:.2f}%" ) global_results[model_name] = (mean_primary, std_primary) continue except Exception: pass input_size = get_model_input_size(model_name) print(f"Input size: {input_size}x{input_size}") fold_results = [] for fold_idx in range(args.n_folds): fold_id = fold_idx + 1 print(f"\n-- Fold {fold_id}/{args.n_folds} --") metrics_json = model_dir / f"fold{fold_id}_best_metrics.json" weight_path = model_dir / f"fold{fold_id}_best.pth" if metrics_json.is_file() and weight_path.is_file(): try: with open(metrics_json, "r", encoding="utf-8") as f: cached = json.load(f) fold_results.append({ "best_epoch": cached["best_epoch"], "metrics": cached["metrics"], "per_class": cached["per_class"], }) print( f"Fold {fold_id}: cached result found " f"(Macro-F1={cached['metrics']['macro_f1']:.2f}%, " f"BA={cached['metrics']['balanced_accuracy']:.2f}%), skipped." ) continue except Exception: pass train_idx, val_idx = fold_splits[fold_idx] train_df = df.iloc[train_idx].reset_index(drop=True) val_df = df.iloc[val_idx].reset_index(drop=True) train_patients = set(train_df[args.group_col].astype(str).tolist()) val_patients = set(val_df[args.group_col].astype(str).tolist()) overlap = train_patients & val_patients if overlap: raise RuntimeError( f"Leakage detected in fold {fold_id}: {len(overlap)} overlapping patients/groups." ) print(f"Train: {len(train_df)} | Validation: {len(val_df)}") print( f"Train patients: {train_df[args.group_col].nunique()} | " f"Validation patients: {val_df[args.group_col].nunique()}" ) print( f"Train class dist: {dict(train_df['label'].value_counts().sort_index())} | " f"Val class dist: {dict(val_df['label'].value_counts().sort_index())}" ) train_loader, val_loader, class_weights = create_fold_loaders( train_df=train_df, val_df=val_df, input_size=input_size, batch_size=args.batch_size, num_classes=args.num_classes, balance_mode=args.balance_mode, num_workers=args.num_workers, ) try: model = model_fn() except Exception as exc: print(f"Model creation failed for {model_name}: {exc}") break model = replace_classifier(model_name, model, args.num_classes) model = patch_vit_for_large_input(model, model_name, input_size) model = configure_small_batch_behavior(model_name, model, args.batch_size) model = model.to(device) dummy = torch.randn(1, 3, input_size, input_size, device=device) model.eval() with torch.no_grad(): out = model(dummy) out = _extract_logits(out) out_dim = out.shape[-1] if out_dim != args.num_classes: raise RuntimeError( f"Fatal: classifier replacement failed for {model_name}. " f"Output dim={out_dim}, expected={args.num_classes}." ) print(f"Forward sanity check passed: output dim={out_dim}") del dummy, out if device.type == "cuda": torch.cuda.empty_cache() result = train_one_fold( model_name=model_name, model=model, train_loader=train_loader, val_loader=val_loader, epochs=args.epochs, num_classes=args.num_classes, backbone_lr=args.backbone_lr, head_lr=args.head_lr, class_weights=class_weights, fold_id=fold_id, save_dir=model_dir, freeze_backbone_epochs=args.freeze_backbone_epochs, max_grad_norm=args.max_grad_norm, primary_metric=PRIMARY_METRIC, loss_type=args.loss_type, focal_gamma=args.focal_gamma, label_smoothing=args.label_smoothing, use_tta=args.use_tta, ) fold_results.append(result) del model if device.type == "cuda": torch.cuda.empty_cache() if len(fold_results) == args.n_folds: mean_primary, std_primary = save_kfold_summary( model_name, fold_results, args.num_classes, model_dir, ) global_results[model_name] = (mean_primary, std_primary) else: print(f"Warning: {model_name} completed only {len(fold_results)}/{args.n_folds} folds.") print("\n" + "=" * 70) print(f"Global leaderboard ({args.n_folds}-Fold CV)") print(f"Sorted by: {PRIMARY_METRIC}") print("=" * 70) sorted_results = sorted(global_results.items(), key=lambda x: x[1][0], reverse=True) print(f"{'Rank':<6} {'Model':<25} {PRIMARY_METRIC:>12} {'Std':>10}") print("-" * 62) for rank, (name, (mean_primary, std_primary)) in enumerate(sorted_results, 1): print(f"{rank:<6} {name:<25} {mean_primary:>11.2f}% {std_primary:>9.2f}%") leaderboard_path = output_root / f"global_leaderboard_{PRIMARY_METRIC}.csv" pd.DataFrame([ { "rank": idx + 1, "model": name, f"mean_{PRIMARY_METRIC}": mean_primary, f"std_{PRIMARY_METRIC}": std_primary, } for idx, (name, (mean_primary, std_primary)) in enumerate(sorted_results) ]).to_csv(leaderboard_path, index=False) print(f"\nLeaderboard saved to: {leaderboard_path}") if __name__ == "__main__": main()