eval.py

import tqdm
import argparse
import math
# import torchac
import sys
import os
import time
import logging
from datetime import datetime

import torch
import torch.nn as nn
import torch.optim as optim

import torchvision
from torch.utils.data import DataLoader
from torchvision import transforms
from torchvision.models import resnet50

import yaml
from pytorch_msssim import ms_ssim
from DISTS_pytorch import DISTS
from util.lpips import LPIPS
from torch.nn import functional as F
from torchvision import utils as vutils
import numpy as np
import glob

import util.misc as misc
import util.lr_sched as lr_sched
from torch.utils.tensorboard import SummaryWriter
import models_mage_codec
import timm.optim.optim_factory as optim_factory
from util.misc import NativeScalerWithGradNormCount as NativeScaler
import json
import PIL.Image as Image
import torch.backends.cudnn as cudnn
from pathlib import Path
import random
import torch.distributed as dist
import lpips
from util.dataloader import MSCOCO, Kodak, prepadding

class CalMetrics(nn.Module):
    """Calculate BPP, PSNR, MS-SSIM, LPIPS and DISTS for the reconstructed image."""

    def __init__(self):
        super().__init__()
        self.mse = nn.MSELoss()

    def bpp_loss(self, ori, out_net):
        b, _, h, w = ori.shape
        num_pixels = b * h * w
        bpp = torch.log(out_net["likelihoods"]).sum() / (-math.log(2) * num_pixels)
        bs_mask_token = out_net['bs_mask_token']
        bytes_length = len(bs_mask_token)
        # 因为每个字节包含8位，所以总位数是字节数 * 8
        total_bits = bytes_length * 8
        # 计算每像素的位数（bpp）
        bpp_mask = total_bits / num_pixels
        return bpp, bpp_mask

    def psnr(self, rec, ori):
        mse = torch.mean((rec - ori) ** 2)
        if(mse == 0):
            return 100
        max_pixel = 1.
        psnr = 10 * torch.log10(max_pixel / mse)
        return torch.mean(psnr)

    def lpips(self, rec, ori):
        loss_fn_alex = lpips.LPIPS(net='alex').cuda()
        lipis_alex = loss_fn_alex(rec, ori)
        return lipis_alex.mean()
    
    def dists(self, rec, ori):
        D = DISTS().cuda()
        dists_value = D(rec, ori)
        return dists_value.mean()
    
    def cal_total_loss(self, lpips, bpp, out_net):
        # task_loss = out_net['task_loss'] + 0.1 * lpips
        task_loss = out_net['task_loss']
        total_loss = bpp + out_net['lambda'] * task_loss
        return total_loss

    def forward(self, ori, out_net, rec=None):
        out = {}
        out["bpp"], out["bpp_mask"] = self.bpp_loss(ori, out_net)
        # out["bpp_loss"] = out["bpp"] + out["bpp_mask"]
        out["bpp_loss"] = out["bpp"]
        if rec is not None:
            out["psnr"] = self.psnr(torch.clamp(rec, 0, 1), ori)
            out["msssim"] = ms_ssim(torch.clamp(rec, 0, 1), ori, data_range=1, size_average=True)
            out["lpips"] = self.lpips(torch.clamp(rec, 0, 1), ori)
            out["dists"] = self.dists(torch.clamp(rec, 0, 1), ori)
            out["total_loss"] = self.cal_total_loss(out["lpips"], out["bpp_loss"], out_net)
        return out


class FeatureHook():
    def __init__(self, module):
        module.register_forward_hook(self.attach)
    
    def attach(self, model, input, output):
        self.feature = output

class Clsloss(nn.Module):
    def __init__(self, device, cls_loss=True) -> None:
        super().__init__()
        self.ce = nn.CrossEntropyLoss()
        self.classifier = resnet50(True)
        self.classifier.requires_grad_(False)
        self.hooks = [FeatureHook(i) for i in [ # for calculating perceptual loss
            self.classifier.layer1,
            self.classifier.layer2,
            self.classifier.layer3,
            self.classifier.layer4,
        ]]
        self.classifier = self.classifier.to(device)
        for k, p in self.classifier.named_parameters():
            p.requires_grad = False
        self.classifier.eval()
        self.cls_loss = cls_loss
        self.transform = transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
        self.perceptual_loss = False

    def accuracy(output, target, topk=(1,)):
        maxk = max(topk)
        batch_size = target.size(0)

        _, pred = output.topk(maxk, 1, True, True)
        pred = pred.t()
        correct = pred.eq(target.view(1, -1).expand_as(pred))

        res = []
        for k in topk:
            correct_k = correct[:k].view(-1).float().sum(0)
            res.append(correct_k.mul_(100.0 / batch_size))
        return res

    def forward(self, d, rec, y_true):
        # output = codec(d), d = original image, y_true = label
        x_hat = torch.clamp(rec,0,1)
        pred = self.classifier(self.transform(x_hat))   # transform is normalization
        cls_loss = self.ce(pred, y_true)
        accu = sum(torch.argmax(pred,-1)==y_true)/pred.shape[0]
        if self.perceptual_loss:
            pred_feat = [i.feature.clone() for i in self.hooks]
            _ = self.classifier(self.transform(d))
            ori_feat = [i.feature.clone() for i in self.hooks]
            perc_loss = torch.stack([nn.functional.mse_loss(p,o, reduction='none').mean((1,2,3)) for p,o in zip(pred_feat, ori_feat)])
            perc_loss = perc_loss.mean()
            return perc_loss

        return cls_loss, accu, None


class AverageMeter:
    """Compute running average."""

    def __init__(self):
        self.val = 0
        self.avg = 0
        self.sum = 0
        self.count = 0

    def update(self, val, n=1):
        self.val = val
        self.sum += val * n
        self.count += n
        self.avg = self.sum / self.count

class CustomDataParallel(nn.DataParallel):
    """Custom DataParallel to access the module methods."""

    def __getattr__(self, key):
        try:
            return super().__getattr__(key)
        except AttributeError:
            return getattr(self.module, key)

def init(args):
    base_dir = f'{args.root}/{args.exp_name}/'
    os.makedirs(base_dir, exist_ok=True)
    return base_dir

def setup_logger(log_dir):
    log_formatter = logging.Formatter("%(asctime)s [%(levelname)-5.5s]  %(message)s")
    root_logger = logging.getLogger()
    root_logger.setLevel(logging.INFO)

    log_file_handler = logging.FileHandler(log_dir, encoding='utf-8')
    log_file_handler.setFormatter(log_formatter)
    root_logger.addHandler(log_file_handler)

    log_stream_handler = logging.StreamHandler(sys.stdout)
    log_stream_handler.setFormatter(log_formatter)
    root_logger.addHandler(log_stream_handler)

    logging.info('Logging file is %s' % log_dir)

def load_img(p, padding=True, factor=64):
    x = Image.open(p)
    x = torch.from_numpy(np.asarray(x))
    if len(x.shape) == 2:
        x = x.unsqueeze(-1).repeat(1,1,3)    # h,w -> h,w,3
    x = x.permute(2, 0, 1).unsqueeze(0).float().div(255)
    h, w = x.shape[2:4]

    if padding:
        dh = factor * math.ceil(h / factor) - h
        dw = factor * math.ceil(w / factor) - w
        x = F.pad(x, (0, dw, 0, dh))
    return x, h, w

def save_img(img: torch.Tensor, vis_path, input_p, mask=False):
    img = img.clone().detach()
    img = img.to(torch.device('cpu'))
    if os.path.isdir(vis_path) is not True:
        os.makedirs(vis_path)
    end = '/'
    if mask:
        img_name = vis_path + 'mask_' + str(input_p[input_p.rfind(end):])
    else:
        img_name = vis_path + str(input_p[input_p.rfind(end):])
    vutils.save_image(img, os.path.join(vis_path, img_name), nrow=8)

def inference(epoch, test_loader, model, metrics_criterion, device, manual_mask_ratio, args, stage='test'):
    model.eval()
    bpp_loss = AverageMeter() 
    bpp_mask = AverageMeter() 
    psnr = AverageMeter()
    msssim = AverageMeter()
    lpips = AverageMeter()
    dists = AverageMeter()
    test_loss = AverageMeter()

    vis_path = os.path.join("./MIM_vbr_eval_lossless/", stage)
    os.makedirs(vis_path, exist_ok=True)
    if stage == 'test':
        test_vis_path = os.path.join("/home/t2vg-a100-G4-10/project/qyp/rec_fid", manual_mask_ratio)
        os.makedirs(test_vis_path, exist_ok=True)

    with torch.no_grad():
        # tqdm_meter = tqdm.tqdm(eval_path, leave=False)
        tqdm_meter = tqdm.tqdm(enumerate(test_loader),leave=False, total=len(test_loader))
        for i, (d, l) in tqdm_meter:
            d = d.to(device)
            # l = l.to(device)
        # for input_p in eval_path:
            # x, hx, wx = load_img(input_p, padding=True, factor=64)
            # x = x.to(device)
            out_net = model(d, is_training=False, manual_mask_rate=manual_mask_ratio)
            rec = model.module.gen_img(out_net['logits'], out_net['token_all_mask'], out_net['token_indices'])

            # x = x[:, :, :hx, :wx]
            # rec = rec[:, :, :hx, :wx]
            rec = rec.to(device)
            out_criterion = metrics_criterion(d, out_net, rec)

            bpp_loss.update(out_criterion["bpp_loss"])
            bpp_mask.update(out_criterion["bpp_mask"])
            psnr.update(out_criterion['psnr'])
            msssim.update(out_criterion['msssim'])
            lpips.update(out_criterion['lpips'])
            dists.update(out_criterion['dists'])
            test_loss.update(out_criterion['total_loss'])
    
            ## ======================= update progress bar & visualization ======================= ##
            if stage == 'val':
                if i % 5 == 0:
                    with torch.no_grad():
                        real_fake_images = torch.cat((d, rec), dim=0)
                        vutils.save_image(real_fake_images, os.path.join(vis_path, f"{epoch}_{i}.jpg"), nrow=8)
                        vutils.save_image(out_net['mask_vis'], os.path.join(vis_path, f"{epoch}_{i}_mask.jpg"), nrow=8)
            if stage == 'test':
                with torch.no_grad():
                    vutils.save_image(rec, os.path.join(test_vis_path, f"{i}.jpg"), nrow=8)
    # txt = f"Rec Loss:{test_loss.avg:.4f}|Bpp:{bpp_loss.avg:.4f}|lpips:{lpips.avg:.4f}|msssim:{msssim.avg:.4f}|dists:{dists.avg:.4f}|psnr:{psnr.avg:.4f}\n"
    # tqdm_meter.set_postfix_str(txt)

    model.train()

    # 假设其它变量和环境已经正确设置
    if torch.distributed.is_initialized():
        rank = dist.get_rank()
    else:
        rank = 0  # 假设未使用DDP，则默认为单进程模式，rank为0

    if rank == 0:
        log_txt = f"{epoch}|bpp:{bpp_loss.avg.item():.5f}|mask:{bpp_mask.avg:.5f}|mask_ratio:{manual_mask_ratio}|psnr:{psnr.avg.item():.5f}|msssim:{msssim.avg.item():.5f}|lpips:{lpips.avg.item():.5f}|dists:{dists.avg.item():.5f}|Test loss:{test_loss.avg.item():.5f}"
        logging.info(log_txt)
    return test_loss.avg

def save_images_by_batch(rec, vis_path, batch_idx):
    # 保存每个图像
    for i in range(rec.size(0)):  # 假设 rec 是 (batch_size, C, H, W)
        image_filename = os.path.join(vis_path, f"batch_{batch_idx}_{i}.jpg")
        # rec[i] 选择单个图像，unsqueeze(0) 添加一个批次维度
        vutils.save_image(rec[i], image_filename)

def inference_with_acc(epoch, test_loader, model, metrics_criterion, cls_criterion, device, manual_mask_ratio, args, stage='test'):
    model.eval()
    bpp_loss = AverageMeter() 
    bpp_mask = AverageMeter() 
    psnr = AverageMeter()
    msssim = AverageMeter()
    lpips = AverageMeter()
    dists = AverageMeter()
    accuracy = AverageMeter()
    test_loss = AverageMeter()

    if stage == 'test':
        test_vis_path = os.path.join("/home/t2vg-a100-G4-10/project/qyp/rec_fid", str(manual_mask_ratio))
        os.makedirs(test_vis_path, exist_ok=True)

    with torch.no_grad():
        # tqdm_meter = tqdm.tqdm(eval_path, leave=False)
        tqdm_meter = tqdm.tqdm(enumerate(test_loader),leave=False, total=len(test_loader))
        for i, (d, l) in tqdm_meter:
            d = d.to(device)
            l = l.to(device)
        # for input_p in eval_path:
            # x, hx, wx = load_img(input_p, padding=True, factor=64)
            # x = x.to(device)
            out_net = model(d, is_training=False, manual_mask_rate=manual_mask_ratio)
            rec = model.module.gen_img(out_net['logits'], out_net['token_all_mask'], out_net['token_indices'])

            # x = x[:, :, :hx, :wx]
            # rec = rec[:, :, :hx, :wx]
            rec = rec.to(device)
            out_criterion = metrics_criterion(d, out_net, rec)
            _, accu, _ = cls_criterion(d, rec, l)

            bpp_loss.update(out_criterion["bpp_loss"])
            bpp_mask.update(out_criterion["bpp_mask"])
            psnr.update(out_criterion['psnr'])
            msssim.update(out_criterion['msssim'])
            lpips.update(out_criterion['lpips'])
            dists.update(out_criterion['dists'])
            test_loss.update(out_criterion['total_loss'])
            accuracy.update(accu)
    
            ## ======================= update progress bar & visualization ======================= ##
            if stage == 'test':
                with torch.no_grad():
                    save_images_by_batch(rec, test_vis_path, i)
                    # vutils.save_image(rec, os.path.join(test_vis_path, f"{i}.jpg"))
    # txt = f"Rec Loss:{test_loss.avg:.4f}|Bpp:{bpp_loss.avg:.4f}|lpips:{lpips.avg:.4f}|msssim:{msssim.avg:.4f}|dists:{dists.avg:.4f}|psnr:{psnr.avg:.4f}\n"
    # tqdm_meter.set_postfix_str(txt)

    model.train()
    # 假设其它变量和环境已经正确设置
    if torch.distributed.is_initialized():
        rank = dist.get_rank()
    else:
        rank = 0  # 假设未使用DDP，则默认为单进程模式，rank为0

    if rank == 0:
        log_txt = f"{epoch}|bpp:{bpp_loss.avg.item():.5f}|mask:{bpp_mask.avg:.5f}|mask_ratio:{manual_mask_ratio}|psnr:{psnr.avg.item():.5f}|msssim:{msssim.avg.item():.5f}|lpips:{lpips.avg.item():.5f}|dists:{dists.avg.item():.5f}|accu:{accuracy.avg:.5f}|Test loss:{test_loss.avg.item():.5f}"
        logging.info(log_txt)
    return test_loss.avg

def inference_with_kodak(epoch, test_loader, model, metrics_criterion, cls_criterion, device, manual_mask_ratio, args, stage='test'):
    model.eval()
    bpp_loss = AverageMeter() 
    bpp_mask = AverageMeter() 
    psnr = AverageMeter()
    msssim = AverageMeter()
    lpips = AverageMeter()
    dists = AverageMeter()
    accuracy = AverageMeter()
    test_loss = AverageMeter()

    if stage == 'test':
        test_vis_path = os.path.join("/home/t2vg-a100-G4-10/project/qyp/mimc_rope/eval_kodak_256", str(manual_mask_ratio))
        os.makedirs(test_vis_path, exist_ok=True)

    with torch.no_grad():
        # tqdm_meter = tqdm.tqdm(eval_path, leave=False)
        tqdm_meter = tqdm.tqdm(enumerate(test_loader), leave=False, total=len(test_loader))
        for i, d in tqdm_meter:
            d = d.to(device)
            d, h_ori, w_ori = prepadding(d, factor=256)
            # l = l.to(device)
        # for input_p in eval_path:
            # x, hx, wx = load_img(input_p, padding=True, factor=64)
            # x = x.to(device)
            out_net = model(d, is_training=False, manual_mask_rate=manual_mask_ratio)
            rec = model.module.gen_img(out_net['logits'], out_net['token_all_mask'], out_net['token_indices'])

            d = d[:, :, :h_ori, :w_ori]
            rec = rec[:, :, :h_ori, :w_ori]
            rec = rec.to(device)
            out_criterion = metrics_criterion(d, out_net, rec)
            # _, accu, _ = cls_criterion(d, rec, l)

            bpp_loss.update(out_criterion["bpp_loss"])
            bpp_mask.update(out_criterion["bpp_mask"])
            psnr.update(out_criterion['psnr'])
            msssim.update(out_criterion['msssim'])
            lpips.update(out_criterion['lpips'])
            dists.update(out_criterion['dists'])
            test_loss.update(out_criterion['total_loss'])
            # accuracy.update(accu)
    
            ## ======================= update progress bar & visualization ======================= ##
            if stage == 'test':
                with torch.no_grad():
                    save_images_by_batch(rec, test_vis_path, i)
                    # vutils.save_image(rec, os.path.join(test_vis_path, f"{i}.jpg"))
    # txt = f"Rec Loss:{test_loss.avg:.4f}|Bpp:{bpp_loss.avg:.4f}|lpips:{lpips.avg:.4f}|msssim:{msssim.avg:.4f}|dists:{dists.avg:.4f}|psnr:{psnr.avg:.4f}\n"
    # tqdm_meter.set_postfix_str(txt)

    model.train()
    # 假设其它变量和环境已经正确设置
    if torch.distributed.is_initialized():
        rank = dist.get_rank()
    else:
        rank = 0  # 假设未使用DDP，则默认为单进程模式，rank为0

    if rank == 0:
        log_txt = f"{epoch}|bpp:{bpp_loss.avg.item():.5f}|mask:{bpp_mask.avg:.5f}|mask_ratio:{manual_mask_ratio}|psnr:{psnr.avg.item():.5f}|msssim:{msssim.avg.item():.5f}|lpips:{lpips.avg.item():.5f}|dists:{dists.avg.item():.5f}|Test loss:{test_loss.avg.item():.5f}"
        logging.info(log_txt)
    return test_loss.avg

def save_checkpoint(state, is_best, base_dir, filename="checkpoint.pth.tar"):
    torch.save(state, base_dir+filename)
    if is_best:
        torch.save(state, base_dir+"checkpoint_best.pth.tar")
        # shutil.copyfile(base_dir+filename, base_dir+"checkpoint_best_loss.pth.tar")

def parse_args(argv):
    parser = argparse.ArgumentParser(description="Example training script.")
    parser.add_argument(
        "-c",
        "--config",
        default="config/vpt_default.yaml",
        help="Path to config file",
    )
    parser.add_argument(
        '--name', 
        default=datetime.now().strftime('%Y-%m-%d_%H_%M_%S'), 
        type=str,
        help='Result dir name', 
    )
    parser.add_argument('--lr', type=float, default=None, metavar='LR',
                        help='learning rate (absolute lr)')
    given_configs, remaining = parser.parse_known_args(argv)
    # distributed training parameters
    parser.add_argument('--world_size', default=1, type=int,
                        help='number of distributed processes')
    parser.add_argument('--local-rank', default=-1, type=int)
    parser.add_argument('--dist_on_itp', action='store_true')
    parser.add_argument('--dist_url', default='env://',
                        help='url used to set up distributed training')
    with open(given_configs.config) as file:
        yaml_data= yaml.safe_load(file)
        parser.set_defaults(**yaml_data)
    
    parser.add_argument(
        "-T",
        "--TEST",
        action='store_true',
        help='Testing'
    )
    args = parser.parse_args(remaining)
    return args

def load_eval_ps(eval_path):
        eval_ps = sorted(glob.glob(os.path.join(eval_path, '*.png')))
        return eval_ps

def main(argv):
    args = parse_args(argv)
    base_dir = init(args)   # create the base dir for saving the results

    if args.output_dir:
        Path(args.output_dir).mkdir(parents=True, exist_ok=True)
    args.log_dir = args.output_dir

    misc.init_distributed_mode(args)

    print('job dir: {}'.format(os.path.dirname(os.path.realpath(__file__))))
    print("{}".format(args).replace(', ', ',\n'))

    device = torch.device(args.device)
     # fix the seed for reproducibility
    seed = args.seed + misc.get_rank()
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    np.random.seed(seed)
    random.seed(seed)

    cudnn.benchmark = True
    
    setup_logger(base_dir + '/' + time.strftime('%Y%m%d_%H%M%S') + '.log')
    msg = f'======================= {args.name} ======================='
    logging.info(msg)
    for k in args.__dict__:
        logging.info(k + ':' + str(args.__dict__[k]))
    logging.info('=' * len(msg))

    ## ======================= prepare dataset ======================= ##
    transform_train = transforms.Compose([
            transforms.RandomResizedCrop(args.input_size, scale=(0.8, 1.0)),
            transforms.RandomHorizontalFlip(),
            transforms.ToTensor()])
    transform_test = transforms.Compose(    # crop to 256x256
        [transforms.Resize(256), transforms.CenterCrop(256), transforms.ToTensor()]
    )
    transform_val = transforms.Compose([
            transforms.Resize(256), 
            transforms.CenterCrop(256),
            transforms.ToTensor()
    ])

    if args.dataset=='imagenet': 
        train_dataset = torchvision.datasets.ImageFolder(os.path.join(args.dataset_path, "train"), transform=transform_train)
        test_dataset = torchvision.datasets.ImageFolder(os.path.join(args.dataset_path, "val"), transform=transform_test)
        # val_dataset, _ = torch.utils.data.random_split(test_dataset, [2000, 48000])
        val_dataset = Kodak(args.kodak_path, transform_val)
        small_train_datasets = torch.utils.data.random_split(train_dataset, [40000]*32+[1167])

    device = "cuda" if args.cuda and torch.cuda.is_available() else "cpu"

    # if args.distributed:
    if True:
        num_tasks = misc.get_world_size()
        global_rank = misc.get_rank()
        sampler_test = torch.utils.data.DistributedSampler(
            test_dataset, num_replicas=num_tasks, rank=global_rank, shuffle=True
        )
        sampler_val = torch.utils.data.DistributedSampler(
            val_dataset, num_replicas=num_tasks, rank=global_rank, shuffle=True
        )
    else:
        sampler_train = torch.utils.data.RandomSampler(train_dataset)

    if global_rank == 0 and args.log_dir is not None:
        os.makedirs(args.log_dir, exist_ok=True)
        log_writer = SummaryWriter(log_dir=args.log_dir)
    else:
        log_writer = None


    test_dataloader = DataLoader(test_dataset, sampler=sampler_test, batch_size=args.test_batch_size,
                                 num_workers=args.num_workers, shuffle=False, pin_memory=args.pin_mem,)
    val_dataloader = DataLoader(val_dataset,  sampler=sampler_val, batch_size=1,
                                num_workers=args.num_workers, shuffle=False, pin_memory=args.pin_mem, drop_last=True)

    ## ======================= prepare model ======================= ##
    vqgan_ckpt_path = '/home/t2vg-a100-G4-10/project/qyp/mage/vqgan_jax_strongaug.ckpt'
    model = models_mage_codec.__dict__[args.model](mask_ratio_mu=args.mask_ratio_mu, mask_ratio_std=args.mask_ratio_std,
                                             mask_ratio_min=args.mask_ratio_min, mask_ratio_max=args.mask_ratio_max,
                                             vqgan_ckpt_path=vqgan_ckpt_path)

    model.to(device)
    model_without_ddp = model
    print("Model = %s" % str(model_without_ddp))
    eff_batch_size = args.batch_size * args.accum_iter * misc.get_world_size()
    if args.lr is None:  # only base_lr is specified
        args.lr = args.blr * eff_batch_size / 256
    print("base lr: %.2e" % (args.lr * 256 / eff_batch_size))
    print("actual lr: %.2e" % args.lr)

    print("accumulate grad iterations: %d" % args.accum_iter)
    print("effective batch size: %d" % eff_batch_size)

    if args.distributed:
        model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.gpu], find_unused_parameters=True)
        model_without_ddp = model.module
    
    # following timm: set wd as 0 for bias and norm layers
    param_groups = optim_factory.add_weight_decay(model_without_ddp, args.weight_decay)
    optimizer = torch.optim.AdamW(param_groups, lr=args.lr, betas=(0.9, 0.95))
    print(optimizer)
    loss_scaler = NativeScaler()

    # resume from a checkpoint
    misc.load_model(args=args, model_without_ddp=model_without_ddp, 
                    optimizer=optimizer, loss_scaler=loss_scaler, strict=False)
    
    metrics_criterion = CalMetrics()
    cls_criterion = Clsloss(device, True)
    ## ======================= Start Training ======================= ##
    last_epoch = args.start_epoch
    
    if True:   # test only 
        best_loss = float("inf")
        tqrange = tqdm.trange(last_epoch, args.epochs)
        # loss = test_epoch(-1, test_dataloader, net, rate_criterion, args.VPT_lmbda, args, 'test')
        for manual_mask_ratio in [0.8, 0.75, 0.7, 0.65, 0.6, 0.55, 0.5, 0.45, 0.4, 0.35, 0.3, 0.25, 0.2, 0.15, 0.1, 0.05]:
            loss = inference_with_kodak(-1, val_dataloader, model, metrics_criterion, cls_criterion, device, manual_mask_ratio, args, 'test')
        return



if __name__ == "__main__":
    main(sys.argv[1:])