train_meta.py

'''
change from osr_closed_set_all_you_need/train_zhou.py and CVPR21-Proser/proser_unknown_detection.py
21CVPR_Learning placeholders for open-set recognition
'''

import os
import argparse
import datetime
import time
import pandas as pd
import importlib

import torch
from torch.utils.data import DataLoader
import torch.nn as nn
import torch.backends.cudnn as cudnn

from methods.ARPL.arpl_models import gan
from methods.ARPL.arpl_models.arpl_models import classifier32ABN
# from methods.ARPL.arpl_models.resnetABN import resnet50ABN,resnet18ABN
from methods.ARPL.arpl_models.wrapper_classes import TimmResNetWrapper
from methods.ARPL.core import train, train_cs, test

from utils.utils import seed_torch, str2bool, get_default_hyperparameters, save_list_to_txt,save_networks
from utils.schedulers import get_scheduler
from data.open_set_datasets import get_class_splits, get_datasets
from models.model_utils import get_model
from models.resnetABN32 import resnet50ABN,resnet18ABN

import yaml
from torch.utils.tensorboard import SummaryWriter
import shutil
import sys
from methods.Proser import core as ProserCore
from methods.Ours import core as OursCore
from methods.Ours import gan as Oursgan
from methods.ConfidentCLS import core as ConfidentCore
from methods.ConfidentCLS import gan as ConfidentModels
from methods.Collective import core as CollectiveCore
from utils.data_utils import CategoriesSampler,CategoriesSamplerV1
from copy import deepcopy
from collections import OrderedDict

parser = argparse.ArgumentParser("Training")
parser.add_argument('-default_args_path', type=str, default='')
parser.add_argument('-method', type=str, default='Proser',
                        choices=['Proser','allYouNeed','Ours','ConfidentCLS','Collective'])
parser.add_argument('-ProserStageIsPre', action='store_true',default=False, help="only for Proser, pretrain or train on target dataset")
parser.add_argument('-ProserPreCP', type=str, default=None, help="")
parser.add_argument('-use_val', action='store_true', help="use val dataset during training process", default=False)
parser.add_argument('-img_path', type=str, default=None, help="") #指定的dataset的图像路径
# parser.add_argument('-resume_path', type=str, default=None, help="") 

# Dataset
parser.add_argument('-dataset', type=str, default='cub', help="")
parser.add_argument('-out-num', type=int, default=10, help='For cifar-10-100')
def to_list(arg):
    str_len = len(arg.split(","))
    if str_len==1 or (str_len==2 and arg.split(",")[-1]==''): #64 64,
        return int(arg.split(",")[0]) # 64
    elif str_len ==2 or (str_len==3 and arg.split(",")[-1]==''): #64,64 64,64,
        return tuple([int(i) for i in arg.split(",")]) # (64,64)
parser.add_argument('-image_size', type=to_list, default=64)
# parser.add_argument('-image_size', type=int, default=64)
#------------------for episode learning----------------
parser.add_argument('-data_base_meta', action='store_true',default=False)
parser.add_argument('-train_episode', type=int, default=100) 
parser.add_argument('-episode_way', type=int, default=4) 
parser.add_argument('-episode_shot', type=int, default=5) 
parser.add_argument('-episode_query', type=int, default=15) 
parser.add_argument('-episode_negative_query', type=int, default=15) 
parser.add_argument('-openness_index', type=int, default=0,choices=[i for i in range(5)]) #-----add 230624, only for mobile
parser.add_argument('-train_two_way', action='store_true',default=True)

# optimization
parser.add_argument('-optim', type=str, default=None, help="Which optimizer to use {adam, sgd}")
parser.add_argument('-lr', type=float, default=0.1, help="learning rate for model")
parser.add_argument('-weight_decay', type=float, default=1e-4, help="LR regularisation on weights")
parser.add_argument('-gan_optim', type=str, default='adam',choices=['adam','sgd'])
parser.add_argument('-gan_lr', type=float, default=0.0002, help="learning rate for gan")
parser.add_argument('-gan_weight_decay', type=float, default=5e-4, help="LR regularisation on weights")
parser.add_argument('-batch_size', type=int, default=128)
parser.add_argument('-max-epoch', type=int, default=100)
parser.add_argument('-gan_start_epoch', type=int, default=10)
parser.add_argument('-scheduler', type=str, default='cosine_warm_restarts')
parser.add_argument('-gan_scheduler', type=str, default='cosine_warm_restarts')
parser.add_argument('-steps', type=list, default=[], help="for scheduler == steps")
parser.add_argument('-temp', type=float, default=1.0, help="temp")
parser.add_argument('-num_restarts', type=int, default=2, help='How many restarts for cosine_warm_restarts schedule')
parser.add_argument('-warmup_epochs', type=int, default=10, help='')
parser.add_argument('-num-centers', type=int, default=1)

# model
parser.add_argument('-loss', type=str, default='Softmax')
parser.add_argument('-weight-pl', type=float, default=0.1, help="weight for center loss")
parser.add_argument('-label_smoothing', type=float, default=None, help="Smoothing constant for label smoothing."
                                                                        "No smoothing if None or 0")
parser.add_argument('-beta', type=float, default=0.1, help="weight for entropy loss")
parser.add_argument('-model', type=str, default='classifier32',
                    choices=['classifier32','WRN-28-10','ResNet50','ResNet18','ResNet34','ResNet101'])
parser.add_argument('-ResNet_pretrained', action='store_true',default=True) # for ResNet50 and ResNet18
parser.add_argument('-KUthr_classWise', action='store_true',default=False) # for ResNet50 and ResNet18
parser.add_argument('-pro_type', type=str, default='inputDependent',
                    choices=['inputDependent','negativeGuided'])
parser.add_argument('-only_use_queries', action='store_true',default=False)
parser.add_argument('-genLoss_type', type=str, default='uniformDist',
                    choices=['uniformDist','negDist'])
parser.add_argument('-resnet50_pretrain', type=str, default='places_moco',
                        help='Which pretraining to use if -model=timm_resnet50_pretrained.'
                             'Options are: {iamgenet_moco, places_moco, places}', metavar='BOOL')
parser.add_argument('-feat_dim', type=int, default=128, help="Feature vector dim, only for classifier32 at the moment")
parser.add_argument('-alpha', default=1.0,type=float,help='alpha value for beta distribution')
parser.add_argument('-lamda1', default=1.0,type=float,help='trade-off between loss')
parser.add_argument('-lamda2', default=1.0,type=float,help='trade-off between loss')

# aug
parser.add_argument('-transform', type=str, default='rand-augment')
parser.add_argument('-rand_aug_m', type=int, default=None)
parser.add_argument('-rand_aug_n', type=int, default=None)

# misc
parser.add_argument('-num_workers', default=4, type=int)
parser.add_argument('-split_train_val', default=False, type=str2bool,
                        help='Subsample training set to create validation set', metavar='BOOL')
parser.add_argument('-use_default_parameters', default=False, type=str2bool,
                    help='Set to True to use optimized hyper-parameters from paper', metavar='BOOL')
parser.add_argument('-nz', type=int, default=100)
parser.add_argument('-ns', type=int, default=1)
parser.add_argument('-eval-freq', type=int, default=1)
parser.add_argument('-print-freq', type=int, default=100)
parser.add_argument('-checkpt_freq', type=int, default=20)
parser.add_argument('-gpu', type=str, default='0')
parser.add_argument('-seed', type=int, default=1)
parser.add_argument('-use-cpu', action='store_true')
parser.add_argument('-eval', action='store_true', help="Eval", default=False)
parser.add_argument('-cs', action='store_true', help="Confusing Sample", default=False)
parser.add_argument('-equal_num', action='store_true', default=False)
parser.add_argument('-aux_fake', action='store_true', default=False)
parser.add_argument('-filter_fake', action='store_true', default=False)
parser.add_argument('-gan_input', type=str,default='None',choices=['all','normal','None'])
parser.add_argument('-gan_input_init', action='store_true', default=True)
parser.add_argument('-train_feat_extractor', default=True, type=str2bool,
                        help='Train feature extractor (only implemented for renset_50_faces)', metavar='BOOL')
parser.add_argument('-split_idx', default=0, type=int, help='0-4 OSR splits for each dataset')
parser.add_argument('-use_softmax_in_eval', default=False, type=str2bool,
                        help='Do we use softmax or logits for evaluation', metavar='BOOL')

def init_save_dirs(args):
    '''
    args.log_dir = args.exp_root+''
    args.writer
    args.model_dir = args.log_dir
    '''
    args.log_dir = os.path.join(args.exp_root,args.dataset+'_'+args.method+'_'+args.model+'_'+args.loss+'_'+time.strftime('%Y%m%d_%H%M%S', time.localtime()))
    if not os.path.exists(args.log_dir):
        os.makedirs(args.log_dir)
    print(f'Experiment saved to: {args.log_dir}')
    args.result_list.append(f'Experiment saved to: {args.log_dir} \n')
    args.writer = SummaryWriter(log_dir=args.log_dir)
    
    hparam_dict = {}
    for k, v in vars(args).items():
        if isinstance(v, (int, float, str, bool, torch.Tensor)):
            hparam_dict[k] = v
    args.writer.add_hparams(hparam_dict=hparam_dict, metric_dict={})
    
    args.model_dir = args.log_dir
    # 将参数文件保存下来
    shutil.copy(args.default_args_path, args.log_dir)
    print(args)
    args.result_list.append(args)
    return args
    

def get_optimizer(args, params_list):

    if args.optim is None:

        if options['dataset'] == 'tinyimagenet':
            optimizer = torch.optim.Adam(params_list, lr=args.lr)
        else:
            optimizer = torch.optim.SGD(params_list, lr=args.lr, momentum=0.9, weight_decay=args.weight_decay)

    elif args.optim == 'sgd':

        optimizer = torch.optim.SGD(params_list, lr=args.lr, momentum=0.9, weight_decay=args.weight_decay)

    elif args.optim == 'adam':

        optimizer = torch.optim.Adam(params_list, lr=args.lr)

    else:

        raise NotImplementedError

    return optimizer


def get_mean_lr(optimizer):
    return torch.mean(torch.Tensor([param_group['lr'] for param_group in optimizer.param_groups])).item()


# TODO: Args and options are largely duplicates: tidy up
def main_worker(options, args):

    # torch.manual_seed(options['seed'])
    os.environ['CUDA_VISIBLE_DEVICES'] = options['gpu']
    use_gpu = torch.cuda.is_available()
    if options['use_cpu']: use_gpu = False

    if use_gpu:
        print("Currently using GPU: {}".format(options['gpu']))
        args.result_list.append("Currently using GPU: {}".format(options['gpu']))
        cudnn.benchmark = False
        torch.cuda.manual_seed_all(options['seed'])
    else:
        print("Currently using CPU")
        args.result_list.append("Currently using CPU")

    # -----------------------------
    # DATALOADERS
    # -----------------------------
    if 'train_meta' in dataloaders.keys():
        trainloader_meta = dataloaders['train_meta']
    trainloader = dataloaders['train']
    valloader = dataloaders['val']
    testloader = dataloaders['test_known']
    outloader = dataloaders['test_unknown']

    # -----------------------------
    # MODEL
    # -----------------------------
    netG = None
    print("Creating model: {}".format(options['model']))
    args.result_list.append("Creating model: {}".format(options['model']))
    if args.method == 'Proser':
        # if args.model == 'WideResnet':
        #     checkpoint='/home/zcf/00AnomalyDetection/Proser_results/Modelof_Epoch59.pth'
        net,args.feat_dim = ProserCore.getmodel(backbone=args.model,known_class=len(args.train_classes),
                       checkpoint=args.ProserPreCP,dummynumber=1,ResNet_pretrained=args.ResNet_pretrained,
                       pretrainONtarget=args.ProserStageIsPre) #TODO dummynumber
        if args.ProserStageIsPre:
            args.best_acc = -1
    elif args.method == 'ConfidentCLS':
        net,args.feat_dim = ProserCore.getmodel(backbone=args.model,known_class=len(args.train_classes),
                       checkpoint=args.ProserPreCP,dummynumber=-1,ResNet_pretrained=args.ResNet_pretrained
                       ) #dummynumber=-1
        if args.cs:
            print("Creating GAN")
            args.result_list.append("Creating GAN") 
            netG = ConfidentModels.Generator(1, args.nz, 64, 3) # ngpu, nz, ngf, nc , same as ARPL
            netD = ConfidentModels.Discriminator(1, 3, 64) # ngpu, nc, ndf, ARPL修改了最后一个Conv2d，增加了avgpool+classifier
            criterionD = nn.BCELoss()
    elif args.method == 'allYouNeed':   
        if args.cs and args.loss == 'ARPLoss':
            if args.model == 'classifier32':
                net = classifier32ABN(num_classes=len(args.train_classes), feat_dim=args.feat_dim)
            elif args.model == 'ResNet50': #zhou
                # methods.ARPL.arpl_models下实现的resnet和常规resnet不一样，比较庞大，缩放8倍（而不是标准的32倍）
                net = resnet50ABN(num_classes=len(args.train_classes),pretrained = args.ResNet_pretrained)  #------------zhou, Image reduction by 32*32 times
                args.feat_dim = 2048  #要不然resnet50对应的ARPLoss的Dist会报错
            elif args.model == 'ResNet18': #zhou
                net = resnet18ABN(num_classes=len(args.train_classes),pretrained = args.ResNet_pretrained) #------------zhou, Image reduction by 32*32 times
                args.feat_dim = 512        
            else:
                raise NotImplementedError
    
        else:
            if args.model == 'timm_resnet50_pretrained':
                wrapper_class = TimmResNetWrapper #让模型存在两种输出，1）fc之前的特征 and logits；2）logits
                args.feat_dim = 2048  #要不然resnet50对应的ARPLoss的Dist会报错
            else:
                wrapper_class = None #模型本身存在两种输出，1）fc之前的特征 and logits；2）logits
            net = get_model(args, wrapper_class=wrapper_class)
        # -----------------------------
        # --CS MODEL AND LOSS
        # -----------------------------
        if args.cs and args.loss == 'ARPLoss':
            print("Creating GAN")
            args.result_list.append("Creating GAN")
            nz, ns = options['nz'], 1
            if args.image_size >= 64:
                netG = gan.Generator(1, nz, 64, 3) #nz=Input channel, 3=output channel = fake images channel
                netD = gan.Discriminator(1, 3, 64) #3=Input channel
            else:
                netG = gan.Generator32(1, nz, 64, 3)
                netD = gan.Discriminator32(1, 3, 64)
            # fixed_noise = torch.FloatTensor(64, nz, 1, 1).normal_(0, 1)
            criterionD = nn.BCELoss()
    elif args.method == 'Ours':
        if args.dataset in ['cifar-10-10']:
            args.ResNet_pretrained = False
        net = OursCore.getmodel(backbone=args.model,known_class=len(args.train_classes),
                                episode_way = args.episode_way,checkpoint=None,
                                dummynumber=1,temperature=16,ResNet_pretrained=args.ResNet_pretrained) #TODO
        if args.cs: #------------gan0510
            netG_init = False
    elif args.method == 'Collective':
        net = CollectiveCore.getmodel(backbone=args.model,known_class=len(args.train_classes),
                       checkpoint=args.ProserPreCP,ResNet_pretrained=args.ResNet_pretrained)
            
    feat_dim = args.feat_dim

    # Loss
    options.update(
        {
            'feat_dim': feat_dim,
            'use_gpu':  use_gpu
        }
    )
    args.use_gpu = use_gpu

    # -----------------------------
    # GET LOSS
    # -----------------------------
    if args.method == 'Proser':
        criterion = nn.CrossEntropyLoss()
    elif args.method == 'allYouNeed':
        Loss = importlib.import_module('methods.ARPL.loss.'+options['loss'])
        criterion = getattr(Loss, options['loss'])(**options)
    elif args.method in ['Ours','ConfidentCLS','Collective']:
        criterion = None
    # -----------------------------
    # PREPARE EXPERIMENT
    # -----------------------------
    if use_gpu:
        net = nn.DataParallel(net).cuda()
        if criterion:
            criterion = criterion.cuda()
        if netG is not None:
            netG = nn.DataParallel(netG, device_ids=[i for i in range(len(options['gpu'].split(',')))]).cuda()
            netD = nn.DataParallel(netD, device_ids=[i for i in range(len(options['gpu'].split(',')))]).cuda()
            # fixed_noise.cuda()

    # model_path = os.path.join(args.log_dir, 'arpl_models', options['dataset'])
    # if not os.path.exists(model_path):
    #     os.makedirs(model_path)
    if criterion:
        params_list = [{'params': net.parameters()},
                        {'params': criterion.parameters()}]
    else:
        params_list = [{'params': net.parameters()}]   
    
    # Get base network and criterion
    optimizer = get_optimizer(args=args, params_list=params_list)

    if netG is not None:
        if args.gan_optim == 'sgd':
            optimizerD = torch.optim.SGD(netD.parameters(), lr=args.gan_lr, momentum=0.9, weight_decay=args.gan_weight_decay)
            optimizerG = torch.optim.SGD(netG.parameters(), lr=args.gan_lr, momentum=0.9, weight_decay=args.gan_weight_decay)
        elif args.gan_optim == 'adam':
            optimizerD = torch.optim.Adam(netD.parameters(), lr=args.gan_lr, betas=(0.5, 0.999))
            optimizerG = torch.optim.Adam(netG.parameters(), lr=args.gan_lr, betas=(0.5, 0.999))

    # -----------------------------
    # GET SCHEDULER
    # ----------------------------
    if args.scheduler == 'None':
        scheduler = None
    else:
        scheduler = get_scheduler(optimizer, args.scheduler, args)
        
    if args.gan_scheduler == 'None' or netG is None:
        schedulerD = None
        schedulerG = None
    else:
        schedulerD = get_scheduler(optimizerD, args.gan_scheduler, args) 
        schedulerG = get_scheduler(optimizerG, args.gan_scheduler, args) 

    start_time = time.time()

    # -----------------------------
    # TRAIN
    # -----------------------------
#     if options['max_epoch'] <=0:
#         model_dir = '/data/zcf/00AnomalyDetection/osr_closed_set_out/mobile_Ours_ResNet50_Softmax_20230502_080906/checkpoints/mobile_19_1_Softmax.pth'
#         model_dict = torch.load(model_dir)
#         net.load_state_dict(model_dict)
#         img_feas1 = OursCore.test_feas(net, criterion, valloader, outloader=None, epoch=0, args=args,save_feas=True)
#         img_feas2 = OursCore.test_feas(net, criterion, testloader, outloader, epoch=0, args=args,save_feas=True)
#         torch.save({'protos':model_dict['module.protos_save.weight'].data.cpu(),'img_feas1':img_feas1,'img_feas2':img_feas2},
#                    model_dir.split(model_dir.split('/')[-1])[0]+'img_feas.pth')
#         print('hello')
#         return
    start_epoch = 0
# 算了，从一半开始训练的，跟从头开始训练完全不一样
#     if args.resume_path is not None:
#         all_save = torch.load(args.resume_path)
#         start_epoch = all_save['epoch']
#         net.load_state_dict(all_save['net'])
#         params_list = [{'params': net.parameters()}]   
#         optimizer = get_optimizer(args=args, params_list=params_list)
#         optimizer.load_state_dict(all_save['optimizer']) #optimizer.state_dict()
#         scheduler.load_state_dict(all_save['scheduler'])
#         scheduler.T_cur = start_epoch
#         optimizer.lr = scheduler.get_lr()
    for epoch in range(start_epoch,args.max_epoch):
        print("==> Epoch {}/{}".format(epoch+1, args.max_epoch))
        args.result_list.append("==> Epoch {}/{}".format(epoch+1, args.max_epoch))

        if args.method == 'Ours':
            if args.cs and ((epoch+1) > args.gan_start_epoch):
                if netG_init==False:
                    # 增加网络而不使用，和不增加网络的结果不一样，只有放在这个地方，两者结果一致 
                    print("Creating GAN")
                    args.result_list.append("Creating GAN")
                    netG = Oursgan.Generator((net.module.pos2neg.d_k, net.module.pos2neg.d_v, net.module.pos2neg.d_model))
                    netD = Oursgan.Discriminator(net.module.pos2neg.d_model[0],num_classes=1)
                    netG = nn.DataParallel(netG, device_ids=[i for i in range(len(options['gpu'].split(',')))]).cuda()
                    netD = nn.DataParallel(netD, device_ids=[i for i in range(len(options['gpu'].split(',')))]).cuda()
                    if args.gan_optim == 'sgd':
                        optimizerD = torch.optim.SGD(netD.parameters(), lr=args.gan_lr, momentum=0.9, weight_decay=args.gan_weight_decay)
                        optimizerG = torch.optim.SGD(netG.parameters(), lr=args.gan_lr, momentum=0.9, weight_decay=args.gan_weight_decay)
                    elif args.gan_optim == 'adam':
                        optimizerD = torch.optim.Adam(netD.parameters(), lr=args.gan_lr, betas=(0.5, 0.999))
                        optimizerG = torch.optim.Adam(netG.parameters(), lr=args.gan_lr, betas=(0.5, 0.999))
                    if args.gan_scheduler == 'None' or netG is None:
                        schedulerD = None
                        schedulerG = None
                    else:
                        schedulerD = get_scheduler(optimizerD, args.gan_scheduler, args) 
                        schedulerG = get_scheduler(optimizerG, args.gan_scheduler, args)                     
                    netG.module.gen.load_state_dict(net.module.pos2neg.state_dict())
                    netG_init = True
                    
                OursCore.train_cs(net, netD, netG, criterion, None,
                    optimizer, optimizerD, optimizerG,
                    trainloader_meta, epoch=epoch, args=args)
                # -----------------------------
                # STEP SCHEDULER
                # ----------------------------
                if schedulerD is not None:
                    if args.gan_scheduler == 'plateau' or args.gan_scheduler == 'warm_restarts_plateau':
                        print('not implement')
                    elif args.gan_scheduler in ['multi_step','cosine','OneCycleLR']:
                        schedulerD.step()
                        schedulerG.step()
                    else:
                        schedulerD.step(epoch=epoch-args.gan_start_epoch)
                        schedulerG.step(epoch=epoch-args.gan_start_epoch)
            if (not args.cs) or args.train_two_way or (epoch < args.gan_start_epoch):
                OursCore.train(net, criterion, optimizer, trainloader_meta, epoch=epoch, args=args) 
        elif args.method == 'allYouNeed': 
            if args.cs and args.loss == 'ARPLoss':
                train_cs(net, netD, netG, criterion, criterionD,
                    optimizer, optimizerD, optimizerG,
                    trainloader, epoch=epoch, args=args)
            train(net, criterion, optimizer, trainloader, epoch=epoch, args=args)
        elif args.method == 'Proser':
            if args.ProserStageIsPre:
                ProserCore.preTrain(net, criterion, optimizer, trainloader, epoch=epoch, args=args) 
                ProserCore.preVal(net, criterion, valloader, epoch=epoch, args=args)
                results = dict()
            else:
                ProserCore.traindummy(net, criterion, optimizer, trainloader, epoch=epoch, args=args)
        elif args.method == 'ConfidentCLS':
            ConfidentCore.train_cs(net, netD, netG, criterion, criterionD,optimizer, optimizerD, optimizerG,
                                   trainloader, epoch=epoch, args=args)
        elif args.method == 'Collective':
            CollectiveCore.train(net, criterion, optimizer, trainloader, epoch=epoch, args=args)
        print('\n')
        args.result_list.append('\n')
        if args.ProserStageIsPre:
            print('skip')
        elif options['eval_freq'] > 0 and (epoch+1) % options['eval_freq'] == 0 or (epoch+1) == options['max_epoch']:
            print("==> Test", options['loss'])
            args.result_list.append("==> Test")
            if args.method == 'Ours':
                KnownProtosLoader = deepcopy(trainloader)
                KnownProtosLoader.dataset.transform = valloader.dataset.transform
                pos_neg_protos = OursCore.save_KnownProtos(net,KnownProtosLoader,args=args)
                pos_neg_protos_start = 1
                if args.dataset in ['mobile','PVELAD']:
                    pos_neg_protos_start = 0
                for i in range(pos_neg_protos_start,len(pos_neg_protos)):
                    net.module.protos_save.weight.data[:pos_neg_protos[i].shape[0]] = pos_neg_protos[i]
                    if args.use_val:
                        results_val = OursCore.test(net, criterion, valloader, outloader=None, epoch=epoch, args=args)
                    results = OursCore.test(net, criterion, testloader, outloader, epoch=epoch, args=args)
                    print("---test---Epoch {}: Acc (%): {:.3f}\t AUROC (%): {:.3f}\t AUROC_normalDefect (%): {:.3f}\t".format(
                        epoch+1,results['ACC'],results['AUROC'],results['auroc_normalDefect']),
                        "OSCR (%): {:.3f}\t f1_0 (%): {:.3f}\t f1_best (%): {:.3f}\t".format(
                            results['OSCR'],results['macro_f1_0'],results['macro_f1_best']),
                        "precisionU (%): {:.3f}\t recallU (%): {:.3f}\t f1U (%): {:.3f}\t thr_best (%): {:.3f}\t".format(
                            results['precisionU'],results['recallU'],results['f1U'],results['best_threshold']))
                    print('\n')
                    args.result_list.append("---test---Epoch {}: Acc (%): {:.3f}\t AUROC (%): {:.3f}\t AUROC_normalDefect (%): {:.3f}\t".format(
                        epoch+1,results['ACC'],results['AUROC'],results['auroc_normalDefect'])+
                        "OSCR (%): {:.3f}\t f1_0 (%): {:.3f}\t f1_best (%): {:.3f}\t".format(
                            results['OSCR'],results['macro_f1_0'],results['macro_f1_best'])+
                        "precisionU (%): {:.3f}\t recallU (%): {:.3f}\t f1U (%): {:.3f}\t thr_best (%): {:.3f}\t".format(
                            results['precisionU'],results['recallU'],results['f1U'],results['best_threshold']))
                    args.result_list.append('\n')
                    if (epoch+1) % options['checkpt_freq'] == 0 or epoch == options['max_epoch'] - 1:
                        save_networks(epoch+1,net,optimizer,scheduler,
                                      args.model_dir, file_name.split('.')[0]+'_lr{}_{}_{}'.format(args.lr,epoch,i))
            else:
                if args.use_val:
                    if args.method == 'Proser':
                        results_val = ProserCore.valdummy(net, criterion, valloader, outloader=None, epoch=epoch, args=args)
                    elif args.method == 'allYouNeed':
                        results_val = test(net, criterion, valloader, outloader=None, epoch=epoch, args=args)
                    elif args.method == 'ConfidentCLS':
                        results_val = ConfidentCore.test(net, criterion, valloader, outloader=None, epoch=epoch, args=args)
                    elif args.method == 'Collective':
                        results_val = CollectiveCore.test(net, criterion, valloader, outloader=None, epoch=epoch, args=args)
                    print("---val---Epoch {}: Acc (%): {:.3f}\n".format(epoch+1,results_val['ACC']))
                if args.method == 'Proser':
                    results = ProserCore.valdummy(net, criterion, testloader, outloader, epoch=epoch, args=args)
                elif args.method == 'allYouNeed':
                    results = test(net, criterion, testloader, outloader, epoch=epoch, args=args)
                elif args.method == 'ConfidentCLS':
                    results = ConfidentCore.test(net, criterion, testloader, outloader, epoch=epoch, args=args)
                elif args.method == 'Collective':
                    results = CollectiveCore.test(net, criterion, testloader, outloader, epoch=epoch, args=args)
                print_list1 = "---test---Epoch {}: Acc (%): {:.3f}\t AUROC (%): {:.3f}\t AUROC_normalDefect (%): {:.3f}\t".format(
                    epoch+1,results['ACC'],results['AUROC'],results['auroc_normalDefect'])
                print_list2 = "OSCR (%): {:.3f}\t f1_0 (%): {:.3f}\t f1_best (%): {:.3f}\t".format(
                    results['OSCR'],results['macro_f1_0'],results['macro_f1_best'])
                print_list3 = "precisionU (%): {:.3f}\t recallU (%): {:.3f}\t f1U (%): {:.3f}\t thr_best (%): {:.3f}\t".format(
                     results['precisionU'],results['recallU'],results['f1U'],results['best_threshold'])
                print(print_list1,print_list2,print_list3)
                args.result_list.append(print_list1+print_list2+print_list3)
                print('\n')
                args.result_list.append('\n')
                if (epoch+1) % options['checkpt_freq'] == 0 or epoch == options['max_epoch'] - 1:
                    save_networks(epoch+1,net,optimizer,scheduler,
                                  args.model_dir, file_name.split('.')[0]+'_{}_{}'.format(args.lr,epoch))

            # ----------------
            # LOG
            # ----------------
            if args.use_val:
                args.writer.add_scalar('Val Acc Top 1', results_val['ACC'], epoch)
                if 'allset_ACC' in results_val.keys():
                    args.writer.add_scalar('Val Acc allset', results_val['allset_ACC'], epoch)
                if 'allset_ACC_per_class' in results_val.keys():
                    args.writer.add_scalar('Val Acc allset per_class', results_val['allset_ACC_per_class'], epoch)
            args.writer.add_scalar('Test Acc Top 1', results['ACC'], epoch)
            if 'allset_ACC' in results.keys():
                args.writer.add_scalar('Test Acc allset', results['allset_ACC'], epoch)
            if 'allset_ACC_per_class' in results.keys():
                args.writer.add_scalar('Test Acc allset per_class', results['allset_ACC_per_class'], epoch)
            args.writer.add_scalar('Test AUROC', results['AUROC'], epoch)
            args.writer.add_scalar('Test AUROC_normalDefect', results['auroc_normalDefect'], epoch)
            args.writer.add_scalar('Test OSCR', results['OSCR'], epoch)
            args.writer.add_scalar('Test macro_f1_0', results['macro_f1_0'], epoch)
            args.writer.add_scalar('Test macro_f1_best', results['macro_f1_best'], epoch)
            args.writer.add_scalar('Test precisionU', results['precisionU'], epoch)
            args.writer.add_scalar('Test recallU', results['recallU'], epoch)
            args.writer.add_scalar('Test f1U', results['f1U'], epoch)
            args.writer.add_scalar('Test best_threshold', results['best_threshold'], epoch)

        args.writer.add_scalar('LR', get_mean_lr(optimizer), epoch)

        # -----------------------------
        # STEP SCHEDULER
        # ----------------------------
        if scheduler is not None:
            if args.scheduler == 'plateau' or args.scheduler == 'warm_restarts_plateau':
                if args.use_val:
                    scheduler.step(results_val['ACC'], epoch)
                else:
                    scheduler.step(results['ACC'], epoch)
            elif args.scheduler in ['multi_step','cosine','OneCycleLR']:
                scheduler.step()
            else:
                scheduler.step(epoch=epoch)
                
        save_list_to_txt(os.path.join(args.log_dir, 'results.txt'), args.result_list)
        args.result_list = []
        
    elapsed = round(time.time() - start_time)
    elapsed = str(datetime.timedelta(seconds=elapsed))
    print("Finished. Total elapsed time (h:m:s): {}".format(elapsed))
    save_list_to_txt(os.path.join(args.log_dir, 'results.txt'), ["Finished. Total elapsed time (h:m:s): {}".format(elapsed),
                                                                 f"Experiment saved to: {args.log_dir}"])

    return results

if __name__ == '__main__':

    args = parser.parse_args()
    #----------yaml文件级别高于parser.parse_args
    if args.default_args_path is not None:
        with open(args.default_args_path, 'r') as f:
            default_args = yaml.safe_load(f)
        if 'exp_root' not in default_args:
            default_args['exp_root'] = os.path.dirname(__file__)
        #说明img_path是相对路径，要更换为真实的路径
        if '.' in default_args.get('img_path','img_path'):
            img_path_root = args.default_args_path.split(args.default_args_path.split('/')[-1])[0]
            img_path_root2 = default_args['img_path'].split(default_args['img_path'].split('/')[-1])[0]
            img_path_root3 = os.path.join(img_path_root,img_path_root2)
            sys.path.append(img_path_root3) #img_path_root3为.py文件所在的路径
            img_path_py = default_args['img_path'].split('/')[-1].split('.')[0] #img_path_py为.py文件的名字data_config
            img_path_name = default_args['img_path'].split('/')[-1].split('.')[1] #img_path_name为.py文件中的变量名字
            exec('from '+img_path_py+' import '+img_path_name)
            default_args['img_path'] = eval(img_path_name)
        if default_args.get('img_path',None) is None:
           default_args['img_path'] = None
        parser.set_defaults(**default_args)
        args = parser.parse_args()
    #----------yaml文件级别高于一切，包括parser.parse_args
    args.result_list = []
    # ------------------------
    # Update parameters with default hyperparameters if specified
    # ------------------------
    if args.use_default_parameters:
        print('NOTE: Using default hyper-parameters...')
        args = get_default_hyperparameters(args)

    results = dict()

    for i in range(1):

        # ------------------------
        # INIT
        # ------------------------
        if args.feat_dim is None:
            args.feat_dim = 128 if args.model == 'classifier32' else 2048

        if args.dataset in ['mobile','PVELAD','tianchiAL']:
            args.train_classes, args.open_set_classes = None,None
        else:
            args.train_classes, args.open_set_classes = get_class_splits(args.dataset, args.split_idx,
                                                                     cifar_plus_n=args.out_num)
        args = init_save_dirs(args)

        # ------------------------
        # SEED
        # ------------------------
        seed_torch(args.seed)

        # ------------------------
        # DATASETS
        # ------------------------
        datasets = get_datasets(args.dataset, transform=args.transform, train_classes=args.train_classes,
                                open_set_classes=args.open_set_classes, balance_open_set_eval=True,
                                split_train_val=args.split_train_val, image_size=args.image_size, seed=args.seed,
                                args=args)
        if args.dataset in ['mobile','PVELAD','tianchiAL']:
            args.train_classes, args.open_set_classes = datasets['train'].unique_classes.tolist(),datasets['test_unknown'].unique_classes.tolist()

        # ------------------------
        # RANDAUG HYPERPARAM SWEEP
        # ------------------------
        if args.transform == 'rand-augment':
            if args.rand_aug_m is not None:
                if args.rand_aug_n is not None:
                    datasets['train'].transform.transforms[0].m = args.rand_aug_m
                    datasets['train'].transform.transforms[0].n = args.rand_aug_n

        # ------------------------
        # DATALOADER
        # ------------------------
        dataloaders = {}
        for k, v, in datasets.items():
            shuffle = True if k == 'train' else False
            if k == 'train' and args.data_base_meta:
                if args.episode_negative_query == -1:
                    args.episode_negative_query = args.episode_query
                if args.dataset in ['mobile','PVELAD','tianchiAL','NEU']:
                    sampler = CategoriesSampler(v.targets, args.train_episode, args.episode_way,
                                    args.episode_shot + args.episode_query, args.episode_negative_query) #label = 0,1,2...0,1,2
                if args.dataset in ['NEU']:
                    sampler = CategoriesSampler(v.targets, args.train_episode, args.episode_way,
                                    args.episode_shot + args.episode_query, args.episode_negative_query) #label = 0,1,2...0,1,2
                elif args.dataset in ['cifar-10-10']:
                    sampler = CategoriesSampler([v.target_transform(i) for i in v.targets], args.train_episode, args.episode_way,
                                    args.episode_shot + args.episode_query, args.episode_negative_query) #label = 0,1,2...0,1,2
                dataloaders['train_meta'] = DataLoader(dataset=v, batch_sampler=sampler,
                                            num_workers=args.num_workers,pin_memory=True)
              
            dataloaders[k] = DataLoader(v, batch_size=args.batch_size,shuffle=shuffle, sampler=None, 
                                            num_workers=args.num_workers)

        # ------------------------
        # SAVE PARAMS
        # ------------------------
        options = vars(args)
        options.update(
            {
                'item':     i,
                'known':    args.train_classes,
                'unknown':  args.open_set_classes,
                'img_size': args.image_size,
                'dataloaders': dataloaders,
                'num_classes': len(args.train_classes)
            }
        )

        if options['dataset'] == 'cifar-10-100':
            file_name = '{}_{}.csv'.format(options['dataset'], options['out_num'])
            if args.cs:
                file_name = '{}_{}_cs.csv'.format(options['dataset'], options['out_num'])
        else:
            file_name = options['dataset'] + '.csv'
            if args.cs:
                file_name = options['dataset'] + '_cs' + '.csv'

        print('result path:', os.path.join(args.log_dir, file_name))
        # ------------------------
        # TRAIN
        # ------------------------
        res = main_worker(options, args)

        # ------------------------
        # LOG
        # ------------------------
        print(f'Experiment saved to: {args.log_dir}')
        res['split_idx'] = args.split_idx
        res['unknown'] = args.open_set_classes
        res['known'] = args.train_classes
        res['ID'] = args.log_dir.split("/")[-1]
        results[str(args.split_idx)] = res
        df = pd.DataFrame(results)
        df.to_csv(os.path.join(args.log_dir, file_name), mode='a', header=False)