Upload train.py with huggingface_hub

train.py

@@ -0,0 +1,277 @@
+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')