图像分类完整实战

import torch import torchvision import torchvision.transforms as transforms # 定义基础转换 transform_train = transforms.Compose([ transforms.RandomCrop(32, padding=4), transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)), ]) transform_test = transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)), ]) # 加载数据集 trainset = torchvision.datasets.CIFAR10( root='./data', train=True, download=True, transform=transform_train) testset = torchvision.datasets.CIFAR10( root='./data', train=False, download=True, transform=transform_test) trainloader = torch.utils.data.DataLoader( trainset, batch_size=128, shuffle=True, num_workers=4, pin_memory=True) testloader = torch.utils.data.DataLoader( testset, batch_size=128, shuffle=False, num_workers=4, pin_memory=True)

dataset/ ├── train/ │ ├── class_0/ │ │ ├── img_0001.jpg │ │ ├── img_0002.jpg │ │ └── ... │ ├── class_1/ │ │ ├── img_0001.jpg │ │ └── ... │ └── ... └── val/ ├── class_0/ │ ├── img_0001.jpg │ └── ... ├── class_1/ │ └── ... └── ...

from torchvision.datasets import ImageFolder from torch.utils.data import DataLoader, random_split # 加载完整数据集 full_dataset = ImageFolder( root='dataset/train', transform=transform_train ) # 按比例划分训练集和验证集 train_size = int(0.85 * len(full_dataset)) val_size = len(full_dataset) - train_size train_dataset, val_dataset = random_split( full_dataset, [train_size, val_size], generator=torch.Generator().manual_seed(42) ) # 类别名称自动从目录结构获取 class_names = full_dataset.classes print(f"类别数: {len(class_names)}") print(f"类别列表: {class_names}") print(f"训练样本数: {len(train_dataset)}") print(f"验证样本数: {len(val_dataset)}") train_loader = DataLoader( train_dataset, batch_size=64, shuffle=True, num_workers=4, pin_memory=True ) val_loader = DataLoader( val_dataset, batch_size=64, shuffle=False, num_workers=4, pin_memory=True )

from torch.utils.data import WeightedRandomSampler import numpy as np # 计算类别权重 labels = [label for _, label in train_dataset] class_counts = np.bincount(labels) class_weights = 1.0 / class_counts sample_weights = class_weights[labels] # 创建加权采样器 sampler = WeightedRandomSampler( weights=sample_weights, num_samples=len(sample_weights), replacement=True ) # 使用加权采样器的 DataLoader balanced_loader = DataLoader( train_dataset, batch_size=64, sampler=sampler, # 使用 sampler 后不能使用 shuffle=True num_workers=4, pin_memory=True )

from torchvision import transforms # 强增强策略（适用于训练集） train_transform = transforms.Compose([ # 几何变换 transforms.RandomResizedCrop( size=224, scale=(0.08, 1.0), ratio=(0.75, 1.333)), transforms.RandomHorizontalFlip(p=0.5), transforms.RandomRotation(degrees=15), # 颜色变换 transforms.ColorJitter( brightness=0.2, contrast=0.2, saturation=0.2, hue=0.1), # 高级增强 transforms.RandomGrayscale(p=0.1), transforms.RandomAffine( degrees=0, translate=(0.1, 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]), ])

from torchvision.transforms import AutoAugment, AutoAugmentPolicy from torchvision.transforms import RandAugment, TrivialAugmentWide # AutoAugment - 基于搜索学习的策略 auto_augment = AutoAugment( policy=AutoAugmentPolicy.IMAGENET ) # 可选策略: CIFAR10, SVHN, IMAGENET # RandAugment - 简化版自动增强 (N, M) # N: 每次应用的变换数量 # M: 变换强度 rand_augment = RandAugment( num_ops=2, # 每次应用2个变换 magnitude=9 # 变换强度 (0-30) ) # TrivialAugmentWide - 更宽的强度范围 trivial_augment = TrivialAugmentWide() # 完整训练流水线（结合 RandAugment） train_transform_advanced = transforms.Compose([ transforms.RandomResizedCrop(224), transforms.RandomHorizontalFlip(), rand_augment, # 自动增强 transforms.ToTensor(), transforms.Normalize( mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), transforms.RandomErasing( # 随机擦除 p=0.25, scale=(0.02, 0.33)) ])

import torch import numpy as np # ---- CutOut: 随机遮挡图像区域 ---- class CutOut: """随机遮挡图像中的一个正方形区域""" def __init__(self, mask_size=16, p=0.5): self.mask_size = mask_size self.p = p def __call__(self, img): if np.random.rand() > self.p: return img h, w = img.shape[1:] # C, H, W mask_size_half = self.mask_size // 2 cx = np.random.randint(mask_size_half, w - mask_size_half) cy = np.random.randint(mask_size_half, h - mask_size_half) x1 = max(0, cx - mask_size_half) x2 = min(w, cx + mask_size_half) y1 = max(0, cy - mask_size_half) y2 = min(h, cy + mask_size_half) img[:, y1:y2, x1:x2] = 0 return img # ---- MixUp: 混合两张图像和标签 ---- def mixup_data(x, y, alpha=1.0): """MixUp 数据增强""" 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) # ---- CutMix: 区域替换混合 ---- def cutmix_data(x, y, alpha=1.0): """CutMix 数据增强""" 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) h, w = x.size()[2], x.size()[3] cx = np.random.randint(w) cy = np.random.randint(h) cut_w = int(w * np.sqrt(1 - lam)) cut_h = int(h * np.sqrt(1 - lam)) x1 = max(0, cx - cut_w // 2) y1 = max(0, cy - cut_h // 2) x2 = min(w, cx + cut_w // 2) y2 = min(h, cy + cut_h // 2) x[:, :, y1:y2, x1:x2] = x[index, :, y1:y2, x1:x2] lam = 1 - ((x2 - x1) * (y2 - y1) / (w * h)) return x, y, y[index], lam # 训练循环中使用 MixUp def train_one_epoch(model, loader, optimizer, criterion, device): model.train() running_loss = 0.0 correct = 0 total = 0 for inputs, targets in loader: inputs, targets = inputs.to(device), targets.to(device) # 应用 MixUp inputs, targets_a, targets_b, lam = mixup_data( inputs, targets, alpha=0.2) optimizer.zero_grad() outputs = model(inputs) loss = mixup_criterion( criterion, outputs, targets_a, targets_b, lam) loss.backward() optimizer.step() running_loss += loss.item() return running_loss / len(loader)

import torchvision.models as models # ---- ResNet 系列（经典 CNN 架构） ---- resnet18 = models.resnet18(weights='IMAGENET1K_V1') resnet50 = models.resnet50(weights='IMAGENET1K_V2') resnet152 = models.resnet152(weights='IMAGENET1K_V2') # ---- EfficientNet 系列（计算效率优化） ---- efficientnet_b0 = models.efficientnet_b0( weights='IMAGENET1K_V1') efficientnet_b7 = models.efficientnet_b7( weights='IMAGENET1K_V1') # ---- Vision Transformer 系列（现代架构） ---- vit_b_16 = models.vit_b_16(weights='IMAGENET1K_V1') vit_l_16 = models.vit_l_16(weights='IMAGENET1K_V1') # ViT-L/16: 大型 ViT，参数量 ~304M # ---- ConvNeXt（现代 CNN） ---- convnext_base = models.convnext_base( weights='IMAGENET1K_V1') convnext_large = models.convnext_large( weights='IMAGENET1K_V1') # ---- 查看模型参数量 ---- def count_parameters(model): total = sum(p.numel() for p in model.parameters()) trainable = sum(p.numel() for p in model.parameters() if p.requires_grad) return total, trainable total, trainable = count_parameters(resnet152) print(f"ResNet-152 总参数量: {total/1e6:.2f}M") print(f"ResNet-152 可训练参数量: {trainable/1e6:.2f}M") # 输出: ResNet-152 总参数量: 60.19M

# ========== 策略一：冻结特征提取器 ========== model = models.resnet50(weights='IMAGENET1K_V2') # 冻结所有参数 for param in model.parameters(): param.requires_grad = False # 替换分类头（新分类头默认 requires_grad=True） num_classes = 10 # 以 CIFAR-10 为例 model.fc = torch.nn.Linear( model.fc.in_features, num_classes) # 只优化分类头 optimizer = torch.optim.Adam( model.fc.parameters(), lr=0.001) # ========== 策略二：全微调 ========== model = models.resnet50(weights='IMAGENET1K_V2') # 所有参数均可训练 for param in model.parameters(): param.requires_grad = True # 替换分类头 model.fc = torch.nn.Linear( model.fc.in_features, num_classes) # 使用较低的学习率 optimizer = torch.optim.SGD( model.parameters(), lr=0.001, momentum=0.9, weight_decay=1e-4) # ========== 策略三：区别学习率 ========== model = models.resnet50(weights='IMAGENET1K_V2') # 替换分类头 model.fc = torch.nn.Linear( model.fc.in_features, num_classes) # 为不同层组设置不同的学习率 def get_optimizer_params(model, lr=0.001): """为不同层设置不同学习率""" # 分类头：使用完整学习率 fc_params = list(model.fc.parameters()) # 最后两个 Block：使用 1/10 学习率 layer4_params = list(model.layer4.parameters()) # 前面层：使用 1/100 学习率 base_params = [ p for name, p in model.named_parameters() if not any(p is fc for fc in fc_params) and not any(p is l4 for l4 in layer4_params) and p.requires_grad ] # 注意：上面的简单区分方式在实际中应使用 # param_group 的 name 进行更精确的控制 return [ {'params': model.fc.parameters(), 'lr': lr}, {'params': model.layer4.parameters(), 'lr': lr * 0.1}, {'params': base_params, 'lr': lr * 0.01, 'weight_decay': 1e-4}, ] optimizer = torch.optim.SGD( get_optimizer_params(model, lr=0.01), momentum=0.9, weight_decay=1e-4) # 学习率调度器 scheduler = torch.optim.lr_scheduler.CosineAnnealingLR( optimizer, T_max=30, eta_min=1e-6 )

# ViT 微调实战 model = models.vit_l_16(weights='IMAGENET1K_V1') # ViT 的分类头在 model.heads 中 num_features = model.heads.head.in_features model.heads.head = torch.nn.Linear( num_features, num_classes) # ViT 需要更大的输入尺寸（默认 224x224） # 同时注意 ViT 对位置编码的依赖 transform_vit = transforms.Compose([ transforms.Resize(256), transforms.CenterCrop(224), # ViT 要求固定尺寸 transforms.ToTensor(), transforms.Normalize( mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), ]) # 使用 AdamW 优化器（ViT 推荐） optimizer = torch.optim.AdamW( model.parameters(), lr=3e-5, # ViT 微调通常使用更小的学习率 weight_decay=0.01 ) # 线性 warmup + cosine 衰减调度 warmup_epochs = 5 total_epochs = 50 def warmup_cosine_scheduler(epoch): if epoch < warmup_epochs: return (epoch + 1) / warmup_epochs return 0.5 * (1 + np.cos( np.pi * (epoch - warmup_epochs) / (total_epochs - warmup_epochs))) scheduler = torch.optim.lr_scheduler.LambdaLR( optimizer, lr_lambda=warmup_cosine_scheduler )

# ---- 策略 1: 阶梯式衰减 ---- scheduler_step = torch.optim.lr_scheduler.MultiStepLR( optimizer, milestones=[30, 60, 80], # 在第30、60、80轮衰减 gamma=0.1 # 衰减因子 ) # ---- 策略 2: Cosine Annealing ---- scheduler_cosine = torch.optim.lr_scheduler.CosineAnnealingLR( optimizer, T_max=100, # 周期长度 eta_min=1e-6 # 最小学习率 ) # ---- 策略 3: Cosine Annealing with Warm Restarts ---- scheduler_restart = torch.optim.lr_scheduler.CosineAnnealingWarmRestarts( optimizer, T_0=20, # 第一个周期长度 T_mult=2, # 周期倍增因子 eta_min=1e-6 ) # ---- 策略 4: OneCycleLR（推荐快速训练） ---- scheduler_onecycle = torch.optim.lr_scheduler.OneCycleLR( optimizer, max_lr=0.01, # 最大学习率 steps_per_epoch=len(train_loader), epochs=50, pct_start=0.3, # 前30% epoch 完成预热 div_factor=25, # 初始 lr = max_lr / 25 final_div_factor=1000 # 最终 lr = max_lr / (25*1000) ) # ---- 学习率查找器（LR Finder） ---- # 在训练开始前运行，找到合适的学习率区间 def lr_finder(model, loader, optimizer, device, init_lr=1e-7, final_lr=10, beta=0.98): model.train() lrs = [] losses = [] best_loss = float('inf') lr = init_lr optimizer.param_groups[0]['lr'] = lr for batch_idx, (inputs, targets) in enumerate(loader): inputs, targets = inputs.to(device), targets.to(device) optimizer.zero_grad() outputs = model(inputs) loss = torch.nn.functional.cross_entropy(outputs, targets) loss.backward() optimizer.step() # 平滑损失 if batch_idx == 0: avg_loss = loss.item() else: avg_loss = beta * avg_loss + (1 - beta) * loss.item() smoothed_loss = avg_loss / (1 - beta ** (batch_idx + 1)) lrs.append(lr) losses.append(smoothed_loss) if smoothed_loss < best_loss: best_loss = smoothed_loss if smoothed_loss > 4 * best_loss: break # 损失爆炸，停止 # 指数增加学习率 lr *= (final_lr / init_lr) ** (1 / len(loader)) optimizer.param_groups[0]['lr'] = lr return lrs, losses # 使用方法：绘制学习率-损失曲线，选择损失下降最快的 LR 区间

# ---- SGD with Momentum（经典选择） ---- optimizer_sgd = torch.optim.SGD( model.parameters(), lr=0.001, momentum=0.9, weight_decay=1e-4, nesterov=True # Nesterov 加速梯度 ) # ---- Adam（自适应学习率） ---- optimizer_adam = torch.optim.Adam( model.parameters(), lr=0.001, betas=(0.9, 0.999), eps=1e-8, weight_decay=1e-4 ) # ---- AdamW（推荐，带解耦权重衰减） ---- optimizer_adamw = torch.optim.AdamW( model.parameters(), lr=0.001, betas=(0.9, 0.999), weight_decay=0.01 ) # ---- 损失函数选择 ---- # 标准分类任务 criterion_ce = torch.nn.CrossEntropyLoss() # 带标签平滑的交叉熵 class LabelSmoothingCrossEntropy(torch.nn.Module): def __init__(self, smoothing=0.1): super().__init__() self.smoothing = smoothing def forward(self, pred, target): n_classes = pred.size(1) log_pred = torch.log_softmax(pred, dim=1) # 构造平滑标签 smooth_target = torch.full_like(log_pred, self.smoothing / (n_classes - 1)) smooth_target.scatter_( 1, target.unsqueeze(1), 1 - self.smoothing) loss = -(smooth_target * log_pred).sum(dim=1).mean() return loss criterion_smooth = LabelSmoothingCrossEntropy(smoothing=0.1) # 焦点损失 Focal Loss（处理难分类样本） class FocalLoss(torch.nn.Module): def __init__(self, gamma=2.0, alpha=None): super().__init__() self.gamma = gamma self.alpha = alpha def forward(self, pred, target): ce_loss = torch.nn.functional.cross_entropy( pred, target, reduction='none') pt = torch.exp(-ce_loss) focal_loss = ((1 - pt) ** self.gamma * ce_loss) if self.alpha is not None: alpha_t = self.alpha[target] focal_loss = alpha_t * focal_loss return focal_loss.mean()

# ---- 梯度累积（解决大 Batch 的 GPU 内存限制） ---- accumulation_steps = 4 effective_batch_size = 64 * accumulation_steps for i, (inputs, targets) in enumerate(train_loader): inputs, targets = inputs.to(device), targets.to(device) outputs = model(inputs) loss = criterion(outputs, targets) # 缩放损失 loss = loss / accumulation_steps loss.backward() if (i + 1) % accumulation_steps == 0: optimizer.step() optimizer.zero_grad() # ---- 混合精度训练（AMP） ---- from torch.cuda.amp import autocast, GradScaler scaler = GradScaler() model.train() for inputs, targets in train_loader: inputs, targets = inputs.to(device), targets.to(device) optimizer.zero_grad() # 自动混合精度上下文 with autocast(): outputs = model(inputs) loss = criterion(outputs, targets) # 缩放梯度并反向传播 scaler.scale(loss).backward() scaler.step(optimizer) scaler.update() # ---- 梯度裁剪（防止梯度爆炸） ---- torch.nn.utils.clip_grad_norm_( model.parameters(), max_norm=1.0) # ---- 完整训练函数 ---- def train_epoch(model, loader, criterion, optimizer, scheduler, device, scaler=None, clip_grad=1.0, accumulation_steps=1): model.train() running_loss = 0.0 correct = 0 total = 0 optimizer.zero_grad() for batch_idx, (inputs, targets) in enumerate(loader): inputs, targets = inputs.to(device), targets.to(device) if scaler is not None: with autocast(): outputs = model(inputs) loss = criterion(outputs, targets) loss = loss / accumulation_steps scaler.scale(loss).backward() else: outputs = model(inputs) loss = criterion(outputs, targets) loss = loss / accumulation_steps loss.backward() if (batch_idx + 1) % accumulation_steps == 0: if scaler is not None: scaler.unscale_(optimizer) torch.nn.utils.clip_grad_norm_( model.parameters(), clip_grad) scaler.step(optimizer) scaler.update() else: torch.nn.utils.clip_grad_norm_( model.parameters(), clip_grad) optimizer.step() optimizer.zero_grad() running_loss += loss.item() * accumulation_steps _, predicted = outputs.max(1) total += targets.size(0) correct += predicted.eq(targets).sum().item() if scheduler is not None: scheduler.step() epoch_loss = running_loss / len(loader) epoch_acc = 100. * correct / total return epoch_loss, epoch_acc

import matplotlib.pyplot as plt import seaborn as sns from sklearn.metrics import confusion_matrix, classification_report import numpy as np def evaluate_model(model, loader, device, class_names): """全面评估模型并生成可视化结果""" model.eval() all_preds = [] all_targets = [] all_probs = [] with torch.no_grad(): for inputs, targets in loader: inputs = inputs.to(device) outputs = model(inputs) probs = torch.softmax(outputs, dim=1) _, preds = outputs.max(1) all_preds.extend(preds.cpu().numpy()) all_targets.extend(targets.numpy()) all_probs.extend(probs.cpu().numpy()) all_preds = np.array(all_preds) all_targets = np.array(all_targets) all_probs = np.array(all_probs) # 1. 混淆矩阵 cm = confusion_matrix(all_targets, all_preds) plt.figure(figsize=(12, 10)) sns.heatmap(cm, annot=True, fmt='d', cmap='Blues', xticklabels=class_names, yticklabels=class_names) plt.title('Confusion Matrix') plt.xlabel('Predicted Label') plt.ylabel('True Label') plt.tight_layout() plt.savefig('confusion_matrix.png', dpi=150) plt.show() # 2. 分类报告（精确率、召回率、F1-score） report = classification_report( all_targets, all_preds, target_names=class_names, digits=4 ) print("=== 分类报告 ===") print(report) # 3. 按类别准确率排序 class_acc = {} for i, name in enumerate(class_names): mask = all_targets == i class_acc[name] = (all_preds[mask] == i).mean() sorted_acc = sorted(class_acc.items(), key=lambda x: x[1]) print("=== 各类别准确率（从低到高）===") for name, acc in sorted_acc: print(f" {name}: {acc*100:.2f}%") return all_preds, all_targets, all_probs def plot_error_analysis(all_preds, all_targets, all_probs, class_names, top_k=5): """错误分析：找出最容易被混淆的类别对""" cm = confusion_matrix(all_targets, all_preds) np.fill_diagonal(cm, 0) # 忽略正确分类 print("=== 最常见的错误分类（Top-5）===") for _ in range(top_k): idx = np.unravel_index(np.argmax(cm), cm.shape) print(f" 真实类别 '{class_names[idx[0]]}' " f"被误分为 '{class_names[idx[1]]}': " f"{cm[idx]} 次") cm[idx] = 0 # 置信度分布图 correct_probs = all_probs[all_preds == all_targets].max(axis=1) wrong_probs = all_probs[all_preds != all_targets].max(axis=1) plt.figure(figsize=(10, 5)) plt.hist(correct_probs, bins=30, alpha=0.6, label='Correct', color='green') plt.hist(wrong_probs, bins=30, alpha=0.6, label='Wrong', color='red') plt.xlabel('Confidence') plt.ylabel('Count') plt.title('Prediction Confidence Distribution') plt.legend() plt.show()

class GradCAM: """Grad-CAM 可视化工具""" def __init__(self, model, target_layer): self.model = model self.target_layer = target_layer self.gradients = None self.activations = None self._register_hooks() def _register_hooks(self): def forward_hook(module, input, output): self.activations = output def backward_hook(module, grad_in, grad_out): self.gradients = grad_out[0] self.target_layer.register_forward_hook(forward_hook) self.target_layer.register_backward_hook(backward_hook) def generate(self, input_image, class_idx=None): """生成 Grad-CAM 热力图""" # 前向传播 model_output = self.model(input_image) if class_idx is None: class_idx = model_output.argmax(dim=1).item() # 反向传播（仅针对目标类别） self.model.zero_grad() one_hot = torch.zeros_like(model_output) one_hot[0][class_idx] = 1 model_output.backward(gradient=one_hot) # 计算 Grad-CAM 权重 gradients = self.gradients[0] # [C, H, W] activations = self.activations[0] # [C, H', W'] # 全局平均池化获得权重 weights = gradients.mean(dim=(1, 2), keepdim=True) # 加权激活图 cam = (weights * activations).sum(dim=0) cam = torch.relu(cam) # 仅保留正值 # 归一化到 [0, 1] cam = cam - cam.min() cam = cam / (cam.max() + 1e-8) # 调整尺寸到原图大小 from torch.nn.functional import interpolate cam = interpolate( cam.unsqueeze(0).unsqueeze(0), size=input_image.shape[2:], mode='bilinear', align_corners=False ).squeeze() return cam.detach().cpu().numpy() def visualize_gradcam(model, image_tensor, class_names, target_layers, device): """可视化多个层的 Grad-CAM""" model.eval() image_tensor = image_tensor.to(device) fig, axes = plt.subplots( 1, len(target_layers) + 1, figsize=(5*(len(target_layers)+1), 5)) # 原始图像 img = image_tensor.cpu().squeeze() img = img * torch.tensor([0.229, 0.224, 0.225])[:, None, None] img = img + torch.tensor([0.485, 0.456, 0.406])[:, None, None] img = img.permute(1, 2, 0).numpy() img = np.clip(img, 0, 1) with torch.no_grad(): output = model(image_tensor) pred_idx = output.argmax(dim=1).item() axes[0].imshow(img) axes[0].set_title(f"Original\nPred: {class_names[pred_idx]}") axes[0].axis('off') # 生成各层的 Grad-CAM for i, (name, layer) in enumerate(target_layers): gradcam = GradCAM(model, layer) cam = gradcam.generate(image_tensor, class_idx=pred_idx) axes[i+1].imshow(img) axes[i+1].imshow(cam, cmap='jet', alpha=0.5) axes[i+1].set_title(f"{name}") axes[i+1].axis('off') plt.tight_layout() plt.show() # 使用示例 # 假设已加载模型和图像 target_layers = [ ("layer1", model.layer1[-1]), ("layer2", model.layer2[-1]), ("layer3", model.layer3[-1]), ("layer4", model.layer4[-1]), ] # visualize_gradcam(model, image_tensor, class_names, # target_layers, device)

class TrainingVisualizer: """训练过程可视化工具""" def __init__(self): self.train_losses = [] self.val_losses = [] self.train_accs = [] self.val_accs = [] self.lr_history = [] def update(self, train_loss, train_acc, val_loss, val_acc, lr): self.train_losses.append(train_loss) self.train_accs.append(train_acc) self.val_losses.append(val_loss) self.val_accs.append(val_acc) self.lr_history.append(lr) def plot(self, save_path='training_curves.png'): fig, axes = plt.subplots(1, 3, figsize=(18, 5)) # 损失曲线 axes[0].plot(self.train_losses, label='Train Loss', color='blue') axes[0].plot(self.val_losses, label='Val Loss', color='red') axes[0].set_xlabel('Epoch') axes[0].set_ylabel('Loss') axes[0].set_title('Training and Validation Loss') axes[0].legend() axes[0].grid(True, alpha=0.3) # 准确率曲线 axes[1].plot(self.train_accs, label='Train Acc', color='blue') axes[1].plot(self.val_accs, label='Val Acc', color='red') axes[1].set_xlabel('Epoch') axes[1].set_ylabel('Accuracy (%)') axes[1].set_title('Training and Validation Accuracy') axes[1].legend() axes[1].grid(True, alpha=0.3) # 学习率曲线 axes[2].plot(self.lr_history, label='Learning Rate', color='green') axes[2].set_xlabel('Epoch') axes[2].set_ylabel('LR') axes[2].set_title('Learning Rate Schedule') axes[2].legend() axes[2].grid(True, alpha=0.3) plt.tight_layout() plt.savefig(save_path, dpi=150) plt.show() def find_best_epoch(self, metric='val_acc'): """找出最佳验证准确率对应的 epoch""" if metric == 'val_acc': best_idx = np.argmax(self.val_accs) best_value = self.val_accs[best_idx] elif metric == 'val_loss': best_idx = np.argmin(self.val_losses) best_value = self.val_losses[best_idx] print(f"Best {metric} at epoch {best_idx+1}: {best_value:.4f}") return best_idx, best_value

""" 图像分类完整实战 - 端到端训练脚本 基于 PyTorch + torchvision 支持: CIFAR-10/100, ImageFolder, 自定义数据集 """ import os import sys import time import json import numpy as np import torch import torch.nn as nn import torch.optim as optim import torchvision import torchvision.transforms as transforms from torchvision.datasets import ImageFolder from torch.utils.data import DataLoader, random_split from torch.cuda.amp import autocast, GradScaler import matplotlib.pyplot as plt # ============ 1. 配置类 ============ class Config: # 数据配置 dataset = 'cifar10' # cifar10 / cifar100 / custom data_path = './data' num_classes = 10 input_size = 224 # 模型配置 model_name = 'resnet50' # resnet50 / resnet152 / vit_b_16 pretrained = True # 训练配置 batch_size = 64 epochs = 100 lr = 0.001 weight_decay = 1e-4 momentum = 0.9 # 优化配置 use_amp = True # 混合精度 label_smoothing = 0.1 grad_clip = 1.0 # 数据增强 use_randaugment = True augmentation_magnitude = 9 # 系统配置 device = 'cuda' if torch.cuda.is_available() else 'cpu' num_workers = 4 seed = 42 save_dir = './checkpoints' @classmethod def setup(cls): os.makedirs(cls.save_dir, exist_ok=True) torch.manual_seed(cls.seed) np.random.seed(cls.seed) if torch.cuda.is_available(): torch.cuda.manual_seed_all(cls.seed) print(f"Using device: {cls.device}") print(f"PyTorch version: {torch.__version__}") # ============ 2. 数据加载 ============ def get_data_loaders(cfg): """根据配置加载数据集""" # ImageNet 归一化参数 mean, std = [0.485, 0.456, 0.406], [0.229, 0.224, 0.225] # 训练集增强 train_transforms = [ transforms.RandomResizedCrop( cfg.input_size, scale=(0.08, 1.0)), transforms.RandomHorizontalFlip(), ] if cfg.use_randaugment: train_transforms.append( transforms.RandAugment( num_ops=2, magnitude=cfg.augmentation_magnitude)) train_transforms.extend([ transforms.ToTensor(), transforms.Normalize(mean, std), ]) train_transform = transforms.Compose(train_transforms) val_transform = transforms.Compose([ transforms.Resize(int(cfg.input_size * 1.14)), transforms.CenterCrop(cfg.input_size), transforms.ToTensor(), transforms.Normalize(mean, std), ]) # 加载 CIFAR-10 if cfg.dataset == 'cifar10': cfg.num_classes = 10 trainset = torchvision.datasets.CIFAR10( root=cfg.data_path, train=True, download=True, transform=train_transform) testset = torchvision.datasets.CIFAR10( root=cfg.data_path, train=False, download=True, transform=val_transform) class_names = trainset.classes elif cfg.dataset == 'cifar100': cfg.num_classes = 100 trainset = torchvision.datasets.CIFAR100( root=cfg.data_path, train=True, download=True, transform=train_transform) testset = torchvision.datasets.CIFAR100( root=cfg.data_path, train=False, download=True, transform=val_transform) class_names = trainset.classes else: # 自定义数据集 full_dataset = ImageFolder( root=cfg.data_path, transform=train_transform) train_size = int(0.85 * len(full_dataset)) val_size = len(full_dataset) - train_size trainset, testset = random_split( full_dataset, [train_size, val_size]) class_names = full_dataset.classes cfg.num_classes = len(class_names) train_loader = DataLoader( trainset, batch_size=cfg.batch_size, shuffle=True, num_workers=cfg.num_workers, pin_memory=True, drop_last=True) test_loader = DataLoader( testset, batch_size=cfg.batch_size, shuffle=False, num_workers=cfg.num_workers, pin_memory=True) print(f"Dataset: {cfg.dataset}") print(f"Classes: {cfg.num_classes}") print(f"Train samples: {len(trainset)}") print(f"Test samples: {len(testset)}") return train_loader, test_loader, class_names # ============ 3. 模型创建 ============ def create_model(cfg): """创建并配置模型""" model_fn = getattr(torchvision.models, cfg.model_name) if cfg.pretrained: model = model_fn(weights='IMAGENET1K_V2') print("Loaded ImageNet pretrained weights") else: model = model_fn(weights=None) # 替换分类头 if 'vit' in cfg.model_name: in_features = model.heads.head.in_features model.heads.head = nn.Linear( in_features, cfg.num_classes) elif 'efficientnet' in cfg.model_name: in_features = model.classifier[-1].in_features model.classifier[-1] = nn.Linear( in_features, cfg.num_classes) elif 'convnext' in cfg.model_name: in_features = model.classifier[-1].in_features model.classifier[-1] = nn.Linear( in_features, cfg.num_classes) else: # ResNet 等 in_features = model.fc.in_features model.fc = nn.Linear(in_features, cfg.num_classes) model = model.to(cfg.device) 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: {total_params/1e6:.2f}M") print(f"Trainable params: {trainable_params/1e6:.2f}M") return model # ============ 4. 训练与验证 ============ class Trainer: def __init__(self, model, train_loader, val_loader, cfg): self.model = model self.train_loader = train_loader self.val_loader = val_loader self.cfg = cfg # 损失函数（标签平滑） self.criterion = nn.CrossEntropyLoss( label_smoothing=cfg.label_smoothing) # 优化器 self.optimizer = optim.SGD( model.parameters(), lr=cfg.lr, momentum=cfg.momentum, weight_decay=cfg.weight_decay, nesterov=True) # 学习率调度器 self.scheduler = optim.lr_scheduler.CosineAnnealingLR( self.optimizer, T_max=cfg.epochs) # 混合精度 self.scaler = GradScaler() if cfg.use_amp else None # 训练记录 self.history = { 'train_loss': [], 'train_acc': [], 'val_loss': [], 'val_acc': [], 'lr': [] } self.best_val_acc = 0.0 def train_epoch(self): self.model.train() running_loss = 0.0 correct = 0 total = 0 for inputs, targets in self.train_loader: inputs = inputs.to(self.cfg.device) targets = targets.to(self.cfg.device) self.optimizer.zero_grad() if self.scaler is not None: with autocast(): outputs = self.model(inputs) loss = self.criterion(outputs, targets) self.scaler.scale(loss).backward() self.scaler.unscale_(self.optimizer) nn.utils.clip_grad_norm_( self.model.parameters(), self.cfg.grad_clip) self.scaler.step(self.optimizer) self.scaler.update() else: outputs = self.model(inputs) loss = self.criterion(outputs, targets) loss.backward() nn.utils.clip_grad_norm_( self.model.parameters(), self.cfg.grad_clip) self.optimizer.step() running_loss += loss.item() _, predicted = outputs.max(1) total += targets.size(0) correct += predicted.eq(targets).sum().item() return (running_loss / len(self.train_loader), 100. * correct / total) @torch.no_grad() def validate(self): self.model.eval() running_loss = 0.0 correct = 0 total = 0 for inputs, targets in self.val_loader: inputs = inputs.to(self.cfg.device) targets = targets.to(self.cfg.device) outputs = self.model(inputs) loss = self.criterion(outputs, targets) running_loss += loss.item() _, predicted = outputs.max(1) total += targets.size(0) correct += predicted.eq(targets).sum().item() return (running_loss / len(self.val_loader), 100. * correct / total) def run(self): print(f"\n{'='*50}") print(f"Training {self.cfg.model_name}") print(f"{'='*50}") start_time = time.time() for epoch in range(1, self.cfg.epochs + 1): train_loss, train_acc = self.train_epoch() val_loss, val_acc = self.validate() self.history['train_loss'].append(train_loss) self.history['train_acc'].append(train_acc) self.history['val_loss'].append(val_loss) self.history['val_acc'].append(val_acc) current_lr = self.optimizer.param_groups[0]['lr'] self.history['lr'].append(current_lr) self.scheduler.step() # 保存最佳模型 if val_acc > self.best_val_acc: self.best_val_acc = val_acc torch.save({ 'epoch': epoch, 'model_state_dict': self.model.state_dict(), 'optimizer_state_dict': self.optimizer.state_dict(), 'val_acc': val_acc, }, os.path.join( self.cfg.save_dir, 'best_model.pth')) if epoch % 5 == 0 or epoch == 1: elapsed = time.time() - start_time print(f"Epoch {epoch:3d}/{self.cfg.epochs} | " f"Train Loss: {train_loss:.4f} | " f"Train Acc: {train_acc:.2f}% | " f"Val Loss: {val_loss:.4f} | " f"Val Acc: {val_acc:.2f}% | " f"LR: {current_lr:.6f} | " f"Time: {elapsed:.0f}s") elapsed = time.time() - start_time print(f"\nTraining completed in {elapsed/60:.1f} minutes") print(f"Best validation accuracy: {self.best_val_acc:.2f}%") return self.history # ============ 5. 主函数 ============ def main(): Config.setup() train_loader, test_loader, class_names = \ get_data_loaders(Config) model = create_model(Config) trainer = Trainer( model, train_loader, test_loader, Config) history = trainer.run() # 保存训练历史 with open(os.path.join(Config.save_dir, 'history.json'), 'w') as f: json.dump(history, f, indent=2) print(f"\nTraining complete. Best accuracy: " f"{trainer.best_val_acc:.2f}%") print(f"Model saved to {Config.save_dir}/") if __name__ == '__main__': main()

# ---- TorchScript 导出（跨平台部署） ---- model.eval() example_input = torch.randn(1, 3, 224, 224).to(device) # Tracing 方式导出 traced_model = torch.jit.trace(model, example_input) traced_model.save('model_traced.pt') # Scripting 方式导出（支持控制流） scripted_model = torch.jit.script(model) scripted_model.save('model_scripted.pt') # 加载导出的模型 loaded_model = torch.jit.load('model_traced.pt') loaded_model.eval() # ---- ONNX 导出（跨框架部署） ---- torch.onnx.export( model, example_input, 'model.onnx', input_names=['input'], output_names=['output'], dynamic_axes={ 'input': {0: 'batch_size'}, 'output': {0: 'batch_size'} }, opset_version=17, do_constant_folding=True ) print("Model exported to ONNX format")

# ---- FP16 推理加速 ---- model.half() # 将模型转换为 FP16 with torch.no_grad(): for inputs in test_loader: inputs = inputs.half().to(device) outputs = model(inputs) # ---- TensorRT 集成（NVIDIA GPU 最优推理） ---- # 需要先安装 torch-tensorrt # pip install torch-tensorrt import torch_tensorrt # 编译为 TensorRT 引擎 trt_model = torch_tensorrt.compile( model, inputs=[torch.randn(1, 3, 224, 224)], enabled_precisions={torch.half}, workspace_size=1 << 30, # 1GB ) # 推理 API class ImageClassifier: """完整推理服务封装""" def __init__(self, model_path, class_names, device='cuda'): self.device = device self.class_names = class_names self.model = torch.jit.load(model_path) self.model.eval().to(device) self.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]), ]) @torch.no_grad() def predict(self, image_path, top_k=5): from PIL import Image img = Image.open(image_path).convert('RGB') img_tensor = self.transform(img).unsqueeze(0).to(self.device) outputs = self.model(img_tensor) probs = torch.softmax(outputs, dim=1) top_probs, top_indices = probs.topk(top_k, dim=1) results = [] for i in range(top_k): results.append({ 'class': self.class_names[top_indices[0][i]], 'probability': top_probs[0][i].item() }) return results def batch_predict(self, image_paths, batch_size=32): """批量推理""" results = [] for i in range(0, len(image_paths), batch_size): batch = image_paths[i:i+batch_size] batch_tensors = [] for path in batch: img = Image.open(path).convert('RGB') batch_tensors.append( self.transform(img)) batch_tensor = torch.stack(batch_tensors).to(self.device) outputs = self.model(batch_tensor) probs = torch.softmax(outputs, dim=1) results.append(probs.cpu()) return torch.cat(results, dim=0)