""" ResNet50蒸馏训练ResNet18实现 学生网络使用ArcFace损失 支持单机双卡训练 """ import os import torch import torch.nn as nn import torch.distributed as dist import torch.multiprocessing as mp from torch.nn.parallel import DistributedDataParallel as DDP from torch.optim.lr_scheduler import CosineAnnealingLR from torch.cuda.amp import GradScaler from model import resnet18, resnet50, ArcFace from tqdm import tqdm import torch.nn.functional as F from tools.dataset import load_data # from config import config as conf import yaml import math def setup(rank, world_size): os.environ['MASTER_ADDR'] = '0.0.0.0' os.environ['MASTER_PORT'] = '12355' dist.init_process_group("nccl", rank=rank, world_size=world_size) def cleanup(): dist.destroy_process_group() class DistillTrainer: def __init__(self, rank, world_size, conf): self.rank = rank self.world_size = world_size self.device = torch.device(f'cuda:{rank}') # 初始化模型 self.teacher = resnet50(pretrained=True, scale=conf['models']['channel_ratio']).to(self.device) self.student = resnet18(pretrained=True, scale=conf['models']['student_channel_ratio']).to(self.device) # 加载预训练教师模型 # teacher_path = os.path.join('checkpoints', 'resnet50_0519', 'best.pth') teacher_path = conf['models']['teacher_model_path'] if os.path.exists(teacher_path): teacher_state = torch.load(teacher_path, map_location=self.device) new_state_dict = {} for k, v in teacher_state.items(): if k.startswith('module.'): new_state_dict[k[7:]] = v # 去除前7个字符'module.' else: new_state_dict[k] = v # 加载处理后的状态字典 self.teacher.load_state_dict(new_state_dict, strict=False) if self.rank == 0: print(f"Successfully loaded teacher model from {teacher_path}") else: raise FileNotFoundError(f"Teacher model weights not found at {teacher_path}") # 数据加载 self.train_loader, num_classes = load_data(training=True, cfg=conf) self.val_loader, _ = load_data(training=False, cfg=conf) # ArcFace损失 self.metric = ArcFace(conf['base']['embedding_size'], num_classes).to(self.device) # 分布式训练 if world_size > 1: self.teacher = DDP(self.teacher, device_ids=[rank]) self.student = DDP(self.student, device_ids=[rank]) self.metric = DDP(self.metric, device_ids=[rank]) # 优化器 self.optimizer = torch.optim.SGD([ {'params': self.student.parameters()}, {'params': self.metric.parameters()} ], lr=conf['training']['lr'], momentum=0.9, weight_decay=5e-4) self.scheduler = CosineAnnealingLR(self.optimizer, T_max=conf['training']['epochs']) self.scaler = GradScaler() # 损失函数 self.arcface_loss = nn.CrossEntropyLoss() self.distill_loss = nn.KLDivLoss(reduction='batchmean') self.conf = conf def cosine_annealing(self, epoch, total_epochs, initial_weight, final_weight=0.1): """ 余弦退火法动态调整蒸馏权重 参数： epoch: 当前训练轮次 total_epochs: 总训练轮次 initial_weight: 初始蒸馏权重（如0.8） final_weight: 最终蒸馏权重（如0.1） 返回： 当前轮次的蒸馏权重 """ return final_weight + 0.5 * (initial_weight - final_weight) * (1 + math.cos(math.pi * epoch / total_epochs)) def train_epoch(self, epoch): self.teacher.eval() self.student.train() if self.rank == 0: print(f"\nTeacher network type: {type(self.teacher)}") print(f"Student network type: {type(self.student)}") total_loss = 0 for data, labels in tqdm(self.train_loader, desc=f"Epoch {epoch}"): data = data.to(self.device) labels = labels.to(self.device) # with autocast(): # 教师输出 with torch.no_grad(): teacher_logits = self.teacher(data) # 学生输出 student_features = self.student(data) student_logits = self.metric(student_features, labels) # 计算损失 arc_loss = self.arcface_loss(student_logits, labels) distill_loss = self.distill_loss( F.log_softmax(student_features / self.conf['training']['temperature'], dim=1), F.softmax(teacher_logits / self.conf['training']['temperature'], dim=1) ) * (self.conf['training']['temperature'] ** 2) # 温度缩放后需要乘以T^2保持梯度规模 current_distill_weight = self.cosine_annealing(epoch, self.conf['training']['epochs'], self.conf['training']['distill_weight']) loss = (1-current_distill_weight) * arc_loss + current_distill_weight * distill_loss self.optimizer.zero_grad() self.scaler.scale(loss).backward() self.scaler.step(self.optimizer) self.scaler.update() total_loss += loss.item() self.scheduler.step() return total_loss / len(self.train_loader) def validate(self): self.student.eval() total_loss = 0 correct = 0 total = 0 with torch.no_grad(): for data, labels in self.val_loader: data = data.to(self.device) labels = labels.to(self.device) features = self.student(data) logits = self.metric(features, labels) loss = self.arcface_loss(logits, labels) total_loss += loss.item() _, predicted = torch.max(logits.data, 1) total += labels.size(0) correct += (predicted == labels).sum().item() return total_loss / len(self.val_loader), correct / total def save_checkpoint(self, epoch, is_best=False): if self.rank != 0: return state = { 'epoch': epoch, 'student_state_dict': self.student.state_dict(), 'metric_state_dict': self.metric.state_dict(), 'optimizer_state_dict': self.optimizer.state_dict(), } filename = 'best.pth' if is_best else f'checkpoint_{epoch}.pth' if not os.path.exists(self.conf['training']['checkpoints']): os.makedirs(self.conf['training']['checkpoints']) if filename != 'best.pth': torch.save(state, os.path.join(self.conf['training']['checkpoints'], filename)) else: torch.save(state['student_state_dict'], os.path.join(self.conf['training']['checkpoints'], filename)) def train(rank, world_size): setup(rank, world_size) with open('configs/distill.yml', 'r') as f: conf = yaml.load(f, Loader=yaml.FullLoader) trainer = DistillTrainer(rank, world_size, conf) best_acc = 0 for epoch in range(conf['training']['epochs']): train_loss = trainer.train_epoch(epoch) val_loss, val_acc = trainer.validate() if rank == 0: print(f"Epoch {epoch}: Train Loss: {train_loss:.4f}, Val Loss: {val_loss:.4f}, Val Acc: {val_acc:.4f}") if val_acc > best_acc: best_acc = val_acc trainer.save_checkpoint(epoch, is_best=True) cleanup() if __name__ == '__main__': world_size = torch.cuda.device_count() if world_size > 1: mp.spawn(train, args=(world_size,), nprocs=world_size, join=True) else: train(0, 1)