test.py

"""
Training a single model (student or teacher)
"""
from __future__ import print_function

import glob
import os

# os.environ['CUDA_VISIBLE_DEVICES'] = '2,3'

import argparse
import time
import importlib
import torch
import torch.nn as nn
import torch.optim as optim
import torch.multiprocessing as mp
import torch.distributed as dist
import torch.nn as nn
import torch.backends.cudnn as cudnn
import tensorboard_logger as tb_logger
import random
import pandas as pd
import numpy as np

from models import model_dict
from dataset.cifar100 import get_cifar100_dataloaders
from dataset.imagenet import get_imagenet_dataloader
from dataset.imagenet_dali import get_dali_data_loader
from helper.loops_RFF import train_vanilla as train, validate_vanilla
from learning.contrast_trainer import BaseTrainer
from helper.util import save_dict_to_json, reduce_tensor, adjust_learning_rate, update_dict_to_json, \
    load_pretrained_weights, load_pretrained_weights_teacher
from dataset.histo_dataset import get_histo_dataloader, histodata_list, get_histo_dataloader_sample, get_histo_testloader
from model_def import load_model, load_model_test

def parse_option():
    parser = argparse.ArgumentParser('argument for training')

    # baisc
    parser.add_argument('--print_freq', type=int, default=50, help='print frequency')
    parser.add_argument('--batch_size', type=int, default=64, help='batch_size')
    parser.add_argument('--num_workers', type=int, default=8, help='num_workers')
    parser.add_argument('--epochs', type=int, default=60, help='number of training epochs')
    parser.add_argument('--gpu_id', type=str, default='0', help='id(s) for CUDA_VISIBLE_DEVICES')
    parser.add_argument('--seed', default=12345, type=int,
                        help='seed for initializing training. choices=[None, 0, 1],')
    parser.add_argument('--trial', type=str, default='1', help='trial id')

    # optimization
    parser.add_argument('--learning_rate', type=float, default=0.05, help='learning rate')
    parser.add_argument('--lr_decay_epochs', type=str, default='30,40,60', help='where to decay lr, can be a list')
    parser.add_argument('--lr_decay_rate', type=float, default=0.1, help='decay rate for learning rate')
    parser.add_argument('--weight_decay', type=float, default=1e-4, help='weight decay')
    parser.add_argument('--momentum', type=float, default=0.9, help='momentum')
    parser.add_argument('--cosine', action='store_true', help='using cosine annealing')

    # dataset
    # dataset and model
    parser.add_argument('--dataset', type=str, default='prostate_hv', help='dataset')
    parser.add_argument('--models_dir', type=str, default='/data1/trinh/code/KD/SimKD_model/KD_teNst_ImageNet_2021_process4/KD_CMO_att', help='dataset')
    parser.add_argument('--model_path', type=str, default='./save/test4/', help='where to save')
    parser.add_argument('--csv_dir', type=str, default='Acc_KD_CMO_att_cmo_kather19_e_.csv', help='dataset')
    parser.add_argument('--model', type=str, default='effiB0',)
    #                         choices=['resnet8', 'resnet14', 'resnet20', 'resnet32', 'resnet44', 'resnet56', 'resnet110',
    #                                  'ResNet18', 'ResNet34',
    #                                  'resnet8x4', 'resnet32x4', 'wrn_16_1', 'wrn_16_2', 'wrn_40_1', 'wrn_40_2',
    #                                  'vgg8', 'vgg11', 'vgg13', 'vgg16', 'vgg19',
    #                                  'MobileNetV2', 'ShuffleV1', 'ShuffleV2', 'ResNet50',
    #                                  'effiB0']
    parser.add_argument('--pretrain', type=str, default='PANDA', help='tma_class, tma_kd, ImageNet')
    parser.add_argument('--pre_strict', action='store_false', help='strict by default')

    # Augment
    parser.add_argument('--aug_train', type=str, default='RA', help='aug_train')
    parser.add_argument('--crop', type=float, default=0.2, help='crop threshold for RandomResizedCrop')
    parser.add_argument('--image_size', type=int, default=512, help='image_size')
    parser.add_argument('--image_resize', action='store_true')
    parser.add_argument('--n_cls', type=int, default=8, help='image_size')
    parser.add_argument('--skip_test', action='store_true', help='strict by default')

    parser.add_argument('--dali', type=str, choices=['cpu', 'gpu'], default=None)

    # multiprocessing
    parser.add_argument('--multiprocessing-distributed', action='store_true',
                        help='Use multi-processing distributed training to launch '
                             'N processes per node, which has N GPUs. This is the '
                             'fastest way to use PyTorch for either single node or '
                             'multi node data parallel training')
    parser.add_argument('--dist-url', default='tcp://127.0.0.1:23451', type=str,
                        help='url used to set up distributed training')
    parser.add_argument('--deterministic', action='store_false', help='Make results reproducible, true by default')
    parser.add_argument('--skip_validation', action='store_false', help='Skip validation of teacher')

    opt = parser.parse_args()

    # set different learning rates for these MobileNet/ShuffleNet models
    # if opt.model_s in ['MobileNetV2', 'MobileNetV2_1_0', 'ShuffleV1', 'ShuffleV2', 'ShuffleV2_1_5']:
    #     opt.learning_rate = 0.01

    # set the path of model and tensorboard
    # opt.model_path = f'./save/test4/'
    opt.tb_path = './save/students/tensorboard_test'

    iterations = opt.lr_decay_epochs.split(',')
    opt.lr_decay_epochs = list([])
    for it in iterations:
        opt.lr_decay_epochs.append(int(it))

    # set the model name
    opt.model_name = f'{opt.csv_dir}_test_{opt.dataset}_aug_{opt.aug_train}'


    opt.save_folder = os.path.join(opt.model_path, opt.model_name)
    if not os.path.isdir(opt.save_folder):
        os.makedirs(opt.save_folder)

    return opt


best_acc = 0
best_f1 = 0
total_time = time.time()


def main():
    opt = parse_option()

    # ASSIGN CUDA_ID
    os.environ['CUDA_VISIBLE_DEVICES'] = opt.gpu_id

    ngpus_per_node = torch.cuda.device_count()
    opt.ngpus_per_node = ngpus_per_node
    if opt.multiprocessing_distributed:
        # Since we have ngpus_per_node processes per node, the total world_size
        # needs to be adjusted accordingly
        world_size = 1
        opt.world_size = ngpus_per_node * world_size
        # Use torch.multiprocessing.spawn to launch distributed processes: the
        # main_worker process function
        mp.spawn(main_worker, nprocs=ngpus_per_node, args=(ngpus_per_node, opt))
    else:
        main_worker(None if ngpus_per_node > 1 else opt.gpu_id, ngpus_per_node, opt)


def main_worker(gpu, ngpus_per_node, opt):

    global best_acc, best_f1, total_time
    opt.gpu = int(gpu)
    opt.gpu_id = int(gpu)
    opt.rank = 0
    opt.dist_backend = 'nccl'

    if opt.gpu is not None:
        print("Use GPU: {} for training".format(opt.gpu))

    trainer = BaseTrainer(opt)
    trainer.init_ddp_environment(gpu, ngpus_per_node)

    # if opt.deterministic:
    if opt.seed is not None:
        print('opt.deterministic',  opt.deterministic)
        random.seed(opt.seed)
        torch.manual_seed(opt.seed)
        cudnn.deterministic = True
        cudnn.benchmark = False
        np.random.seed(opt.seed)

        torch.backends.cudnn.benchmark = False
        torch.backends.cudnn.deterministic = True
        torch.cuda.manual_seed(opt.seed)
        torch.cuda.manual_seed_all(opt.seed)


    print('opt.n_cls:', opt.n_cls)


    criterion = nn.CrossEntropyLoss()


    # dataloader

    test_loader = get_histo_testloader(
        opt=opt,
        batch_size=opt.batch_size, num_workers=opt.num_workers,
        multiprocessing_distributed=opt.multiprocessing_distributed)
    df = pd.read_csv(f'./save/model_stat/{opt.csv_dir}')
    name_list = list(df['name_list'])
    max_idx = list(df['idx'])
    models_dict = dict(zip(name_list, max_idx))
    # routine
    out_stat = {}
    model_tested = 0
    for name in name_list:
        print(name)

        # if 'Class_ce_model_effiB0_kather19_IS_224_BS256_StdPre_PANDA_CPU8_GPU2_seed12345_epoch50_trial_7' != name:
        #     continue
        if 'cmo_crc_tp_folder1' in name:
            continue
        # if 'trial_1' not in name:
        #     continue
        # if 'd1.0' in name:
        #     continue
        model_idx = models_dict[name]
        pretrain_path = glob.glob(f'{opt.models_dir}/*/{name}/model_{model_idx}.pth')
        if len(pretrain_path) < 1:
            # continue
            if 'f1' in  opt.csv_dir:
                print('f1 in ',  opt.csv_dir)
                pretrain_path = glob.glob(f'{opt.models_dir}/*/{name}/net_best_f1.pth')
            else:
                pretrain_path = glob.glob(f'{opt.models_dir}/*/{name}/net_best_acc.pth')

            if len(pretrain_path) < 1:
                print("There is no model at:", f'{opt.models_dir}/*/{name}')
                continue
            pretrain_path = pretrain_path[0]
        else:
            pretrain_path = pretrain_path[0]

        print(pretrain_path)
        model = load_model_test(opt.model, pretrain_path, opt.n_cls, opt.pre_strict, opt.gpu, opt.multiprocessing_distributed)



        if torch.cuda.is_available():
            # For multiprocessing distributed, DistributedDataParallel constructor
            # should always set the single device scope, otherwise,
            # DistributedDataParallel will use all available devices.
            if opt.multiprocessing_distributed:
                if opt.gpu is not None:

                    torch.cuda.set_device(opt.gpu)
                    model.cuda(opt.gpu)
                    DDP = torch.nn.parallel.DistributedDataParallel
                    model = DDP(model, device_ids=[opt.gpu])
                else:
                    print('multiprocessing_distributed must be with a specifiec gpu id')
            else:
                criterion = criterion.cuda()
                if torch.cuda.device_count() > 1:
                    model = nn.DataParallel(model).cuda()
                else:
                    model = model.cuda()

            if not opt.deterministic:
                cudnn.benchmark = True

        print("==> training...")




        print('GPU %d validating' % (opt.gpu))
        test_acc, test_loss, test_output_stat = validate_vanilla(test_loader, model, criterion, opt, prefix='Val')
        # val_acc, val_loss, val_output_stat = validate_distill(val_loader, module_list, criterion_cls, opt, prefix='Val')

        for i in test_output_stat.keys():
            print(i, [test_output_stat[i]])



        test_merics = {
            'name': name,

            'test_cf': pd.Series({'conf_mat': test_output_stat['conf_mat']}).to_json(orient='records'),
            'test_loss': test_loss,
            'test_acc': test_acc
        }


        # update_dict_to_json(model_tested, test_merics, os.path.join(opt.save_folder, "stat.json"))
        model_tested += 1
        out_stat[model_tested] = test_merics
    import json
    with open(os.path.join(opt.save_folder, "stat.json"), 'w') as json_file:
        json.dump(out_stat, json_file, indent=4)

if __name__ == '__main__':
    main()