main_OPDA.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
import torch.nn as nn
from models.OPDA.basenet import Prototypical
from models.OPDA.loss import bce_loss
from dataset import OSHeDADataset, NUSDataset, MRTDataset
from utils import get_config, save_model, save_result
from sklearn.metrics import roc_auc_score
import random
import torch.multiprocessing
from pathlib import Path
torch.multiprocessing.set_sharing_strategy('file_system')
warnings.filterwarnings('ignore')


def test(model, configuration, srctar, src_dataset, tgt_l_dataset, tgt_u_dataset):
    model.eval()
    if srctar == 'source':
        N = configuration['ns']
        feature, label = torch.from_numpy(src_dataset.data['ft']).float(), torch.from_numpy(src_dataset.data['lb']).long()
    elif srctar == 'labeled_target':
        N = configuration['nl']
        feature, label = torch.from_numpy(tgt_l_dataset.data['ft']).float(), torch.from_numpy(tgt_l_dataset.data['lb']).long()
    elif srctar == 'unlabeled_target':
        N = configuration['nu']
        feature, label = torch.from_numpy(tgt_u_dataset.data['ft']).float(), torch.from_numpy(tgt_u_dataset.data['lb']).long()
    else:
        raise Exception('Parameter srctar invalid! ')

    with torch.no_grad():
        if torch.cuda.is_available():
            feature, label = feature.cuda(), label.cuda()
        classifier_output, _ = model(input_feature=feature)
        if srctar in ['source', 'labeled_target']:
            _ , pred = torch.max(classifier_output.data, 1)
            acc = 0.
            for i in range(configuration['class_number'] - 1):
                idx = (label == i)
                n_correct = (pred[idx] == label[idx]).sum().item()
                acc += (float(n_correct) / sum(idx) * 100.).item()
            acc /= configuration['class_number'] - 1
            return acc
        else:
            _, pred = torch.max(classifier_output, 1)
            unk_prob = torch.softmax(classifier_output, dim=1)[:, -1]
            sel_idx = torch.argsort(unk_prob)
            idx = int(len(sel_idx) * configuration['lamda'])
            sel_idx = sel_idx[-idx:]
            pred[sel_idx] = configuration['class_number'] - 1
            os = 0.
            for i in range(configuration['class_number'] - 1):
                idx = (label == i)
                n_correct = (pred[idx] == label[idx]).sum().item()
                os += (float(n_correct) / sum(idx) * 100.).item()
            os /= configuration['class_number'] - 1
            idx = (label == configuration['class_number'] - 1)
            # print(torch.softmax(classifier_output, dim=1)[idx][:, -1])
            n_correct = (pred[idx] == label[idx]).sum().item()
            unk = (float(n_correct) / sum(idx) * 100.).item()
            hos = 2 * os * unk / (os + unk)
            unk_label = (label == configuration['class_number'] - 1).int()
            unk_auc = roc_auc_score(unk_label.cpu().numpy(), unk_prob.cpu().numpy())
            return hos, os, unk, unk_auc


def train(model, optimizer, configuration, src_dataset, tgt_l_dataset, tgt_u_dataset):
    best_hos = -float('inf')
    best_val_score = -float('inf')
    acc_src_list = []
    acc_labeled_tar_list = []
    hos_unlabeled_tar_list = []
    criterion = nn.CrossEntropyLoss().to(device)
    # training
    for epoch in range(args.nepoch):
        start_time = time.time()
        model.train()
        optimizer['G'].zero_grad()
        optimizer['C'].zero_grad()
        # prepare data
        source_feature, source_label = torch.from_numpy(src_dataset.data['ft']).float(), torch.from_numpy(src_dataset.data['lb']).long()
        l_target_feature, l_target_label = torch.from_numpy(tgt_l_dataset.data['ft']).float(), torch.from_numpy(tgt_l_dataset.data['lb']).long()
        u_target_feature, u_target_label = torch.from_numpy(tgt_u_dataset.data['ft']).float(), torch.from_numpy(tgt_u_dataset.data['lb']).long()
        if torch.cuda.is_available():
            source_feature, source_label = source_feature.cuda(), source_label.cuda()
            l_target_feature, l_target_label = l_target_feature.cuda(), l_target_label.cuda()
            u_target_feature = u_target_feature.cuda()
        # forward propagation
        source_output, _ = model(input_feature=source_feature)
        l_target_output, _ = model(input_feature=l_target_feature)
        # s_output = torch.cat([source_output, l_target_output])
        # s_label = torch.cat([source_label, l_target_label])
        s_output = l_target_output
        s_label = l_target_label

        # ========================source data loss============================
        # labeled source data
        # CrossEntropy loss
        cls_loss = criterion(s_output, s_label)
        cls_loss.backward()
        if epoch % 1 == 0:
            print('Use source CE loss: %.4f' % (cls_loss.item()))
        # ========================unknown loss============================
        # Calculate unknown loss
        target_funk = torch.FloatTensor(len(u_target_feature), 2).fill_(0.5).to(u_target_feature.device)
        out_t, _ = model(input_feature=u_target_feature, reverse=True)
        out_t = F.softmax(out_t)
        prob1 = torch.sum(out_t[:, :configuration['class_number'] - 1], 1).view(-1, 1)
        prob2 = out_t[:, configuration['class_number'] - 1].contiguous().view(-1, 1)

        prob = torch.cat((prob1, prob2), 1)
        unk_loss = bce_loss(prob, target_funk)
        unk_loss.backward()
        if epoch % 1 == 0:
            print('Use unknown loss: %.4f' % (unk_loss.item()))

        # backward propagation
        optimizer['G'].step()
        optimizer['C'].step()
        optimizer['G'].zero_grad()
        optimizer['C'].zero_grad()

        # Testing Phase
        acc_src = test(model, configuration, 'source', src_dataset, tgt_l_dataset, tgt_u_dataset)
        acc_labeled_tar = test(model, configuration, 'labeled_target', src_dataset, tgt_l_dataset, tgt_u_dataset)
        hos_unlabeled_tar, os_unlabeled_tar, unk_unlabeled_tar, unk_auc = test(model, configuration, 'unlabeled_target', src_dataset, tgt_l_dataset, tgt_u_dataset)
        end_time = time.time()
        print('ACC -> ', end='')
        print('Epoch: [{}/{}], {:.1f}s, Src acc: {:.4f}%, LTar acc: {:.4f}%, UTar HOS: {:.4f}%, OS: {:.4f}%, UNK: {:.4f}%, AUC:{:.4f}'.format(
            epoch, args.nepoch, end_time - start_time, acc_src, acc_labeled_tar, hos_unlabeled_tar, os_unlabeled_tar, unk_unlabeled_tar, unk_auc))

        acc_src_list.append(acc_src)
        acc_labeled_tar_list.append(acc_labeled_tar)
        hos_unlabeled_tar_list.append((hos_unlabeled_tar, os_unlabeled_tar, unk_unlabeled_tar, unk_auc))

        val_score = acc_src + acc_labeled_tar
        if best_val_score < val_score:
            best_val_score = val_score
            best_val_score_idx = epoch
            save_model(model, model_dir, 'val', args.setting, args.seed)

        if best_hos < hos_unlabeled_tar:
            best_hos = hos_unlabeled_tar
            best_hos_idx = epoch
            save_model(model, model_dir, 'test', args.setting, args.seed)

    # end for max_epoch
    print('Best Test HOS at epoch %d: %.4f - %.4f- %.4f - %.4f' % (best_hos_idx, *hos_unlabeled_tar_list[best_hos_idx]))
    save_result(hos_unlabeled_tar_list[best_val_score_idx], hos_unlabeled_tar_list[best_hos_idx], score_dir, args.setting, args.seed)


if __name__ == '__main__':
    parser = argparse.ArgumentParser(
    description='Open Set Domain Adaptation by Back-propagation')
    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.nepoch = 1000
    args.d_common = 256
    args.optimizer = 'Adam'
    args.lr = 0.001
    args.time_string = datetime.datetime.strftime(datetime.datetime.now(), '%Y-%m-%d %H-%M-%S')
    args.layer = 'double'

    # 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/OPDA'
    score_dir = 'output/OPDA'
    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['n_labels'], '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']}
    

    model = Prototypical(configuration['d_source'], configuration['d_target'], args.d_common, configuration['class_number'], args.layer)
    if torch.cuda.is_available():
        model = model.cuda()
    optimizer = {}
    if args.optimizer == 'SGD':
        optimizer['G'] = optim.SGD(list(model.projector_source.parameters()) + list(model.projector_target.parameters()), lr=args.lr)
        optimizer['C'] = optim.SGD(model.classifier.parameters(), lr=args.lr)
    elif args.optimizer == 'mSGD':
        optimizer['G'] = optim.SGD(list(model.projector_source.parameters()) + list(model.projector_target.parameters()), 
                                   lr=args.lr, momentum=0.9, weight_decay=0.001, nesterov=True)
        optimizer['C'] = optim.SGD(model.classifier.parameters(), lr=args.lr, momentum=0.9,
                                weight_decay=0.001, nesterov=True)
    elif args.optimizer == 'Adam':
        optimizer['G'] = optim.Adam(list(model.projector_source.parameters()) + list(model.projector_target.parameters()), 
                                    lr=args.lr, betas=(0.9, 0.99))
        optimizer['C'] = optim.Adam(model.classifier.parameters(), lr=args.lr, betas=(0.9, 0.99))

    train(model, optimizer, configuration, src_dataset, tgt_l_dataset, tgt_u_dataset)