train.py

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms
import os
import numpy as np
import wandb
from PIL import Image
from models.resnet import resnet18, resnet34, resnet50
from models.openmax import OpenMax
# from models.metamax import MetaMax
from utils.data_stats import calculate_dataset_stats, load_dataset_stats
from utils.eval_utils import evaluate_known_classes, evaluate_openmax, evaluate_metamax
from pprint import pprint
import math


class GameDataset(Dataset):
    def __init__(self, data_dir, num_labels=20, transform=None):
        self.data_dir = data_dir
        self.transform = transform
        self.images = []
        self.labels = []
        self.image_paths = []
        
        if not os.path.exists(data_dir):
            raise ValueError(f"Data directory {data_dir} does not exist")
            
        # 遍历数据目录加载图片和标签
        for class_dir in range(num_labels):  # 训练集为0-19类,验证集为0-20类
            class_path = os.path.join(data_dir, f"{class_dir:02d}")
            if os.path.exists(class_path):
                for img_name in os.listdir(class_path):
                    if img_name.endswith('.png'):
                        img_path = os.path.join(class_path, img_name)
                        try:
                            # 读取PNG图片，只保留RGB通道
                            img = np.array(Image.open(img_path))[:, :, :3]  # 只取前3个通道
                            if img.shape != (50, 50, 3):
                                print(f"Skipping {img_path} due to invalid shape: {img.shape}")
                                continue
                                
                            self.images.append(img)
                            self.labels.append(class_dir)
                            self.image_paths.append(img_path)
                        except Exception as e:
                            print(f"Error loading {img_path}: {e}")
                            continue
        
        self.images = np.array(self.images)
        self.labels = np.array(self.labels)
        print(f"Loaded {len(self.images)} images from {data_dir}")
    
    def __len__(self):
        return len(self.images)
    
    def __getitem__(self, idx):
        image = self.images[idx]
        label = self.labels[idx]
        path = self.image_paths[idx]
        
        if self.transform:
            image = self.transform(image)
        
        return image, label, path



def train(num_epochs = 20, batch_size = 256, learning_rate = 0.001, dropout_rate = 0.3, patience = 10, model_type='resnet34'):
    from post_train import collect_features
    os.makedirs('models', exist_ok=True)
    os.makedirs('wandb_logs', exist_ok=True)
    images_path = os.path.join('jk_zfls', 'round0_train')
    # 尝试加载已保存的数据集统计信息，如果不存在则重新计算
    try:
        mean, std = load_dataset_stats()
        print("Loaded pre-calculated dataset statistics")
    except FileNotFoundError:
        print("FileNotFound, Calculating dataset statistics...")
        mean, std = calculate_dataset_stats(images_path)
        
    wandb.init(
        project="jk_zfls",
        name=f"{model_type}-training",
        config={
            "learning_rate": learning_rate,
            "batch_size": batch_size,
            "epochs": num_epochs,
            "model": f"{model_type}",
            "num_classes": 20
        },
        dir="./wandb_logs"
    )
    
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    
    # 计算填充值 (将均值从[0,1]转换为[0,255])
    fill_value = tuple(int(x * 255) for x in mean)
    
    # 增加数据增强
    transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.RandomAffine(
            degrees=15,
            translate=(0.1, 0.1),
            scale=(0.9, 1.1),
            fill=fill_value  # 使用数据集的均值作为填充值
        ),
        transforms.Normalize(mean=mean, std=std)
    ])
    
    # 验证集不需要数据增强
    val_transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize(mean=mean, std=std)
    ])
    
    # 加载数据集
    train_dataset = GameDataset('jk_zfls/round0_train', num_labels=20, transform=transform)
    val_dataset = GameDataset('jk_zfls/round0_eval', num_labels=21, transform=val_transform)
    
    train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=4, pin_memory=True)
    val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False, num_workers=4, pin_memory=True)
    
    # 根据选择加载不同的模型
    if model_type == 'resnet18':
        model = resnet18(num_classes=20, dropout_rate=dropout_rate)
    elif model_type == 'resnet34':
        model = resnet34(num_classes=20, dropout_rate=dropout_rate)
    elif model_type == 'resnet50':
        model = resnet50(num_classes=20, dropout_rate=dropout_rate)
    else:
        raise ValueError(f"Unsupported model type: {model_type}")
        
    # 加载模型（和已有参数）
    # checkpoint = torch.load('models/best_model_99.75.pth')
    # model.load_state_dict(checkpoint['model_state_dict'])
    model = model.to(device)
    
    # 定义损失函数和优化器，使用更小的学习率
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.Adam(model.parameters(), lr=learning_rate * 0.1, weight_decay=1e-3)
    
    # optimizer = optim.Adam(model.parameters(), lr=learning_rate, weight_decay=1e-4)
    # scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=10)
    # 使用带 warmup 的 cosine 调度器
    num_training_steps = len(train_loader) * num_epochs
    num_warmup_steps = len(train_loader) * 5      # 5个epoch的warmup
    
    # 定义warmup调度器和ReduceLROnPlateau调度器
    warmup_scheduler = optim.lr_scheduler.LinearLR(
        optimizer,
        start_factor=0.1,  # 从0.1倍的学习率开始
        end_factor=1.0,    # 最终达到设定的学习率
        total_iters=num_warmup_steps
    )

    reduce_scheduler = optim.lr_scheduler.ReduceLROnPlateau(
        optimizer,
        mode='max',
        factor=0.5,
        patience=5,
        verbose=True,
        min_lr=1e-6
    )

    patience_counter = 0  # 计数器，记录连续没有提升的轮数
    best_params = {
        'epoch': None,
        'model_state_dict': None,
        'optimizer_state_dict': None,
        'loss': None,
        'best_val_acc': 0
    }
    for epoch in range(num_epochs):
        # 训练阶段
        model.train()
        total_loss = 0
        
        for batch_idx, (images, labels, paths) in enumerate(train_loader):
            images, labels = images.to(device), labels.to(device)
            
            optimizer.zero_grad()
            logits = model(images)
            loss = criterion(logits, labels)
            loss.backward()
            optimizer.step()
            
            total_loss += loss.item()
            
            if batch_idx % 10 == 0:
                print(f'Epoch: {epoch}, Batch: {batch_idx}, Loss: {loss.item():.4f}')
            
            # 在warmup阶段更新学习率
            if epoch * len(train_loader) + batch_idx < num_warmup_steps:
                warmup_scheduler.step()
        
        train_loss = total_loss / len(train_loader)
        
        # 验证阶段（只验证已知类别）
        val_loss, val_acc, val_errors = evaluate_known_classes(model, val_loader, criterion, device)
        
        # 记录到wandb
        wandb.log({
            'epoch': epoch,
            'train_loss': train_loss,
            'val_loss': val_loss,
            'val_accuracy': val_acc
        })
        
        print(f'Epoch {epoch}:')
        print(f'Train Loss = {train_loss:.4f}, Val Loss = {val_loss:.4f}, Val Accuracy = {val_acc:.2f}%')
        
        # 验证阶段后更新ReduceLROnPlateau
        reduce_scheduler.step(val_acc)
        
        # 打印当前学习率
        current_lr = optimizer.param_groups[0]['lr']
        print(f'Current learning rate: {current_lr:.2e}')
        
        # 记录最佳模型（基于验证集准确率）
        if val_acc > best_params['best_val_acc']:
            patience_counter = 0  # 重置计数器
            best_params.update({
                'epoch': epoch,
                'model_state_dict': model.state_dict(),
                'optimizer_state_dict': optimizer.state_dict(),
                'loss': val_loss,
                'best_val_acc': val_acc
            })
        else:
            patience_counter += 1  # 增加计数器
            print(f'Validation accuracy did not improve. Patience: {patience_counter}/{patience}')
        
        # 早停检查
        if patience_counter >= patience:
            print(f"\nEarly stopping triggered! No improvement for {patience} consecutive epochs.")
            break
            
        if val_acc == 100:
            print(f'Achieved 100% accuracy at epoch {epoch}')
            break
            
    
    # 训练完成后，保存最佳模型的参数
    print("Saving best model parameters...")
    torch.save(best_params, f'models/{model_type}_{best_params["best_val_acc"]:.2f}.pth')
    
    # 使用最佳模型收集features
    print("Collecting features from best model for OpenMax/MetaMax training...")
    model.load_state_dict(best_params['model_state_dict'])
    model.eval()
    features, labels = collect_features(model, train_loader, device, return_logits=False)
    
    # 训练OpenMax/MetaMax
    openmax = OpenMax(num_classes=20)
    openmax.fit(features, labels)
    
    # metamax = MetaMax(num_classes=20)
    # metamax.fit(features, labels)
    
    # 保存模型
    torch.save(openmax, 'models/openmax.pth')
    # torch.save(metamax, 'models/metamax.pth')
    print("OpenMax and MetaMax models saved")
    # 在训练完OpenMax后添加评估
    print("Evaluating OpenMax and MetaMax...")
    val_features, val_logits, val_labels = collect_features(model, val_loader, device, return_logits=True)

    overall_acc, known_acc, unknown_acc = evaluate_openmax(openmax, val_features, val_logits, val_labels, multiplier=0.5)
    print(f"Multiplier: 0.5, Overall Accuracy: {overall_acc:.2f}%")
    # evaluate_metamax(metamax, val_features, val_labels, device)
    wandb.finish()

if __name__ == '__main__':
    train(num_epochs=100, batch_size=64, learning_rate=0.001, dropout_rate=0.3, patience=20, model_type='resnet50')