utils.py

import os
import argparse
import time
from config import get_config

import torch
import torch.nn as nn
import torch.distributed as dist
import matplotlib.pyplot as plt
import numpy as np
from PIL import Image
import torchvision.transforms as transforms

from timm.utils import accuracy, AverageMeter
from data.build import DATASET_STATS


def add_noise_to_images(images, noise_type='gaussian', noise_level=0.1):
    if noise_type == 'gaussian':
        noise = torch.randn_like(images) * noise_level
        noisy_images = torch.clamp(images + noise, 0, 1)
    elif noise_type == 'uniform':
        noise = (torch.rand_like(images) - 0.5) * 2 * noise_level
        noisy_images = torch.clamp(images + noise, 0, 1)
    elif noise_type == 'salt_pepper':
        noisy_images = images.clone()
        mask = torch.rand_like(images) < noise_level
        salt_pepper_values = torch.randint(0, 2, noisy_images[mask].shape, device=images.device).float()
        noisy_images[mask] = salt_pepper_values
    else:
        raise ValueError(f"不支持的噪声类型: {noise_type}")
    
    return noisy_images


def add_noise_to_text(tokens, noise_level=0.1, mask_token_id=0):
    """
    为文本token添加mask噪声
    
    Args:
        tokens: 输入的token序列 [B, L]
        noise_level: 噪声比例,表示要mask的token比例
        mask_token_id: 用于mask的token id,默认为0(PAD token)
    
    Returns:
        noisy_tokens: 添加噪声后的token序列
    """
    noisy_tokens = tokens.clone()
    mask = torch.rand(tokens.shape, device=tokens.device) < noise_level
    noisy_tokens[mask] = mask_token_id
    
    return noisy_tokens


def load_checkpoint(config, model, optimizer, lr_scheduler, logger):
    logger.info(f"==============> Resuming form {config.MODEL.RESUME}....................")
    if config.MODEL.RESUME.startswith('https'):
        checkpoint = torch.hub.load_state_dict_from_url(
            config.MODEL.RESUME, map_location='cpu', check_hash=True)
    else:
        checkpoint = torch.load(config.MODEL.RESUME, map_location='cpu', weights_only=False)
    msg = model.load_state_dict(checkpoint['model'], strict=False)
    logger.info(msg)
    max_accuracy = 0.0
    if not config.EVAL_MODE and 'optimizer' in checkpoint and 'lr_scheduler' in checkpoint and 'epoch' in checkpoint:
        optimizer.load_state_dict(checkpoint['optimizer'])
        lr_scheduler.load_state_dict(checkpoint['lr_scheduler'])
        config.defrost()
        config.TRAIN.START_EPOCH = checkpoint['epoch']
        config.freeze()
        logger.info(f"=> loaded successfully '{config.MODEL.RESUME}' (epoch {checkpoint['epoch']})")
    if 'max_accuracy' in checkpoint:
        max_accuracy = checkpoint['max_accuracy']

    del checkpoint
    torch.cuda.empty_cache()
    return max_accuracy

def load_pretrained(ckpt_path, model, logger):
    logger.info(f"==============> Loading pretrained form {ckpt_path}....................")
    checkpoint = torch.load(ckpt_path, map_location='cpu')
    # msg = model.load_pretrained(checkpoint['model'])
    # logger.info(msg)
    # logger.info(f"=> Loaded successfully {ckpt_path} ")
    # del checkpoint
    # torch.cuda.empty_cache()
    state_dict = checkpoint['model'] if 'model' in checkpoint.keys() else checkpoint

    # delete relative_position_index since we always re-init it
    relative_position_index_keys = [k for k in state_dict.keys() if "relative_position_index" in k]
    for k in relative_position_index_keys:
        del state_dict[k]

    # delete relative_coords_table since we always re-init it
    relative_position_index_keys = [k for k in state_dict.keys() if "relative_coords_table" in k]
    for k in relative_position_index_keys:
        del state_dict[k]

    # delete attn_mask since we always re-init it
    attn_mask_keys = [k for k in state_dict.keys() if "attn_mask" in k]
    for k in attn_mask_keys:
        del state_dict[k]

    # linear interpolate fsa bias if not match
    fsa_bias_keys = [k for k in state_dict.keys() if ("ah_bias" in k) or ("aw_bias" in k)
                                or ("ha_bias" in k) or ("wa_bias" in k)]
    for k in fsa_bias_keys:
        if "ah_bias" in k:
            squeeze_dim, permute = -1, False
        elif "aw_bias" in k:
            squeeze_dim, permute = -2, False
        elif "ha_bias" in k:
            squeeze_dim, permute = -2, True
        else:
            squeeze_dim, permute = -3, True
        fsa_bias_pretrained = state_dict[k].squeeze(dim=0).squeeze(dim=squeeze_dim)
        fsa_bias_current = model.state_dict()[k].squeeze(dim=0).squeeze(dim=squeeze_dim)
        if permute:
            fsa_bias_pretrained = fsa_bias_pretrained.permute(0, 2, 1)
            fsa_bias_current = fsa_bias_current.permute(0, 2, 1)
        num_heads1, fsa_num1, hw1 = fsa_bias_pretrained.size()
        num_heads2, fsa_num2, hw2 = fsa_bias_current.size()
        if (num_heads1 != num_heads2) or (fsa_num1 != fsa_num2):
            logger.warning(f"Error in loading {k}, passing......")
        else:
            if hw1 != hw2:
                # linear interpolate fsa bias if not match
                fsa_bias_pretrained_resized = torch.nn.functional.interpolate(
                    fsa_bias_pretrained, size=hw2, mode='linear')
                if permute:
                    fsa_bias_pretrained_resized = fsa_bias_pretrained_resized.permute(0, 2, 1)
                state_dict[k] = fsa_bias_pretrained_resized.unsqueeze(dim=0).unsqueeze(dim=squeeze_dim)

    # bicubic interpolate patch_embed.proj if not match
    patch_embed_keys = [k for k in state_dict.keys() if ("patch_embed" in k) and (".proj.weight" in k)]
    for k in patch_embed_keys:
        patch_embed_pretrained = state_dict[k]
        patch_embed_current = model.state_dict()[k]
        out1, in1, h1, w1 = patch_embed_pretrained.size()
        out2, in2, h2, w2 = patch_embed_current.size()
        if (out1 != out2) or (in1 != in2):
            logger.warning(f"Error in loading {k}, passing......")
        else:
            if (h1 != h2) or (w1 != w2):
                # bicubic interpolate patch_embed.proj if not match
                patch_embed_pretrained_resized = torch.nn.functional.interpolate(
                    patch_embed_pretrained, size=(h2, w2), mode='bicubic')
                state_dict[k] = patch_embed_pretrained_resized

    # bicubic interpolate relative_position_bias_table if not match
    relative_position_bias_table_keys = [k for k in state_dict.keys() if "relative_position_bias_table" in k]
    for k in relative_position_bias_table_keys:
        relative_position_bias_table_pretrained = state_dict[k]
        relative_position_bias_table_current = model.state_dict()[k]
        L1, nH1 = relative_position_bias_table_pretrained.size()
        L2, nH2 = relative_position_bias_table_current.size()
        if nH1 != nH2:
            logger.warning(f"Error in loading {k}, passing......")
        else:
            if L1 != L2:
                # bicubic interpolate relative_position_bias_table if not match
                S1 = int(L1 ** 0.5)
                S2 = int(L2 ** 0.5)
                relative_position_bias_table_pretrained_resized = torch.nn.functional.interpolate(
                    relative_position_bias_table_pretrained.permute(1, 0).view(1, nH1, S1, S1), size=(S2, S2),
                    mode='bicubic')
                state_dict[k] = relative_position_bias_table_pretrained_resized.view(nH2, L2).permute(1, 0)

    # bicubic interpolate absolute_pos_embed if not match
    absolute_pos_embed_keys = [k for k in state_dict.keys() if "pos_embed" in k]
    for k in absolute_pos_embed_keys:
        # dpe
        absolute_pos_embed_pretrained = state_dict[k]
        absolute_pos_embed_current = model.state_dict()[k]
        _, L1, C1 = absolute_pos_embed_pretrained.size()
        _, L2, C2 = absolute_pos_embed_current.size()
        if C1 != C1:
            logger.warning(f"Error in loading {k}, passing......")
        else:
            if L1 != L2:
                S1 = int(L1 ** 0.5)
                S2 = int(L2 ** 0.5)
                i, j = L1 - S1 ** 2, L2 - S2 ** 2
                absolute_pos_embed_pretrained_ = absolute_pos_embed_pretrained[:, i:, :].reshape(-1, S1, S1, C1)
                absolute_pos_embed_pretrained_ = absolute_pos_embed_pretrained_.permute(0, 3, 1, 2)
                absolute_pos_embed_pretrained_resized = torch.nn.functional.interpolate(
                    absolute_pos_embed_pretrained_, size=(S2, S2), mode='bicubic')
                absolute_pos_embed_pretrained_resized = absolute_pos_embed_pretrained_resized.permute(0, 2, 3, 1)
                absolute_pos_embed_pretrained_resized = absolute_pos_embed_pretrained_resized.flatten(1, 2)
                state_dict[k] = torch.cat([absolute_pos_embed_pretrained[:, :j, :],
                                           absolute_pos_embed_pretrained_resized], dim=1)

    # check classifier, if not match, then re-init classifier to zero
    head_bias_pretrained = state_dict['head.bias']
    Nc1 = head_bias_pretrained.shape[0]
    Nc2 = model.head.bias.shape[0]
    if (Nc1 != Nc2):
        if Nc1 == 21841 and Nc2 == 1000:
            logger.info("loading ImageNet-22K weight to ImageNet-1K ......")
            map22kto1k_path = f'data/map22kto1k.txt'
            with open(map22kto1k_path) as f:
                map22kto1k = f.readlines()
            map22kto1k = [int(id22k.strip()) for id22k in map22kto1k]
            state_dict['head.weight'] = state_dict['head.weight'][map22kto1k, :]
            state_dict['head.bias'] = state_dict['head.bias'][map22kto1k]
        else:
            torch.nn.init.constant_(model.head.bias, 0.)
            torch.nn.init.constant_(model.head.weight, 0.)
            del state_dict['head.weight']
            del state_dict['head.bias']
            logger.warning(f"Error in loading classifier head, re-init classifier head to 0")

    msg = model.load_state_dict(state_dict, strict=False)
    logger.warning(msg)

    logger.info(f"=> loaded successfully '{ckpt_path}'")

    del checkpoint
    torch.cuda.empty_cache()


def save_checkpoint(config, epoch, model, max_accuracy, optimizer, lr_scheduler, logger):
    save_state = {'model': model.state_dict(),
                  'optimizer': optimizer.state_dict(),
                  'lr_scheduler': lr_scheduler.state_dict(),
                  'max_accuracy': max_accuracy,
                  'epoch': epoch,
                  'config': config}

    save_path = os.path.join(config.OUTPUT, f'ckpt_epoch_{epoch}.pth')
    logger.info(f"{save_path} saving......")
    torch.save(save_state, save_path)
    logger.info(f"{save_path} saved !!!")
    
def save_checkpoint_new(config, epoch, model, max_accuracy, optimizer, lr_scheduler, logger, name=None):
    save_state = {'model': model.state_dict(),
                  'optimizer': optimizer.state_dict(),
                  'lr_scheduler': lr_scheduler.state_dict(),
                  'max_accuracy': max_accuracy,
                  'epoch': epoch,
                  'config': config}
    if name == None:
        old_ckpt = os.path.join(config.OUTPUT, f'ckpt_epoch_{epoch-3}.pth') 
        if os.path.exists(old_ckpt):
            os.remove(old_ckpt)

    if name!=None:
        save_path = os.path.join(config.OUTPUT, f'{name}.pth')
        logger.info(f"{save_path} saving......")
        torch.save(save_state, save_path)
        logger.info(f"{save_path} saved !!!")
    else:
        save_path = os.path.join(config.OUTPUT, f'ckpt_epoch_{epoch}.pth')
        logger.info(f"{save_path} saving......")
        torch.save(save_state, save_path)
        logger.info(f"{save_path} saved !!!")

def get_grad_norm(parameters, norm_type=2):
    if isinstance(parameters, torch.Tensor):
        parameters = [parameters]
    parameters = list(filter(lambda p: p.grad is not None, parameters))
    norm_type = float(norm_type)
    total_norm = 0
    for p in parameters:
        param_norm = p.grad.data.norm(norm_type)
        total_norm += param_norm.item() ** norm_type
    total_norm = total_norm ** (1. / norm_type)
    return total_norm


def auto_resume_helper(output_dir):
    checkpoints = os.listdir(output_dir)
    checkpoints = [ckpt for ckpt in checkpoints if ckpt.endswith('pth')]
    print(f"All checkpoints founded in {output_dir}: {checkpoints}")
    if len(checkpoints) > 0:
        latest_checkpoint = max([os.path.join(output_dir, d) for d in checkpoints], key=os.path.getmtime)
        print(f"The latest checkpoint founded: {latest_checkpoint}")
        resume_file = latest_checkpoint
    else:
        resume_file = None
    return resume_file


def reduce_tensor(tensor):
    rt = tensor.clone()
    dist.all_reduce(rt, op=dist.ReduceOp.SUM)
    rt /= dist.get_world_size()
    return rt


def parse_option():
    parser = argparse.ArgumentParser('training and evaluation script', add_help=False)
    parser.add_argument('--cfg', type=str, required=True, metavar="FILE", help='path to config file', )
    parser.add_argument('--opts', default=None, nargs='+')
    parser.add_argument('--batch-size', type=int, help="batch size for single GPU")
    parser.add_argument('--data-path', type=str, help='path to dataset')
    parser.add_argument('--resume', help='path to checkpoint you want to resume')
    parser.add_argument('--amp', action='store_true')
    parser.add_argument('--zip', action='store_true')
    parser.add_argument('--cache-mode', type=str, default='part', choices=['no', 'full', 'part'],
                        help='no: no cache, '
                             'full: cache all data, '
                             'part: sharding the dataset into nonoverlapping pieces and only cache one piece')
    parser.add_argument('--use-checkpoint', action='store_true',
                        help="whether to use gradient checkpointing to save memory")
    parser.add_argument('--output', type=str, metavar='PATH', help='root of output folder, the full path is <output>/<model_name>/<tag> (default: output)')
    parser.add_argument('--tag', help='tag of experiment')
    parser.add_argument('--eval', action='store_true', help='Perform evaluation only')
    parser.add_argument('--test', action='store_true', help='Perform test only')
    parser.add_argument('--get_attention_maps', action='store_true', help='Get attention maps')
    parser.add_argument('--image_path', type=str, help='Path to specific image for attention visualization')
    parser.add_argument('--save_image_path', type=str, help='Path to save attention maps')
    parser.add_argument('--throughput', action='store_true', help='Test throughput only')
    parser.add_argument('--pretrained', type=str, help='Finetune 384 initial checkpoint.', default='')
    parser.add_argument('--find_unused_params', action='store_true', default=False)
    parser.add_argument('--seed', type=int, help='random seed', default=0)
    parser.add_argument('--local-rank', type=int)
    
    args, _ = parser.parse_known_args()
    config = get_config(args)
    return args, config


@torch.no_grad()
def validate(config, data_loader, model, logger, noise_type=None, noise_level=None):
    criterion = nn.CrossEntropyLoss()
    model.eval()

    loss_meter = AverageMeter()
    acc1_meter = AverageMeter()
    
    is_text_dataset = config.DATA.DATASET.lower() in ['rotten_tomatoes', 'imdb', '20_newsgroups', 'ag_news']

    for idx, (samples, target) in enumerate(data_loader):
        samples = samples.cuda(non_blocking=True)
        target = target.cuda(non_blocking=True)

        if noise_type and noise_level:
            if is_text_dataset:
                samples = add_noise_to_text(samples, noise_level=noise_level, mask_token_id=0)
            else:
                samples = add_noise_to_images(samples, noise_type, noise_level)

        output = model(samples)

        loss = criterion(output, target)
        acc1 = accuracy(output, target, topk=(1,))[0]
        
        acc1 = reduce_tensor(acc1)
        loss = reduce_tensor(loss)

        loss_meter.update(loss.item(), target.size(0))
        acc1_meter.update(acc1.item(), target.size(0))

    return acc1_meter.avg, loss_meter.avg


@torch.no_grad()
def throughput(data_loader, model, logger):
    model.eval()
    images, _ = data_loader[0]
    images = images.cuda(non_blocking=True)
    batch_size = images.shape[0]
    for _ in range(50):
        model(images)
    torch.cuda.synchronize()
    logger.info(f"throughput averaged with 30 times")
    tic1 = time.time()
    for _ in range(30):
        model(images)
    torch.cuda.synchronize()
    tic2 = time.time()
    logger.info(f"batch_size {batch_size} throughput {30 * batch_size / (tic2 - tic1)}")


def load_image_from_path(image_path, config):
    img = Image.open(image_path).convert('RGB')

    mean = DATASET_STATS[config.DATA.DATASET]['mean']
    std = DATASET_STATS[config.DATA.DATASET]['std']

    transform = transforms.Compose([
        transforms.Resize((config.DATA.IMG_SIZE, config.DATA.IMG_SIZE)),
        transforms.ToTensor(),
        transforms.Normalize(mean, std)
    ])
    
    img_tensor = transform(img)  # [C, H, W]
    
    img_tensor = img_tensor.unsqueeze(0)  # [1, C, H, W]
    
    return img_tensor


def visualize_attention_maps(config, images, attention_maps, feature_maps, save_path="visualization/vit/attention_map.png"):
    """
    images: Tensor [B, C, H, W]
    attention_maps: List of torch.Tensor，每个 [1, num_heads, N, N] 或 [1, N, N]
    """

    image = images[0].cpu() # [C, H, W]
    image_np = image.permute(1, 2, 0).numpy() # [H, W, C]
    image_np = np.clip(image_np, 0, 1)

    num_layers = len(attention_maps)
    fig, axes = plt.subplots(1, num_layers + 1, figsize=(8 * (num_layers + 1), 8))

    axes[0].imshow(image_np)
    axes[0].axis('off')

    for idx, amap in enumerate(attention_maps):
        attn = amap.squeeze(0).mean(dim=0) # [N, N]
        cls_attention = attn[0, 1:].detach().cpu()
        H = W = int(cls_attention.shape[0] ** 0.5)
        attn_2d = cls_attention.reshape(H, W).unsqueeze(0).unsqueeze(0) # [1, 1, H, W]
        attn_resized = torch.nn.functional.interpolate(attn_2d, size=image_np.shape[:2], mode='bilinear', align_corners=False).squeeze().numpy()
        # axes[idx + 1].imshow(image_np)
        im = axes[idx + 1].imshow(attn_resized, cmap='jet', alpha=0.6)
        plt.colorbar(im, ax=axes[idx + 1], fraction=0.046, pad=0.04)
        axes[idx + 1].set_title(f"Attn-{idx}")
        axes[idx + 1].axis('off')

    plt.tight_layout()
    plt.savefig(save_path)
    plt.close()
    print(f"可视化已保存到: {save_path}")