AI模型训练优化实战指南：从损失函数到分布式训练

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

# 1. Focal Loss（解决类别不平衡）
class FocalLoss(nn.Module):
    def __init__(self, alpha=0.25, gamma=2.0):
        super().__init__()
        self.alpha = alpha
        self.gamma = gamma

    def forward(self, pred, target):
        ce_loss = F.cross_entropy(pred, target, reduction='none')
        pt = torch.exp(-ce_loss)
        loss = self.alpha * (1 - pt) ** self.gamma * ce_loss
        return loss.mean()

# 2. 类别加权交叉熵
class WeightedCrossEntropy(nn.Module):
    def __init__(self, class_weights):
        super().__init__()
        self.class_weights = torch.tensor(class_weights, dtype=torch.float32)

    def forward(self, pred, target):
        return F.cross_entropy(pred, target, weight=self.class_weights.to(pred.device))

# 3. Label Smoothing（防止过自信）
def label_smoothing_loss(pred, target, smoothing=0.1):
    n_classes = pred.size(-1)
    log_preds = F.log_softmax(pred, dim=-1)
    nll_loss = F.nll_loss(log_preds, target, reduction='none')
    smooth_loss = -log_preds.mean(dim=-1)
    return (1 - smoothing) * nll_loss + smoothing * smooth_loss

from torch.optim.lr_scheduler import (
    StepLR, CosineAnnealingLR, CosineAnnealingWarmRestarts,
    OneCycleLR, ReduceLROnPlateau, LambdaLR
)

# 1. Warmup + Cosine（Transformer 标配）
def get_cosine_schedule_with_warmup(optimizer, num_warmup_steps, num_training_steps):
    def lr_lambda(current_step):
        if current_step < num_warmup_steps:
            return float(current_step) / float(max(1, num_warmup_steps))
        progress = float(current_step - num_warmup_steps) / float(max(1, num_training_steps - num_warmup_steps))
        return max(0.0, 0.5 * (1.0 + math.cos(math.pi * progress)))
    return LambdaLR(optimizer, lr_lambda)

# 2. OneCycleLR（快速收敛）
scheduler = OneCycleLR(
    optimizer,
    max_lr=1e-3,
    total_steps=num_training_steps,
    pct_start=0.3,  # 30% 时间升学习率
    anneal_strategy='cos'
)

# 3. ReduceLROnPlateau（验证集不降就降学习率）
scheduler = ReduceLROnPlateau(
    optimizer,
    mode='min',
    factor=0.5,   # 学习率乘以 0.5
    patience=3,   # 3 个 epoch 不降就触发
    min_lr=1e-7
)

# 在优化器中加权重衰减（等价于 L2 正则化）
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3, weight_decay=1e-4)

# 自定义 L1 正则化
def l1_regularization(model, lambda_l1=0.01):
    l1_loss = sum(p.abs().sum() for p in model.parameters())
    return lambda_l1 * l1_loss

loss = task_loss + l1_regularization(model)

class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(768, 256)
        self.dropout = nn.Dropout(0.5)
        self.fc2 = nn.Linear(256, 10)

    def forward(self, x):
        x = F.relu(self.fc1(x))
        x = self.dropout(x)  # 训练时随机失活 50%
        return self.fc2(x)

# Min-Max：线性映射到 [0, 1]，适合分布均匀、无明显离群值
def minmax_normalize(X, feature_range=(0, 1)):
    min_val, max_val = X.min(axis=0), X.max(axis=0)
    normalized = (X - min_val) / (max_val - min_val)
    return normalized * (feature_range[1] - feature_range[0]) + feature_range[0]

# Z-Score：减均值除标准差，对离群值更鲁棒（基于分布统计量而非极值）
def zscore_normalize(X):
    return (X - X.mean(axis=0)) / X.std(axis=0)

# 训练时保存统计量，测试时直接复用（不要重新算）
train_mean, train_std = X_train.mean(axis=0), X_train.std(axis=0)
X_train_norm = (X_train - train_mean) / (train_std + 1e-8)
X_test_norm  = (X_test  - train_mean) / (train_std + 1e-8)  # 用训练集统计量

def batch_norm(x, gamma, beta, running_mean, running_var, eps=1e-5, momentum=0.1, training=True):
    if training:
        mean = x.mean(dim=(0, 2, 3), keepdim=True)   # 跨 (B, H, W) 按通道求均值
        var = x.var(dim=(0, 2, 3), keepdim=True)
        running_mean = momentum * running_mean + (1 - momentum) * mean.data  # 更新全局统计量
        running_var  = momentum * running_var  + (1 - momentum) * var.data
    else:
        mean, var = running_mean, running_var       # 推理用固定的全局统计量
    x_norm = (x - mean) / torch.sqrt(var + eps)
    return gamma * x_norm + beta

model.train()  # 训练：统计量会更新
model.eval()   # 推理：统计量固定

# 多 GPU 训练时把普通 BN 换成跨卡同步版本
from torch.nn import SyncBatchNorm
model = SyncBatchNorm.convert_sync_batchnorm(model)

def layer_norm(x, gamma, beta, eps=1e-5):
    mean = x.mean(dim=-1, keepdim=True)
    var = x.var(dim=-1, keepdim=True)
    return gamma * (x - mean) / torch.sqrt(var + eps) + beta

# GroupNorm：64 通道分 8 组，每组 8 个通道一起归一化
self.gn = nn.GroupNorm(num_groups=8, num_channels=64)

import albumentations as A

transform = A.Compose([
    A.RandomResizedCrop(224, 224),
    A.HorizontalFlip(p=0.5),
    A.ColorJitter(brightness=0.2, contrast=0.2),
    A.GaussianBlur(p=0.3),
    A.Normalize(),
    A.CoarseDropout(max_holes=8, max_height=32, max_width=32, p=0.5),
])

def mixup_data(x, y, alpha=0.2):
    lam = np.random.beta(alpha, alpha)
    batch_size = x.size(0)
    index = torch.randperm(batch_size)
    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):
    return lam * criterion(pred, y_a) + (1 - lam) * criterion(pred, y_b)

class EarlyStopping:
    def __init__(self, patience=5, min_delta=0):
        self.patience = patience
        self.min_delta = min_delta
        self.counter = 0
        self.best_loss = float('inf')
        self.early_stop = False

    def __call__(self, val_loss):
        if val_loss < self.best_loss - self.min_delta:
            self.best_loss = val_loss
            self.counter = 0
        else:
            self.counter += 1
            if self.counter >= self.patience:
                self.early_stop = True

# 使用
early_stopping = EarlyStopping(patience=5)

for epoch in range(num_epochs):
    train_loss = train_one_epoch()
    val_loss = validate()

    early_stopping(val_loss)
    if early_stopping.early_stop:
        print(f"Early stopping at epoch {epoch}")
        break

accumulation_steps = 4  # 等效 batch size = batch_size * 4

optimizer.zero_grad()

for i, batch in enumerate(dataloader):
    outputs = model(batch)
    loss = criterion(outputs, targets) / accumulation_steps  # 关键：除以累积步数
    loss.backward()

    if (i + 1) % accumulation_steps == 0:
        optimizer.step()
        optimizer.zero_grad()

# 防止梯度爆炸
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)

# 或者按值裁剪
torch.nn.utils.clip_grad_value_(model.parameters(), clip_value=0.5)

# SGD（带动量）
optimizer = torch.optim.SGD(model.parameters(), lr=0.01, momentum=0.9)

# Adam（自适应）
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

# AdamW（推荐 Transformer 用）
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4, weight_decay=0.01)

from torch.cuda.amp import autocast, GradScaler

scaler = GradScaler()

for batch in dataloader:
    optimizer.zero_grad()

    # 前向传播用混合精度
    with autocast():
        outputs = model(batch)
        loss = criterion(outputs, targets)

    # 反向传播用梯度缩放
    scaler.scale(loss).backward()
    scaler.unscale_(optimizer)
    torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
    scaler.step(optimizer)
    scaler.update()

# BF16 更稳定，推荐 A100/H100 用
model = model.to(torch.bfloat16)

# FP16 需要 GradScaler
model = model.to(torch.float16)

import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP

def setup(rank, world_size):
    dist.init_process_group("nccl", rank=rank, world_size=world_size)

def train(rank, world_size):
    setup(rank, world_size)
    torch.cuda.set_device(rank)

    model = Model().to(rank)
    model = DDP(model, device_ids=[rank])

    optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4)

    for epoch in range(num_epochs):
        for batch in dataloader:
            outputs = model(batch)
            loss = criterion(outputs, targets)
            loss.backward()
            optimizer.step()
            optimizer.zero_grad()

    dist.destroy_process_group()

# 启动
torchrun --nproc_per_node=4 train.py

{
    "zero_optimization": {
        "stage": 2,
        "offload_optimizer": {
            "device": "cpu"
        },
        "allgather_partitions": true,
        "allgather_bucket_size": 5e8,
        "overlap_comm": true,
        "contiguous_gradients": true
    },
    "fp16": {
        "enabled": true,
        "loss_scale": 0,
        "loss_scale_window": 1000
    }
}

def save_checkpoint(model, optimizer, scheduler, epoch, path):
    torch.save({
        'epoch': epoch,
        'model_state_dict': model.state_dict(),
        'optimizer_state_dict': optimizer.state_dict(),
        'scheduler_state_dict': scheduler.state_dict(),
    }, path)

def load_checkpoint(model, optimizer, scheduler, path):
    checkpoint = torch.load(path)
    model.load_state_dict(checkpoint['model_state_dict'])
    optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
    scheduler.load_state_dict(checkpoint['scheduler_state_dict'])
    return checkpoint['epoch']

# 训练循环
start_epoch = 0
if resume:
    start_epoch = load_checkpoint(model, optimizer, scheduler, 'checkpoint.pth')

for epoch in range(start_epoch, num_epochs):
    train()
    if epoch % save_interval == 0:
        save_checkpoint(model, optimizer, scheduler, epoch, 'checkpoint.pth')

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

        self.qkv = nn.Linear(embed_dim, embed_dim * 3)
        self.proj = nn.Linear(embed_dim, embed_dim)

    def forward(self, x):
        B, N, C = x.shape
        qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, self.head_dim).permute(2, 0, 3, 1, 4)
        q, k, v = qkv[0], qkv[1], qkv[2]

        attn = (q @ k.transpose(-2, -1)) * (self.head_dim ** -0.5)
        attn = attn.softmax(dim=-1)
        x = (attn @ v).transpose(1, 2).reshape(B, N, C)
        return self.proj(x)

# SimCLR（对比学习）
class SimCLR(nn.Module):
    def __init__(self, backbone, projection_dim=128):
        super().__init__()
        self.backbone = backbone
        self.projector = nn.Sequential(
            nn.Linear(backbone.dim, backbone.dim),
            nn.ReLU(),
            nn.Linear(backbone.dim, projection_dim)
        )

    def forward(self, x1, x2):
        h1 = self.backbone(x1)
        h2 = self.backbone(x2)
        z1 = self.projector(h1)
        z2 = self.projector(h2)
        return z1, z2

def nt_xent_loss(z1, z2, temperature=0.5):
    batch_size = z1.shape[0]
    z = torch.cat([z1, z2], dim=0)
    z = F.normalize(z, dim=-1)

    sim = torch.matmul(z, z.T) / temperature
    sim_i_j = torch.diag(sim, batch_size)
    sim_j_i = torch.diag(sim, -batch_size)

    positive_samples = torch.cat([sim_i_j, sim_j_i], dim=0)
    mask = (~torch.eye(2 * batch_size, dtype=torch.bool, device=z.device)).float()
    negative_samples = sim * mask

    logits = torch.cat([positive_samples.unsqueeze(1), negative_samples], dim=1)
    labels = torch.zeros(2 * batch_size, dtype=torch.long, device=z.device)
    return F.cross_entropy(logits, labels)

# FGSM（Fast Gradient Sign Method）
def fgsm_attack(model, x, y, epsilon):
    x.requires_grad = True
    output = model(x)
    loss = F.cross_entropy(output, y)
    loss.backward()

    perturbation = epsilon * x.grad.sign()
    adversarial_x = x + perturbation
    return adversarial_x.detach()

# PGD（Project Gradient Descent，更强的攻击）
def pgd_attack(model, x, y, epsilon, alpha, num_steps):
    x_adv = x.clone().detach()
    for _ in range(num_steps):
        x_adv.requires_grad = True
        output = model(x_adv)
        loss = F.cross_entropy(output, y)
        loss.backward()

        perturbation = alpha * x_adv.grad.sign()
        x_adv = x_adv + perturbation
        x_adv = torch.clamp(x_adv, x - epsilon, x + epsilon)
        x_adv = x_adv.detach()

    return x_adv

# 对抗训练
for x, y in dataloader:
    x_adv = pgd_attack(model, x, y, epsilon=0.03, alpha=0.01, num_steps=7)
    # 用对抗样本训练
    outputs = model(x_adv)
    loss = criterion(outputs, y)

def distillation_loss(student_logits, teacher_logits, temperature=4.0):
    soft_targets = F.softmax(teacher_logits / temperature, dim=-1)
    soft_prob = F.log_softmax(student_logits / temperature, dim=-1)
    return F.kl_div(soft_prob, soft_targets, reduction='batchmean') * (temperature ** 2)

# 训练
for x, y in dataloader:
    with torch.no_grad():
        teacher_logits = teacher_model(x)
    student_logits = student_model(x)

    # 硬标签 + 软标签
    hard_loss = F.cross_entropy(student_logits, y)
    soft_loss = distillation_loss(student_logits, teacher_logits)
    loss = 0.5 * hard_loss + 0.5 * soft_loss

class MultiTaskModel(nn.Module):
    def __init__(self, backbone):
        super().__init__()
        self.backbone = backbone
        self.classification_head = nn.Linear(backbone.dim, num_classes)
        self.regression_head = nn.Linear(backbone.dim, 1)

    def forward(self, x):
        features = self.backbone(x)
        return {
            'classification': self.classification_head(features),
            'regression': self.regression_head(features)
        }

# 动态权重
class UncertaintyWeighting(nn.Module):
    def __init__(self, num_tasks):
        super().__init__()
        self.log_vars = nn.Parameter(torch.zeros(num_tasks))

    def forward(self, losses):
        total = 0
        for i, loss in enumerate(losses):
            precision = torch.exp(-self.log_vars[i])
            total += precision * loss + self.log_vars[i]
        return total

class CurriculumSampler:
    """从简单到困难的课程学习"""
    def __init__(self, dataset, difficulty_fn):
        self.difficulties = [difficulty_fn(item) for item in dataset]
        self.indices = list(range(len(dataset)))

    def get_subset(self, epoch, total_epochs):
        """根据 epoch 返回不同难度的子集"""
        progress = epoch / total_epochs
        max_difficulty = self.difficulties.quantile(progress)
        return [i for i, d in zip(self.indices, self.difficulties) if d <= max_difficulty]

def uncertainty_sampling(model, unlabeled_pool, n_samples):
    """选择模型最不确定的样本标注"""
    model.eval()
    uncertainties = []

    with torch.no_grad():
        for x in unlabeled_pool:
            output = model(x)
            prob = F.softmax(output, dim=-1)
            entropy = -(prob * torch.log(prob + 1e-8)).sum()
            uncertainties.append(entropy)

    # 选不确定性最高的 n 个
    top_indices = torch.topk(torch.tensor(uncertainties), n_samples).indices
    return [unlabeled_pool[i] for i in top_indices]

import torch_geometric.nn as gnn

class GCN(nn.Module):
    def __init__(self, in_channels, hidden_channels, num_classes):
        super().__init__()
        self.conv1 = gnn.GCNConv(in_channels, hidden_channels)
        self.conv2 = gnn.GCNConv(hidden_channels, num_classes)

    def forward(self, x, edge_index):
        x = F.relu(self.conv1(x, edge_index))
        x = F.dropout(x, p=0.5, training=self.training)
        x = self.conv2(x, edge_index)
        return x

class GAT(nn.Module):
    def __init__(self, in_channels, hidden_channels, num_classes, heads=8):
        super().__init__()
        self.conv1 = gnn.GATConv(in_channels, hidden_channels, heads=heads)
        self.conv2 = gnn.GATConv(hidden_channels * heads, num_classes, heads=1)

    def forward(self, x, edge_index):
        x = F.elu(self.conv1(x, edge_index))
        x = self.conv2(x, edge_index)
        return x

import gensim
from gensim.models import Word2Vec

# 训练 Word2Vec
sentences = [['hello', 'world'], ['foo', 'bar']]
model = Word2Vec(
    sentences,
    vector_size=100,
    window=5,
    min_count=1,
    workers=4,
    sg=0  # 0=CBOW, 1=Skip-gram
)

# 使用
vector = model.wv['hello']
similar = model.wv.most_similar('hello')

class MoELayer(nn.Module):
    def __init__(self, experts, gate):
        super().__init__()
        self.experts = nn.ModuleList(experts)
        self.gate = gate  # 路由网络

    def forward(self, x):
        gate_scores = self.gate(x)  # [batch, num_experts]
        top_k_scores, top_k_indices = torch.topk(gate_scores, k=2, dim=-1)

        output = torch.zeros_like(x)
        for i, expert in enumerate(self.experts):
            mask = (top_k_indices == i).any(dim=-1)
            if mask.any():
                expert_input = x[mask]
                expert_output = expert(expert_input)
                weights = top_k_scores[mask].max(dim=-1).values.unsqueeze(-1)
                output[mask] += expert_output * weights

        return output

# 完整训练配置模板
config = {
    # 模型
    'model_name': 'resnet50',
    'pretrained': True,
    'num_classes': 1000,

    # 数据
    'batch_size': 64,
    'num_workers': 8,
    'pin_memory': True,

    # 训练
    'epochs': 100,
    'learning_rate': 1e-4,
    'weight_decay': 0.01,
    'warmup_steps': 1000,
    'max_grad_norm': 1.0,

    # 优化
    'optimizer': 'AdamW',
    'scheduler': 'cosine_with_warmup',
    'mixed_precision': True,
    'gradient_accumulation_steps': 1,

    # 正则化
    'dropout': 0.1,
    'label_smoothing': 0.1,

    # 早停
    'early_stopping_patience': 5,

    # 保存
    'save_interval': 5,
    'save_top_k': 3,
}