深度学习调试完全指南：从训练失败诊断到性能优化

def diagnose_training(model, train_loader, val_loader, optimizer, loss_fn, device):
    """
    训练开始前执行快速诊断的函数
    """
    print("=== 深度学习训练诊断检查清单 ===\n")

    # 1. 数据验证
    print("[1] 正在验证数据...")
    batch = next(iter(train_loader))
    X, y = batch
    print(f"  输入形状: {X.shape}")
    print(f"  标签形状: {y.shape}")
    print(f"  输入范围: [{X.min():.4f}, {X.max():.4f}]")
    print(f"  输入中是否存在 NaN: {torch.isnan(X).any()}")
    print(f"  输入中是否存在 Inf: {torch.isinf(X).any()}")

    # 2. 模型参数验证
    print("\n[2] 正在验证模型参数...")
    total_params = sum(p.numel() for p in model.parameters())
    trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
    print(f"  总参数量: {total_params:,}")
    print(f"  可训练参数量: {trainable_params:,}")

    # 3. 前向传播测试
    print("\n[3] 正在测试前向传播...")
    model.eval()
    with torch.no_grad():
        try:
            output = model(X.to(device))
            print(f"  输出形状: {output.shape}")
            print(f"  输出中是否存在 NaN: {torch.isnan(output).any()}")
            loss = loss_fn(output, y.to(device))
            print(f"  初始 Loss: {loss.item():.4f}")
        except Exception as e:
            print(f"  前向传播失败: {e}")

    # 4. 反向传播测试
    print("\n[4] 正在测试反向传播...")
    model.train()
    optimizer.zero_grad()
    output = model(X.to(device))
    loss = loss_fn(output, y.to(device))
    loss.backward()

    # 梯度验证
    grad_norms = []
    for name, param in model.named_parameters():
        if param.grad is not None:
            grad_norms.append((name, param.grad.norm().item()))

    if grad_norms:
        print("  各层梯度范数（前 5 名）:")
        for name, norm in sorted(grad_norms, key=lambda x: x[1], reverse=True)[:5]:
            print(f"    {name}: {norm:.6f}")

    print("\n诊断完成！")

import torch
import torch.nn as nn
import torch.optim as optim
import numpy as np
import matplotlib.pyplot as plt


def find_learning_rate(model, train_loader, loss_fn, device,
                        start_lr=1e-7, end_lr=10, num_iter=100):
    """
    通过 LR Range Test 寻找最优学习率范围。
    """
    optimizer = optim.SGD(model.parameters(), lr=start_lr)

    lr_multiplier = (end_lr / start_lr) ** (1 / num_iter)

    lrs = []
    losses = []
    best_loss = float('inf')

    model.train()
    data_iter = iter(train_loader)

    for i in range(num_iter):
        try:
            X, y = next(data_iter)
        except StopIteration:
            data_iter = iter(train_loader)
            X, y = next(data_iter)

        X, y = X.to(device), y.to(device)

        optimizer.zero_grad()
        output = model(X)
        loss = loss_fn(output, y)

        if torch.isnan(loss) or loss.item() > best_loss * 4:
            print(f"检测到 Loss 爆炸 at lr={optimizer.param_groups[0]['lr']:.2e}")
            break

        if loss.item() < best_loss:
            best_loss = loss.item()

        lrs.append(optimizer.param_groups[0]['lr'])
        losses.append(loss.item())

        loss.backward()
        optimizer.step()

        for pg in optimizer.param_groups:
            pg['lr'] *= lr_multiplier

    plt.figure(figsize=(10, 4))
    plt.plot(lrs, losses)
    plt.xscale('log')
    plt.xlabel('Learning Rate')
    plt.ylabel('Loss')
    plt.title('LR Range Test')
    plt.grid(True)
    plt.savefig('lr_range_test.png')
    plt.show()

    return lrs, losses


def safe_training_step(model, X, y, optimizer, loss_fn, scaler=None):
    """
    检测并跳过 NaN Loss 的安全训练步骤
    """
    optimizer.zero_grad()

    # 输入验证
    if torch.isnan(X).any() or torch.isinf(X).any():
        print("警告: 输入中存在 NaN/Inf，跳过该步骤")
        return None

    if scaler is not None:
        with torch.cuda.amp.autocast():
            output = model(X)
            loss = loss_fn(output, y)

        if torch.isnan(loss) or torch.isinf(loss):
            print(f"警告: Loss 为 {loss.item()}，跳过该步骤")
            return None

        scaler.scale(loss).backward()
        scaler.unscale_(optimizer)
        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
        scaler.step(optimizer)
        scaler.update()
    else:
        output = model(X)
        loss = loss_fn(output, y)

        if torch.isnan(loss) or torch.isinf(loss):
            print(f"警告: Loss 为 {loss.item()}，跳过该步骤")
            return None

        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
        optimizer.step()

    return loss.item()

# 错误做法: 可能出现 log(0)
def bad_cross_entropy(pred, target):
    return -torch.sum(target * torch.log(pred))

# 正确做法: 用 eps 保证数值稳定性
def safe_cross_entropy(pred, target, eps=1e-8):
    pred = torch.clamp(pred, min=eps, max=1-eps)
    return -torch.sum(target * torch.log(pred))

# 更好的做法: 使用 PyTorch 内置函数（内部应用了 log-sum-exp 技巧）
loss_fn = nn.CrossEntropyLoss()  # 数值稳定
log_softmax = nn.LogSoftmax(dim=1)  # log 与 softmax 的结合

# 自定义对数损失（数值稳定）
def numerically_stable_log_loss(logits, targets):
    """
    利用 log-sum-exp 技巧实现的数值稳定交叉熵
    """
    # F.cross_entropy 内部就应用了这一技巧
    import torch.nn.functional as F
    return F.cross_entropy(logits, targets)

import torch

# 启用异常检测模式（仅在开发/调试时使用）
# 会拖慢性能，生产环境中应关闭
with torch.autograd.detect_anomaly():
    output = model(X)
    loss = loss_fn(output, y)
    loss.backward()  # 一旦此处出现 NaN/Inf，会打印出准确的位置

# 或者全局设置
torch.autograd.set_detect_anomaly(True)

# 应用到训练循环中的示例
def train_with_anomaly_detection(model, loader, optimizer, loss_fn, device, epochs=5):
    model.train()
    for epoch in range(epochs):
        for batch_idx, (X, y) in enumerate(loader):
            X, y = X.to(device), y.to(device)

            with torch.autograd.detect_anomaly():
                optimizer.zero_grad()
                output = model(X)
                loss = loss_fn(output, y)

                if torch.isnan(loss):
                    print(f"NaN Loss at epoch {epoch}, batch {batch_idx}")
                    # 打印输入状态
                    print(f"Input stats: mean={X.mean():.4f}, std={X.std():.4f}")
                    print(f"Output stats: mean={output.mean():.4f}, std={output.std():.4f}")
                    break

                loss.backward()
                optimizer.step()

import torch
import torch.nn as nn
import matplotlib.pyplot as plt


def check_gradient_flow(model):
    """
    通过可视化各层的梯度大小来诊断消失/爆炸
    """
    ave_grads = []
    max_grads = []
    layers = []

    for name, param in model.named_parameters():
        if param.requires_grad and param.grad is not None:
            layers.append(name)
            ave_grads.append(param.grad.abs().mean().item())
            max_grads.append(param.grad.abs().max().item())

    plt.figure(figsize=(12, 6))
    plt.bar(range(len(ave_grads)), ave_grads, alpha=0.5, label='平均梯度')
    plt.bar(range(len(max_grads)), max_grads, alpha=0.5, label='最大梯度')
    plt.xticks(range(len(layers)), layers, rotation=90)
    plt.xlabel("层")
    plt.ylabel("梯度大小")
    plt.title("各层梯度流动情况")
    plt.legend()
    plt.yscale('log')
    plt.tight_layout()
    plt.savefig('gradient_flow.png')

    # 检测梯度消失
    for name, avg_grad in zip(layers, ave_grads):
        if avg_grad < 1e-6:
            print(f"警告: {name} 层可能存在梯度消失（avg={avg_grad:.2e}）")

    return layers, ave_grads, max_grads


def register_gradient_hooks(model):
    """
    注册梯度钩子以进行实时监控
    """
    gradient_stats = {}

    def make_hook(name):
        def hook(grad):
            gradient_stats[name] = {
                'mean': grad.abs().mean().item(),
                'max': grad.abs().max().item(),
                'std': grad.std().item(),
                'has_nan': torch.isnan(grad).any().item(),
                'has_inf': torch.isinf(grad).any().item()
            }
            if torch.isnan(grad).any():
                print(f"检测到 NaN 梯度: {name}")
            return grad
        return hook

    hooks = []
    for name, param in model.named_parameters():
        if param.requires_grad:
            hook = param.register_hook(make_hook(name))
            hooks.append(hook)

    return gradient_stats, hooks


# 解决梯度消失: He 初始化 + BatchNorm + 残差连接
class ResidualBlock(nn.Module):
    def __init__(self, dim):
        super().__init__()
        self.block = nn.Sequential(
            nn.Linear(dim, dim),
            nn.BatchNorm1d(dim),
            nn.ReLU(),
            nn.Linear(dim, dim),
            nn.BatchNorm1d(dim)
        )
        self.relu = nn.ReLU()

    def forward(self, x):
        return self.relu(self.block(x) + x)  # 残差连接


# 解决梯度消失: Xavier/He 初始化
def init_weights(m):
    if isinstance(m, nn.Linear):
        nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
        if m.bias is not None:
            nn.init.zeros_(m.bias)
    elif isinstance(m, nn.Conv2d):
        nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
        if m.bias is not None:
            nn.init.zeros_(m.bias)
    elif isinstance(m, nn.BatchNorm2d):
        nn.init.ones_(m.weight)
        nn.init.zeros_(m.bias)

# 通过 model.apply(init_weights) 应用

import torch
import torch.nn as nn


def train_with_gradient_clipping(model, loader, optimizer, loss_fn, device,
                                   max_norm=1.0, epochs=10):
    """
    应用了梯度裁剪的安全训练循环
    """
    model.train()
    history = {'train_loss': [], 'grad_norm': []}

    for epoch in range(epochs):
        epoch_loss = 0
        epoch_grad_norms = []

        for X, y in loader:
            X, y = X.to(device), y.to(device)

            optimizer.zero_grad()
            output = model(X)
            loss = loss_fn(output, y)
            loss.backward()

            # 计算梯度范数（裁剪前）
            total_norm = 0
            for p in model.parameters():
                if p.grad is not None:
                    param_norm = p.grad.data.norm(2)
                    total_norm += param_norm.item() ** 2
            total_norm = total_norm ** 0.5
            epoch_grad_norms.append(total_norm)

            # 应用梯度裁剪
            torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=max_norm)

            # 或者基于数值的裁剪
            # torch.nn.utils.clip_grad_value_(model.parameters(), clip_value=0.5)

            optimizer.step()
            epoch_loss += loss.item()

        avg_loss = epoch_loss / len(loader)
        avg_grad_norm = sum(epoch_grad_norms) / len(epoch_grad_norms)

        history['train_loss'].append(avg_loss)
        history['grad_norm'].append(avg_grad_norm)

        print(f"Epoch {epoch+1}: Loss={avg_loss:.4f}, Grad Norm={avg_grad_norm:.4f}")

        if avg_grad_norm > max_norm * 10:
            print(f"警告: 梯度范数非常大（{avg_grad_norm:.4f}）。请降低学习率。")

    return history

import matplotlib.pyplot as plt
import numpy as np


def plot_learning_curves(train_losses, val_losses, train_accs=None, val_accs=None):
    """
    通过训练/验证 Loss 与准确率曲线诊断过拟合
    """
    fig, axes = plt.subplots(1, 2 if train_accs else 1, figsize=(14, 5))

    if not isinstance(axes, np.ndarray):
        axes = [axes]

    # Loss 曲线
    axes[0].plot(train_losses, label='Train Loss', color='blue')
    axes[0].plot(val_losses, label='Val Loss', color='red', linestyle='--')
    axes[0].set_xlabel('Epoch')
    axes[0].set_ylabel('Loss')
    axes[0].set_title('训练/验证 Loss')
    axes[0].legend()
    axes[0].grid(True)

    # 检测过拟合发生的时间点
    min_val_idx = np.argmin(val_losses)
    axes[0].axvline(x=min_val_idx, color='green', linestyle=':', label=f'最优 epoch: {min_val_idx}')
    axes[0].legend()

    # 准确率曲线（如果有的话）
    if train_accs and val_accs:
        axes[1].plot(train_accs, label='Train Acc', color='blue')
        axes[1].plot(val_accs, label='Val Acc', color='red', linestyle='--')
        axes[1].set_xlabel('Epoch')
        axes[1].set_ylabel('Accuracy')
        axes[1].set_title('训练/验证准确率')
        axes[1].legend()
        axes[1].grid(True)

    # 计算过拟合差距
    final_gap = val_losses[-1] - train_losses[-1]
    print(f"最终过拟合差距（Val-Train Loss）: {final_gap:.4f}")
    if final_gap > 0.1:
        print("警告: 过拟合非常严重！")

    plt.tight_layout()
    plt.savefig('learning_curves.png')
    plt.show()

class EarlyStopping:
    """
    监控验证 Loss，在出现过拟合时提前终止训练
    """
    def __init__(self, patience=10, min_delta=0.001, restore_best=True, verbose=True):
        self.patience = patience
        self.min_delta = min_delta
        self.restore_best = restore_best
        self.verbose = verbose

        self.best_loss = float('inf')
        self.best_epoch = 0
        self.counter = 0
        self.best_weights = None
        self.stopped_epoch = 0

    def __call__(self, val_loss, model, epoch):
        if val_loss < self.best_loss - self.min_delta:
            self.best_loss = val_loss
            self.best_epoch = epoch
            self.counter = 0
            if self.restore_best:
                import copy
                self.best_weights = copy.deepcopy(model.state_dict())
            if self.verbose:
                print(f"验证 Loss 有所改善: {val_loss:.6f}（epoch {epoch}）")
            return False  # 继续训练
        else:
            self.counter += 1
            if self.verbose:
                print(f"EarlyStopping counter: {self.counter}/{self.patience}")
            if self.counter >= self.patience:
                self.stopped_epoch = epoch
                if self.restore_best and self.best_weights:
                    model.load_state_dict(self.best_weights)
                    print(f"恢复最优权重（epoch {self.best_epoch}）")
                return True  # 停止训练
        return False


def train_with_regularization(model, train_loader, val_loader,
                               optimizer, loss_fn, device, epochs=100):
    """
    应用了多种正则化技巧的训练循环
    """
    early_stopping = EarlyStopping(patience=15, min_delta=0.001)

    # 学习率调度器
    scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
        optimizer, mode='min', factor=0.5, patience=5, verbose=True
    )

    train_losses = []
    val_losses = []

    for epoch in range(epochs):
        # 训练
        model.train()
        train_loss = 0
        for X, y in train_loader:
            X, y = X.to(device), y.to(device)
            optimizer.zero_grad()
            output = model(X)
            loss = loss_fn(output, y)
            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
            optimizer.step()
            train_loss += loss.item()

        # 验证
        model.eval()
        val_loss = 0
        with torch.no_grad():
            for X, y in val_loader:
                X, y = X.to(device), y.to(device)
                output = model(X)
                val_loss += loss_fn(output, y).item()

        train_loss /= len(train_loader)
        val_loss /= len(val_loader)

        train_losses.append(train_loss)
        val_losses.append(val_loss)

        scheduler.step(val_loss)

        print(f"Epoch {epoch+1}: Train={train_loss:.4f}, Val={val_loss:.4f}")

        if early_stopping(val_loss, model, epoch):
            print(f"在 epoch {epoch+1} 提前终止训练")
            break

    return train_losses, val_losses


# Dropout 与 L2 正则化示例
class RegularizedModel(nn.Module):
    def __init__(self, input_dim, hidden_dim, output_dim, dropout_rate=0.3):
        super().__init__()
        self.network = nn.Sequential(
            nn.Linear(input_dim, hidden_dim),
            nn.BatchNorm1d(hidden_dim),
            nn.ReLU(),
            nn.Dropout(p=dropout_rate),  # Dropout
            nn.Linear(hidden_dim, hidden_dim),
            nn.BatchNorm1d(hidden_dim),
            nn.ReLU(),
            nn.Dropout(p=dropout_rate),
            nn.Linear(hidden_dim, output_dim)
        )

    def forward(self, x):
        return self.network(x)

# L2 正则化（Weight Decay）在 optimizer 中设置
optimizer = torch.optim.AdamW(
    model.parameters(),
    lr=1e-3,
    weight_decay=1e-4  # L2 正则化强度
)

import torchvision.transforms as transforms
from torch.utils.data import DataLoader, Dataset


# 用于图像分类的数据增强
train_transform = transforms.Compose([
    transforms.RandomHorizontalFlip(p=0.5),
    transforms.RandomRotation(degrees=15),
    transforms.RandomResizedCrop(224, scale=(0.8, 1.0)),
    transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2, hue=0.1),
    transforms.RandomGrayscale(p=0.1),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

# 用于验证/测试（不做增强，只做归一化）
val_transform = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

# Mixup 数据增强
def mixup_data(x, y, alpha=0.2, device='cuda'):
    """
    Mixup 增强: 对两个样本做线性插值，生成新样本
    """
    if alpha > 0:
        lam = np.random.beta(alpha, alpha)
    else:
        lam = 1

    batch_size = x.size(0)
    index = torch.randperm(batch_size).to(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):
    return lam * criterion(pred, y_a) + (1 - lam) * criterion(pred, y_b)

import torch
from torch.utils.data import DataLoader
import time


def profile_dataloader(dataset, batch_size=32, num_workers_list=[0, 2, 4, 8]):
    """
    比较不同 num_workers 设置下的数据加载速度
    """
    results = {}

    for num_workers in num_workers_list:
        loader = DataLoader(
            dataset,
            batch_size=batch_size,
            num_workers=num_workers,
            pin_memory=True,  # 提升向 GPU 传输的速度
            prefetch_factor=2 if num_workers > 0 else None,
            persistent_workers=True if num_workers > 0 else False
        )

        start = time.time()
        for i, batch in enumerate(loader):
            if i >= 10:  # 仅测量 10 个 batch
                break
        elapsed = time.time() - start

        results[num_workers] = elapsed
        print(f"num_workers={num_workers}: {elapsed:.3f} 秒（10 个 batch）")

    best_workers = min(results, key=results.get)
    print(f"\n最优 num_workers: {best_workers}")
    return results


# 优化后的 DataLoader 设置
def create_optimized_dataloader(dataset, batch_size, is_train=True):
    return DataLoader(
        dataset,
        batch_size=batch_size,
        shuffle=is_train,
        num_workers=4,           # 建议设为 CPU 核心数的一半
        pin_memory=True,         # 用 CUDA pinned memory 加速向 GPU 传输
        prefetch_factor=2,       # 预加载的 batch 数
        persistent_workers=True, # 保持 worker 进程存活（消除重建开销）
        drop_last=is_train       # 丢弃不完整的最后一个 batch
    )

import torch
from torch.cuda.amp import autocast, GradScaler


def train_mixed_precision(model, loader, optimizer, loss_fn, device, epochs=10):
    """
    使用 FP16 混合精度训练，速度提升 2-3 倍
    """
    scaler = GradScaler()
    model.train()

    for epoch in range(epochs):
        for X, y in loader:
            X, y = X.to(device), y.to(device)
            optimizer.zero_grad()

            # 在 autocast 上下文中以 FP16 执行前向传播
            with autocast(device_type='cuda', dtype=torch.float16):
                output = model(X)
                loss = loss_fn(output, y)

            # 将 FP16 梯度放大为 FP32
            scaler.scale(loss).backward()

            # 梯度裁剪（考虑缩放）
            scaler.unscale_(optimizer)
            torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)

            # 更新参数
            scaler.step(optimizer)
            scaler.update()

        print(f"Epoch {epoch+1} 完成，scaler scale: {scaler.get_scale()}")

import torch

# 使用 PyTorch 2.0+ 的 torch.compile 提升速度
# 首次运行会有编译耗时，之后会变快

model = MyModel().to(device)

# 基础编译
compiled_model = torch.compile(model)

# 最大性能模式（编译时间更长）
compiled_model = torch.compile(model, mode='max-autotune')

# 适合输入尺寸经常变化的场景
compiled_model = torch.compile(model, dynamic=True)

# 比较训练时间
import time

def benchmark_model(model, inputs, n_iters=100):
    # 预热(warmup)
    for _ in range(10):
        _ = model(inputs)

    torch.cuda.synchronize()
    start = time.time()
    for _ in range(n_iters):
        _ = model(inputs)
    torch.cuda.synchronize()
    elapsed = time.time() - start

    return elapsed / n_iters

import torch
import gc


def print_gpu_memory_summary(device=0):
    """
    详细打印 GPU 显存使用情况
    """
    if not torch.cuda.is_available():
        print("无法使用 CUDA。")
        return

    print(f"=== GPU {device} 显存摘要 ===")
    print(f"总显存: {torch.cuda.get_device_properties(device).total_memory / 1e9:.2f} GB")
    print(f"已预留显存: {torch.cuda.memory_reserved(device) / 1e9:.2f} GB")
    print(f"已使用显存: {torch.cuda.memory_allocated(device) / 1e9:.2f} GB")
    print(f"缓存中的显存: {(torch.cuda.memory_reserved(device) - torch.cuda.memory_allocated(device)) / 1e9:.2f} GB")
    print()
    print(torch.cuda.memory_summary(device=device, abbreviated=False))


def find_memory_leaks(model, loader, device):
    """
    显存泄漏检测工具
    """
    import tracemalloc

    initial_memory = torch.cuda.memory_allocated(device)

    for i, (X, y) in enumerate(loader):
        X, y = X.to(device), y.to(device)
        output = model(X)

        current_memory = torch.cuda.memory_allocated(device)
        diff = (current_memory - initial_memory) / 1e6

        if i % 10 == 0:
            print(f"Batch {i}: GPU Memory Delta = {diff:.2f} MB")

        if i >= 50:
            break

    # 清理显存
    del X, y, output
    gc.collect()
    torch.cuda.empty_cache()


def clear_gpu_memory():
    """
    清理 GPU 显存缓存
    """
    gc.collect()
    torch.cuda.empty_cache()
    torch.cuda.synchronize()
    print(f"清理后 GPU 显存: {torch.cuda.memory_allocated() / 1e9:.2f} GB")

import torch
import torch.nn as nn
from torch.utils.checkpoint import checkpoint, checkpoint_sequential


class MemoryEfficientModel(nn.Module):
    """
    应用了 Gradient Checkpointing 的显存高效模型
    以少量重计算带来的速度损失，换取显存节省
    """
    def __init__(self, layers):
        super().__init__()
        self.layers = nn.ModuleList(layers)

    def forward(self, x):
        # checkpoint_sequential: 对连续的层自动应用 checkpointing
        # segments: 划分为多少个 chunk（越多越省显存，速度越慢）
        return checkpoint_sequential(self.layers, segments=4, input=x)

    def forward_with_manual_checkpoints(self, x):
        # 只对特定层选择性应用
        x = self.layers[0](x)  # 第一层直接执行
        for layer in self.layers[1:-1]:
            x = checkpoint(layer, x)  # 对中间层应用 checkpointing
        x = self.layers[-1](x)  # 最后一层直接执行
        return x


# 在 Transformer 中启用 Gradient Checkpointing
from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained("gpt2")
model.gradient_checkpointing_enable()  # 在 Hugging Face 模型中一行代码即可启用

def find_optimal_batch_size(model, loss_fn, device,
                              start_batch=8, max_batch=512):
    """
    在不发生 OOM 的前提下，寻找可用的最大 batch size
    """
    batch_size = start_batch
    optimal_batch_size = start_batch

    while batch_size <= max_batch:
        try:
            # 用 dummy 数据测试
            dummy_input = torch.randn(batch_size, 3, 224, 224).to(device)
            dummy_target = torch.randint(0, 1000, (batch_size,)).to(device)

            output = model(dummy_input)
            loss = loss_fn(output, dummy_target)
            loss.backward()

            optimal_batch_size = batch_size
            print(f"batch size {batch_size}: 成功")

            batch_size *= 2

            # 清理显存
            del dummy_input, dummy_target, output, loss
            torch.cuda.empty_cache()

        except RuntimeError as e:
            if "out of memory" in str(e):
                print(f"batch size {batch_size}: OOM")
                torch.cuda.empty_cache()
                break
            else:
                raise e

    print(f"\n推荐 batch size: {optimal_batch_size}（含安全余量: {optimal_batch_size // 2}）")
    return optimal_batch_size

import matplotlib.pyplot as plt
import numpy as np
import torch
from collections import Counter


def visualize_batch(loader, num_samples=16, class_names=None):
    """
    可视化数据 batch 样本，检查预处理结果
    """
    X, y = next(iter(loader))

    fig, axes = plt.subplots(4, 4, figsize=(12, 12))
    axes = axes.flatten()

    for i in range(min(num_samples, len(X))):
        img = X[i].numpy()

        # 反归一化（基于 ImageNet 标准）
        mean = np.array([0.485, 0.456, 0.406])
        std = np.array([0.229, 0.224, 0.225])
        img = std[:, None, None] * img + mean[:, None, None]
        img = np.clip(img, 0, 1)

        axes[i].imshow(img.transpose(1, 2, 0))
        label = y[i].item()
        title = class_names[label] if class_names else f"Label: {label}"
        axes[i].set_title(title)
        axes[i].axis('off')

    plt.tight_layout()
    plt.savefig('data_samples.png')
    plt.show()


def check_label_distribution(dataset):
    """
    检查标签分布——检测类别不平衡
    """
    labels = [dataset[i][1] for i in range(len(dataset))]
    counter = Counter(labels)

    classes = sorted(counter.keys())
    counts = [counter[c] for c in classes]
    total = sum(counts)

    print("标签分布:")
    for cls, count in zip(classes, counts):
        pct = count / total * 100
        bar = '#' * int(pct / 2)
        print(f"  类别 {cls}: {count:5d} ({pct:.1f}%) {bar}")

    # 计算不平衡比例
    max_count = max(counts)
    min_count = min(counts)
    imbalance_ratio = max_count / min_count

    if imbalance_ratio > 10:
        print(f"\n警告: 类别不平衡非常严重！（比例: {imbalance_ratio:.1f}:1）")
        print("解决方法: 建议使用加权采样或类别加权损失函数。")

    return counter


def create_weighted_sampler(dataset):
    """
    为解决类别不平衡而创建加权采样器
    """
    labels = [dataset[i][1] for i in range(len(dataset))]
    class_counts = Counter(labels)

    # 计算每个类别的采样权重
    weights = [1.0 / class_counts[label] for label in labels]
    weights = torch.DoubleTensor(weights)

    sampler = torch.utils.data.WeightedRandomSampler(
        weights=weights,
        num_samples=len(weights),
        replacement=True
    )
    return sampler


def check_normalization(loader, expected_mean=None, expected_std=None):
    """
    验证数据归一化数值
    """
    all_data = []
    for X, _ in loader:
        all_data.append(X)
        if len(all_data) >= 10:  # 仅采样 10 个 batch
            break

    all_data = torch.cat(all_data, dim=0)
    actual_mean = all_data.mean(dim=[0, 2, 3]) if all_data.dim() == 4 else all_data.mean()
    actual_std = all_data.std(dim=[0, 2, 3]) if all_data.dim() == 4 else all_data.std()

    print(f"实际均值: {actual_mean.tolist()}")
    print(f"实际标准差: {actual_std.tolist()}")

    if expected_mean:
        mean_diff = abs(actual_mean - torch.tensor(expected_mean)).max().item()
        print(f"与期望均值的差异: {mean_diff:.4f}")
        if mean_diff > 0.1:
            print("警告: 归一化数值与期望不符！")

from torchinfo import summary
import torch
import torch.nn as nn


def analyze_model(model, input_size):
    """
    分析模型结构并定位瓶颈层
    """
    # 基础摘要（形状、参数、显存）
    model_stats = summary(
        model,
        input_size=input_size,
        col_names=["input_size", "output_size", "num_params", "kernel_size",
                   "mult_adds"],
        verbose=1
    )

    # 检查各层参数量
    print("\n各层参数分布:")
    for name, module in model.named_modules():
        num_params = sum(p.numel() for p in module.parameters(recurse=False))
        if num_params > 0:
            print(f"  {name}: {num_params:,} 个参数")

    return model_stats


def monitor_activations(model, X):
    """
    监控中间激活值以检测死神经元（Dead Neurons）
    """
    activations = {}

    def make_activation_hook(name):
        def hook(module, input, output):
            activations[name] = output.detach()
        return hook

    # 注册钩子
    hooks = []
    for name, module in model.named_modules():
        if isinstance(module, (nn.ReLU, nn.GELU, nn.Tanh, nn.Sigmoid)):
            hook = module.register_forward_hook(make_activation_hook(name))
            hooks.append(hook)

    # 前向传播
    with torch.no_grad():
        model(X)

    # 分析激活值
    print("\n激活统计:")
    for name, act in activations.items():
        dead_neurons = (act == 0).float().mean().item()
        print(f"  {name}:")
        print(f"    均值: {act.mean():.4f}，标准差: {act.std():.4f}")
        print(f"    死神经元比例: {dead_neurons:.2%}")
        if dead_neurons > 0.5:
            print(f"    警告: {dead_neurons:.0%} 的神经元处于失活状态！")

    # 移除钩子
    for hook in hooks:
        hook.remove()

    return activations


def visualize_weight_distribution(model):
    """
    通过可视化权重分布检测初始化问题
    """
    import matplotlib.pyplot as plt

    fig, axes = plt.subplots(2, 3, figsize=(15, 10))
    axes = axes.flatten()

    linear_layers = [(name, m) for name, m in model.named_modules()
                     if isinstance(m, (nn.Linear, nn.Conv2d))]

    for i, (name, layer) in enumerate(linear_layers[:6]):
        if i >= len(axes):
            break
        weight_data = layer.weight.data.cpu().numpy().flatten()
        axes[i].hist(weight_data, bins=50, color='blue', alpha=0.7)
        axes[i].set_title(f"{name}\n(mean={weight_data.mean():.4f}, std={weight_data.std():.4f})")
        axes[i].set_xlabel('权重值')
        axes[i].set_ylabel('频次')
        axes[i].grid(True, alpha=0.3)

    plt.suptitle("各层权重分布")
    plt.tight_layout()
    plt.savefig('weight_distribution.png')
    plt.show()

from torch.utils.tensorboard import SummaryWriter
import torch
import numpy as np


class TensorBoardLogger:
    def __init__(self, log_dir='runs/experiment'):
        self.writer = SummaryWriter(log_dir)
        self.step = 0

    def log_scalars(self, metrics: dict, epoch: int):
        """记录标量指标"""
        for name, value in metrics.items():
            self.writer.add_scalar(name, value, epoch)

    def log_model_gradients(self, model, epoch: int):
        """记录梯度直方图"""
        for name, param in model.named_parameters():
            if param.grad is not None:
                self.writer.add_histogram(f'gradients/{name}',
                                           param.grad, epoch)
                self.writer.add_histogram(f'weights/{name}',
                                           param.data, epoch)

    def log_images(self, images: torch.Tensor, tag: str, epoch: int, n=8):
        """记录图像 batch"""
        self.writer.add_images(tag, images[:n], epoch)

    def log_learning_rate(self, optimizer, epoch: int):
        """记录学习率"""
        for i, pg in enumerate(optimizer.param_groups):
            self.writer.add_scalar(f'lr/group_{i}', pg['lr'], epoch)

    def log_confusion_matrix(self, cm, class_names, epoch: int):
        """记录混淆矩阵"""
        import matplotlib.pyplot as plt
        import seaborn as sns

        fig, ax = plt.subplots(figsize=(10, 8))
        sns.heatmap(cm, annot=True, fmt='d', cmap='Blues',
                    xticklabels=class_names, yticklabels=class_names)
        ax.set_xlabel('预测')
        ax.set_ylabel('实际')
        ax.set_title('混淆矩阵')

        self.writer.add_figure('confusion_matrix', fig, epoch)
        plt.close()

    def close(self):
        self.writer.close()


def train_with_tensorboard(model, train_loader, val_loader,
                             optimizer, loss_fn, device, epochs=50):
    logger = TensorBoardLogger(log_dir='runs/debug_session')

    for epoch in range(epochs):
        # 训练
        model.train()
        train_loss, train_correct = 0, 0
        for X, y in train_loader:
            X, y = X.to(device), y.to(device)
            optimizer.zero_grad()
            output = model(X)
            loss = loss_fn(output, y)
            loss.backward()
            optimizer.step()

            train_loss += loss.item()
            train_correct += (output.argmax(1) == y).sum().item()

        train_loss /= len(train_loader)
        train_acc = train_correct / len(train_loader.dataset)

        # 验证
        model.eval()
        val_loss, val_correct = 0, 0
        with torch.no_grad():
            for X, y in val_loader:
                X, y = X.to(device), y.to(device)
                output = model(X)
                val_loss += loss_fn(output, y).item()
                val_correct += (output.argmax(1) == y).sum().item()

        val_loss /= len(val_loader)
        val_acc = val_correct / len(val_loader.dataset)

        # TensorBoard 记录
        logger.log_scalars({
            'Loss/Train': train_loss,
            'Loss/Val': val_loss,
            'Accuracy/Train': train_acc,
            'Accuracy/Val': val_acc,
        }, epoch)
        logger.log_model_gradients(model, epoch)
        logger.log_learning_rate(optimizer, epoch)

    logger.close()
    print("TensorBoard: 运行 tensorboard --logdir=runs 命令即可启动")

import wandb
import torch


def train_with_wandb(model, train_loader, val_loader, optimizer, loss_fn,
                      device, config=None):
    """
    使用 W&B 追踪实验
    """
    if config is None:
        config = {
            'learning_rate': 1e-3,
            'batch_size': 32,
            'epochs': 50,
            'optimizer': 'AdamW',
            'weight_decay': 1e-4,
            'model': model.__class__.__name__
        }

    # 初始化 W&B
    run = wandb.init(
        project="deep-learning-debug",
        config=config,
        tags=["debugging", "experiment"]
    )

    # 监视模型（自动记录梯度、参数）
    wandb.watch(model, log='all', log_freq=100)

    for epoch in range(config['epochs']):
        model.train()
        train_loss = 0

        for batch_idx, (X, y) in enumerate(train_loader):
            X, y = X.to(device), y.to(device)
            optimizer.zero_grad()
            output = model(X)
            loss = loss_fn(output, y)
            loss.backward()

            # 梯度裁剪
            grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
            optimizer.step()

            train_loss += loss.item()

            # batch 级别的记录
            if batch_idx % 10 == 0:
                wandb.log({
                    'batch/loss': loss.item(),
                    'batch/grad_norm': grad_norm.item(),
                    'batch': epoch * len(train_loader) + batch_idx
                })

        # epoch 级别的记录
        wandb.log({
            'epoch/train_loss': train_loss / len(train_loader),
            'epoch/learning_rate': optimizer.param_groups[0]['lr'],
            'epoch': epoch
        })

    run.finish()

import random
import numpy as np
import torch
import os


def set_seed(seed: int = 42):
    """
    为实现完全可复现性而固定随机种子
    为所有随机数生成器设置相同的种子
    """
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)  # 多 GPU

    # cuDNN 确定性模式
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False  # 性能会略有下降

    # PyTorch 确定性运算
    torch.use_deterministic_algorithms(True)

    # 设置环境变量
    os.environ['PYTHONHASHSEED'] = str(seed)
    os.environ['CUBLAS_WORKSPACE_CONFIG'] = ':4096:8'

    print(f"已用种子 {seed} 完成所有随机数生成器的初始化")


def save_experiment_config(config: dict, save_path: str = 'experiment_config.json'):
    """
    记录完整的实验环境
    """
    import json
    import subprocess

    full_config = config.copy()

    # 记录 Python 及各依赖包的版本
    full_config['environment'] = {
        'python': subprocess.getoutput('python --version'),
        'torch': torch.__version__,
        'cuda': torch.version.cuda,
        'cudnn': str(torch.backends.cudnn.version()),
        'gpu': torch.cuda.get_device_name(0) if torch.cuda.is_available() else 'CPU'
    }

    # 记录 Git commit hash
    try:
        full_config['git_hash'] = subprocess.getoutput('git rev-parse HEAD')
    except Exception:
        full_config['git_hash'] = 'unknown'

    with open(save_path, 'w', encoding='utf-8') as f:
        json.dump(full_config, f, indent=2, ensure_ascii=False)

    print(f"实验配置已保存: {save_path}")
    return full_config


# 可复现性测试
def test_reproducibility(model_fn, train_fn, seed=42, n_runs=3):
    """
    用相同种子多次运行以验证可复现性
    """
    results = []

    for run in range(n_runs):
        set_seed(seed)
        model = model_fn()
        loss = train_fn(model)
        results.append(loss)
        print(f"Run {run+1}: Final Loss = {loss:.6f}")

    max_diff = max(results) - min(results)
    print(f"\n最大差异: {max_diff:.8f}")

    if max_diff < 1e-5:
        print("可复现性验证通过！")
    else:
        print("警告: 存在可复现性问题。")

    return results

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


def setup_distributed(rank, world_size, backend='nccl'):
    """
    初始化分布式训练环境
    """
    os.environ['MASTER_ADDR'] = 'localhost'
    os.environ['MASTER_PORT'] = '12355'

    dist.init_process_group(
        backend=backend,
        rank=rank,
        world_size=world_size
    )

    torch.cuda.set_device(rank)
    print(f"进程 {rank}/{world_size} 初始化完成")


def cleanup_distributed():
    dist.destroy_process_group()


def debug_ddp_training(rank, world_size, model, dataset):
    """
    DDP 训练调试示例
    """
    setup_distributed(rank, world_size)

    device = torch.device(f'cuda:{rank}')
    model = model.to(device)

    # DDP 包装
    model = DDP(model, device_ids=[rank], find_unused_parameters=True)

    # 分布式采样器（为每个 rank 分配不同数据）
    sampler = torch.utils.data.distributed.DistributedSampler(
        dataset,
        num_replicas=world_size,
        rank=rank,
        shuffle=True
    )

    loader = torch.utils.data.DataLoader(
        dataset,
        batch_size=32,
        sampler=sampler,
        num_workers=4
    )

    optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3)
    loss_fn = torch.nn.CrossEntropyLoss()

    for epoch in range(10):
        # 重要: 每个 epoch 都要更新采样器的种子
        sampler.set_epoch(epoch)

        for X, y in loader:
            X, y = X.to(device), y.to(device)
            optimizer.zero_grad()
            output = model(X)
            loss = loss_fn(output, y)
            loss.backward()
            optimizer.step()

        # 仅在 rank 0 上记录
        if rank == 0:
            print(f"Epoch {epoch+1} 完成")

            # 模型保存也只在 rank 0 上进行
            torch.save(model.module.state_dict(), f'checkpoint_epoch{epoch}.pt')

        # 同步所有 rank
        dist.barrier()

    cleanup_distributed()


def check_gradient_sync(model):
    """
    验证 DDP 中的梯度同步
    """
    for name, param in model.named_parameters():
        if param.grad is not None:
            # 收集所有 rank 的梯度总和
            grad_sum = param.grad.data.clone()
            dist.all_reduce(grad_sum, op=dist.ReduceOp.SUM)

            # 检查各 rank 的梯度是否已同步
            world_size = dist.get_world_size()
            expected = grad_sum / world_size

            if not torch.allclose(param.grad.data, expected, atol=1e-5):
                print(f"警告: {name} 层的梯度同步不一致！")

import mlflow
import mlflow.pytorch
import torch
import optuna


def train_with_mlflow(model, train_loader, val_loader, optimizer,
                       loss_fn, device, params: dict):
    """
    使用 MLflow 追踪实验并管理模型版本
    """
    mlflow.set_tracking_uri("http://localhost:5000")
    mlflow.set_experiment("deep-learning-debug")

    with mlflow.start_run():
        # 记录超参数
        mlflow.log_params(params)

        best_val_loss = float('inf')

        for epoch in range(params['epochs']):
            # 训练
            model.train()
            train_loss = 0
            for X, y in train_loader:
                X, y = X.to(device), y.to(device)
                optimizer.zero_grad()
                output = model(X)
                loss = loss_fn(output, y)
                loss.backward()
                optimizer.step()
                train_loss += loss.item()

            train_loss /= len(train_loader)

            # 验证
            model.eval()
            val_loss = 0
            with torch.no_grad():
                for X, y in val_loader:
                    X, y = X.to(device), y.to(device)
                    output = model(X)
                    val_loss += loss_fn(output, y).item()
            val_loss /= len(val_loader)

            # 记录指标
            mlflow.log_metrics({
                'train_loss': train_loss,
                'val_loss': val_loss
            }, step=epoch)

            # 保存最优模型
            if val_loss < best_val_loss:
                best_val_loss = val_loss
                mlflow.pytorch.log_model(model, "best_model")

        # 最终指标
        mlflow.log_metric("best_val_loss", best_val_loss)

    return best_val_loss


def hyperparameter_optimization_with_optuna(model_fn, train_loader,
                                              val_loader, device, n_trials=50):
    """
    使用 Optuna 进行超参数优化
    """
    def objective(trial):
        # 定义要搜索的超参数范围
        lr = trial.suggest_float('lr', 1e-5, 1e-1, log=True)
        weight_decay = trial.suggest_float('weight_decay', 1e-6, 1e-2, log=True)
        dropout = trial.suggest_float('dropout', 0.0, 0.5)
        batch_size = trial.suggest_categorical('batch_size', [16, 32, 64, 128])

        model = model_fn(dropout=dropout).to(device)
        optimizer = torch.optim.AdamW(model.parameters(),
                                       lr=lr, weight_decay=weight_decay)
        loss_fn = torch.nn.CrossEntropyLoss()

        # 用短程训练做快速评估
        val_loss = train_with_mlflow(
            model, train_loader, val_loader, optimizer, loss_fn, device,
            params={'lr': lr, 'weight_decay': weight_decay,
                    'dropout': dropout, 'batch_size': batch_size, 'epochs': 10}
        )

        return val_loss

    study = optuna.create_study(
        direction='minimize',
        sampler=optuna.samplers.TPESampler(seed=42),
        pruner=optuna.pruners.MedianPruner()
    )

    study.optimize(objective, n_trials=n_trials)

    print("\n最优超参数:")
    for key, value in study.best_params.items():
        print(f"  {key}: {value}")
    print(f"最优验证 Loss: {study.best_value:.4f}")

    return study.best_params

# 深度学习调试的最小检查清单
def minimum_debug_checklist(model, train_loader, device):
    """
    训练开始前必须确认的最小检查清单
    """
    print("深度学习训练前检查清单")
    print("=" * 50)

    # 1. 单 batch 过拟合测试
    print("[1] 正在进行单 batch 过拟合测试...")
    X, y = next(iter(train_loader))
    X, y = X.to(device), y.to(device)

    optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
    loss_fn = torch.nn.CrossEntropyLoss()

    initial_loss = None
    for step in range(100):
        optimizer.zero_grad()
        output = model(X)
        loss = loss_fn(output, y)
        if initial_loss is None:
            initial_loss = loss.item()
        loss.backward()
        optimizer.step()

    final_loss = loss.item()
    overfit_ratio = initial_loss / final_loss if final_loss > 0 else float('inf')

    if overfit_ratio > 10:
        print(f"  通过: Loss 从 {initial_loss:.4f} 降至 {final_loss:.4f}（比例: {overfit_ratio:.1f}x）")
    else:
        print(f"  警告: 单 batch 未能被过拟合（比例: {overfit_ratio:.1f}x）")
        print("  → 请检查模型容量、学习率、数据错误")

    print("\n检查清单完成！")