第5章:现代CNN架构
    ↓
理解残差连接、轻量化设计
    ↓
第6章:Attention与Transformer
    ↓
掌握注意力机制、ViT架构
    ↓
第7章:数据增强与训练技巧
    ↓
完整训练流程实战

# 安装核心依赖
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118
pip install timm  # PyTorch Image Models
pip install transformers  # Hugging Face Transformers
pip install albumentations  # 数据增强
pip install tensorboard  # 可视化
pip install opencv-python pillow matplotlib

# 可选:实验追踪
pip install wandb

F(x) = H(x) - x
H(x) = F(x) + x  # 输出 = 残差 + 输入

计算量: 标准卷积 vs 深度可分离卷积
标准: D_K × D_K × M × N × D_F × D_F
深度可分离: D_K × D_K × M × D_F × D_F + M × N × D_F × D_F
压缩比: 1/N + 1/D_K²

Attention(Q, K, V) = softmax(QK^T / √d_k) × V

# 小数据集:冻结大部分层
for param in model.parameters():
    param.requires_grad = False
# 只训练分类头

# 中等数据集:冻结早期层
for name, param in model.named_parameters():
    if 'layer1' in name or 'layer2' in name:
        param.requires_grad = False

# 大数据集:全模型fine-tune
for param in model.parameters():
    param.requires_grad = True

# 迁移学习典型设置
lr_backbone = 1e-4  # 预训练层
lr_head = 1e-3      # 新增层

现象: 56层网络训练误差 > 20层网络训练误差
原因: 不是过拟合(训练误差更高),而是优化困难

浅层网络(20层) → 训练误差: 15%
深层网络(56层) → 训练误差: 20% (理论上应该≤15%)

传统: H(x) = 目标映射
ResNet: F(x) = H(x) - x
        H(x) = F(x) + x

输入x
  ↓
3×3 Conv → BN → ReLU
  ↓
3×3 Conv → BN
  ↓
  + ← x (跳跃连接)
  ↓
ReLU
  ↓
输出

输入x (256维)
  ↓
1×1 Conv(64) → BN → ReLU  # 降维
  ↓
3×3 Conv(64) → BN → ReLU  # 特征提取
  ↓
1×1 Conv(256) → BN         # 升维
  ↓
  + ← x
  ↓
ReLU
  ↓
输出(256维)

输入: 224×224×3
  ↓
Conv1: 7×7, 64, stride=2 → 112×112×64
  ↓
MaxPool: 3×3, stride=2 → 56×56×64
  ↓
Layer1: [1×1,64; 3×3,64; 1×1,256] × 3 → 56×56×256
  ↓
Layer2: [1×1,128; 3×3,128; 1×1,512] × 4 → 28×28×512
  ↓
Layer3: [1×1,256; 3×3,256; 1×1,1024] × 6 → 14×14×1024
  ↓
Layer4: [1×1,512; 3×3,512; 1×1,2048] × 3 → 7×7×2048
  ↓
AvgPool: 7×7 → 1×1×2048
  ↓
FC: 1000 classes

class ResidualBlock(nn.Module):
    def __init__(self, in_channels, out_channels, stride=1):
        super().__init__()
        self.conv1 = nn.Conv2d(in_channels, out_channels, 3, stride, 1)
        self.bn1 = nn.BatchNorm2d(out_channels)
        self.conv2 = nn.Conv2d(out_channels, out_channels, 3, 1, 1)
        self.bn2 = nn.BatchNorm2d(out_channels)

        # 如果维度不匹配,使用1×1卷积调整
        self.shortcut = nn.Sequential()
        if stride != 1 or in_channels != out_channels:
            self.shortcut = nn.Sequential(
                nn.Conv2d(in_channels, out_channels, 1, stride),
                nn.BatchNorm2d(out_channels)
            )

    def forward(self, x):
        out = F.relu(self.bn1(self.conv1(x)))
        out = self.bn2(self.conv2(out))
        out += self.shortcut(x)  # 残差连接
        out = F.relu(out)
        return out

反向传播:
∂L/∂x = ∂L/∂H × (∂F/∂x + I)
      = ∂L/∂H × ∂F/∂x + ∂L/∂H

关键: ∂L/∂H 直接传播,不经过中间层!

输入: H × W × M
卷积核: D_K × D_K × M × N
输出: H × W × N

计算量: D_K × D_K × M × N × H × W

输入: H × W × M
卷积核: D_K × D_K × 1 × M (每个通道独立卷积)
输出: H × W × M

计算量: D_K × D_K × M × H × W

输入: H × W × M
卷积核: 1 × 1 × M × N
输出: H × W × N

计算量: M × N × H × W

深度可分离: D_K² × M × HW + M × N × HW
标准卷积: D_K² × M × N × HW

压缩比: (D_K² × M + M × N) / (D_K² × M × N)
       = 1/N + 1/D_K²

例: D_K=3, N=128
压缩比 = 1/128 + 1/9 ≈ 0.119 (压缩8.4倍)

输入: 224×224×3
  ↓
Conv: 3×3, 32, stride=2 → 112×112×32
  ↓
DW: 3×3, 32
PW: 1×1, 64 → 112×112×64
  ↓
DW: 3×3, 64, stride=2
PW: 1×1, 128 → 56×56×128
  ↓
DW: 3×3, 128
PW: 1×1, 128 → 56×56×128
  ↓
... (重复若干次)
  ↓
AvgPool → 1×1×1024
  ↓
FC: 1000

传统残差块: 高维 → 低维 → 高维
倒残差块: 低维 → 高维 → 低维

输入: H × W × C (低维, 如32)
  ↓
1×1 Conv, C×t (扩展, t=6) → H × W × C×6 (高维, 192)
  ↓
3×3 DW Conv → H × W × C×6
  ↓
1×1 Conv, C → H × W × C (压缩回低维)
  ↓
+ ← 输入 (残差连接)

输入 → Global Pool → FC → ReLU → FC → Sigmoid → 乘回输入

h-swish(x) = x × ReLU6(x + 3) / 6
优点: 比ReLU性能好,比swish计算快

深度: d = α^φ
宽度: w = β^φ
分辨率: r = γ^φ

约束: α × β² × γ² ≈ 2
      α ≥ 1, β ≥ 1, γ ≥ 1

同样83%准确率:
- ResNet-152: 60M参数, 11.6G FLOPs
- EfficientNet-B3: 12M参数, 1.8G FLOPs
压缩比: 5倍参数, 6.4倍FLOPs

import torchvision.models as models

# 方式1: 直接加载
model = models.resnet50(weights=models.ResNet50_Weights.IMAGENET1K_V2)

# 方式2: 使用timm(更多模型)
import timm
model = timm.create_model('resnet50', pretrained=True)

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

# 冻结特定层
for name, param in model.named_parameters():
    if 'layer1' in name or 'layer2' in name:
        param.requires_grad = False

# ResNet
num_classes = 10
model.fc = nn.Linear(model.fc.in_features, num_classes)

# EfficientNet (timm)
model.classifier = nn.Linear(model.classifier.in_features, num_classes)

# 方式1: 参数分组
params_to_update = []
params_backbone = []

for name, param in model.named_parameters():
    if param.requires_grad:
        if 'fc' in name or 'classifier' in name:
            params_to_update.append(param)
        else:
            params_backbone.append(param)

optimizer = torch.optim.Adam([
    {'params': params_backbone, 'lr': 1e-4},  # 小学习率
    {'params': params_to_update, 'lr': 1e-3}  # 大学习率
])

# Epoch 1-5: 只训练分类头
# Epoch 6-10: 解冻layer4
# Epoch 11-15: 解冻layer3
# ...

def unfreeze_layer(model, layer_name):
    for name, param in model.named_parameters():
        if layer_name in name:
            param.requires_grad = True

# 越深的层学习率越大
lr_mult = [0.1, 0.3, 0.6, 1.0]  # layer1-4
optimizer = torch.optim.SGD([
    {'params': model.layer1.parameters(), 'lr': base_lr * lr_mult[0]},
    {'params': model.layer2.parameters(), 'lr': base_lr * lr_mult[1]},
    {'params': model.layer3.parameters(), 'lr': base_lr * lr_mult[2]},
    {'params': model.layer4.parameters(), 'lr': base_lr * lr_mult[3]},
    {'params': model.fc.parameters(), 'lr': base_lr * 10}
], momentum=0.9)

from torch.optim.lr_scheduler import LinearLR, SequentialLR, CosineAnnealingLR

warmup_epochs = 5
total_epochs = 100

# 预热阶段
warmup_scheduler = LinearLR(
    optimizer, start_factor=0.1, end_factor=1.0, total_iters=warmup_epochs
)

# 主训练阶段
main_scheduler = CosineAnnealingLR(
    optimizer, T_max=total_epochs - warmup_epochs
)

# 组合调度器
scheduler = SequentialLR(
    optimizer,
    schedulers=[warmup_scheduler, main_scheduler],
    milestones=[warmup_epochs]
)

准确率排序:
EfficientNet-B7 (84.3%) > ResNet-152 (78.3%) > ResNet-50 (76.2%) > MobileNetV3 (75.2%)

速度排序(GPU):
MobileNetV3-Small (5ms) > MobileNetV3-Large (8ms) > ResNet-50 (22ms) > EfficientNet-B7 (80ms)

参数量排序:
MobileNetV3-Small (2.5M) < MobileNetV3-Large (5.4M) < ResNet-50 (25.6M) < EfficientNet-B7 (66M)

α_i = exp(e_i) / Σ exp(e_j)  # 注意力权重
output = Σ α_i × v_i           # 加权求和

从序列自身计算注意力
不依赖外部信息

多个注意力头
捕获不同类型的关系

Query (查询): Q = XW_Q ∈ ℝⁿˣᵈₖ
Key (键): K = XW_K ∈ ℝⁿˣᵈₖ
Value (值): V = XW_V ∈ ℝⁿˣᵈᵥ

参数: W_Q, W_K ∈ ℝᵈˣᵈₖ, W_V ∈ ℝᵈˣᵈᵥ

Attention(Q, K, V) = softmax(QK^T / √d_k) × V

步骤:
1. 计算相似度: S = QK^T ∈ ℝⁿˣⁿ
2. 缩放: S = S / √d_k (防止梯度消失)
3. 归一化: A = softmax(S) (注意力权重)
4. 加权求和: Output = A × V

假设 Q, K 的元素独立同分布 N(0, 1)
则 QK^T 的元素方差约为 d_k
除以 √d_k 使方差归一化到 1
避免 softmax 进入饱和区(梯度消失)

1. 将每个词映射到Q, K, V:
   Q_我 = [0.2, 0.5, ...], K_我 = [0.3, 0.1, ...], V_我 = [0.8, 0.2, ...]
   ...

2. 计算"我"对所有词的注意力分数:
   score(我, 我) = Q_我 · K_我 = 0.8
   score(我, 爱) = Q_我 · K_爱 = 0.6
   score(我, 自然语言处理) = Q_我 · K_自然语言处理 = 0.2

3. 归一化(softmax):
   α_我→我 = 0.4
   α_我→爱 = 0.5
   α_我→自然语言处理 = 0.1

4. 加权求和:
   Output_我 = 0.4 × V_我 + 0.5 × V_爱 + 0.1 × V_自然语言处理

import torch
import torch.nn as nn
import torch.nn.functional as F

class SelfAttention(nn.Module):
    def __init__(self, embed_dim: int, head_dim: int):
        super().__init__()
        self.embed_dim = embed_dim
        self.head_dim = head_dim

        # Q, K, V投影
        self.q_proj = nn.Linear(embed_dim, head_dim)
        self.k_proj = nn.Linear(embed_dim, head_dim)
        self.v_proj = nn.Linear(embed_dim, head_dim)

        self.scale = head_dim ** -0.5  # 1/√d_k

    def forward(self, x):
        """
        Args:
            x: (batch_size, seq_len, embed_dim)

        Returns:
            (batch_size, seq_len, head_dim)
        """
        B, N, C = x.shape

        # 投影
        Q = self.q_proj(x)  # (B, N, head_dim)
        K = self.k_proj(x)  # (B, N, head_dim)
        V = self.v_proj(x)  # (B, N, head_dim)

        # 计算注意力分数
        attn = (Q @ K.transpose(-2, -1)) * self.scale  # (B, N, N)

        # 归一化
        attn = F.softmax(attn, dim=-1)  # (B, N, N)

        # 加权求和
        out = attn @ V  # (B, N, head_dim)

        return out

输入序列
  ↓
Embedding + Positional Encoding
  ↓
Encoder (N层):
  - Multi-Head Attention
  - Add & Norm
  - Feed Forward
  - Add & Norm
  ↓
Decoder (N层):
  - Masked Multi-Head Attention
  - Add & Norm
  - Cross Attention
  - Add & Norm
  - Feed Forward
  - Add & Norm
  ↓
输出

输入 x
  ↓
Multi-Head Attention → Add & Norm
  ↓
Feed Forward Network → Add & Norm
  ↓
输出

输入 X ∈ ℝⁿˣᵈ
  ↓
分成h个头,每个头维度 d_k = d/h
  ↓
Head₁ = Attention(XW_Q1, XW_K1, XW_V1)
Head₂ = Attention(XW_Q2, XW_K2, XW_V2)
...
Headₕ = Attention(XW_Qh, XW_Kh, XW_Vh)
  ↓
Concat(Head₁, Head₂, ..., Headₕ)
  ↓
线性投影 W_O
  ↓
输出 ∈ ℝⁿˣᵈ

MultiHead(Q, K, V) = Concat(head₁, ..., headₕ)W_O

其中:
head_i = Attention(QW_Qi, KW_Ki, VW_Vi)

每个头的维度: d_k = 512/8 = 64
Head 1: 关注局部纹理
Head 2: 关注全局形状
Head 3: 关注颜色
...

输入: [A, B, C]
输入: [C, B, A]
注意力矩阵相同(如果Q, K, V投影相同)

PE(pos, 2i) = sin(pos / 10000^(2i/d))
PE(pos, 2i+1) = cos(pos / 10000^(2i/d))

其中:
- pos: 位置索引
- i: 维度索引
- d: 嵌入维度

position_embedding = nn.Parameter(torch.randn(1, num_patches + 1, embed_dim))

FFN(x) = max(0, xW₁ + b₁)W₂ + b₂
      = GELU(xW₁ + b₁)W₂ + b₂  (现代变体)

其中:
- W₁: d → d_ff (通常 d_ff = 4d)
- W₂: d_ff → d

LN(x) = γ × (x - μ) / σ + β

其中:
- μ, σ: 沿特征维度计算的均值和标准差
- γ, β: 可学习参数

Post-LN (原始):
x = x + Attention(LN(x))

Pre-LN (更稳定):
x = LN(x + Attention(x))

class TransformerEncoder(nn.Module):
    def __init__(
        self,
        embed_dim: int = 768,
        num_heads: int = 12,
        mlp_ratio: int = 4,
        dropout: float = 0.1
    ):
        super().__init__()
        self.norm1 = nn.LayerNorm(embed_dim)
        self.attn = nn.MultiheadAttention(
            embed_dim, num_heads, dropout=dropout, batch_first=True
        )
        self.norm2 = nn.LayerNorm(embed_dim)
        self.mlp = nn.Sequential(
            nn.Linear(embed_dim, embed_dim * mlp_ratio),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(embed_dim * mlp_ratio, embed_dim),
            nn.Dropout(dropout)
        )

    def forward(self, x):
        # Multi-Head Attention
        x = x + self.attn(self.norm1(x), self.norm1(x), self.norm1(x))[0]

        # FFN
        x = x + self.mlp(self.norm2(x))

        return x

输入图像: H × W × C (如 224×224×3)
  ↓
分割成patch: (H/P) × (W/P) × (P²C)
  例: 224/16 × 224/16 × (16²×3) = 14×14×768
  ↓
展平: N × (P²C), N = (H/W)² = 196
  ↓
线性投影: N × D (D=768)
  ↓
添加 [CLS] token: (N+1) × D
  ↓
添加位置编码: (N+1) × D
  ↓
Transformer Encoder × L层
  ↓
提取 [CLS] token
  ↓
MLP Head → 分类

# 方式1: 展平 + 线性投影
patch_size = 16
num_patches = (224 // 16) ** 2  # 196

# 输入: (B, 3, 224, 224)
x = x.unfold(2, patch_size, patch_size).unfold(3, patch_size, patch_size)
# (B, 3, 14, 14, 16, 16)

x = x.permute(0, 2, 3, 4, 5, 1).reshape(B, num_patches, -1)
# (B, 196, 768)

x = linear_projection(x)
# (B, 196, embed_dim)

# 方式2: 卷积(等价且更快)
self.patch_embed = nn.Conv2d(
    in_channels=3,
    out_channels=embed_dim,
    kernel_size=patch_size,
    stride=patch_size
)

# 可学习的分类token
self.cls_token = nn.Parameter(torch.randn(1, 1, embed_dim))

# 添加到序列开头
cls_tokens = self.cls_token.expand(B, -1, -1)  # (B, 1, embed_dim)
x = torch.cat([cls_tokens, x], dim=1)  # (B, N+1, embed_dim)

# 可学习位置编码
self.pos_embed = nn.Parameter(
    torch.randn(1, num_patches + 1, embed_dim)
)

x = x + self.pos_embed

小数据集(<100k):
ResNet-50 (76.2%) > ViT-B/16 (75.0%)

大数据集(ImageNet-21K预训练):
ViT-B/16 (84.0%) > ResNet-50 (80.4%)

超大数据集(JFT-300M预训练):
ViT-H/14 (88.5%) >> ResNet-152 (82.0%)

import timm

# 查看可用模型
models = timm.list_models('vit*', pretrained=True)
print(models[:5])
# ['vit_base_patch16_224', 'vit_base_patch32_224', ...]

# 加载模型
model = timm.create_model('vit_base_patch16_224', pretrained=True, num_classes=10)

# 查看模型结构
print(model)

# ImageNet-1K预训练
model = timm.create_model('vit_base_patch16_224', pretrained=True)

# ImageNet-21K预训练(更好)
model = timm.create_model('vit_base_patch16_224.augreg_in21k', pretrained=True)

# 预训练: 224×224
# 微调: 可以用更大尺寸(384×384)
# 需要插值位置编码

def interpolate_pos_embed(pos_embed, orig_size, new_size):
    # pos_embed: (1, N+1, D)
    cls_pos = pos_embed[:, :1, :]  # CLS token
    patch_pos = pos_embed[:, 1:, :]  # Patch positions

    # 插值
    patch_pos = F.interpolate(
        patch_pos.reshape(1, orig_size, orig_size, -1).permute(0, 3, 1, 2),
        size=(new_size, new_size),
        mode='bicubic'
    )
    patch_pos = patch_pos.permute(0, 2, 3, 1).reshape(1, -1, pos_embed.shape[-1])

    return torch.cat([cls_pos, patch_pos], dim=1)

from torch.cuda.amp import autocast, GradScaler

scaler = GradScaler()

for images, labels in dataloader:
    with autocast():
        outputs = model(images)
        loss = criterion(outputs, labels)

    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()

def visualize_attention(model, image, layer_idx=11, head_idx=0):
    """
    可视化ViT的注意力图

    Args:
        model: ViT模型
        image: 输入图像
        layer_idx: Transformer层索引
        head_idx: 注意力头索引
    """
    model.eval()

    # 注册hook获取注意力权重
    attentions = []

    def hook_fn(module, input, output):
        # output: (attn_weights, attn_output)
        attentions.append(output[0])

    hook = model.blocks[layer_idx].attn.attn_drop.register_forward_hook(hook_fn)

    # 前向传播
    with torch.no_grad():
        _ = model(image)

    hook.remove()

    # 获取注意力权重: (B, num_heads, N+1, N+1)
    attn = attentions[0][0, head_idx, 0, 1:]  # CLS对所有patch的注意力

    # 重塑为2D
    size = int(attn.shape[0] ** 0.5)
    attn_map = attn.reshape(size, size).cpu().numpy()

    # 可视化
    plt.imshow(attn_map, cmap='viridis')
    plt.colorbar()
    plt.title(f'Attention Map - Layer {layer_idx}, Head {head_idx}')
    plt.show()

Self-Attention:
Attention(Q, K, V) = softmax(QK^T / √d_k) × V

Multi-Head:
MultiHead(Q, K, V) = Concat(head₁, ..., headₕ)W_O

Positional Encoding:
PE(pos, 2i) = sin(pos / 10000^(2i/d))
PE(pos, 2i+1) = cos(pos / 10000^(2i/d))

from torchvision import transforms

# 基础增强组合
basic_transforms = transforms.Compose([
    transforms.RandomHorizontalFlip(p=0.5),  # 水平翻转
    transforms.RandomVerticalFlip(p=0.1),    # 垂直翻转（适用于特定任务）
    transforms.RandomRotation(degrees=15),    # 随机旋转
    transforms.RandomResizedCrop(224, scale=(0.8, 1.0)),  # 随机裁剪缩放
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406],
                        std=[0.229, 0.224, 0.225])
])

color_transforms = transforms.Compose([
    transforms.ColorJitter(
        brightness=0.2,   # 亮度变化 ±20%
        contrast=0.2,     # 对比度变化 ±20%
        saturation=0.2,   # 饱和度变化 ±20%
        hue=0.1          # 色调变化 ±10%
    ),
    transforms.RandomGrayscale(p=0.1),  # 10%概率转灰度
    transforms.GaussianBlur(kernel_size=3, sigma=(0.1, 2.0)),  # 高斯模糊
])

import albumentations as A
from albumentations.pytorch import ToTensorV2

# 高级增强组合
train_transforms = A.Compose([
    # 几何变换
    A.HorizontalFlip(p=0.5),
    A.ShiftScaleRotate(
        shift_limit=0.1,
        scale_limit=0.2,
        rotate_limit=15,
        p=0.5
    ),
    A.RandomResizedCrop(height=224, width=224, scale=(0.8, 1.0)),

    # 色彩变换
    A.OneOf([
        A.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2),
        A.HueSaturationValue(hue_shift_limit=20, sat_shift_limit=30),
        A.RandomBrightnessContrast(),
    ], p=0.5),

    # 噪声与模糊
    A.OneOf([
        A.GaussianBlur(blur_limit=3),
        A.GaussNoise(var_limit=(10.0, 50.0)),
        A.ISONoise(),
    ], p=0.3),

    # 标准化
    A.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
    ToTensorV2()
])

# 验证集增强（只做标准化）
val_transforms = A.Compose([
    A.Resize(224, 224),
    A.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
    ToTensorV2()
])

# 使用方式
class CustomDataset(torch.utils.data.Dataset):
    def __init__(self, image_paths, labels, transform=None):
        self.image_paths = image_paths
        self.labels = labels
        self.transform = transform

    def __getitem__(self, idx):
        image = cv2.imread(self.image_paths[idx])
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        label = self.labels[idx]

        if self.transform:
            augmented = self.transform(image=image)
            image = augmented['image']

        return image, label

def mixup_data(x, y, alpha=0.2):
    """
    Mixup数据增强

    Args:
        x: 输入图像 (B, C, H, W)
        y: 标签 (B,)
        alpha: Beta分布参数

    Returns:
        mixed_x: 混合后的图像
        y_a, y_b: 原始标签对
        lam: 混合比例
    """
    if alpha > 0:
        lam = np.random.beta(alpha, alpha)
    else:
        lam = 1

    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(criterion, pred, y_a, y_b, lam):
    """Mixup损失函数"""
    return lam * criterion(pred, y_a) + (1 - lam) * criterion(pred, y_b)

# 训练循环中使用
for images, labels in train_loader:
    images, labels = images.cuda(), labels.cuda()

    # 应用Mixup
    images, labels_a, labels_b, lam = mixup_data(images, labels, alpha=0.2)

    outputs = model(images)
    loss = mixup_criterion(criterion, outputs, labels_a, labels_b, lam)

    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

def cutmix_data(x, y, alpha=1.0):
    """
    CutMix数据增强

    Args:
        x: 输入图像 (B, C, H, W)
        y: 标签 (B,)
        alpha: Beta分布参数
    """
    lam = np.random.beta(alpha, alpha)
    batch_size = x.size(0)
    index = torch.randperm(batch_size).to(x.device)

    # 计算剪切区域
    _, _, H, W = x.shape
    cut_ratio = np.sqrt(1 - lam)
    cut_h = int(H * cut_ratio)
    cut_w = int(W * cut_ratio)

    # 随机中心点
    cx = np.random.randint(W)
    cy = np.random.randint(H)

    # 边界框
    bbx1 = np.clip(cx - cut_w // 2, 0, W)
    bby1 = np.clip(cy - cut_h // 2, 0, H)
    bbx2 = np.clip(cx + cut_w // 2, 0, W)
    bby2 = np.clip(cy + cut_h // 2, 0, H)

    # 替换区域
    x[:, :, bby1:bby2, bbx1:bbx2] = x[index, :, bby1:bby2, bbx1:bbx2]

    # 调整lambda（基于实际裁剪面积）
    lam = 1 - ((bbx2 - bbx1) * (bby2 - bby1) / (W * H))

    y_a, y_b = y, y[index]

    return x, y_a, y_b, lam

from torchvision.transforms import AutoAugment, AutoAugmentPolicy

# 使用预定义的AutoAugment策略
auto_augment = transforms.Compose([
    transforms.RandomResizedCrop(224),
    transforms.AutoAugment(policy=AutoAugmentPolicy.IMAGENET),  # ImageNet策略
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406],
                        std=[0.229, 0.224, 0.225])
])

# 可选策略
# - AutoAugmentPolicy.IMAGENET  # ImageNet优化
# - AutoAugmentPolicy.CIFAR10   # CIFAR-10优化
# - AutoAugmentPolicy.SVHN      # SVHN优化

from torchvision.transforms import RandAugment

# N=2: 随机选择2个变换
# M=9: 变换强度（0-30）
rand_augment = transforms.Compose([
    transforms.RandomResizedCrop(224),
    transforms.RandAugment(num_ops=2, magnitude=9),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406],
                        std=[0.229, 0.224, 0.225])
])

from torch.optim.lr_scheduler import StepLR

optimizer = torch.optim.SGD(model.parameters(), lr=0.1)
scheduler = StepLR(optimizer, step_size=30, gamma=0.1)

# 每30个epoch学习率乘以0.1
# 0.1 -> 0.01 -> 0.001 -> ...

from torch.optim.lr_scheduler import CosineAnnealingLR

scheduler = CosineAnnealingLR(optimizer, T_max=100, eta_min=1e-6)

# 学习率从初始值余弦衰减到eta_min
# lr = eta_min + 0.5 * (lr_init - eta_min) * (1 + cos(π * t / T_max))

from torch.optim.lr_scheduler import CosineAnnealingWarmRestarts

scheduler = CosineAnnealingWarmRestarts(optimizer, T_0=10, T_mult=2)

# 周期性重启：10, 20, 40, 80... epochs

from torch.optim.lr_scheduler import OneCycleLR

scheduler = OneCycleLR(
    optimizer,
    max_lr=0.1,
    steps_per_epoch=len(train_loader),
    epochs=100,
    pct_start=0.3,  # 前30%用于warmup
    anneal_strategy='cos'
)

# 每个batch后更新
for batch in train_loader:
    # ... 训练 ...
    scheduler.step()

class WarmupCosineScheduler:
    def __init__(self, optimizer, warmup_epochs, total_epochs, min_lr=1e-6):
        self.optimizer = optimizer
        self.warmup_epochs = warmup_epochs
        self.total_epochs = total_epochs
        self.min_lr = min_lr
        self.base_lr = optimizer.param_groups[0]['lr']

    def step(self, epoch):
        if epoch < self.warmup_epochs:
            # 线性warmup
            lr = self.base_lr * (epoch + 1) / self.warmup_epochs
        else:
            # 余弦衰减
            progress = (epoch - self.warmup_epochs) / (self.total_epochs - self.warmup_epochs)
            lr = self.min_lr + 0.5 * (self.base_lr - self.min_lr) * (1 + math.cos(math.pi * progress))

        for param_group in self.optimizer.param_groups:
            param_group['lr'] = lr

        return lr

# 使用
scheduler = WarmupCosineScheduler(optimizer, warmup_epochs=5, total_epochs=100)

for epoch in range(100):
    lr = scheduler.step(epoch)
    print(f"Epoch {epoch}, LR: {lr:.6f}")
    # ... 训练 ...

学习率变化曲线示意:

StepLR:           ___________
                           \___________
                                      \___________

CosineAnnealing:  ‾‾‾‾‾‾‾‾‾‾‾\_____________________/

OneCycle:              /\
                      /  \
                     /    \_______

Warmup+Cosine:       /‾‾‾‾‾‾‾‾‾‾\_____________________

optimizer = torch.optim.AdamW(
    model.parameters(),
    lr=1e-3,
    weight_decay=0.01  # L2正则化系数
)

# 注意：AdamW是解耦的weight decay，推荐使用
# SGD + weight_decay = L2正则化
# Adam + weight_decay ≠ 严格的L2正则化
# AdamW = 正确的weight decay实现

class ModelWithDropout(nn.Module):
    def __init__(self, num_classes=1000):
        super().__init__()
        self.features = nn.Sequential(
            # ... 卷积层 ...
        )
        self.classifier = nn.Sequential(
            nn.Dropout(p=0.5),  # 训练时随机丢弃50%
            nn.Linear(2048, 512),
            nn.ReLU(),
            nn.Dropout(p=0.3),
            nn.Linear(512, num_classes)
        )

    def forward(self, x):
        x = self.features(x)
        x = x.flatten(1)
        x = self.classifier(x)
        return x

class LabelSmoothingLoss(nn.Module):
    def __init__(self, num_classes, smoothing=0.1):
        super().__init__()
        self.num_classes = num_classes
        self.smoothing = smoothing
        self.confidence = 1.0 - smoothing

    def forward(self, pred, target):
        """
        pred: (B, num_classes) - 未归一化的logits
        target: (B,) - 类别索引
        """
        pred = pred.log_softmax(dim=-1)

        with torch.no_grad():
            # 创建平滑标签
            true_dist = torch.zeros_like(pred)
            true_dist.fill_(self.smoothing / (self.num_classes - 1))
            true_dist.scatter_(1, target.unsqueeze(1), self.confidence)

        return torch.mean(torch.sum(-true_dist * pred, dim=-1))

# 使用
criterion = LabelSmoothingLoss(num_classes=1000, smoothing=0.1)

# 或使用PyTorch内置
criterion = nn.CrossEntropyLoss(label_smoothing=0.1)

class StochasticDepthResBlock(nn.Module):
    """随机深度残差块"""

    def __init__(self, in_channels, out_channels, survival_prob=0.8):
        super().__init__()
        self.survival_prob = survival_prob
        self.conv1 = nn.Conv2d(in_channels, out_channels, 3, padding=1)
        self.bn1 = nn.BatchNorm2d(out_channels)
        self.conv2 = nn.Conv2d(out_channels, out_channels, 3, padding=1)
        self.bn2 = nn.BatchNorm2d(out_channels)

    def forward(self, x):
        residual = x

        if self.training:
            # 训练时随机跳过
            if torch.rand(1) > self.survival_prob:
                return residual

        out = F.relu(self.bn1(self.conv1(x)))
        out = self.bn2(self.conv2(out))

        if self.training:
            out = out / self.survival_prob  # 缩放补偿

        return F.relu(out + residual)

from torch.cuda.amp import autocast, GradScaler

# 初始化
model = Model().cuda()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
scaler = GradScaler()

for epoch in range(num_epochs):
    for images, labels in train_loader:
        images, labels = images.cuda(), labels.cuda()

        optimizer.zero_grad()

        # 自动混合精度前向传播
        with autocast():
            outputs = model(images)
            loss = criterion(outputs, labels)

        # 缩放反向传播
        scaler.scale(loss).backward()

        # 优化器步骤
        scaler.step(optimizer)
        scaler.update()

import torch
import torch.nn as nn
from torch.cuda.amp import autocast, GradScaler
from torch.utils.data import DataLoader
import albumentations as A
from albumentations.pytorch import ToTensorV2
from tqdm import tqdm
import wandb

class Trainer:
    def __init__(
        self,
        model,
        train_loader,
        val_loader,
        optimizer,
        criterion,
        scheduler,
        config
    ):
        self.model = model.cuda()
        self.train_loader = train_loader
        self.val_loader = val_loader
        self.optimizer = optimizer
        self.criterion = criterion
        self.scheduler = scheduler
        self.config = config

        self.scaler = GradScaler() if config.use_amp else None
        self.best_acc = 0.0

    def train_epoch(self, epoch):
        self.model.train()
        total_loss = 0
        correct = 0
        total = 0

        pbar = tqdm(self.train_loader, desc=f"Epoch {epoch}")
        for images, labels in pbar:
            images, labels = images.cuda(), labels.cuda()

            # Mixup (可选)
            if self.config.use_mixup:
                images, labels_a, labels_b, lam = mixup_data(images, labels)

            self.optimizer.zero_grad()

            # 混合精度前向
            if self.config.use_amp:
                with autocast():
                    outputs = self.model(images)
                    if self.config.use_mixup:
                        loss = mixup_criterion(self.criterion, outputs, labels_a, labels_b, lam)
                    else:
                        loss = self.criterion(outputs, labels)

                self.scaler.scale(loss).backward()
                self.scaler.step(self.optimizer)
                self.scaler.update()
            else:
                outputs = self.model(images)
                if self.config.use_mixup:
                    loss = mixup_criterion(self.criterion, outputs, labels_a, labels_b, lam)
                else:
                    loss = self.criterion(outputs, labels)
                loss.backward()
                self.optimizer.step()

            total_loss += loss.item()
            _, predicted = outputs.max(1)
            total += labels.size(0)
            correct += predicted.eq(labels).sum().item()

            pbar.set_postfix({
                'loss': f'{loss.item():.4f}',
                'acc': f'{100.*correct/total:.2f}%'
            })

        return total_loss / len(self.train_loader), correct / total

    @torch.no_grad()
    def validate(self):
        self.model.eval()
        total_loss = 0
        correct = 0
        total = 0

        for images, labels in self.val_loader:
            images, labels = images.cuda(), labels.cuda()

            outputs = self.model(images)
            loss = self.criterion(outputs, labels)

            total_loss += loss.item()
            _, predicted = outputs.max(1)
            total += labels.size(0)
            correct += predicted.eq(labels).sum().item()

        return total_loss / len(self.val_loader), correct / total

    def train(self):
        for epoch in range(self.config.epochs):
            # 训练
            train_loss, train_acc = self.train_epoch(epoch)

            # 验证
            val_loss, val_acc = self.validate()

            # 更新学习率
            self.scheduler.step()
            current_lr = self.optimizer.param_groups[0]['lr']

            # 日志
            print(f"Epoch {epoch}: "
                  f"Train Loss: {train_loss:.4f}, Train Acc: {train_acc*100:.2f}%, "
                  f"Val Loss: {val_loss:.4f}, Val Acc: {val_acc*100:.2f}%, "
                  f"LR: {current_lr:.6f}")

            # W&B记录
            if self.config.use_wandb:
                wandb.log({
                    'epoch': epoch,
                    'train_loss': train_loss,
                    'train_acc': train_acc,
                    'val_loss': val_loss,
                    'val_acc': val_acc,
                    'lr': current_lr
                })

            # 保存最佳模型
            if val_acc > self.best_acc:
                self.best_acc = val_acc
                torch.save({
                    'epoch': epoch,
                    'model_state_dict': self.model.state_dict(),
                    'optimizer_state_dict': self.optimizer.state_dict(),
                    'best_acc': self.best_acc,
                }, 'best_model.pth')
                print(f"✓ Best model saved (Acc: {val_acc*100:.2f}%)")

# 配置
class Config:
    epochs = 100
    batch_size = 32
    lr = 0.001
    weight_decay = 0.01
    warmup_epochs = 5
    use_amp = True
    use_mixup = True
    use_wandb = True

# 使用
config = Config()

# 数据增强
train_transforms = A.Compose([
    A.RandomResizedCrop(224, 224, scale=(0.8, 1.0)),
    A.HorizontalFlip(p=0.5),
    A.RandAugment(num_ops=2, magnitude=9),
    A.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
    ToTensorV2()
])

# 模型
model = timm.create_model('resnet50', pretrained=True, num_classes=num_classes)

# 优化器
optimizer = torch.optim.AdamW(model.parameters(), lr=config.lr, weight_decay=config.weight_decay)

# 调度器
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=config.epochs)

# 损失函数
criterion = nn.CrossEntropyLoss(label_smoothing=0.1)

# 训练
trainer = Trainer(model, train_loader, val_loader, optimizer, criterion, scheduler, config)
trainer.train()