import os import sys import time import json import math import logging from dataclasses import dataclass from datetime import datetime from typing import Optional, Tuple import torch import torch.nn as nn import torch.nn.functional as F import torch.optim as optim from torch.utils.data import DataLoader from torchvision import datasets, transforms from tqdm import tqdm import matplotlib.pyplot as plt # 将项目根目录加入路径，便于导入 sys.path.append(os.path.join(os.path.dirname(__file__), '..')) from net.resnet_embedding import create_resnet50_embedding from settings import settings # 日志 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s') logger = logging.getLogger(__name__) device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') print(f'使用设备: {device}') @dataclass class TaskConfig: name: str train_dir: str val_dir: str embedding_dim: int batch_size: int lr: float aug_strength: str # "strong" | "medium" | "shape" TASKS = { 'dish': TaskConfig( name='DishClassification', train_dir=os.path.join(settings.BASE_DIR, 'dataset', 'DishClassification', 'train'), val_dir=os.path.join(settings.BASE_DIR, 'dataset', 'DishClassification', 'val'), embedding_dim=512, batch_size=32, lr=5e-4, aug_strength='medium', ), 'whole_ingredient': TaskConfig( name='WholeIngredientRecognition', train_dir=os.path.join(settings.BASE_DIR, 'dataset', 'WholeIngredientRecognition', 'train'), val_dir=os.path.join(settings.BASE_DIR, 'dataset', 'WholeIngredientRecognition', 'val'), embedding_dim=512, batch_size=64, lr=8e-4, aug_strength='medium', ), 'processed_ingredient': TaskConfig( name='ProcessedIngredientRecognition', train_dir=os.path.join(settings.BASE_DIR, 'dataset', 'ProcessedIngredientRecognition', 'train'), val_dir=os.path.join(settings.BASE_DIR, 'dataset', 'ProcessedIngredientRecognition', 'val'), embedding_dim=512, batch_size=32, lr=5e-4, aug_strength='medium', ), } def build_transforms(aug_strength: str): if aug_strength == 'strong': return transforms.Compose([ transforms.Resize((224, 224)), transforms.RandomHorizontalFlip(p=0.5), transforms.RandomRotation(20), transforms.ColorJitter(0.3, 0.3, 0.3, 0.1), transforms.RandomAffine(degrees=0, translate=(0.12, 0.12)), transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), ]), transforms.Compose([ transforms.Resize((224, 224)), transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), ]) if aug_strength == 'medium': return transforms.Compose([ transforms.Resize((224, 224)), transforms.RandomHorizontalFlip(p=0.5), transforms.RandomRotation(15), transforms.ColorJitter(0.2, 0.2, 0.2, 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]), ]), transforms.Compose([ transforms.Resize((224, 224)), transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), ]) # shape: 强调几何与尺度，弱化强色抖动 return transforms.Compose([ transforms.Resize((224, 224)), transforms.RandomHorizontalFlip(p=0.5), transforms.RandomAffine(degrees=15, translate=(0.1, 0.1), scale=(0.9, 1.1)), transforms.RandomPerspective(distortion_scale=0.3, p=0.3), transforms.GaussianBlur(kernel_size=3, sigma=(0.1, 1.0)), transforms.ColorJitter(0.1, 0.1, 0.1, 0.03), transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), ]), transforms.Compose([ transforms.Resize((224, 224)), transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), ]) class CosFaceHead(nn.Module): def __init__(self, in_features: int, num_classes: int, s: float = 64.0, m: float = 0.35): super().__init__() self.s = s self.m = m self.weight = nn.Parameter(torch.randn(num_classes, in_features)) nn.init.xavier_normal_(self.weight) def forward(self, features: torch.Tensor, labels: Optional[torch.Tensor] = None): # 假设输入features已L2归一化（backbone已做），此处仍做一次以保证数值稳健 x = F.normalize(features, dim=1) W = F.normalize(self.weight, dim=1) cos_theta = torch.mm(x, W.t()) # [B, C] if labels is None: return self.s * cos_theta one_hot = F.one_hot(labels, num_classes=W.size(0)).float() cos_theta_m = cos_theta - one_hot * self.m logits = self.s * cos_theta_m return logits, cos_theta def accuracy_top1(logits: torch.Tensor, labels: torch.Tensor) -> float: pred = logits.argmax(dim=1) return (pred == labels).float().mean().item() * 100.0 def evaluate_open_set(cos_scores_known: torch.Tensor, cos_scores_unknown: torch.Tensor, far: float = 0.05) -> Tuple[float, float, float]: """ 基于max_cos分数的简单阈值估计：给定未知集FAR，返回阈值和两侧TPR/FPR。 返回: (threshold, known_accept_rate, unknown_reject_rate) """ # 阈值取未知集分布的(1 - FAR)分位数（使未知中约 FAR 被错误接受） threshold = torch.quantile(cos_scores_unknown, 1 - far).item() if len(cos_scores_unknown) > 0 else 0.5 known_accept = (cos_scores_known >= threshold).float().mean().item() if len(cos_scores_known) > 0 else 0.0 unknown_reject = (cos_scores_unknown < threshold).float().mean().item() if len(cos_scores_unknown) > 0 else 0.0 return threshold, known_accept, unknown_reject def plot_training_curves(train_losses, val_losses, train_accuracies, val_accuracies, save_path: str): epochs = range(1, len(train_losses) + 1) plt.figure(figsize=(12, 5)) plt.subplot(1, 2, 1) plt.plot(epochs, train_losses, label='Train Loss') plt.plot(epochs, val_losses, label='Val Loss') plt.xlabel('Epoch'); plt.ylabel('Loss'); plt.title('Loss Curve'); plt.grid(True); plt.legend() plt.subplot(1, 2, 2) plt.plot(epochs, train_accuracies, label='Train Acc') plt.plot(epochs, val_accuracies, label='Val Acc') plt.xlabel('Epoch'); plt.ylabel('Accuracy (%)'); plt.title('Accuracy Curve'); plt.grid(True); plt.legend() plt.tight_layout(); plt.savefig(save_path, dpi=300, bbox_inches='tight'); plt.close() def run_epoch(model, head, loader, criterion, optimizer=None, train: bool = True): if train: model.train(); head.train() else: model.eval(); head.eval() running_loss = 0.0 running_acc = 0.0 n_batches = 0 with torch.set_grad_enabled(train): pbar = tqdm(loader, desc='训练中' if train else '验证中') for images, labels in pbar: images = images.to(device); labels = labels.to(device) feats = model(images) # 已L2归一化 if train: logits, _ = head(feats, labels) loss = criterion(logits, labels) optimizer.zero_grad(); loss.backward(); optimizer.step() else: logits = head(feats) # 不加边距 loss = criterion(logits, labels) acc = accuracy_top1(logits, labels) running_loss += loss.item() running_acc += acc n_batches += 1 pbar.set_postfix({ 'Loss': f'{running_loss / n_batches:.4f}', 'Acc': f'{running_acc / n_batches:.2f}%' }) return running_loss / max(n_batches, 1), running_acc / max(n_batches, 1) def collect_max_cos_scores(model, head, loader) -> torch.Tensor: model.eval(); head.eval() scores = [] with torch.no_grad(): for images, _ in loader: images = images.to(device) feats = model(images) _, cos = head(feats, labels=None), None # head(feats)返回logits= s*cos # 还需要原始cos分数：logits/s logits = head(feats) # s*cos cos_scores = logits / getattr(head, 's', 64.0) max_cos = cos_scores.max(dim=1).values scores.append(max_cos.cpu()) return torch.cat(scores, dim=0) if scores else torch.tensor([]) def main(task_key: str = 'dish', s: float = 64.0, m: float = 0.35, num_epochs: int = 60, unknown_dir: Optional[str] = None, far: float = 0.05): cfg = TASKS[task_key] timestamp = datetime.now().strftime('%Y%m%d_%H%M%S') save_dir = os.path.join(settings.BASE_DIR, 'model', cfg.name, f'cosface_{timestamp}') os.makedirs(save_dir, exist_ok=True) logger.info(f'[{cfg.name}] 模型保存目录: {save_dir}') transform_train, transform_val = build_transforms(cfg.aug_strength) # 数据集 train_ds = datasets.ImageFolder(cfg.train_dir, transform=transform_train) val_ds = datasets.ImageFolder(cfg.val_dir, transform=transform_val) train_loader = DataLoader(train_ds, batch_size=cfg.batch_size, shuffle=True, num_workers=0, drop_last=True) val_loader = DataLoader(val_ds, batch_size=cfg.batch_size, shuffle=False, num_workers=0) class_names = train_ds.classes num_classes = len(class_names) logger.info(f'类别数: {num_classes}; 类别: {class_names}') # 模型与头 model = create_resnet50_embedding(embedding_dim=cfg.embedding_dim, pretrained=True, use_internal_preprocess=False).to(device) head = CosFaceHead(in_features=cfg.embedding_dim, num_classes=num_classes, s=s, m=m).to(device) # 损失与优化器 criterion = nn.CrossEntropyLoss() optimizer = optim.AdamW(list(model.parameters()) + list(head.parameters()), lr=cfg.lr, weight_decay=1e-4) scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=num_epochs) best_val_acc = 0.0 best_model_path = os.path.join(save_dir, 'best_cosface_model.pth') train_losses, val_losses, train_accs, val_accs = [], [], [], [] start = time.time() logger.info('开始训练...') for epoch in range(num_epochs): logger.info(f'Epoch {epoch+1}/{num_epochs}') tr_loss, tr_acc = run_epoch(model, head, train_loader, criterion, optimizer, train=True) va_loss, va_acc = run_epoch(model, head, val_loader, criterion, optimizer=None, train=False) scheduler.step() train_losses.append(tr_loss); train_accs.append(tr_acc) val_losses.append(va_loss); val_accs.append(va_acc) logger.info(f' Train Loss: {tr_loss:.4f}, Acc: {tr_acc:.2f}%') logger.info(f' Val Loss: {va_loss:.4f}, Acc: {va_acc:.2f}%') logger.info(f' LR: {optimizer.param_groups[0]["lr"]:.6f}') if va_acc > best_val_acc: best_val_acc = va_acc torch.save({ 'epoch': epoch, 'backbone_state_dict': model.state_dict(), 'head_state_dict': head.state_dict(), 'optimizer_state_dict': optimizer.state_dict(), 'val_acc': va_acc, 'val_loss': va_loss, 'class_names': class_names, 'embedding_dim': cfg.embedding_dim, 's': s, 'm': m, }, best_model_path) logger.info(f'保存最佳模型: {best_model_path} (Val Acc: {va_acc:.2f}%)') elapsed = time.time() - start logger.info(f'训练完成，总耗时: {elapsed/3600:.2f} 小时') # 曲线 curves_path = os.path.join(save_dir, 'training_curves.png') plot_training_curves(train_losses, val_losses, train_accs, val_accs, curves_path) # 开放集评估（可选） threshold = None known_accept = None unknown_reject = None if unknown_dir is not None and os.path.isdir(unknown_dir): unknown_ds = datasets.ImageFolder(unknown_dir, transform=transform_val) if os.path.isdir(os.path.join(unknown_dir, os.listdir(unknown_dir)[0])) else datasets.ImageFolder(unknown_dir, transform=transform_val) unknown_loader = DataLoader(unknown_ds, batch_size=cfg.batch_size, shuffle=False, num_workers=0) # 已知集用验证集 known_scores = collect_max_cos_scores(model, head, val_loader) unknown_scores = collect_max_cos_scores(model, head, unknown_loader) if len(unknown_scores) > 0 and len(known_scores) > 0: threshold, known_accept, unknown_reject = evaluate_open_set(known_scores, unknown_scores, far=far) logger.info(f'开放集阈值@FAR={far*100:.1f}%: T={threshold:.4f}, 已知接受率={known_accept*100:.2f}%, 未知拒识率={unknown_reject*100:.2f}%') else: logger.warning('开放集评估数据不足，跳过阈值估计。') # 结果保存 results = { 'training_time': f'{elapsed/3600:.2f} 小时', 'total_epochs': len(train_losses), 'best_val_accuracy': best_val_acc, 'final_train_loss': train_losses[-1] if train_losses else None, 'final_val_loss': val_losses[-1] if val_losses else None, 'final_train_accuracy': train_accs[-1] if train_accs else None, 'final_val_accuracy': val_accs[-1] if val_accs else None, 'model_parameters': sum(p.numel() for p in list(model.parameters()) + list(head.parameters()) if p.requires_grad), 'embedding_dim': cfg.embedding_dim, 'learning_rate': cfg.lr, 'batch_size': cfg.batch_size, 's': s, 'm': m, 'device': str(device), 'threshold_max_cos': threshold, 'known_accept_rate_at_threshold': known_accept, 'unknown_reject_rate_at_threshold': unknown_reject, } with open(os.path.join(save_dir, 'training_results.txt'), 'w', encoding='utf-8') as f: f.write('=== ResNet50 + CosFace 训练结果 ===\n\n') for k, v in results.items(): f.write(f'{k}: {v}\n') # 保存类别信息 with open(os.path.join(save_dir, 'class_info.json'), 'w', encoding='utf-8') as f: json.dump({'class_names': class_names, 'num_classes': num_classes, 'embedding_dim': cfg.embedding_dim}, f, ensure_ascii=False, indent=2) logger.info(f'训练结果已保存到: {save_dir}') if __name__ == '__main__': import argparse parser = argparse.ArgumentParser() # parser.add_argument('task', choices=list(TASKS.keys()), nargs='?', default='dish') parser.add_argument('task', choices=list(TASKS.keys()), nargs='?', default='whole_ingredient') parser.add_argument('--s', type=float, default=64.0) parser.add_argument('--m', type=float, default=0.35) parser.add_argument('--epochs', type=int, default=60) parser.add_argument('--unknown_dir', type=str, default=None, help='开放集评估用未知类目录（可选）') parser.add_argument('--far', type=float, default=0.05, help='未知集允许的FAR，用于阈值估计') args = parser.parse_args() main(task_key=args.task, s=args.s, m=args.m, num_epochs=args.epochs, unknown_dir=args.unknown_dir, far=args.far)