|
|
import logging |
|
|
import copy |
|
|
import os |
|
|
import random |
|
|
import torch |
|
|
import torch.nn.functional as F |
|
|
import numpy as np |
|
|
import math |
|
|
import pandas as pd |
|
|
|
|
|
from .pretrain import PretrainDTCManager |
|
|
from sklearn.cluster import KMeans |
|
|
from sklearn.metrics import confusion_matrix |
|
|
from tqdm import trange, tqdm |
|
|
from utils.metrics import clustering_score, clustering_accuracy_score |
|
|
from utils.functions import save_model, restore_model, set_seed |
|
|
from utils.faster_mix_k_means_pytorch import K_Means |
|
|
from sklearn.metrics import silhouette_score |
|
|
from scipy.optimize import linear_sum_assignment |
|
|
from collections import Counter |
|
|
|
|
|
class DTCManager: |
|
|
|
|
|
def __init__(self, args, data, model, logger_name = 'Discovery'): |
|
|
|
|
|
pretrain_manager = PretrainDTCManager(args, data, model) |
|
|
|
|
|
set_seed(args.seed) |
|
|
self.logger = logging.getLogger(logger_name) |
|
|
|
|
|
loader = data.dataloader |
|
|
self.train_dataloader, self.eval_dataloader, self.test_dataloader = \ |
|
|
loader.train_unlabeled_outputs['loader'], loader.eval_outputs['loader'], loader.test_outputs['loader'] |
|
|
|
|
|
if args.pretrain: |
|
|
|
|
|
self.pretrained_model = pretrain_manager.model |
|
|
self.set_model_optimizer(args, data, model, pretrain_manager) |
|
|
self.load_pretrained_model(self.pretrained_model) |
|
|
else: |
|
|
self.pretrained_model = restore_model(pretrain_manager.model, os.path.join(args.method_output_dir, 'pretrain')) |
|
|
self.set_model_optimizer(args, data, model, pretrain_manager) |
|
|
|
|
|
if args.train: |
|
|
self.load_pretrained_model(self.pretrained_model) |
|
|
else: |
|
|
self.model = restore_model(self.model, args.model_output_dir) |
|
|
|
|
|
def set_model_optimizer(self, args, data, model, pretrain_manager): |
|
|
|
|
|
if args.cluster_num_factor > 1: |
|
|
args.num_labels = self.num_labels = pretrain_manager.num_labels |
|
|
else: |
|
|
args.num_labels = self.num_labels = data.num_labels |
|
|
|
|
|
num_train_examples = len(data.dataloader.train_unlabeled_examples) |
|
|
|
|
|
self.model = model.set_model(args, data, 'bert', args.freeze_bert_parameters) |
|
|
|
|
|
self.warmup_optimizer, self.warmup_scheduler = model.set_optimizer(self.model, num_train_examples, args.train_batch_size, \ |
|
|
args.num_warmup_train_epochs, args.lr, args.warmup_proportion) |
|
|
|
|
|
self.optimizer, self.scheduler = model.set_optimizer(self.model, num_train_examples, args.train_batch_size, \ |
|
|
args.num_train_epochs, args.lr, args.warmup_proportion) |
|
|
|
|
|
self.device = model.device |
|
|
self.model.to(self.device) |
|
|
|
|
|
def initialize_centroids(self, args): |
|
|
|
|
|
self.logger.info("Initialize centroids...") |
|
|
|
|
|
feats = self.get_outputs(args, mode = 'train', get_feats = True) |
|
|
km = KMeans(n_clusters=args.num_labels, n_jobs=-1, random_state=args.seed) |
|
|
km.fit(feats) |
|
|
self.logger.info("Initialization finished...") |
|
|
self.model.cluster_layer.data = torch.tensor(km.cluster_centers_).to(self.device) |
|
|
|
|
|
def warmup_train(self, args): |
|
|
|
|
|
probs = self.get_outputs(args, mode = 'train', get_probs = True) |
|
|
p_target = target_distribution(probs) |
|
|
|
|
|
for epoch in trange(int(args.num_warmup_train_epochs), desc="Warmup_Epoch"): |
|
|
|
|
|
tr_loss, nb_tr_examples, nb_tr_steps = 0, 0, 0 |
|
|
self.model.train() |
|
|
|
|
|
for step, batch in enumerate(tqdm(self.train_dataloader, desc="Warmup_Training")): |
|
|
|
|
|
batch = tuple(t.to(self.device) for t in batch) |
|
|
input_ids, input_mask, segment_ids, label_ids = batch |
|
|
logits, q = self.model(input_ids, segment_ids, input_mask) |
|
|
loss = F.kl_div(q.log(),torch.Tensor(p_target[step * args.train_batch_size: (step+1) * args.train_batch_size]).to(self.device)) |
|
|
|
|
|
loss.backward() |
|
|
tr_loss += loss.item() |
|
|
nb_tr_examples += input_ids.size(0) |
|
|
nb_tr_steps += 1 |
|
|
|
|
|
self.warmup_optimizer.step() |
|
|
self.warmup_scheduler.step() |
|
|
self.warmup_optimizer.zero_grad() |
|
|
|
|
|
eval_true, eval_pred = self.get_outputs(args, mode = 'eval') |
|
|
eval_score = clustering_score(eval_true, eval_pred)['NMI'] |
|
|
eval_results = { |
|
|
'loss': tr_loss, |
|
|
'eval_score': round(eval_score, 2) |
|
|
} |
|
|
self.logger.info("***** Epoch: %s: Eval results *****", str(epoch)) |
|
|
for key in sorted(eval_results.keys()): |
|
|
self.logger.info(" %s = %s", key, str(eval_results[key])) |
|
|
|
|
|
return p_target |
|
|
|
|
|
def get_outputs(self, args, mode = 'eval', get_feats = False, get_probs = False): |
|
|
|
|
|
if mode == 'eval': |
|
|
dataloader = self.eval_dataloader |
|
|
elif mode == 'test': |
|
|
dataloader = self.test_dataloader |
|
|
elif mode == 'train': |
|
|
dataloader = self.train_dataloader |
|
|
|
|
|
self.model.eval() |
|
|
total_labels = torch.empty(0,dtype=torch.long).to(self.device) |
|
|
total_preds = torch.empty(0,dtype=torch.long).to(self.device) |
|
|
|
|
|
total_features = torch.empty((0, args.num_labels)).to(self.device) |
|
|
total_probs = torch.empty((0, args.num_labels)).to(self.device) |
|
|
|
|
|
for batch in tqdm(dataloader, desc="Iteration"): |
|
|
|
|
|
batch = tuple(t.to(self.device) for t in batch) |
|
|
input_ids, input_mask, segment_ids, label_ids = batch |
|
|
|
|
|
with torch.set_grad_enabled(False): |
|
|
logits, probs = self.model(input_ids, segment_ids, input_mask) |
|
|
total_labels = torch.cat((total_labels, label_ids)) |
|
|
total_features = torch.cat((total_features, logits)) |
|
|
total_probs = torch.cat((total_probs, probs)) |
|
|
|
|
|
if get_feats: |
|
|
feats = total_features.cpu().numpy() |
|
|
return feats |
|
|
|
|
|
elif get_probs: |
|
|
return total_probs.cpu().numpy() |
|
|
|
|
|
else: |
|
|
total_preds = total_probs.argmax(1) |
|
|
y_pred = total_preds.cpu().numpy() |
|
|
|
|
|
y_true = total_labels.cpu().numpy() |
|
|
|
|
|
return y_true, y_pred |
|
|
|
|
|
def train(self, args, data): |
|
|
|
|
|
self.initialize_centroids(args) |
|
|
|
|
|
self.logger.info('WarmUp Training start...') |
|
|
self.p_target = self.warmup_train(args) |
|
|
self.logger.info('WarmUp Training finished...') |
|
|
|
|
|
ntrain = len(data.dataloader.train_unlabeled_examples) |
|
|
Z = torch.zeros(ntrain, args.num_labels).float().to(self.device) |
|
|
z_ema = torch.zeros(ntrain, args.num_labels).float().to(self.device) |
|
|
z_epoch = torch.zeros(ntrain, args.num_labels).float().to(self.device) |
|
|
|
|
|
best_model = None |
|
|
best_eval_score = 0 |
|
|
for epoch in trange(int(args.num_train_epochs), desc="Epoch"): |
|
|
|
|
|
|
|
|
tr_loss, nb_tr_examples, nb_tr_steps = 0, 0, 0 |
|
|
self.model.train() |
|
|
|
|
|
for step, batch in enumerate(self.train_dataloader): |
|
|
|
|
|
batch = tuple(t.to(self.device) for t in batch) |
|
|
input_ids, input_mask, segment_ids, label_ids = batch |
|
|
logits, q = self.model(input_ids, segment_ids, input_mask) |
|
|
z_epoch[step * args.train_batch_size: (step+1) * args.train_batch_size, :] = q |
|
|
kl_loss = F.kl_div(q.log(), torch.tensor(self.p_target[step * args.train_batch_size: (step+1) * args.train_batch_size]).to(self.device)) |
|
|
|
|
|
kl_loss.backward() |
|
|
tr_loss += kl_loss.item() |
|
|
nb_tr_examples += input_ids.size(0) |
|
|
nb_tr_steps += 1 |
|
|
|
|
|
self.optimizer.step() |
|
|
self.scheduler.step() |
|
|
self.optimizer.zero_grad() |
|
|
|
|
|
z_epoch = torch.tensor(self.get_outputs(args, mode = 'train', get_probs = True)).to(self.device) |
|
|
Z = args.alpha * Z + (1. - args.alpha) * z_epoch |
|
|
z_ema = Z * (1. / (1. - args.alpha ** (epoch + 1))) |
|
|
|
|
|
if epoch % args.update_interval == 0: |
|
|
self.logger.info('updating target ...') |
|
|
self.p_target = target_distribution(z_ema).float().to(self.device) |
|
|
self.logger.info('updating finished ...') |
|
|
|
|
|
eval_true, eval_pred = self.get_outputs(args, mode = 'eval') |
|
|
eval_score = clustering_score(eval_true, eval_pred)['NMI'] |
|
|
train_loss = tr_loss / nb_tr_steps |
|
|
|
|
|
eval_results = { |
|
|
'train_loss': train_loss, |
|
|
'best_eval_score': best_eval_score, |
|
|
'eval_score': round(eval_score, 2), |
|
|
} |
|
|
self.logger.info("***** Epoch: %s: Eval results *****", str(epoch + 1)) |
|
|
for key in sorted(eval_results.keys()): |
|
|
self.logger.info(" %s = %s", key, str(eval_results[key])) |
|
|
|
|
|
if eval_score > best_eval_score: |
|
|
best_model = copy.deepcopy(self.model) |
|
|
wait = 0 |
|
|
best_eval_score = eval_score |
|
|
|
|
|
elif eval_score > 0: |
|
|
wait += 1 |
|
|
if wait >= args.wait_patient: |
|
|
break |
|
|
|
|
|
self.model = best_model |
|
|
|
|
|
if args.save_model: |
|
|
save_model(self.model, args.model_output_dir) |
|
|
|
|
|
def test(self, args, data): |
|
|
|
|
|
y_true, y_pred = self.get_outputs(args,mode = 'test') |
|
|
test_results = clustering_score(y_true, y_pred) |
|
|
cm = confusion_matrix(y_true,y_pred) |
|
|
|
|
|
self.logger.info |
|
|
self.logger.info("***** Test: Confusion Matrix *****") |
|
|
self.logger.info("%s", str(cm)) |
|
|
self.logger.info("***** Test results *****") |
|
|
|
|
|
for key in sorted(test_results.keys()): |
|
|
self.logger.info(" %s = %s", key, str(test_results[key])) |
|
|
|
|
|
test_results['y_true'] = y_true |
|
|
test_results['y_pred'] = y_pred |
|
|
|
|
|
if args.cluster_num_factor > 1: |
|
|
test_results['estimate_k'] = args.num_labels |
|
|
|
|
|
return test_results |
|
|
|
|
|
def load_pretrained_model(self, pretrained_model): |
|
|
|
|
|
pretrained_dict = pretrained_model.state_dict() |
|
|
classifier_params = ['cluster_layer', 'classifier.weight','classifier.bias'] |
|
|
pretrained_dict = {k: v for k, v in pretrained_dict.items() if k not in classifier_params} |
|
|
self.model.load_state_dict(pretrained_dict, strict=False) |
|
|
|
|
|
|
|
|
def target_distribution(q): |
|
|
weight = q ** 2 / q.sum(0) |
|
|
return (weight.T / weight.sum(1)).T |
|
|
|
