FSL-SSL/legacy/test.py

import torch
import numpy as np
from torch.autograd import Variable
import torch.nn as nn
import torch.optim
import json
import torch.utils.data.sampler
import os
import glob
import random
import time

import config.configs as configs
import models.backbone as backbone
import data.feature_loader as feat_loader
# from data.datamgr import SetDataManager
# from methods.baselinetrain import BaselineTrain
# from methods.baselinefinetune import BaselineFinetune
# from methods.protonet import ProtoNet
# from methods.matchingnet import MatchingNet
# from methods.relationnet import RelationNet
# from methods.maml import MAML
from utils.io_utils import model_dict, parse_args, get_resume_file, get_best_file , get_assigned_file

def feature_evaluation(cl_data_file, model, n_way = 5, n_support = 5, n_query = 15, adaptation = False, semi_inputs=None):
    class_list = cl_data_file.keys()

    select_class = random.sample(class_list,n_way)
    z_all  = []
    for cl in select_class:
        img_feat = cl_data_file[cl]
        perm_ids = np.random.permutation(len(img_feat)).tolist()
        # print(len(img_feat),n_support+n_query)
        z_all.append( [ np.squeeze( img_feat[perm_ids[i]]) for i in range(n_support+n_query) ] )     # stack each batch

    z_all = torch.from_numpy(np.array(z_all) )

    model.n_query = n_query
    if adaptation:
        scores  = model.set_forward_adaptation(z_all, is_feature = True, semi_inputs=semi_inputs)
    else:
        scores  = model.set_forward_test(z_all, is_feature = True, semi_inputs=semi_inputs)
    pred = scores.data.cpu().numpy().argmax(axis = 1)
    y = np.repeat(range( n_way ), n_query )
    acc = np.mean(pred == y)*100
    return acc

def feature_evaluation_depth(cl_data_file, model, n_way = 5, n_support = 5, n_query = 15, adaptation = False):
    class_list = cl_data_file.keys()

    select_class = random.sample(class_list,n_way)
    z_all  = []
    # import ipdb; ipdb.set_trace()
    for cl in select_class:
        img_feat = cl_data_file[cl]
        perm_ids = np.random.permutation(len(img_feat)).tolist()
        # print(len(img_feat),n_support+n_query)
        z_all.append( [ np.squeeze( np.concatenate((img_feat[perm_ids[i]][0],img_feat[perm_ids[i]][1])) ) for i in range(n_support+n_query) ] )     # stack each batch

    z_all = torch.from_numpy(np.array(z_all) )

    model.n_query = n_query
    if adaptation:
        ## not implemented yet
        scores  = model.set_forward_adaptation(z_all, is_feature = True)
    else:
        scores  = model.set_forward_depth(z_all, is_feature = True)
    pred = scores.data.cpu().numpy().argmax(axis = 1)
    y = np.repeat(range( n_way ), n_query )
    acc = np.mean(pred == y)*100
    return acc

if __name__ == '__main__':
    params = parse_args('test')

    isAircraft = (params.dataset == 'aircrafts')

    acc_all = []

    iter_num = 600

    few_shot_params = dict(n_way = params.test_n_way , n_support = params.n_shot)

    if params.dataset in ['omniglot', 'cross_char']:
        assert params.model == 'Conv4' and not params.train_aug ,'omniglot only support Conv4 without augmentation'
        params.model = 'Conv4S'

    if params.method == 'baseline':
        model           = BaselineFinetune( model_dict[params.model], **few_shot_params )
    elif params.method == 'baseline++':
        model           = BaselineFinetune( model_dict[params.model], loss_type = 'dist', **few_shot_params )
    elif params.method == 'protonet':
        model           = ProtoNet( model_dict[params.model], **few_shot_params )
    elif params.method == 'matchingnet':
        model           = MatchingNet( model_dict[params.model], **few_shot_params )
    elif params.method in ['relationnet', 'relationnet_softmax']:
        if params.model == 'Conv4':
            feature_model = backbone.Conv4NP
        elif params.model == 'Conv6':
            feature_model = backbone.Conv6NP
        elif params.model == 'Conv4S':
            feature_model = backbone.Conv4SNP
        else:
            feature_model = lambda: model_dict[params.model]( flatten = False )
        loss_type = 'mse' if params.method == 'relationnet' else 'softmax'
        model           = RelationNet( feature_model, loss_type = loss_type , **few_shot_params )
    elif params.method in ['maml' , 'maml_approx']:
        backbone.ConvBlock.maml = True
        backbone.SimpleBlock.maml = True
        backbone.BottleneckBlock.maml = True
        backbone.ResNet.maml = True
        model = MAML(  model_dict[params.model], approx = (params.method == 'maml_approx') , **few_shot_params )
        if params.dataset in ['omniglot', 'cross_char']: #maml use different parameter in omniglot
            model.n_task     = 32
            model.task_update_num = 1
            model.train_lr = 0.1
    else:
       raise ValueError('Unknown method')

    model = model.cuda()
    model.feature = model.feature.cuda()

    if params.json_seed is not None:
        checkpoint_dir = '%s/checkpoints/%s_%s/%s_%s_%s' %(configs.save_dir, params.dataset, params.json_seed, params.date, params.model, params.method)
    else:
        checkpoint_dir = '%s/checkpoints/%s/%s_%s_%s' %(configs.save_dir, params.dataset, params.date, params.model, params.method)
    if params.train_aug:
        checkpoint_dir += '_aug'
    if not params.method in ['baseline', 'baseline++'] :
        checkpoint_dir += '_%dway_%dshot_%dquery' %( params.train_n_way, params.n_shot, params.n_query)

    checkpoint_dir += '_%d'%params.image_size

    ## Use another dataset (dataloader) for unlabeled data
    if params.dataset_unlabel is not None:
        checkpoint_dir += params.dataset_unlabel
        checkpoint_dir += str(params.bs)

    ## Use grey image
    if params.grey:
        checkpoint_dir += '_grey'

    ## Add jigsaw
    if params.jigsaw:
        checkpoint_dir += '_jigsawonly_alldata_lbda%.2f'%(params.lbda)
        checkpoint_dir += params.optimization
    ## Add rotation
    if params.rotation:
        checkpoint_dir += '_rotation_lbda%.2f'%(params.lbda)
        checkpoint_dir += params.optimization

    checkpoint_dir += '_lr%.4f'%(params.lr)
    if params.finetune:
        checkpoint_dir += '_finetune'

    if params.random:
        checkpoint_dir = 'checkpoints/CUB/random'

    #modelfile   = get_resume_file(checkpoint_dir)

    if params.loadfile != '':
        # modelfile   = params.loadfile
        checkpoint_dir = params.loadfile
    else:
        if not params.method in ['baseline', 'baseline++'] :
            if params.save_iter != -1:
                modelfile   = get_assigned_file(checkpoint_dir,params.save_iter)
            else:
                modelfile   = get_best_file(checkpoint_dir)
            # if modelfile is not None:
            #     tmp = torch.load(modelfile)
            #     model.load_state_dict(tmp['state'], strict=False)

    # if not params.method in ['baseline', 'baseline++'] :
    if params.method in ['maml', 'maml_approx']:
        if modelfile is not None:
            tmp = torch.load(modelfile)
            state = tmp['state']
            state_keys = list(state.keys())
            for i, key in enumerate(state_keys):
                if "feature." in key:
                    newkey = key.replace("feature.","")  # an architecture model has attribute 'feature', load architecture feature to backbone by casting name from 'feature.trunk.xx' to 'trunk.xx'
                    state[newkey] = state.pop(key)
                else:
                    state.pop(key)
            # model.load_state_dict(tmp['state'], strict=False)
            # import ipdb; ipdb.set_trace()
            # model.feature.load_state_dict(tmp['state'])
            model.feature.load_state_dict(tmp['state'], strict=False)##added 0912
        print('modelfile:',modelfile)

    split = params.split
    if params.save_iter != -1:
        split_str = split + "_" +str(params.save_iter)
    else:
        split_str = split
    if params.method in ['maml', 'maml_approx']: #maml do not support testing with feature
        if 'Conv' in params.model:
            if params.dataset in ['omniglot', 'cross_char']:
                image_size = 28
            else:
                image_size = 84
        else:
            image_size = 224

        datamgr         = SetDataManager(image_size, n_eposide = iter_num, n_query = 15 , **few_shot_params, isAircraft=isAircraft, grey=params.grey)

        if params.dataset == 'cross':
            if split == 'base':
                loadfile = configs.data_dir['miniImagenet'] + 'all.json'
            else:
                loadfile   = configs.data_dir['CUB'] + split +'.json'
        elif params.dataset == 'cross_char':
            if split == 'base':
                loadfile = configs.data_dir['omniglot'] + 'noLatin.json'
            else:
                loadfile  = configs.data_dir['emnist'] + split +'.json'
        else:
            if params.json_seed is not None:
                loadfile = configs.data_dir[params.dataset] + split + params.json_seed + '.json'
            else:
                if '_' in params.dataset:
                    loadfile = configs.data_dir[params.dataset.split('_')[0]] + split + '.json'
                else:
                    loadfile = configs.data_dir[params.dataset] + split + '.json'

        novel_loader     = datamgr.get_data_loader( loadfile, aug = False)
        if params.adaptation:
            model.task_update_num = 100 #We perform adaptation on MAML simply by updating more times.
        model.eval()
        acc_mean, acc_std = model.test_loop( novel_loader, return_std = True)

    else:
        novel_file = os.path.join( checkpoint_dir.replace("checkpoints","features"), split_str +".hdf5") #defaut split = novel, but you can also test base or val classes
        print('novel_file',novel_file)
        cl_data_file = feat_loader.init_loader(novel_file)

        for i in range(iter_num):
            # import ipdb; ipdb.set_trace()
            # acc = feature_evaluation(cl_data_file, model, n_query = 15, adaptation = params.adaptation, **few_shot_params)
            acc = feature_evaluation(cl_data_file, model, n_query = params.n_query, adaptation = params.adaptation, **few_shot_params)
            acc_all.append(acc)

        acc_all  = np.asarray(acc_all)
        acc_mean = np.mean(acc_all)
        acc_std  = np.std(acc_all)
        print('%d Test Acc = %4.2f%% +- %4.2f%%' %(iter_num, acc_mean, 1.96* acc_std/np.sqrt(iter_num)))
    #with open('./record/results.txt' , 'a') as f:
    if params.method in ['maml', 'maml_approx']:
    	if not os.path.isdir(checkpoint_dir.replace("checkpoints","features")):
		    os.mkdir(checkpoint_dir.replace("checkpoints","features"))

    with open(os.path.join( checkpoint_dir.replace("checkpoints","features"), split_str +"_test.txt") , 'a') as f:
        timestamp = time.strftime("%Y%m%d-%H%M%S", time.localtime())
        aug_str = '-aug' if params.train_aug else ''
        aug_str += '-adapted' if params.adaptation else ''
        if params.method in ['baseline', 'baseline++'] :
            exp_setting = '%s-%s-%s-%s%s %sshot %sway_test' %(params.dataset, split_str, params.model, params.method, aug_str, params.n_shot, params.test_n_way )
        else:
            exp_setting = '%s-%s-%s-%s%s %sshot %sway_train %sway_test' %(params.dataset, split_str, params.model, params.method, aug_str , params.n_shot , params.train_n_way, params.test_n_way )
        acc_str = '%d Test Acc = %4.2f%% +- %4.2f%%' %(iter_num, acc_mean, 1.96* acc_std/np.sqrt(iter_num))
        f.write( 'Time: %s, Setting: %s, Acc: %s \n' %(timestamp,exp_setting,acc_str)  )