main_KPG.py

import os
import time
import datetime
import argparse
import warnings
import torch
import torch.nn.functional as F
import torch.optim as optim
import torch.utils.data
import numpy as np
from dataset import OSHeDADataset, NUSDataset, MRTDataset
from utils import get_config, save_model, save_result
from models.KPG.utils import cost_matrix, structure_metrix_relation
# from models.KPG.keypointguide_POT.linearprog import lp
from models.KPG.keypointguide_POT.sinkhorn import sinkhorn_log_domain
from sklearn.svm import SVC
import random
import torch.multiprocessing
from pathlib import Path
torch.multiprocessing.set_sharing_strategy('file_system')
warnings.filterwarnings('ignore')


def test(clf, configuration, srctar, src_data, tgt_l_dataset, tgt_u_dataset):
    if srctar == 'source':
        N = configuration['ns']
        feature, label = src_data['ft'], src_data['lb']
    elif srctar == 'labeled_target':
        N = configuration['nl']
        feature, label = tgt_l_dataset.data['ft'], tgt_l_dataset.data['lb']
    elif srctar == 'unlabeled_target':
        N = configuration['nu']
        feature, label = tgt_u_dataset.data['ft'], tgt_u_dataset.data['lb']
    else:
        raise Exception('Parameter srctar invalid! ')

    if srctar in ['source', 'labeled_target']:
        pred = clf.predict(feature)
        acc = 0.
        for i in range(configuration['class_number']):
            idx = (label == i)
            n_correct = (pred[idx] == label[idx]).sum().item()
            acc += (float(n_correct) / sum(idx) * 100.).item()
        acc /= configuration['class_number']
        return acc
    else:
        pass
        # classifier_output = clf.decision_function(feature)
        # logit = np.max(classifier_output, axis=1)
        # pred = np.argmax(classifier_output, axis=1)
        # indices = np.argsort(logit)
        # idx = int(len(indices) * configuration['lamda'])
        # indices = indices[:idx]
        # pred[indices] = configuration['class_number']
        # os = 0.
        # for i in range(configuration['class_number']):
        #     idx = (label == i)
        #     n_correct = (pred[idx] == label[idx]).sum().item()
        #     os += (float(n_correct) / sum(idx) * 100.)
        # os /= configuration['class_number']
        # idx = (label == configuration['class_number'])
        # n_correct = (pred[idx] == label[idx]).sum().item()
        # unk = float(n_correct) / sum(idx) * 100.
        # hos = 2 * os * unk / (os + unk)
        # return hos, os, unk


def train(configuration, src_dataset, tgt_l_dataset, tgt_u_dataset):
    # training
    start_time = time.time()
    # prepare data
    feat_s, label_s = src_dataset.data['ft'], src_dataset.data['lb']
    feat_tl, label_tl = tgt_l_dataset.data['ft'], tgt_l_dataset.data['lb']
    feat_tu, label_tu = tgt_u_dataset.data['ft'], tgt_u_dataset.data['lb']
    
    ####key point
    I = []
    J = []
    t = 0
    feat_sl = []
    for l in label_tl:
        I.append(t)
        J.append(t)
        fl = feat_s[label_s==l]
        feat_sl.append(np.mean(fl,axis=0))
        t += 1
    feat_sl = np.vstack(feat_sl)
    feat_s_ = np.vstack((feat_sl,feat_s))
    feat_t_ = np.vstack((feat_tl,feat_tu))
    Cs = cost_matrix(feat_s_,feat_s_)
    Cs /= Cs.max()
    Ct = cost_matrix(feat_t_,feat_t_)
    Ct /= Ct.max()
    p = np.ones(len(Cs))/len(Cs)
    q = np.ones(len(Ct))/len(Ct)
    C = structure_metrix_relation(Cs,Ct,I,J)
    C = C/C.max()
    ###mask
    M = np.ones_like(C)
    M[I,:] = 0
    M[:,J] = 0
    M[I,J] = 1

    ###key point OT
    print("solving partial kpg-ot...")
    # s = len(p)/len(q)
    # thr = 1e-3
    s = 1 - configuration['lamda']
    p = p*s
    xi = C.max()
    C_ = np.vstack((C, xi * np.ones((1,len(q)))))
    b = np.ones(len(q))
    b[J] = 0
    M_ = np.vstack((M,b.reshape((1,-1))))
    p_ = np.hstack((p,np.sum(q)-s))

    # try:
    #     pi_ = lp(p_,q,C_,M_)
    # except RuntimeError as e:
    #     print(e)
    #     return
    pi_ = sinkhorn_log_domain(p_,q,C_,M_)

    # calculate threshold
    pi_tu = pi_[-1, len(feat_tl):].reshape((-1,))
    num_unknown = int(len(label_tu) * configuration['lamda'])
    thr = np.sort(pi_tu)[-num_unknown-1:-num_unknown+1].mean()

    select_index = np.argwhere(pi_[-1, :] < thr).reshape((-1,))
    pi = pi_[:, select_index]
    pi = pi[:-1, :]
    feat_s_trans = pi@feat_t_[select_index]/p.reshape((-1,1))
    feat_train = np.vstack((feat_tl,feat_s_trans[len(feat_tl):]))
    label_train = np.hstack((label_tl,label_s))
    src_data = {'ft': feat_s_trans[len(feat_tl):], 'lb': label_s}

    print("train svm...")
    clf = SVC(gamma='auto')
    clf.fit(feat_train,label_train)
    print('SVM score: ', clf.score(feat_train, label_train))

    # Testing Phase
    acc_src = test(clf, configuration, 'source', src_data, tgt_l_dataset, tgt_u_dataset)
    acc_labeled_tar = test(clf, configuration, 'labeled_target', src_data, tgt_l_dataset, tgt_u_dataset)
    # hos_unlabeled_tar, os_unlabeled_tar, unk_unlabeled_tar = test(clf, configuration, 'unlabeled_target', src_data, tgt_l_dataset, tgt_u_dataset)
    
    num_class = configuration['class_number']
    out_index = np.argwhere(pi_[-1, len(feat_tl):] > thr).reshape((-1,))
    pred_out = np.ones_like(label_tu)
    pred_out[out_index] = num_class
    label_tu[label_tu>=num_class] =num_class
    acc_unk = np.mean(label_tu[label_tu==num_class]==pred_out[label_tu==num_class])*100
    
    in_index = np.argwhere(pi_[-1, len(feat_tl):] <= thr).reshape((-1,))
    if in_index.sum() != 0:
        pred = clf.predict(feat_tu[in_index])
        pred_out[in_index] = pred
    # acc_kno= np.mean(label_tu[label_tu != num_class] == pred_out[label_tu != num_class])*100
    acc_kno = 0
    for i in range(num_class):
        idx = (label_tu == i)
        n_correct = (pred_out[idx] == label_tu[idx]).sum()
        acc_kno += (float(n_correct) / sum(idx) * 100.)
    acc_kno /= num_class
    h_score = 2 * acc_kno * acc_unk / (acc_kno + acc_unk)
    
    hos_unlabeled_tar, os_unlabeled_tar, unk_unlabeled_tar = h_score, acc_kno, acc_unk
    hos_unlabeled_tar_list = (hos_unlabeled_tar, os_unlabeled_tar, unk_unlabeled_tar)
    end_time = time.time()
    print('ACC -> ', end='')
    print('{:.1f}s, Src acc: {:.4f}%, LTar acc: {:.4f}%, UTar HOS: {:.4f}%, OS: {:.4f}%, UNK: {:.4f}%'.format(
        end_time - start_time, acc_src, acc_labeled_tar, hos_unlabeled_tar, os_unlabeled_tar, unk_unlabeled_tar))

    save_model(clf, model_dir, 'svm', args.setting, args.seed, svc=True)

    save_result(hos_unlabeled_tar_list, hos_unlabeled_tar_list, score_dir, args.setting, args.seed)


if __name__ == '__main__':
    parser = argparse.ArgumentParser(
    description='Keypoint-Guided Optimal Transport')
    parser.add_argument('--seed', type=int, default=0)
    parser.add_argument('--gpu', type=str, default='0')
    parser.add_argument('--setting', type=str)
    args = parser.parse_args()
    args.time_string = datetime.datetime.strftime(datetime.datetime.now(), '%Y-%m-%d %H-%M-%S')

    # parameter initialization
    random.seed(args.seed)
    np.random.seed(args.seed)
    torch.manual_seed(args.seed)
    torch.cuda.manual_seed(args.seed)
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False

    if args.gpu != 'osc':
        os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

    model_dir = 'saved_model/KPG'
    score_dir = 'output/KPG'
    Path(model_dir).mkdir(parents=True, exist_ok=True)
    Path(score_dir).mkdir(parents=True, exist_ok=True)

    config = get_config()[args.setting]

    dataset = args.setting.split('_')[0]

    if dataset in ['BT', 'ImageCLEF', 'OfficeCaltech', 'Wikipedia']:
        src_dataset = OSHeDADataset(config, 'src', args.seed, 'l2')
        tgt_l_dataset = OSHeDADataset(config, 'tgt_l', args.seed, 'l2')
        tgt_u_dataset = OSHeDADataset(config, 'tgt_u', args.seed, 'l2')
    elif dataset == 'MRT':
        src_dataset = MRTDataset(config, 'src', args.seed, 'l2')
        tgt_l_dataset = MRTDataset(config, 'tgt_l', args.seed, 'l2')
        tgt_u_dataset = MRTDataset(config, 'tgt_u', args.seed, 'l2')
    elif dataset == 'NUSIMAGE':
        src_dataset = NUSDataset(config, 'src', args.seed, 'l2')
        tgt_l_dataset = NUSDataset(config, 'tgt_l', args.seed, 'l2')
        tgt_u_dataset = NUSDataset(config, 'tgt_u', args.seed, 'l2')
    
    configuration = {'ns': len(src_dataset.data['lb']), 'nl': len(tgt_l_dataset.data['lb']), 
                    'nu': len(tgt_u_dataset.data['lb']), 'nt': len(tgt_l_dataset.data['lb']) + len(tgt_u_dataset.data['lb']), 
                    'class_number': config['unk_idx'], 'labeled_amount': len(tgt_l_dataset.data['lb']) // config['unk_idx'], 
                    'd_source': src_dataset.data['ft'].shape[1], 'd_target': tgt_l_dataset.data['ft'].shape[1], 'lamda': config['lamda']}

    train(configuration, src_dataset, tgt_l_dataset, tgt_u_dataset)