main.py

from logging import debug
import os
import time
import argparse
import random

from utils.utils import get_logger
from utils.cli_utils import *
from dataset.selectedRotateImageFolder import prepare_test_data

import torch    
import numpy as np

import tent

import models.Res as Resnet
import models.ViT as ViT

import nni


def validate(val_loader, model, criterion, args, mode='eval', logger=None):
    batch_time = AverageMeter('Time', ':6.3f')
    top1 = AverageMeter('Acc@1', ':6.2f')
    top5 = AverageMeter('Acc@5', ':6.2f')
    progress = ProgressMeter(
        len(val_loader),
        [batch_time, top1, top5],
        prefix='Test: ')

    with torch.no_grad():
        end = time.time()
        for i, dl in enumerate(val_loader):
            images, target = dl[0], dl[1]
            if args.gpu is not None:
                images = images.cuda()
            if torch.cuda.is_available():
                target = target.cuda()
            output = model(images)
            # measure accuracy and record loss
            acc1, acc5 = accuracy(output, target, topk=(1, 5))
            if logger:
                logger.info(f"[{i}/{len(val_loader)}] Top-1 Accuracy: {acc1[0]:.5f} and Top-5 Accuracy: {acc5[0]:.5f}")
            top1.update(acc1[0], images.size(0))
            top5.update(acc5[0], images.size(0))

            # measure elapsed time
            batch_time.update(time.time() - end)
            end = time.time()

            if i % 50 == 0:
                progress.display(i)
    return top1.avg, top5.avg


def get_args():

    parser = argparse.ArgumentParser(description='PyTorch ImageNet-C Testing')

    # path of data, output dir
    parser.add_argument('--data', default='/path/to/ImageNet', help='path to dataset')
    parser.add_argument('--data_corruption', default='/path/to/ImageNet-C', help='path to corruption dataset')
    parser.add_argument('--output', default='output', help='the output directory of this experiment')

    # general parameters, dataloader parameters
    parser.add_argument('--seed', default=1, type=int, help='seed for initializing training. ')
    parser.add_argument('--gpu', default=0, type=int, help='GPU id to use.')
    parser.add_argument('--debug', default=False, type=bool, help='debug or not.')
    parser.add_argument('--workers', default=2, type=int, help='number of data loading workers (default: 4)')
    parser.add_argument('--batch_size', default=64, type=int, help='mini-batch size (default: 64)')
    parser.add_argument('--if_shuffle', default=True, type=bool, help='if shuffle the test set.')

    # dataset settings
    parser.add_argument('--level', default=5, type=int, help='corruption level of test(val) set.')
    parser.add_argument('--corruption', default='gaussian_noise', type=str, help='corruption type of test(val) set.')

    # model name, support resnets
    parser.add_argument('--arch', default='resnet50', type=str, help='resnet50 or ViT_B16')

    # overall experimental settings
    parser.add_argument('--exp_type', default='continual', type=str, help='continual or single') 
    parser.add_argument('--lr', type=float, default=0.01, help='learning rate')

    # DEM*
    parser.add_argument('--dem', action='store_true', default=False, help='use dem or not')
    parser.add_argument('--tau', type=float, default=1.0)
    parser.add_argument('--alpha', type=float, default=1.0)
    parser.add_argument('--nni', action='store_true', default=False, help='use nni or not')

    # AdaDEM
    parser.add_argument('--adadem', action='store_true', default=False, help='use adadem or not')
    parser.add_argument('--pi', type=float, default=0.1, help='momentum for MEC')

    return parser.parse_args()


if __name__ == '__main__':

    args = get_args()

    if args.nni:
        nxt_params = nni.get_next_parameter()
        args = nni.utils.merge_parameter(args, nxt_params)

    # set random seeds
    if args.seed is not None:
        random.seed(args.seed)
        np.random.seed(args.seed)
        torch.manual_seed(args.seed)

    if args.arch in ['ViT_B16']:
        subnet = ViT.__dict__[args.arch](pretrained=True)
    else:
        subnet = Resnet.__dict__[args.arch](pretrained=True)

    subnet = subnet.cuda()

    if not os.path.exists(args.output):
        os.makedirs(args.output, exist_ok=True)

    if not args.nni:
        if args.dem:
            log_name = time.strftime("%Y-%m-%d-%H-%M-%S", time.localtime()) + f"-{args.exp_type}-{args.arch}-seed{args.seed}-dem-tau{args.tau}-alpha{args.alpha}-lr{args.lr}-log.txt"
        elif args.adadem:
            log_name = time.strftime("%Y-%m-%d-%H-%M-%S", time.localtime()) + f"-{args.exp_type}-{args.arch}-seed{args.seed}-adadem-pi{args.pi}-lr{args.lr}-log.txt"
        else:
            log_name = time.strftime("%Y-%m-%d-%H-%M-%S", time.localtime()) + f"-{args.exp_type}-{args.arch}-seed{args.seed}-em-lr{args.lr}-log.txt"
        logger = get_logger(name="project", output_directory=args.output, log_name=log_name, debug=False)
    else:
        trial_id = nni.get_trial_id()
        logger = get_logger(name="project", output_directory=args.output, log_name=f"trial_{trial_id}.txt", debug=False)

    if args.nni:
        logger.info(f"Experiment ID: {nni.get_experiment_id()}")
        logger.info(f"Trial ID: {nni.get_trial_id()}")
        logger.info(f"NNI Params: {nxt_params}")
    
    common_corruptions = ['gaussian_noise', 'shot_noise', 'impulse_noise', 'defocus_blur', 'glass_blur', 'motion_blur', 'zoom_blur', 'snow', 'frost', 'fog', 'brightness', 'contrast', 'elastic_transform', 'pixelate', 'jpeg_compression']
    
    logger.info(args)
    logger.info(common_corruptions)

    subnet = tent.configure_model(subnet)
    params, param_names = tent.collect_params(subnet)
    optimizer = torch.optim.SGD(params, args.lr)
    adapt_model = tent.Tent(subnet, optimizer, args=args)

    accs = []
    for corrupt in common_corruptions:
        if args.exp_type == 'single':
            adapt_model.reset()
        elif args.exp_type == 'continual':
            print("continue")
        else:
            assert False, NotImplementedError

        args.corruption = corrupt
        logger.info(args.corruption)

        val_dataset, val_loader = prepare_test_data(args)
        val_dataset.switch_mode(True, False)

        top1, top5 = validate(val_loader, adapt_model, None, args, mode='eval', logger=logger)
        logger.info(f"Under shift type {args.corruption} Top-1 Accuracy: {top1:.5f} and Top-5 Accuracy: {top5:.5f}")
        accs.append(float(top1))

        if args.nni:
            nni.report_intermediate_result(float(top1))

    accs = np.array(accs)
    logger.info(f"Average accuracy is {accs.mean():.5f}")

    if args.nni:
        nni.report_final_result(float(accs.mean()))