PyTorch 1.7.1 戴眼镜分类模型训练：3种骨干网络对比与98.6%准确率复现

PyTorch 1.7.1 戴眼镜分类模型训练：3种骨干网络对比与98.6%准确率复现

在计算机视觉任务中，戴眼镜检测是一个具有实际应用价值的分类问题。本文将基于PyTorch 1.7.1框架，对比Mobilenet_v2、Googlenet和Resnet18三种骨干网络在戴眼镜分类任务上的表现，并提供完整的训练代码和参数配置，帮助开发者快速复现98.6%的高准确率。

1. 环境准备与数据加载

首先需要配置PyTorch 1.7.1环境。建议使用Python 3.7或3.8版本以避免兼容性问题：

conda create -n glasses_cls python=3.7 conda activate glasses_cls pip install torch==1.7.1+cu110 torchvision==0.8.2+cu110 -f https://download.pytorch.org/whl/torch_stable.html pip install opencv-python pillow pandas pyyaml tqdm tensorboard

戴眼镜数据集包含约20,000张标注图像，分为训练集和测试集：

• 训练集：戴眼镜10,475张，不戴眼镜12,841张
• 测试集：戴眼镜和不戴眼镜各1,000张

数据目录结构应如下：

eyeglasses-dataset/ ├── train/ │ ├── face/ # 不戴眼镜图像 │ └── face-eyeglasses/ # 戴眼镜图像 └── test/ ├── face/ └── face-eyeglasses/

使用自定义Dataset类加载数据：

from torch.utils.data import Dataset from PIL import Image import os class EyeglassesDataset(Dataset): def __init__(self, root_dir, transform=None): self.root_dir = root_dir self.transform = transform self.samples = [] for label, class_name in enumerate(['face', 'face-eyeglasses']): class_dir = os.path.join(root_dir, class_name) for img_name in os.listdir(class_dir): self.samples.append((os.path.join(class_dir, img_name), label)) def __len__(self): return len(self.samples) def __getitem__(self, idx): img_path, label = self.samples[idx] image = Image.open(img_path).convert('RGB') if self.transform: image = self.transform(image) return image, label

2. 模型架构与训练配置

我们对比三种经典CNN架构的性能表现：

训练配置参数如下：

# config.yaml input_size: [112, 112] batch_size: 32 lr: 0.01 momentum: 0.9 weight_decay: 0.0005 num_epochs: 100 milestones: [20, 50, 80] # 学习率调整时机

数据增强策略对模型性能至关重要：

from torchvision import transforms train_transform = transforms.Compose([ transforms.RandomHorizontalFlip(), transforms.RandomRotation(10), transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2), transforms.Resize(128), transforms.RandomCrop(112), transforms.ToTensor(), transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]) ]) val_transform = transforms.Compose([ transforms.Resize(112), transforms.ToTensor(), transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]) ])

3. 训练过程实现

完整的训练流程包括模型初始化、损失函数定义和训练循环：

import torch import torch.nn as nn import torch.optim as optim from torch.utils.tensorboard import SummaryWriter def train_model(model, train_loader, val_loader, config): device = torch.device("cuda" if torch.cuda.is_available() else "cpu") model = model.to(device) criterion = nn.CrossEntropyLoss() optimizer = optim.SGD(model.parameters(), lr=config['lr'], momentum=config['momentum'], weight_decay=config['weight_decay']) scheduler = optim.lr_scheduler.MultiStepLR(optimizer, milestones=config['milestones'], gamma=0.1) writer = SummaryWriter() best_acc = 0.0 for epoch in range(config['num_epochs']): model.train() running_loss = 0.0 for inputs, labels in train_loader: inputs, labels = inputs.to(device), labels.to(device) optimizer.zero_grad() outputs = model(inputs) loss = criterion(outputs, labels) loss.backward() optimizer.step() running_loss += loss.item() scheduler.step() # 验证集评估 val_loss, val_acc = evaluate(model, val_loader, criterion, device) # 记录TensorBoard日志 writer.add_scalar('Loss/train', running_loss/len(train_loader), epoch) writer.add_scalar('Loss/val', val_loss, epoch) writer.add_scalar('Accuracy/val', val_acc, epoch) # 保存最佳模型 if val_acc > best_acc: best_acc = val_acc torch.save(model.state_dict(), 'best_model.pth') print(f'Epoch {epoch+1}/{config["num_epochs"]} | ' f'Train Loss: {running_loss/len(train_loader):.4f} | ' f'Val Loss: {val_loss:.4f} | Val Acc: {val_acc:.2%}') writer.close() return model

评估函数实现：

def evaluate(model, data_loader, criterion, device): model.eval() total_loss = 0.0 correct = 0 total = 0 with torch.no_grad(): for inputs, labels in data_loader: inputs, labels = inputs.to(device), labels.to(device) outputs = model(inputs) loss = criterion(outputs, labels) total_loss += loss.item() _, predicted = torch.max(outputs.data, 1) total += labels.size(0) correct += (predicted == labels).sum().item() return total_loss/len(data_loader), correct/total

4. 三种骨干网络性能对比

我们使用相同的训练配置对三种模型进行训练，结果如下：

关键训练曲线对比：

1. 损失函数下降趋势：

   • Mobilenet_v2收敛最快，20个epoch后趋于稳定
   • Resnet18初期下降较慢但最终达到最低loss

2. 验证集准确率：

   • 三种模型在50个epoch后准确率均超过98%
   • Googlenet和Resnet18在后期仍有小幅提升

3. 资源消耗：

   • Mobilenet_v2显存占用仅为Resnet18的55%
   • 训练时间方面Mobilenet_v2优势明显

# 模型初始化示例 import torchvision.models as models def init_model(model_name, pretrained=True): if model_name == 'mobilenet_v2': model = models.mobilenet_v2(pretrained=pretrained) model.classifier[1] = nn.Linear(model.last_channel, 2) elif model_name == 'googlenet': model = models.googlenet(pretrained=pretrained) model.fc = nn.Linear(1024, 2) elif model_name == 'resnet18': model = models.resnet18(pretrained=pretrained) model.fc = nn.Linear(512, 2) return model

5. 高准确率复现关键技巧

要达到98%以上的准确率，需注意以下关键点：

1. 数据清洗：

   • 检查并移除标注错误的样本
   • 确保人脸区域裁剪准确
   • 平衡正负样本比例

2. 训练技巧：

   • 使用学习率warmup：前3个epoch线性增加学习率
   • 启用混合精度训练：减少显存占用，加快训练速度
   • 添加Label Smoothing：缓解过拟合

# 混合精度训练示例 from torch.cuda.amp import autocast, GradScaler scaler = GradScaler() with autocast(): outputs = model(inputs) loss = criterion(outputs, labels) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update()

3. 模型微调：
   • 不同层设置不同学习率
   • 冻结底层特征提取层，只训练分类头

# 分层学习率设置示例 optim_params = [ {'params': model.features.parameters(), 'lr': config['lr']*0.1}, {'params': model.classifier.parameters(), 'lr': config['lr']} ] optimizer = optim.SGD(optim_params, momentum=0.9)

6. 模型部署与推理

训练完成后，可以使用以下代码进行单张图像预测：

import cv2 from PIL import Image def predict(image_path, model_path, transform): device = torch.device("cuda" if torch.cuda.is_available() else "cpu") model = init_model('mobilenet_v2') model.load_state_dict(torch.load(model_path)) model.eval() image = Image.open(image_path).convert('RGB') image = transform(image).unsqueeze(0).to(device) with torch.no_grad(): output = model(image) prob = torch.softmax(output, dim=1) pred = torch.argmax(prob).item() return '戴眼镜' if pred == 1 else '不戴眼镜', prob[0][1].item()

对于视频流实时检测，建议采用以下优化策略：

1. 使用OpenCV的DNN模块加载TorchScript格式模型
2. 采用异步处理避免阻塞主线程
3. 添加帧缓存机制减少重复计算

# TorchScript模型导出 model = init_model('mobilenet_v2') model.load_state_dict(torch.load('best_model.pth')) model.eval() example = torch.rand(1, 3, 112, 112).to(device) traced_script = torch.jit.trace(model, example) traced_script.save('eyeglasses_detection.pt')

7. 常见问题与解决方案

在实际项目中可能会遇到以下典型问题：

1. CUDA内存不足：

   • 减小batch_size
   • 使用梯度累积
   • 启用混合精度训练

2. 过拟合：

   • 增加数据增强
   • 添加Dropout层
   • 使用更小的学习率

3. 准确率波动大：

   • 检查数据标注一致性
   • 调整学习率衰减策略
   • 增加训练epoch

# 梯度累积实现 accum_steps = 4 for i, (inputs, labels) in enumerate(train_loader): outputs = model(inputs) loss = criterion(outputs, labels) loss = loss / accum_steps loss.backward() if (i+1) % accum_steps == 0: optimizer.step() optimizer.zero_grad()