diff --git a/train/train_embedding.py b/train/train_embedding.py index e69de29..5953f56 100644 --- a/train/train_embedding.py +++ b/train/train_embedding.py @@ -0,0 +1,716 @@ +import torch +import torch.nn as nn +import torch.optim as optim +import torch.nn.functional as F +from torch.utils.data import DataLoader, Dataset +from torchvision import datasets, transforms +import matplotlib.pyplot as plt +import matplotlib +import numpy as np +from tqdm import tqdm +import os +import sys +import time +import json +import random +from collections import defaultdict +from typing import List, Tuple, Dict, Optional +import logging +from datetime import datetime + +# 添加项目根目录到路径 +sys.path.append(os.path.join(os.path.dirname(__file__), '..')) +from net.resnet_embedding import create_resnet50_embedding +from settings import settings + +# 设置matplotlib支持中文显示 +plt.rcParams['font.sans-serif'] = ['SimHei', 'Microsoft YaHei', 'DejaVu Sans'] +plt.rcParams['axes.unicode_minus'] = False + +# 设置日志 +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}") + + +class TripletDataset(Dataset): + """ + 三元组数据集,用于三元组损失训练 + 每个样本包含:锚点(anchor)、正样本(positive)、负样本(negative) + """ + + def __init__(self, dataset_path: str, transform=None, samples_per_class: int = 100): + """ + 初始化三元组数据集 + + Args: + dataset_path: 数据集路径 + transform: 数据变换 + samples_per_class: 每个类别最多使用的样本数 + """ + self.dataset_path = dataset_path + self.transform = transform + self.samples_per_class = samples_per_class + + # 加载数据集 + self.class_to_idx = {} + self.idx_to_class = {} + self.samples_by_class = defaultdict(list) + self.all_samples = [] + + self._load_dataset() + + def _load_dataset(self): + """加载数据集并按类别组织""" + class_names = sorted(os.listdir(self.dataset_path)) + + for idx, class_name in enumerate(class_names): + if not os.path.isdir(os.path.join(self.dataset_path, class_name)): + continue + + self.class_to_idx[class_name] = idx + self.idx_to_class[idx] = class_name + + class_path = os.path.join(self.dataset_path, class_name) + image_files = [f for f in os.listdir(class_path) + if f.lower().endswith(('.png', '.jpg', '.jpeg'))] + + # 限制每个类别的样本数 + if len(image_files) > self.samples_per_class: + image_files = random.sample(image_files, self.samples_per_class) + + for img_file in image_files: + img_path = os.path.join(class_path, img_file) + self.samples_by_class[idx].append(img_path) + self.all_samples.append((img_path, idx)) + + logger.info(f"加载数据集完成:") + for class_name, class_idx in self.class_to_idx.items(): + logger.info(f" {class_name}: {len(self.samples_by_class[class_idx])} 张图片") + + def __len__(self): + return len(self.all_samples) + + def __getitem__(self, idx): + """ + 获取三元组样本 + + Returns: + tuple: (anchor, positive, negative, anchor_label) + """ + # 获取锚点样本 + anchor_path, anchor_label = self.all_samples[idx] + anchor_img = self._load_image(anchor_path) + + # 获取正样本(同类别的不同图片) + positive_candidates = [path for path in self.samples_by_class[anchor_label] + if path != anchor_path] + if positive_candidates: + positive_path = random.choice(positive_candidates) + else: + positive_path = anchor_path # 如果只有一张图片,使用自己作为正样本 + positive_img = self._load_image(positive_path) + + # 获取负样本(不同类别的图片) + negative_classes = [cls for cls in self.samples_by_class.keys() if cls != anchor_label] + negative_class = random.choice(negative_classes) + negative_path = random.choice(self.samples_by_class[negative_class]) + negative_img = self._load_image(negative_path) + + return anchor_img, positive_img, negative_img, anchor_label + + def _load_image(self, image_path: str): + """加载并预处理图片""" + from PIL import Image + + try: + image = Image.open(image_path).convert('RGB') + if self.transform: + image = self.transform(image) + return image + except Exception as e: + logger.error(f"加载图片失败: {image_path}, 错误: {e}") + # 返回一个黑色图片作为备用 + if self.transform: + return self.transform(Image.new('RGB', (224, 224), (0, 0, 0))) + else: + return Image.new('RGB', (224, 224), (0, 0, 0)) + + def get_class_names(self): + """获取类别名称列表""" + return [self.idx_to_class[i] for i in sorted(self.idx_to_class.keys())] + + +class TripletLoss(nn.Module): + """ + 三元组损失函数 + """ + + def __init__(self, margin: float = 0.3): + """ + 初始化三元组损失 + + Args: + margin: 边界值,正负样本距离的最小差值 + """ + super(TripletLoss, self).__init__() + self.margin = margin + + def forward(self, anchor, positive, negative): + """ + 计算三元组损失 + + Args: + anchor: 锚点特征向量 [batch_size, embedding_dim] + positive: 正样本特征向量 [batch_size, embedding_dim] + negative: 负样本特征向量 [batch_size, embedding_dim] + + Returns: + torch.Tensor: 三元组损失值 + """ + # 计算距离 + pos_dist = F.pairwise_distance(anchor, positive, p=2) + neg_dist = F.pairwise_distance(anchor, negative, p=2) + + # 计算三元组损失 + loss = F.relu(pos_dist - neg_dist + self.margin) + + return loss.mean() + + +class CenterLoss(nn.Module): + """ + 中心损失函数,用于增强类内聚合 + """ + + def __init__(self, num_classes: int, embedding_dim: int, alpha: float = 0.5): + """ + 初始化中心损失 + + Args: + num_classes: 类别数量 + embedding_dim: 特征向量维度 + alpha: 中心更新的学习率 + """ + super(CenterLoss, self).__init__() + self.num_classes = num_classes + self.embedding_dim = embedding_dim + self.alpha = alpha + + # 初始化类别中心 + self.centers = nn.Parameter(torch.randn(num_classes, embedding_dim)) + + def forward(self, embeddings, labels): + """ + 计算中心损失 + + Args: + embeddings: 特征向量 [batch_size, embedding_dim] + labels: 标签 [batch_size] + + Returns: + torch.Tensor: 中心损失值 + """ + batch_size = embeddings.size(0) + + # 获取对应的中心向量 + centers_batch = self.centers[labels] # [batch_size, embedding_dim] + + # 计算特征向量与中心的距离 + loss = F.mse_loss(embeddings, centers_batch) + + return loss + + def update_centers(self, embeddings, labels): + """ + 更新类别中心 + + Args: + embeddings: 特征向量 [batch_size, embedding_dim] + labels: 标签 [batch_size] + """ + with torch.no_grad(): + for label in torch.unique(labels): + mask = labels == label + if mask.sum() > 0: + # 计算该类别的平均特征向量 + class_embeddings = embeddings[mask] + class_center = class_embeddings.mean(dim=0) + + # 更新中心 + self.centers[label] = (1 - self.alpha) * self.centers[label] + self.alpha * class_center + + +class EarlyStopping: + """早停机制""" + + def __init__(self, patience: int = 10, min_delta: float = 0.001): + """ + 初始化早停机制 + + Args: + patience: 容忍的epoch数 + min_delta: 最小改善幅度 + """ + self.patience = patience + self.min_delta = min_delta + self.counter = 0 + self.best_loss = float('inf') + + def __call__(self, val_loss: float) -> bool: + """ + 检查是否应该早停 + + Args: + val_loss: 验证损失 + + Returns: + bool: 是否应该早停 + """ + if val_loss < self.best_loss - self.min_delta: + self.best_loss = val_loss + self.counter = 0 + else: + self.counter += 1 + + return self.counter >= self.patience + + +def train_epoch(model, train_loader, triplet_criterion, center_criterion, + optimizer, center_optimizer, device, epoch): + """ + 训练一个epoch + + Args: + model: 模型 + train_loader: 训练数据加载器 + triplet_criterion: 三元组损失函数 + center_criterion: 中心损失函数 + optimizer: 模型优化器 + center_optimizer: 中心损失优化器 + device: 设备 + epoch: 当前epoch + + Returns: + tuple: (平均损失, 三元组损失, 中心损失) + """ + model.train() + total_loss = 0.0 + total_triplet_loss = 0.0 + total_center_loss = 0.0 + num_batches = 0 + + train_bar = tqdm(train_loader, desc=f'Epoch {epoch+1} 训练中') + + for batch_idx, (anchor, positive, negative, labels) in enumerate(train_bar): + anchor = anchor.to(device) + positive = positive.to(device) + negative = negative.to(device) + labels = labels.to(device) + + # 前向传播 + anchor_emb = model(anchor) + positive_emb = model(positive) + negative_emb = model(negative) + + # 计算三元组损失 + triplet_loss = triplet_criterion(anchor_emb, positive_emb, negative_emb) + + # 计算中心损失 + center_loss = center_criterion(anchor_emb, labels) + + # 总损失 + loss = triplet_loss + 0.1 * center_loss # 中心损失权重为0.1 + + # 反向传播 + optimizer.zero_grad() + center_optimizer.zero_grad() + + loss.backward() + + optimizer.step() + center_optimizer.step() + + # 更新中心 + center_criterion.update_centers(anchor_emb.detach(), labels) + + # 统计 + total_loss += loss.item() + total_triplet_loss += triplet_loss.item() + total_center_loss += center_loss.item() + num_batches += 1 + + # 更新进度条 + train_bar.set_postfix({ + 'Loss': f'{total_loss/num_batches:.4f}', + 'Triplet': f'{total_triplet_loss/num_batches:.4f}', + 'Center': f'{total_center_loss/num_batches:.4f}' + }) + + return (total_loss/num_batches, + total_triplet_loss/num_batches, + total_center_loss/num_batches) + + +def validate_epoch(model, val_loader, triplet_criterion, center_criterion, device): + """ + 验证一个epoch + + Args: + model: 模型 + val_loader: 验证数据加载器 + triplet_criterion: 三元组损失函数 + center_criterion: 中心损失函数 + device: 设备 + + Returns: + tuple: (平均损失, 三元组损失, 中心损失, 准确率) + """ + model.eval() + total_loss = 0.0 + total_triplet_loss = 0.0 + total_center_loss = 0.0 + correct_predictions = 0 + total_predictions = 0 + num_batches = 0 + + with torch.no_grad(): + val_bar = tqdm(val_loader, desc='验证中') + + for anchor, positive, negative, labels in val_bar: + anchor = anchor.to(device) + positive = positive.to(device) + negative = negative.to(device) + labels = labels.to(device) + + # 前向传播 + anchor_emb = model(anchor) + positive_emb = model(positive) + negative_emb = model(negative) + + # 计算损失 + triplet_loss = triplet_criterion(anchor_emb, positive_emb, negative_emb) + center_loss = center_criterion(anchor_emb, labels) + loss = triplet_loss + 0.1 * center_loss + + # 计算准确率(基于最近邻分类) + # 这里简化为检查正样本距离是否小于负样本距离 + pos_dist = F.pairwise_distance(anchor_emb, positive_emb, p=2) + neg_dist = F.pairwise_distance(anchor_emb, negative_emb, p=2) + correct = (pos_dist < neg_dist).sum().item() + + # 统计 + total_loss += loss.item() + total_triplet_loss += triplet_loss.item() + total_center_loss += center_loss.item() + correct_predictions += correct + total_predictions += anchor.size(0) + num_batches += 1 + + # 更新进度条 + val_bar.set_postfix({ + 'Loss': f'{total_loss/num_batches:.4f}', + 'Acc': f'{100.*correct_predictions/total_predictions:.2f}%' + }) + + accuracy = 100. * correct_predictions / total_predictions + + return (total_loss/num_batches, + total_triplet_loss/num_batches, + total_center_loss/num_batches, + accuracy) + + +def plot_training_curves(train_losses, val_losses, train_accuracies, val_accuracies, save_path): + """ + 绘制训练曲线 + + Args: + train_losses: 训练损失列表 + val_losses: 验证损失列表 + train_accuracies: 训练准确率列表 + val_accuracies: 验证准确率列表 + save_path: 保存路径 + """ + fig, ((ax1, ax2), (ax3, ax4)) = plt.subplots(2, 2, figsize=(15, 10)) + + epochs = range(1, len(train_losses) + 1) + + # 总损失 + ax1.plot(epochs, train_losses, 'b-', label='训练损失') + ax1.plot(epochs, val_losses, 'r-', label='验证损失') + ax1.set_title('总损失变化') + ax1.set_xlabel('Epoch') + ax1.set_ylabel('Loss') + ax1.legend() + ax1.grid(True) + + # 准确率 + ax2.plot(epochs, train_accuracies, 'b-', label='训练准确率') + ax2.plot(epochs, val_accuracies, 'r-', label='验证准确率') + ax2.set_title('准确率变化') + ax2.set_xlabel('Epoch') + ax2.set_ylabel('Accuracy (%)') + ax2.legend() + ax2.grid(True) + + # 三元组损失(如果有记录的话) + ax3.set_title('三元组损失变化') + ax3.set_xlabel('Epoch') + ax3.set_ylabel('Triplet Loss') + ax3.grid(True) + + # 中心损失(如果有记录的话) + ax4.set_title('中心损失变化') + ax4.set_xlabel('Epoch') + ax4.set_ylabel('Center Loss') + ax4.grid(True) + + plt.tight_layout() + plt.savefig(save_path, dpi=300, bbox_inches='tight') + plt.close() + + +def save_training_results(results, save_path): + """ + 保存训练结果 + + Args: + results: 训练结果字典 + save_path: 保存路径 + """ + with open(save_path, 'w', encoding='utf-8') as f: + f.write("=== ResNet50 Embedding 训练结果 ===\n\n") + f.write(f"训练时间: {results['training_time']}\n") + f.write(f"总epoch数: {results['total_epochs']}\n") + f.write(f"最佳验证损失: {results['best_val_loss']:.4f}\n") + f.write(f"最佳验证准确率: {results['best_val_accuracy']:.2f}%\n") + f.write(f"最终训练损失: {results['final_train_loss']:.4f}\n") + f.write(f"最终验证损失: {results['final_val_loss']:.4f}\n") + f.write(f"最终训练准确率: {results['final_train_accuracy']:.2f}%\n") + f.write(f"最终验证准确率: {results['final_val_accuracy']:.2f}%\n") + f.write(f"模型参数数量: {results['model_parameters']:,}\n") + f.write(f"Embedding维度: {results['embedding_dim']}\n") + + f.write("\n=== 训练配置 ===\n") + f.write(f"学习率: {results['learning_rate']}\n") + f.write(f"批次大小: {results['batch_size']}\n") + f.write(f"三元组损失边界: {results['triplet_margin']}\n") + f.write(f"中心损失权重: {results['center_loss_weight']}\n") + f.write(f"设备: {results['device']}\n") + + +def main(): + """主训练函数""" + + # 创建保存目录 + timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") + save_dir = os.path.join(settings.BASE_DIR, 'model', f'embedding_{timestamp}') + os.makedirs(save_dir, exist_ok=True) + + logger.info(f"模型保存目录: {save_dir}") + + # 训练参数 + EMBEDDING_DIM = 512 + BATCH_SIZE = 16 # 三元组训练通常使用较小的batch size + LEARNING_RATE = 0.001 + NUM_EPOCHS = 50 + TRIPLET_MARGIN = 0.3 + CENTER_LOSS_WEIGHT = 0.1 + PATIENCE = 10 + + # 数据预处理 + transform_train = transforms.Compose([ + transforms.Resize((224, 224)), + transforms.RandomHorizontalFlip(p=0.5), + transforms.RandomRotation(15), + transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2, hue=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]) + ]) + + transform_val = transforms.Compose([ + transforms.Resize((224, 224)), + transforms.ToTensor(), + transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) + ]) + + # 创建数据集 + train_dataset = TripletDataset( + dataset_path=settings.TRAIN_DATA_DIR, + transform=transform_train, + samples_per_class=200 + ) + + val_dataset = TripletDataset( + dataset_path=settings.VAL_DATA_DIR, + transform=transform_val, + samples_per_class=50 + ) + + # 创建数据加载器 + train_loader = DataLoader( + train_dataset, + batch_size=BATCH_SIZE, + shuffle=True, + num_workers=0, + drop_last=True + ) + + val_loader = DataLoader( + val_dataset, + batch_size=BATCH_SIZE, + shuffle=False, + num_workers=0, + drop_last=True + ) + + logger.info(f"训练集大小: {len(train_dataset)}") + logger.info(f"验证集大小: {len(val_dataset)}") + logger.info(f"类别数量: {len(train_dataset.get_class_names())}") + logger.info(f"类别名称: {train_dataset.get_class_names()}") + + # 创建模型 + model = create_resnet50_embedding( + embedding_dim=EMBEDDING_DIM, + pretrained=True, + use_internal_preprocess=False + ) + model = model.to(device) + + # 创建损失函数 + triplet_criterion = TripletLoss(margin=TRIPLET_MARGIN) + center_criterion = CenterLoss( + num_classes=len(train_dataset.get_class_names()), + embedding_dim=EMBEDDING_DIM, + alpha=0.5 + ).to(device) + + # 创建优化器 + optimizer = optim.Adam(model.parameters(), lr=LEARNING_RATE, weight_decay=1e-4) + center_optimizer = optim.Adam(center_criterion.parameters(), lr=LEARNING_RATE) + + # 学习率调度器 + scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=15, gamma=0.5) + + # 早停机制 + early_stopping = EarlyStopping(patience=PATIENCE, min_delta=0.001) + + # 训练记录 + train_losses = [] + val_losses = [] + train_accuracies = [] + val_accuracies = [] + best_val_loss = float('inf') + best_val_accuracy = 0.0 + + logger.info("开始训练...") + start_time = time.time() + + for epoch in range(NUM_EPOCHS): + # 训练 + train_loss, train_triplet_loss, train_center_loss = train_epoch( + model, train_loader, triplet_criterion, center_criterion, + optimizer, center_optimizer, device, epoch + ) + + # 验证 + val_loss, val_triplet_loss, val_center_loss, val_accuracy = validate_epoch( + model, val_loader, triplet_criterion, center_criterion, device + ) + + # 学习率调度 + scheduler.step() + + # 记录结果 + train_losses.append(train_loss) + val_losses.append(val_loss) + train_accuracies.append(0) # 训练准确率暂时设为0,因为三元组训练中不容易计算 + val_accuracies.append(val_accuracy) + + # 保存最佳模型 + if val_loss < best_val_loss: + best_val_loss = val_loss + best_val_accuracy = val_accuracy + + # 保存模型 + model_save_path = os.path.join(save_dir, 'best_embedding_model.pth') + torch.save({ + 'epoch': epoch, + 'model_state_dict': model.state_dict(), + 'center_state_dict': center_criterion.state_dict(), + 'optimizer_state_dict': optimizer.state_dict(), + 'center_optimizer_state_dict': center_optimizer.state_dict(), + 'val_loss': val_loss, + 'val_accuracy': val_accuracy, + 'class_names': train_dataset.get_class_names(), + 'embedding_dim': EMBEDDING_DIM, + }, model_save_path) + + logger.info(f"保存最佳模型: Epoch {epoch+1}, Val Loss: {val_loss:.4f}, Val Acc: {val_accuracy:.2f}%") + + # 打印epoch结果 + logger.info(f"Epoch {epoch+1}/{NUM_EPOCHS}:") + logger.info(f" Train Loss: {train_loss:.4f} (Triplet: {train_triplet_loss:.4f}, Center: {train_center_loss:.4f})") + logger.info(f" Val Loss: {val_loss:.4f} (Triplet: {val_triplet_loss:.4f}, Center: {val_center_loss:.4f})") + logger.info(f" Val Accuracy: {val_accuracy:.2f}%") + logger.info(f" Learning Rate: {scheduler.get_last_lr()[0]:.6f}") + + # 早停检查 + if early_stopping(val_loss): + logger.info(f"早停触发,在第 {epoch+1} epoch停止训练") + break + + # 训练完成 + training_time = time.time() - start_time + logger.info(f"训练完成,总用时: {training_time/3600:.2f} 小时") + + # 绘制训练曲线 + curves_save_path = os.path.join(save_dir, 'training_curves.png') + plot_training_curves(train_losses, val_losses, train_accuracies, val_accuracies, curves_save_path) + + # 保存训练结果 + results = { + 'training_time': f"{training_time/3600:.2f} 小时", + 'total_epochs': len(train_losses), + 'best_val_loss': best_val_loss, + 'best_val_accuracy': best_val_accuracy, + 'final_train_loss': train_losses[-1], + 'final_val_loss': val_losses[-1], + 'final_train_accuracy': train_accuracies[-1], + 'final_val_accuracy': val_accuracies[-1], + 'model_parameters': sum(p.numel() for p in model.parameters() if p.requires_grad), + 'embedding_dim': EMBEDDING_DIM, + 'learning_rate': LEARNING_RATE, + 'batch_size': BATCH_SIZE, + 'triplet_margin': TRIPLET_MARGIN, + 'center_loss_weight': CENTER_LOSS_WEIGHT, + 'device': str(device), + } + + results_save_path = os.path.join(save_dir, 'training_results.txt') + save_training_results(results, results_save_path) + + # 保存类别信息 + class_info = { + 'class_names': train_dataset.get_class_names(), + 'class_to_idx': train_dataset.class_to_idx, + 'idx_to_class': train_dataset.idx_to_class, + 'embedding_dim': EMBEDDING_DIM, + } + + class_info_path = os.path.join(save_dir, 'class_info.json') + with open(class_info_path, 'w', encoding='utf-8') as f: + json.dump(class_info, f, ensure_ascii=False, indent=2) + + logger.info(f"训练结果已保存到: {save_dir}") + logger.info(f"最佳验证损失: {best_val_loss:.4f}") + logger.info(f"最佳验证准确率: {best_val_accuracy:.2f}%") + + +if __name__ == "__main__": + main() \ No newline at end of file