main.py

import os
import time
import datetime
import numpy as np
import shutil

import torch
import torch.nn as nn
import torch.backends.cudnn as cudnn
import torch.distributed as dist
import torch.distributed as dist
from torch.amp import autocast, GradScaler

from timm.loss import LabelSmoothingCrossEntropy, SoftTargetCrossEntropy
from timm.utils import accuracy, AverageMeter

from models import build_model
from data import build_loader
from lr_scheduler import build_scheduler
from optimizer import build_optimizer
from logger import create_logger
from utils import load_checkpoint, save_checkpoint, save_checkpoint_new, get_grad_norm, auto_resume_helper, \
    reduce_tensor, load_pretrained, parse_option, throughput, validate, visualize_attention_maps, add_noise_to_images, add_noise_to_text, load_image_from_path

import warnings
warnings.filterwarnings('ignore')



def count_parameters(model):
    """计算模型参数量"""
    total_params = sum(p.numel() for p in model.parameters())
    trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
    return total_params, trainable_params


def count_flops(model, config, logger, device=None):
    """计算模型FLOPs (仅使用fvcore)"""
    model.eval()
    
    if device is None:
        device = next(model.parameters()).device
    
    # 统一使用 LongTensor
    is_text_dataset = config.DATA.DATASET.lower() in ['rotten_tomatoes', 'imdb', '20_newsgroups', 'ag_news'] or config.DATA.DATASET.lower().startswith('glue_')
    
    # 创建dummy input
    if is_text_dataset:
        # 文本输入: [1, seq_len]
        dummy_input = torch.randint(0, config.DATA.VOCAB_SIZE, (1, config.DATA.MAX_SEQ_LEN), dtype=torch.long, device=device)
    else:
        # 图像输入: [1, channels, height, width]
        dummy_input = torch.randn(1, config.DATA.IN_CHANS, config.DATA.IMG_SIZE, config.DATA.IMG_SIZE, device=device)
    
    try:
        # 使用 fvcore.nn.FlopCountAnalysis 计算 FLOPs
        from fvcore.nn import FlopCountAnalysis
        flops = FlopCountAnalysis(model, dummy_input)
        return flops.total()
    except ImportError as e:
        logger.warning(f"fvcore import failed: {e}. Please ensure it is installed correctly (pip install fvcore).")
        return None
    except Exception as e:
        logger.warning(f"Failed to calculate FLOPs: {e}")
        return None




def main():
    os.environ["TORCH_NCCL_ASYNC_ERROR_HANDLING"] = "1"
    args, config = parse_option()

    rank = int(os.environ["RANK"])
    world_size = int(os.environ["WORLD_SIZE"])
    local_rank = int(os.environ["LOCAL_RANK"])
    torch.cuda.set_device(local_rank)
    dist.init_process_group(backend='nccl', init_method='env://', world_size=world_size, rank=rank)

    seed = config.SEED + dist.get_rank()
    torch.manual_seed(seed)
    np.random.seed(seed)
    cudnn.enabled = True
    cudnn.benchmark = True

    os.makedirs(config.OUTPUT, exist_ok=True)
    logger = create_logger(output_dir=config.OUTPUT, dist_rank=dist.get_rank(), name=f"{config.MODEL.NAME}")

    # Save config and model
    if dist.get_rank() == 0:
        path = os.path.join(config.OUTPUT, "config.json")
        with open(path, "w") as f:
            f.write(config.dump())
        logger.info(f"Full config saved to {path}")
        
        if config.MODEL.PATH:
            path = os.path.join(config.OUTPUT, "model.py")
            shutil.copy2(config.MODEL.PATH, path)
            logger.info(f"Used model saved to {path}")
    
    # load data
    _, dataset_val, data_loader_train, data_loader_val, dataset_test, data_loader_test, mixup_fn = build_loader(config)

    logger.info(f"Creating model:{config.MODEL.TYPE}/{config.MODEL.NAME}")
    model = build_model(config)
    device = torch.device(f'cuda:{local_rank}')
    model = model.to(device)

    # 计算模型参数量和FLOPs
    if dist.get_rank() == 0:
        total_params, trainable_params = count_parameters(model)
        logger.info(f"Total parameters: {total_params:,} ({total_params / 1e6:.2f}M)")
        logger.info(f"Trainable parameters: {trainable_params:,} ({trainable_params / 1e6:.2f}M)")
        
        flops = count_flops(model, config, logger, device)
        if flops is not None:
            logger.info(f"FLOPs: {flops:,.0f} ({flops / 1e9:.2f}G)")
        

    optimizer = build_optimizer(config, model)

    # 确保所有参数和缓冲区都在正确的设备上
    for param in model.parameters():
        if param.device.type != 'cuda':
            param.data = param.data.to(device)
    for buffer in model.buffers():
        if buffer.device.type != 'cuda':
            buffer.data = buffer.data.to(device)

    model = nn.parallel.DistributedDataParallel(model, device_ids=[local_rank], broadcast_buffers=True, find_unused_parameters=args.find_unused_params)
    model_without_ddp = model.module
    
    lr_scheduler = build_scheduler(config, optimizer, len(data_loader_train))
    total_epochs = config.TRAIN.EPOCHS + config.TRAIN.COOLDOWN_EPOCHS

    if config.DATA.DATASET.lower() in ['rotten_tomatoes', 'imdb', 'ag_news', '20_newsgroups'] or config.DATA.DATASET.lower().startswith('glue_'):
        criterion = nn.CrossEntropyLoss()
    elif config.AUG.MIXUP > 0.:
        criterion = SoftTargetCrossEntropy()
    elif config.MODEL.LABEL_SMOOTHING > 0.:
        criterion = LabelSmoothingCrossEntropy(smoothing=config.MODEL.LABEL_SMOOTHING)
    else:
        criterion = nn.CrossEntropyLoss()

    max_val_accuracy = max_test_accuracy = 0.0
    
    if args.test:
        load_checkpoint(config, model_without_ddp, optimizer, lr_scheduler, logger)
        
        is_text_dataset = config.DATA.DATASET.lower() in ['rotten_tomatoes', 'imdb', '20_newsgroups', 'ag_news'] or config.DATA.DATASET.lower().startswith('glue_')
        
        # 只在验证集上验证，跳过测试集（因为GLUE测试集没有标签）
        # 或者如果有标签才验证
        max_val_accuracy, _ = validate(config, data_loader_val, model, logger)
        if is_text_dataset:
            logger.info(f"Accuracy of the loaded model on the {len(dataset_val)} val samples: {max_val_accuracy:.2f}%")
        else:
            logger.info(f"Accuracy of the loaded model on the {len(dataset_val)} val images: {max_val_accuracy:.2f}%")
            
        # 对于 GLUE 任务，跳过测试集评估，除非我们确定测试集有标签
        is_glue = 'glue' in config.DATA.DATASET.lower()
        if not is_text_dataset or not is_glue:
            max_test_accuracy, _ = validate(config, data_loader_test, model, logger)
            if is_text_dataset:
                logger.info(f"Accuracy of the loaded model on the {len(dataset_test)} test samples: {max_test_accuracy:.2f}%")
            else:
                logger.info(f"Accuracy of the loaded model on the {len(dataset_test)} test images: {max_test_accuracy:.2f}%")
        else:
             # GLUE 测试集通常没有标签，我们可以进行预测但无法计算准确率
             # 如果需要生成预测结果提交，可以修改 validate 函数来支持保存预测结果
             # 这里我们简单地尝试运行验证，如果因为没有标签报错则忽略，或者仅输出日志
             try:
                 # 尝试验证，前提是 validate 函数能处理 -1 标签（我们之前在 dataset 里处理了）
                 # 但注意 validate 内部计算 accuracy 会用到 label，如果全为 -1，acc 计算无意义
                 # 暂时保持跳过，或者输出一个提示
                 logger.info("GLUE test set evaluation skipped (no ground truth labels).")
                 max_test_accuracy = 0.0
             except Exception as e:
                 logger.warning(f"Failed to evaluate on GLUE test set: {e}")
                 max_test_accuracy = 0.0
        
        if is_text_dataset:
            logger.info("test mask noise with noise_level=0.05")
            val_acc_noise, _ = validate(config, data_loader_val, model, logger, 'mask', 0.05)
            logger.info(f"Accuracy of the loaded model on the {len(dataset_val)} val samples with mask noise (noise_level=0.05): {val_acc_noise:.2f}%")
            test_acc_noise, _ = validate(config, data_loader_test, model, logger, 'mask', 0.05)
            logger.info(f"Accuracy of the loaded model on the {len(dataset_test)} test samples with mask noise (noise_level=0.05): {test_acc_noise:.2f}%")
            
            logger.info("test mask noise with noise_level=0.1")
            val_acc_noise, _ = validate(config, data_loader_val, model, logger, 'mask', 0.1)
            logger.info(f"Accuracy of the loaded model on the {len(dataset_val)} val samples with mask noise (noise_level=0.1): {val_acc_noise:.2f}%")
            test_acc_noise, _ = validate(config, data_loader_test, model, logger, 'mask', 0.1)
            logger.info(f"Accuracy of the loaded model on the {len(dataset_test)} test samples with mask noise (noise_level=0.1): {test_acc_noise:.2f}%")
            
            logger.info("test mask noise with noise_level=0.2")
            val_acc_noise, _ = validate(config, data_loader_val, model, logger, 'mask', 0.2)
            logger.info(f"Accuracy of the loaded model on the {len(dataset_val)} val samples with mask noise (noise_level=0.2): {val_acc_noise:.2f}%")
            test_acc_noise, _ = validate(config, data_loader_test, model, logger, 'mask', 0.2)
            logger.info(f"Accuracy of the loaded model on the {len(dataset_test)} test samples with mask noise (noise_level=0.2): {test_acc_noise:.2f}%")
        else:
            # gaussian noise
            logger.info("test gaussian noise with noise_level=0.05")
            val_acc_noise, _ = validate(config, data_loader_val, model, logger, 'gaussian', 0.05)
            logger.info(f"Accuracy of the loaded model on the {len(dataset_val)} val images with gaussian noise (noise_level=0.05): {val_acc_noise:.2f}%")
            test_acc_noise, _ = validate(config, data_loader_test, model, logger, 'gaussian', 0.05)
            logger.info(f"Accuracy of the loaded model on the {len(dataset_test)} test images with gaussian noise (noise_level=0.05): {test_acc_noise:.2f}%")
            logger.info("test gaussian noise with noise_level=0.1")
            val_acc_noise, _ = validate(config, data_loader_val, model, logger, 'gaussian', 0.1)
            logger.info(f"Accuracy of the loaded model on the {len(dataset_val)} val images with gaussian noise (noise_level=0.1): {val_acc_noise:.2f}%")
            test_acc_noise, _ = validate(config, data_loader_test, model, logger, 'gaussian', 0.1)
            logger.info(f"Accuracy of the loaded model on the {len(dataset_test)} test images with gaussian noise (noise_level=0.1): {test_acc_noise:.2f}%")
            logger.info("test gaussian noise with noise_level=0.2")
            val_acc_noise, _ = validate(config, data_loader_val, model, logger, 'gaussian', 0.2)
            logger.info(f"Accuracy of the loaded model on the {len(dataset_val)} val images with gaussian noise (noise_level=0.2): {val_acc_noise:.2f}%")
            test_acc_noise, _ = validate(config, data_loader_test, model, logger, 'gaussian', 0.2)
            logger.info(f"Accuracy of the loaded model on the {len(dataset_test)} test images with gaussian noise (noise_level=0.2): {test_acc_noise:.2f}%")
            
            # test salt_pepper noise
            logger.info("test salt_pepper noise with noise_level=0.1")
            val_acc_sp, _ = validate(config, data_loader_val, model, logger, 'salt_pepper', 0.1)
            logger.info(f"Accuracy of the loaded model on the {len(dataset_val)} val images with salt_pepper noise (noise_level=0.1): {val_acc_sp:.2f}%")
            test_acc_sp, _ = validate(config, data_loader_test, model, logger, 'salt_pepper', 0.1)
            logger.info(f"Accuracy of the loaded model on the {len(dataset_test)} test images with salt_pepper noise (noise_level=0.1): {test_acc_sp:.2f}%")

            # test uniform noise
            logger.info("test uniform noise with noise_level=0.1")
            val_acc_uniform, _ = validate(config, data_loader_val, model, logger, 'uniform', 0.1)
            logger.info(f"Accuracy of the loaded model on the {len(dataset_val)} val images with uniform noise (noise_level=0.1): {val_acc_uniform:.2f}%")
            test_acc_uniform, _ = validate(config, data_loader_test, model, logger, 'uniform', 0.1)
            logger.info(f"Accuracy of the loaded model on the {len(dataset_test)} test images with uniform noise (noise_level=0.1): {test_acc_uniform:.2f}%")
        return
    
    if args.get_attention_maps:
        # TODO: how to draw a better attention map?
        load_checkpoint(config, model_without_ddp, optimizer, lr_scheduler, logger)
        
        # 如果指定了图片路径，使用该图片；否则使用数据集中的图片
        if args.image_path:
            logger.info(f"使用指定图片: {args.image_path}")
            images = load_image_from_path(args.image_path, config)
            images = images.cuda(non_blocking=True)
            images = add_noise_to_images(images, noise_type='gaussian', noise_level=0.1)
        else:
            logger.info("使用数据集中的图片")
            images, _ = next(iter(data_loader_test))
            images = images.cuda(non_blocking=True)
            images = add_noise_to_images(images, noise_type='gaussian', noise_level=0.1)

        _, attention_maps, feature_maps = model(images, track=True)

        visualize_attention_maps(config, images, attention_maps, feature_maps, save_path=args.save_image_path)
        return

    if args.pretrained:
        load_pretrained(args.pretrained, model_without_ddp, logger)

    if config.TRAIN.AUTO_RESUME:
        resume_file = auto_resume_helper(config.OUTPUT)
        if resume_file:
            if config.MODEL.RESUME:
                logger.warning(f"auto-resume changing resume file from {config.MODEL.RESUME} to {resume_file}")
            config.defrost()
            config.MODEL.RESUME = resume_file
            config.freeze()
            logger.info(f'auto resuming from {resume_file}')
        else:
            logger.info(f'no checkpoint found in {config.OUTPUT}, ignoring auto resume')

    if config.MODEL.RESUME:
        load_checkpoint(config, model_without_ddp, optimizer, lr_scheduler, logger)
        max_val_accuracy, _ = validate(config, data_loader_val, model, logger)
        logger.info(f"Accuracy of the loaded model on the {len(dataset_val)} val images: {max_val_accuracy:.2f}%")
        max_test_accuracy, _ = validate(config, data_loader_test, model, logger)
        logger.info(f"Accuracy of the loaded model on the {len(dataset_test)} test images: {max_test_accuracy:.2f}%")
    
    if config.THROUGHPUT_MODE:
        throughput(data_loader_val, model, logger)
        return

    logger.info("Start training")
    start_time = time.time()
    for epoch in range(config.TRAIN.START_EPOCH, total_epochs):
        data_loader_train.sampler.set_epoch(epoch)

        train_acc1 = train_one_epoch(config, model, criterion, data_loader_train, optimizer, epoch, mixup_fn, lr_scheduler, logger, total_epochs)

        val_acc1, _ = validate(config, data_loader_val, model, logger)
        is_text_dataset = config.DATA.DATASET.lower() in ['rotten_tomatoes', 'imdb', '20_newsgroups', 'ag_news'] or config.DATA.DATASET.lower().startswith('glue_')
        if is_text_dataset:
            logger.info(f"Train accuracy: {train_acc1:.2f}%, Val accuracy: {val_acc1:.2f}% on {len(dataset_val)} val samples")
        else:
            logger.info(f"Train accuracy: {train_acc1:.2f}%, Val accuracy: {val_acc1:.2f}% on {len(dataset_val)} val images")
        max_val_accuracy = max(max_val_accuracy, val_acc1)

        # 对于 GLUE 任务，跳过测试集评估
        is_glue = 'glue' in config.DATA.DATASET.lower()
        if not is_glue:
            test_acc1, _ = validate(config, data_loader_test, model, logger)
            logger.info(f"Above is Accuracy of the network on the {len(dataset_test)} test images: {test_acc1:.2f}%")
        else:
            test_acc1 = 0.0 # Placeholder
            logger.info("Skipping test set evaluation for GLUE datasets (labels are not available)")
            
        if dist.get_rank() == 0 and ((epoch + 1) % config.SAVE_FREQ == 0 or (epoch + 1) == (total_epochs)):
            save_checkpoint_new(config, epoch + 1, model_without_ddp, max(max_val_accuracy, val_acc1), optimizer, lr_scheduler, logger)

        if dist.get_rank() == 0 and ((epoch + 1) % config.SAVE_FREQ == 0 or (epoch + 1) == (total_epochs)) and val_acc1 >= max_val_accuracy:
            max_test_accuracy = test_acc1
            save_checkpoint_new(config, epoch + 1, model_without_ddp, max(max_val_accuracy, val_acc1), optimizer, lr_scheduler, logger, name='max_acc')

        logger.info(f'Train accuracy: {train_acc1:.2f}%, max val accuracy: {max_val_accuracy:.2f}%; the test accuracy corresponding to the max val accuracy: {max_test_accuracy:.2f}%')

    total_time = time.time() - start_time
    total_time_str = str(datetime.timedelta(seconds=int(total_time)))
    logger.info(f"Training time {total_time_str}")


def train_one_epoch(config, model, criterion, data_loader, optimizer, epoch, mixup_fn, lr_scheduler, logger, total_epochs):
    model.train()
    optimizer.zero_grad()

    num_steps = len(data_loader)
    batch_time = AverageMeter()
    loss_meter = AverageMeter()
    norm_meter = AverageMeter()
    acc1_meter = AverageMeter()

    start = time.time()
    end = time.time()

    scaler = GradScaler(device="cuda")

    for idx, (samples, targets) in enumerate(data_loader):
        

        optimizer.zero_grad()
        samples = samples.cuda(non_blocking=True)
        targets_orig = targets.cuda(non_blocking=True)
        targets = targets_orig

        if mixup_fn is not None:
            samples, targets = mixup_fn(samples, targets)
        
        if config.AMP: 
            with autocast(device_type="cuda"):
                outputs = model(samples)
                loss = criterion(outputs, targets)
            scaler.scale(loss).backward()
            if config.TRAIN.CLIP_GRAD:
                scaler.unscale_(optimizer)
                grad_norm = nn.utils.clip_grad_norm_(model.parameters(), config.TRAIN.CLIP_GRAD)
                scaler.step(optimizer)
                scaler.update()
            else:
                grad_norm = get_grad_norm(model.parameters())
                scaler.step(optimizer)
                scaler.update()
        else:
            outputs = model(samples)
            if isinstance(outputs, tuple):
                loss = 0.5 * (criterion(outputs[0], targets) + criterion(outputs[1], targets))
            else:
                loss = criterion(outputs, targets)
            loss.backward()
            if config.TRAIN.CLIP_GRAD:
                grad_norm = nn.utils.clip_grad_norm_(model.parameters(), config.TRAIN.CLIP_GRAD)
            else:
                grad_norm = get_grad_norm(model.parameters())
            optimizer.step()

        lr_scheduler.step_update(epoch * num_steps + idx)

        torch.cuda.synchronize()
        
        # 计算准确率：始终使用原始标签 targets_orig，无论是否使用 mixup
        # 处理 outputs 可能是 tuple 的情况
        if isinstance(outputs, tuple):
            outputs_for_acc = outputs[0]
        else:
            outputs_for_acc = outputs
        
        acc1 = accuracy(outputs_for_acc, targets_orig, topk=(1,))[0]
        acc1 = reduce_tensor(acc1)
        acc1_meter.update(acc1.item(), targets_orig.size(0))

        loss_meter.update(loss.item(), targets.size(0))
        norm_meter.update(grad_norm)
        batch_time.update(time.time() - end)
        end = time.time()

        if (idx + 1) % config.PRINT_FREQ == 0:
            lr = optimizer.param_groups[0]['lr']
            memory_used = torch.cuda.max_memory_allocated() / (1024.0 * 1024.0)
            etas = batch_time.avg * (num_steps - idx)
            acc_str = f'acc1 {acc1_meter.val:.3f} ({acc1_meter.avg:.3f})'
            logger.info(
                f'Train: [{epoch + 1}/{total_epochs}][{idx + 1}/{num_steps}]\t'
                f'eta {datetime.timedelta(seconds=int(etas))} lr {lr:.6f}\t'
                f'time {batch_time.val:.4f} ({batch_time.avg:.4f})\t'
                f'loss {loss_meter.val:.4f} ({loss_meter.avg:.4f})\t'
                f'{acc_str}\t'
                f'grad_norm {norm_meter.val:.4f} ({norm_meter.avg:.4f})\t'
                f'mem {memory_used:.0f}MB')
    epoch_time = time.time() - start
    train_acc = acc1_meter.avg if acc1_meter.count > 0 else 0.0
    logger.info(f"EPOCH {epoch + 1} training takes {datetime.timedelta(seconds=int(epoch_time))}, train acc: {train_acc:.2f}%")
    return train_acc


if __name__ == '__main__':
    main()