新增 ResNet34 ImageNet 预训练 + 10% 数据微调模型 (baseline)

.gitignore

@@ -16,6 +16,7 @@
 !/baseline/__init__.py
 !/baseline/VGG_KNN.py
 !/baseline/compare_models.py
+!/baseline/ResNet34_Pretrained_10pct.py
 !/training_log.csv
 !/confusion_matrix.png
 !/roc_curve.png

baseline/ResNet34_Pretrained_10pct.py

@@ -0,0 +1,278 @@
+"""
+baseline/ResNet34_Pretrained_10pct.py
+ResNet-34 ImageNet 预训练权重 + 10% 训练集微调
+可独立运行训练，也可被 compare_models.py 导入
+author: yukun-hh
+date: 2026-5-14
+"""
+import sys, os
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+
+import random
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.optim as optim
+from torch.utils.data import DataLoader, Subset
+from torchvision import models, transforms
+from tqdm import tqdm
+import csv
+import matplotlib.pyplot as plt
+import matplotlib
+
+from Dataloader import RobustImageFolder
+
+matplotlib.rcParams['font.sans-serif'] = ['SimHei', 'DejaVu Sans']
+matplotlib.rcParams['axes.unicode_minus'] = False
+
+# ============================================================
+# ★★★ 可配置参数 ★★★
+# ============================================================
+DATA_ROOT   = '../../trash_division_data/ultimate_4_class/'
+BATCH_SIZE  = 32
+IMAGE_SIZE  = 256
+NUM_WORKERS = 4
+EPOCHS      = 30
+LR          = 0.001
+TRAIN_PCT   = 0.1
+SEED        = 42
+DROPOUT     = 0.3
+MODEL_SAVE_PATH = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'resnet34_10pct.pth')
+LOG_PATH        = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'resnet34_10pct_log.csv')
+# ============================================================
+
+NUM_CLASSES = 4
+CLASS_NAMES = ['厨余垃圾', '可回收物', '其他垃圾', '有害垃圾']
+
+
+class PretrainedResNet34(nn.Module):
+    def __init__(self, num_classes=NUM_CLASSES, dropout=DROPOUT):
+        super().__init__()
+        self.backbone = models.resnet34(weights='IMAGENET1K_V1')
+        in_features = self.backbone.fc.in_features
+        self.backbone.fc = nn.Identity()
+        self.dropout = nn.Dropout(dropout)
+        self.fc = nn.Linear(in_features, num_classes)
+
+    def forward(self, x):
+        x = self.backbone(x)
+        x = self.dropout(x)
+        x = self.fc(x)
+        return x
+
+    def freeze_early_layers(self):
+        for param in self.backbone.conv1.parameters():
+            param.requires_grad = False
+        for param in self.backbone.bn1.parameters():
+            param.requires_grad = False
+        for param in self.backbone.layer1.parameters():
+            param.requires_grad = False
+        for param in self.backbone.layer2.parameters():
+            param.requires_grad = False
+
+    def print_trainable_info(self):
+        frozen = sum(p.numel() for p in self.parameters() if not p.requires_grad)
+        trainable = sum(p.numel() for p in self.parameters() if p.requires_grad)
+        total = frozen + trainable
+        print(f"  冻结参数: {frozen:,}  可训练参数: {trainable:,}  ({100.*trainable/total:.1f}%)")
+
+
+def compute_macro_f1(predicted, targets, num_classes=NUM_CLASSES):
+    tp = torch.zeros(num_classes, device=predicted.device)
+    fp = torch.zeros(num_classes, device=predicted.device)
+    fn = torch.zeros(num_classes, device=predicted.device)
+    for c in range(num_classes):
+        tp[c] = ((predicted == c) & (targets == c)).sum()
+        fp[c] = ((predicted == c) & (targets != c)).sum()
+        fn[c] = ((predicted != c) & (targets == c)).sum()
+    precision = tp / (tp + fp + 1e-8)
+    recall = tp / (tp + fn + 1e-8)
+    f1 = 2 * precision * recall / (precision + recall + 1e-8)
+    return f1.mean().item()
+
+
+def train_one_epoch(model, loader, criterion, optimizer, device, epoch):
+    model.train()
+    running_loss, correct, total = 0.0, 0, 0
+    all_preds, all_labels = [], []
+    pbar = tqdm(loader, desc=f'Epoch {epoch+1} [Train]')
+    for images, labels in pbar:
+        images, labels = images.to(device), labels.to(device)
+        outputs = model(images)
+        loss = criterion(outputs, labels)
+        optimizer.zero_grad()
+        loss.backward()
+        optimizer.step()
+        running_loss += loss.item() * images.size(0)
+        _, predicted = outputs.max(1)
+        total += labels.size(0)
+        correct += predicted.eq(labels).sum().item()
+        all_preds.append(predicted)
+        all_labels.append(labels)
+        batch_f1 = compute_macro_f1(predicted, labels)
+        pbar.set_postfix({'loss': loss.item(), 'F1': f'{batch_f1:.4f}',
+                          'Acc': f'{100.*correct/total:.2f}%'})
+    epoch_loss = running_loss / total
+    epoch_f1 = compute_macro_f1(torch.cat(all_preds), torch.cat(all_labels))
+    epoch_acc = 100. * correct / total
+    return epoch_loss, epoch_f1, epoch_acc
+
+
+@torch.no_grad()
+def validate(model, loader, criterion, device):
+    model.eval()
+    running_loss, correct, total = 0.0, 0, 0
+    all_preds, all_labels = [], []
+    for images, labels in tqdm(loader, desc='[Validate]'):
+        images, labels = images.to(device), labels.to(device)
+        outputs = model(images)
+        loss = criterion(outputs, labels)
+        running_loss += loss.item() * images.size(0)
+        _, predicted = outputs.max(1)
+        total += labels.size(0)
+        correct += predicted.eq(labels).sum().item()
+        all_preds.append(predicted)
+        all_labels.append(labels)
+    epoch_loss = running_loss / total
+    epoch_f1 = compute_macro_f1(torch.cat(all_preds), torch.cat(all_labels))
+    epoch_acc = 100. * correct / total
+    return epoch_loss, epoch_f1, epoch_acc
+
+
+def train_model(model, train_loader, val_loader, device, epochs=EPOCHS, lr=LR):
+    criterion = nn.CrossEntropyLoss()
+    optimizer = optim.SGD(filter(lambda p: p.requires_grad, model.parameters()),
+                          lr=lr, momentum=0.9, weight_decay=1e-4)
+    scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=epochs)
+
+    history = {'train_loss': [], 'train_f1': [], 'train_acc': [],
+               'val_loss': [], 'val_f1': [], 'val_acc': []}
+    best_val_f1 = 0.0
+
+    log_file = open(LOG_PATH, 'w', newline='')
+    log_writer = csv.writer(log_file)
+    log_writer.writerow(['epoch', 'train_loss', 'train_f1', 'train_acc',
+                         'val_loss', 'val_f1', 'val_acc', 'lr', 'best'])
+
+    for epoch in range(epochs):
+        print(f'\n{"="*50}')
+        print(f'Epoch {epoch+1}/{epochs}')
+
+        train_loss, train_f1, train_acc = train_one_epoch(
+            model, train_loader, criterion, optimizer, device, epoch)
+        val_loss, val_f1, val_acc = validate(model, val_loader, criterion, device)
+        scheduler.step()
+
+        history['train_loss'].append(train_loss)
+        history['train_f1'].append(train_f1)
+        history['train_acc'].append(train_acc)
+        history['val_loss'].append(val_loss)
+        history['val_f1'].append(val_f1)
+        history['val_acc'].append(val_acc)
+
+        print(f'Train Loss: {train_loss:.4f} | Train Acc: {train_acc:.2f}% | Train Macro-F1: {train_f1:.4f}')
+        print(f'Val   Loss: {val_loss:.4f} | Val   Acc: {val_acc:.2f}% | Val   Macro-F1: {val_f1:.4f}')
+        print(f'Learning Rate: {optimizer.param_groups[0]["lr"]:.6f}')
+
+        best_mark = ''
+        if val_f1 > best_val_f1:
+            best_val_f1 = val_f1
+            torch.save(model.state_dict(), MODEL_SAVE_PATH)
+            best_mark = 'best'
+            print(f'✓ 保存最佳模型 (Macro-F1: {val_f1:.4f})')
+
+        lr_val = optimizer.param_groups[0]['lr']
+        log_writer.writerow([epoch+1, train_loss, train_f1, train_acc,
+                             val_loss, val_f1, val_acc, lr_val, best_mark])
+        log_file.flush()
+
+    log_file.close()
+    print(f'\n训练完成！最佳验证 Macro-F1: {best_val_f1:.4f}')
+    return history
+
+
+# ============================================================
+# compare_models.py 导入接口
+# ============================================================
+
+def get_resnet34_10pct_preds(train_loader, val_loader, device):
+    model = PretrainedResNet34(num_classes=NUM_CLASSES)
+    model.load_state_dict(torch.load(MODEL_SAVE_PATH, map_location='cpu'))
+    model = model.to(device).eval()
+
+    y_true, y_preds, y_probs = [], [], []
+    with torch.no_grad():
+        for images, labels in tqdm(val_loader, desc='ResNet-34 (10%)'):
+            images, labels = images.to(device), labels
+            logits = model(images)
+            probs = torch.softmax(logits, dim=1)
+            preds = probs.argmax(dim=1)
+            y_true.append(labels.numpy())
+            y_preds.append(preds.cpu().numpy())
+            y_probs.append(probs.cpu().numpy())
+    return np.concatenate(y_true), np.concatenate(y_preds), np.concatenate(y_probs)
+
+
+# ============================================================
+# 独立训练入口
+# ============================================================
+
+if __name__ == '__main__':
+    random.seed(SEED)
+    np.random.seed(SEED)
+    torch.manual_seed(SEED)
+
+    device = torch.device('cuda' if torch.cuda.is_available()
+                          else 'xpu' if hasattr(torch, 'xpu') and torch.xpu.is_available()
+                          else 'cpu')
+    print(f"Device: {device}")
+
+    val_transform = transforms.Compose([
+        transforms.Resize((IMAGE_SIZE, IMAGE_SIZE)),
+        transforms.ToTensor(),
+        transforms.Normalize(mean=[0.485, 0.456, 0.406],
+                             std=[0.229, 0.224, 0.225]),
+    ])
+    train_transform = transforms.Compose([
+        transforms.RandomResizedCrop(IMAGE_SIZE, scale=(0.8, 1.0)),
+        transforms.RandomHorizontalFlip(p=0.5),
+        transforms.RandomRotation(degrees=15),
+        transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2),
+        transforms.ToTensor(),
+        transforms.Normalize(mean=[0.485, 0.456, 0.406],
+                             std=[0.229, 0.224, 0.225]),
+    ])
+
+    full_train_dataset = RobustImageFolder(
+        root=os.path.join(DATA_ROOT, 'train'),
+        transform=train_transform,
+    )
+    val_dataset = RobustImageFolder(
+        root=os.path.join(DATA_ROOT, 'val'),
+        transform=val_transform,
+    )
+
+    n_train = len(full_train_dataset)
+    n_subset = max(1, int(n_train * TRAIN_PCT))
+    indices = random.sample(range(n_train), n_subset)
+    train_dataset = Subset(full_train_dataset, indices)
+    print(f"训练集: {len(train_dataset)} / {n_train} ({TRAIN_PCT*100:.0f}%)")
+    print(f"验证集: {len(val_dataset)}")
+
+    train_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE,
+                              shuffle=True, num_workers=NUM_WORKERS,
+                              pin_memory=True, drop_last=True)
+    val_loader = DataLoader(val_dataset, batch_size=BATCH_SIZE,
+                            shuffle=False, num_workers=NUM_WORKERS,
+                            pin_memory=True, drop_last=False)
+
+    model = PretrainedResNet34(num_classes=NUM_CLASSES, dropout=DROPOUT)
+    model.freeze_early_layers()
+    model.print_trainable_info()
+    model = model.to(device)
+
+    history = train_model(model, train_loader, val_loader, device, epochs=EPOCHS, lr=LR)
+
+    model.load_state_dict(torch.load(MODEL_SAVE_PATH, map_location='cpu'))
+    print(f"模型已保存: {MODEL_SAVE_PATH}")
+    print(f"训练日志已保存: {LOG_PATH}")

baseline/compare_models.py

@@ -22,6 +22,7 @@ from sklearn.metrics import roc_curve, auc, accuracy_score
 from Model import Net
 from Dataloader import RobustImageFolder
 from baseline.VGG_KNN import VGGKNNBaseline
+from baseline.ResNet34_Pretrained_10pct import get_resnet34_10pct_preds
 
 matplotlib.rcParams['font.sans-serif'] = ['SimHei', 'DejaVu Sans']
 matplotlib.rcParams['axes.unicode_minus'] = False
@@ -77,8 +78,9 @@ def get_vgg_knn_preds(train_loader, val_loader, device):
 # ============================================================
 
 MODELS = [
-    ('ResNet-34',           get_resnet34_preds),
+    ('ResNet-34',                get_resnet34_preds),
-    ('VGG16 + KNN  (K=5)',  get_vgg_knn_preds),
+    ('ResNet-34 (10% Fine-tune)', get_resnet34_10pct_preds),
+    ('VGG16 + KNN  (K=5)',       get_vgg_knn_preds),
     # 未来轻松扩展示例：
     # ('ResNet-18 (pretrained)', get_resnet18_preds),
     # ('ResNet-50 (pretrained)', get_resnet50_preds),