迁移学习（二）：预训练 + 微调范式

# 对比：从零训练 vs 迁移学习的效率差异
import torch
import torchvision.models as models

# 从零训练：需要大量数据和时间
# 预训练模型：加载即具备丰富特征表示
pretrained_model = models.resnet50(
    weights=models.ResNet50_Weights.IMAGENET1K_V2
)
print(f"预训练权重加载完成")
print(f"参数量: {sum(p.numel() for p in pretrained_model.parameters()):,}")

# 迁移学习的核心优势量化
import time
from torchvision import datasets, transforms

# 场景模拟：小样本目标域训练
train_time_scratch = "约 72 小时 (从零训练 ResNet-50)"
train_time_transfer = "约 2 小时 (微调预训练模型)"

print(f"从零训练: {train_time_scratch}")
print(f"迁移学习: {train_time_transfer}")
print(f"时间节省: {((72 - 2) / 72 * 100):.0f}%")

# 策略一：特征提取（冻结全部层）
import torch.nn as nn
import torchvision.models as models

model = models.resnet50(weights="IMAGENET1K_V2")

# 冻结所有预训练参数
for param in model.parameters():
    param.requires_grad = False

# 替换分类头
num_features = model.fc.in_features
model.fc = nn.Sequential(
    nn.Linear(num_features, 256),
    nn.ReLU(),
    nn.Dropout(0.5),
    nn.Linear(256, 10)  # 10类目标
)

# 只有分类头的参数可训练
trainable = [p.numel() for p in model.parameters() if p.requires_grad]
print(f"可训练参数: {sum(trainable):,}")

# 策略二：微调（部分解冻 + 低学习率）
def setup_fine_tune(model, unfreeze_from: str = "layer3"):
    """从指定层开始解冻"""
    freeze = True
    for name, param in model.named_parameters():
        if name.startswith(unfreeze_from):
            freeze = False
        param.requires_grad = not freeze
    return model

model = setup_fine_tune(model, unfreeze_from="layer3")

# 不同层使用不同学习率
optimizer = torch.optim.Adam([
    {"params": model.layer3.parameters(), "lr": 1e-5},
    {"params": model.layer4.parameters(), "lr": 1e-4},
    {"params": model.fc.parameters(), "lr": 1e-3},
])

# 渐进式解冻 + 差异化学习率
import torch.nn as nn

def progressive_unfreeze(model, epochs_per_stage=3):
    """分阶段解冻网络层"""
    stage_groups = [
        model.fc,              # 阶段1: 仅分类头
        model.layer4,          # 阶段2: 解冻layer4
        model.layer3,          # 阶段3: 解冻layer3
        [model.layer1, model.layer2]  # 阶段4: 解冻浅层
    ]
    return stage_groups

# 每个阶段使用递减学习率
lr_stages = [1e-3, 5e-4, 1e-4, 5e-5]
print(f"共 {len(lr_stages)} 个训练阶段")
for i, lr in enumerate(lr_stages):
    print(f"阶段 {i+1}: lr={lr:.0e}")

# 学习率调度组合拳
from torch.optim.lr_scheduler import (
    OneCycleLR, CosineAnnealingWarmRestarts
)

# 方案A: OneCycleLR（适合快速微调）
scheduler_one_cycle = OneCycleLR(
    optimizer,
    max_lr=1e-3,
    total_steps=1000,
    pct_start=0.3,       # 30%时间预热
    div_factor=25,       # 初始lr = max_lr / 25
    final_div_factor=1e4 # 最终lr = max_lr / 1e4
)

# 方案B: 余弦退火 + 热重启（适合精细微调）
scheduler_cosine = CosineAnnealingWarmRestarts(
    optimizer, T_0=10, T_mult=2, eta_min=1e-6
)

# 领域适配型数据增强流水线
import albumentations as A
from albumentations.pytorch import ToTensorV2

def get_train_augmentations(img_size=224):
    """针对领域偏移的鲁棒增强策略"""
    return A.Compose([
        A.RandomResizedCrop(img_size, img_size,
                           scale=(0.6, 1.0)),
        A.HorizontalFlip(p=0.5),
        A.ColorJitter(
            brightness=0.3, contrast=0.3,
            saturation=0.3, hue=0.1, p=0.8
        ),
        A.GaussianBlur(blur_limit=7, p=0.3),
        A.Normalize(mean=[0.485, 0.456, 0.406],
                    std=[0.229, 0.224, 0.225]),
        ToTensorV2()
    ])

# MixUp 数据增强（缓解领域偏移）
import torch.nn.functional as F

def mixup_data(x, y, alpha=0.2):
    """MixUp: 线性插值样本对"""
    if alpha > 0:
        lam = torch.distributions.Beta(
            alpha, alpha
        ).sample().item()
    else:
        lam = 1.0

    batch_size = x.size(0)
    index = torch.randperm(batch_size).to(x.device)
    
    mixed_x = lam * x + (1 - lam) * x[index]
    y_a, y_b = y, y[index]
    return mixed_x, y_a, y_b, lam

# 损失函数调整
def mixup_criterion(pred, y_a, y_b, lam):
    return lam * F.cross_entropy(pred, y_a) + \
           (1 - lam) * F.cross_entropy(pred, y_b)

# 多模型基准对比
import torchvision.models as models
from timm import list_models, create_model

def compare_models():
    """对比候选预训练模型"""
    candidates = {
        "resnet50": models.resnet50(
            weights="IMAGENET1K_V2"),
        "efficientnet_b0": models.efficientnet_b0(
            weights="IMAGENET1K_V1"),
        "convnext_tiny": models.convnext_tiny(
            weights="IMAGENET1K_V1"),
    }
    
    results = []
    for name, model in candidates.items():
        params = sum(p.numel() for p in model.parameters())
        results.append((name, params))
    
    for name, params in sorted(results, key=lambda x: x[1]):
        print(f"{name:20s} {params:>10,} params")

compare_models()

# TIMM 库：一站式预训练模型仓库
import timm

# 列出所有可用预训练模型
all_models = timm.list_models(pretrained=True)
print(f"可用预训练模型: {len(all_models)} 个")

# 加载特定预训练权重
model = timm.create_model(
    "convnext_large",
    pretrained=True,
    num_classes=10,        # 自定义类别数
    drop_path_rate=0.2     # 随机深度正则化
)

# 查看模型默认配置
data_cfg = timm.data.resolve_model_data_config(model)
print(f"输入尺寸: {data_cfg['input_size']}")
print(f"归一化均值: {data_cfg['mean']}")
print(f"归一化标准差: {data_cfg['std']}")

# 原型网络（Prototypical Network）
import torch
import torch.nn as nn
import torch.nn.functional as F

class PrototypicalNetwork(nn.Module):
    def __init__(self, encoder):
        super().__init__()
        self.encoder = encoder  # 冻结的预训练特征提取器
    
    def forward(self, support_x, support_y, query_x):
        # 提取支持集特征并计算类原型
        support_feat = self.encoder(support_x)
        prototypes = self.compute_prototypes(
            support_feat, support_y
        )
        # 计算查询样本到原型的距离
        query_feat = self.encoder(query_x)
        dists = self.euclidean_dist(query_feat, prototypes)
        return -dists  # 负距离作为 logits
    
    def compute_prototypes(self, feat, labels):
        classes = torch.unique(labels)
        prototypes = torch.stack([
            feat[labels == c].mean(dim=0) for c in classes
        ])
        return prototypes
    
    def euclidean_dist(self, x, y):
        return torch.cdist(x, y, p=2)

# 小样本微调的超参数配置
few_shot_config = {
    # 冻结策略
    "unfreeze_layers": 1,        # 仅解冻最后1层
    
    # 激进正则化
    "weight_decay": 1e-3,        # 比常规大10倍
    "dropout_rate": 0.6,         # 高Dropout
    "label_smoothing": 0.15,     # 标签平滑
    
    # 学习率
    "lr": 5e-5,                  # 极低学习率
    "warmup_epochs": 5,          # 较长预热
    
    # 训练轮次
    "max_epochs": 50,            # 多轮次精细训练
    "early_stop_patience": 10,   # 耐心早停
}

for k, v in few_shot_config.items():
    print(f"{k}: {v}")

# 完整迁移学习训练流程
import torch
import torch.nn as nn
from torchvision import datasets, transforms, models
from torch.utils.data import DataLoader
import os

# 1. 数据加载
train_transform = transforms.Compose([
    transforms.RandomResizedCrop(224),
    transforms.RandomHorizontalFlip(),
    transforms.ToTensor(),
    transforms.Normalize(
        [0.485, 0.456, 0.406],
        [0.229, 0.224, 0.225]
    )
])

train_dataset = datasets.ImageFolder(
    "data/train", transform=train_transform
)
train_loader = DataLoader(
    train_dataset, batch_size=32, shuffle=True, num_workers=4
)

# 2. 模型构建
model = models.resnet50(weights="IMAGENET1K_V2")
num_ftrs = model.fc.in_features
model.fc = nn.Linear(num_ftrs, 10)  # 自定义10类

# 3. 差异化学习率
optimizer = torch.optim.AdamW([
    {"params": model.conv1.parameters(), "lr": 1e-5},
    {"params": model.layer1.parameters(), "lr": 1e-5},
    {"params": model.layer2.parameters(), "lr": 3e-5},
    {"params": model.layer3.parameters(), "lr": 1e-4},
    {"params": model.layer4.parameters(), "lr": 3e-4},
    {"params": model.fc.parameters(), "lr": 1e-3},
], weight_decay=1e-4)

# 4. 训练循环 + 早停机制
def train_model(model, train_loader, val_loader, epochs=30):
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model = model.to(device)
    
    criterion = nn.CrossEntropyLoss(label_smoothing=0.1)
    scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
        optimizer, mode="min", factor=0.5, patience=3
    )
    
    best_acc = 0.0
    patience_counter = 0
    max_patience = 7
    
    for epoch in range(epochs):
        # 训练阶段
        model.train()
        train_loss = 0.0
        for inputs, labels in train_loader:
            inputs, labels = inputs.to(device), labels.to(device)
            optimizer.zero_grad()
            outputs = model(inputs)
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()
            train_loss += loss.item()
        
        # 验证阶段
        model.eval()
        correct = 0
        total = 0
        with torch.no_grad():
            for inputs, labels in val_loader:
                inputs, labels = inputs.to(device), labels.to(device)
                outputs = model(inputs)
                _, predicted = outputs.max(1)
                total += labels.size(0)
                correct += predicted.eq(labels).sum().item()
        
        val_acc = correct / total
        scheduler.step(val_acc)
        
        # 保存最佳模型
        if val_acc > best_acc:
            best_acc = val_acc
            patience_counter = 0
            torch.save(model.state_dict(), "best_model.pth")
        else:
            patience_counter += 1
        
        print(f"Epoch {epoch+1}: val_acc={val_acc:.4f}, best={best_acc:.4f}")
        
        # 早停检查
        if patience_counter >= max_patience:
            print(f"Early stopping at epoch {epoch+1}")
            break
    
    return best_acc

best_acc = train_model(model, train_loader, val_loader)
print(f"最佳验证准确率: {best_acc:.4f}")