|
|
import os |
|
|
import torch |
|
|
import numpy as np |
|
|
import pandas as pd |
|
|
import random |
|
|
import copy |
|
|
import matplotlib.pyplot as plt |
|
|
import itertools |
|
|
import torch.nn.functional as F |
|
|
import tensorflow as tf |
|
|
from tqdm import tqdm |
|
|
from transformers import WEIGHTS_NAME, CONFIG_NAME |
|
|
|
|
|
def set_seed(seed): |
|
|
random.seed(seed) |
|
|
np.random.seed(seed) |
|
|
torch.manual_seed(seed) |
|
|
tf.random.set_seed(seed) |
|
|
os.environ['TF_DETERMINISTIC_OPS'] = '1' |
|
|
os.environ['TF_CUDNN_DETERMINISTIC'] = '1' |
|
|
|
|
|
def save_npy(npy_file, path, file_name): |
|
|
npy_path = os.path.join(path, file_name) |
|
|
np.save(npy_path, npy_file) |
|
|
|
|
|
def load_npy(path, file_name): |
|
|
npy_path = os.path.join(path, file_name) |
|
|
npy_file = np.load(npy_path) |
|
|
return npy_file |
|
|
|
|
|
def save_model(model, model_dir): |
|
|
|
|
|
save_model = model.module if hasattr(model, 'module') else model |
|
|
model_file = os.path.join(model_dir, WEIGHTS_NAME) |
|
|
model_config_file = os.path.join(model_dir, CONFIG_NAME) |
|
|
torch.save(save_model.state_dict(), model_file) |
|
|
|
|
|
if hasattr(save_model, 'config'): |
|
|
with open(model_config_file, "w") as f: |
|
|
f.write(save_model.config.to_json_string()) |
|
|
|
|
|
def restore_model(model, model_dir): |
|
|
output_model_file = os.path.join(model_dir, WEIGHTS_NAME) |
|
|
model.load_state_dict(torch.load(output_model_file)) |
|
|
return model |
|
|
|
|
|
def save_results(args, test_results, debug_args = None): |
|
|
|
|
|
if 'y_pred' in test_results.keys(): |
|
|
pred_labels_path = os.path.join(args.method_output_dir, 'y_pred.npy') |
|
|
|
|
|
del test_results['y_pred'] |
|
|
|
|
|
if 'y_true' in test_results.keys(): |
|
|
true_labels_path = os.path.join(args.method_output_dir, 'y_true.npy') |
|
|
|
|
|
del test_results['y_true'] |
|
|
|
|
|
if not os.path.exists(args.result_dir): |
|
|
os.makedirs(args.result_dir) |
|
|
|
|
|
var = [args.dataset, args.method, args.backbone, args.known_cls_ratio, args.labeled_ratio, args.cluster_num_factor, args.logger_file_name, args.seed] |
|
|
names = ['dataset', 'method', 'backbone', 'known_cls_ratio', 'labeled_ratio', 'cluster_num_factor', 'logger_file_name', 'seed'] |
|
|
|
|
|
if debug_args is not None: |
|
|
var.extend([args[key] for key in debug_args.keys()]) |
|
|
names.extend(debug_args.keys()) |
|
|
|
|
|
vars_dict = {k:v for k,v in zip(names, var) } |
|
|
results = dict(test_results,**vars_dict) |
|
|
keys = list(results.keys()) |
|
|
values = list(results.values()) |
|
|
|
|
|
results_path = os.path.join(args.result_dir, args.results_file_name) |
|
|
|
|
|
if not os.path.exists(results_path) or os.path.getsize(results_path) == 0: |
|
|
ori = [] |
|
|
ori.append(values) |
|
|
df1 = pd.DataFrame(ori,columns = keys) |
|
|
df1.to_csv(results_path,index=False) |
|
|
else: |
|
|
df1 = pd.read_csv(results_path) |
|
|
new = pd.DataFrame(results,index=[1]) |
|
|
df1 = df1.append(new,ignore_index=True) |
|
|
df1.to_csv(results_path,index=False) |
|
|
data_diagram = pd.read_csv(results_path) |
|
|
|
|
|
print('test_results', data_diagram) |
|
|
|
|
|
def class_count(labels): |
|
|
class_data_num = [] |
|
|
for l in np.unique(labels): |
|
|
num = len(labels[labels == l]) |
|
|
class_data_num.append(num) |
|
|
return class_data_num |
|
|
|
|
|
def centroids_cal(model, args, data, train_dataloader, device): |
|
|
|
|
|
model.eval() |
|
|
centroids = torch.zeros(args.num_labels, args.feat_dim).to(device) |
|
|
total_labels = torch.empty(0, dtype=torch.long).to(device) |
|
|
total_features = torch.empty((0,args.feat_dim)).to(device) |
|
|
|
|
|
with torch.set_grad_enabled(False): |
|
|
|
|
|
for batch in tqdm(train_dataloader, desc="Calculate centroids"): |
|
|
|
|
|
batch = tuple(t.to(device) for t in batch) |
|
|
input_ids, input_mask, segment_ids, label_ids = batch |
|
|
|
|
|
features, logits = model(input_ids, segment_ids, input_mask, feature_ext=True) |
|
|
|
|
|
total_labels = torch.cat((total_labels, label_ids)) |
|
|
total_features = torch.cat((total_features, features)) |
|
|
|
|
|
for i in range(len(label_ids)): |
|
|
label = label_ids[i] |
|
|
centroids[label] += features[i] |
|
|
|
|
|
y_true = total_labels.cpu().numpy() |
|
|
|
|
|
centroids /= torch.tensor(class_count(y_true)).float().unsqueeze(1).to(device) |
|
|
|
|
|
return centroids, total_features, total_labels |
|
|
|
|
|
def plot_confusion_matrix(cm, classes, save_name, normalize=False, title='Confusion matrix', figsize=(12, 10), |
|
|
cmap=plt.cm.Blues, save=False): |
|
|
""" |
|
|
This function prints and plots the confusion matrix. |
|
|
Normalization can be applied by setting `normalize=True`. |
|
|
""" |
|
|
if normalize: |
|
|
cm = cm.astype('float') / cm.sum(axis=1)[:, np.newaxis] |
|
|
print("Normalized confusion matrix") |
|
|
else: |
|
|
print('Confusion matrix, without normalization') |
|
|
plt.switch_backend('agg') |
|
|
|
|
|
np.set_printoptions(precision=2) |
|
|
|
|
|
plt.figure(figsize=figsize) |
|
|
plt.imshow(cm, interpolation='nearest', cmap=cmap) |
|
|
plt.title(title) |
|
|
plt.colorbar() |
|
|
tick_marks = np.arange(len(classes)) |
|
|
plt.xticks(tick_marks, classes, rotation=45) |
|
|
plt.yticks(tick_marks, classes) |
|
|
|
|
|
fmt = '.2f' if normalize else 'd' |
|
|
thresh = cm.max() / 1.2 |
|
|
for i, j in itertools.product(range(cm.shape[0]), range(cm.shape[1])): |
|
|
plt.text(j, i, format(cm[i, j], fmt), |
|
|
horizontalalignment="center", |
|
|
color="white" if cm[i, j] > thresh else "black") |
|
|
|
|
|
plt.ylabel('True label') |
|
|
plt.xlabel('Predicted label') |
|
|
plt.tight_layout() |
|
|
if save: |
|
|
plt.savefig(save_name) |
|
|
|
|
|
def mask_tokens(inputs, tokenizer,\ |
|
|
special_tokens_mask=None, mlm_probability=0.15): |
|
|
""" |
|
|
Prepare masked tokens inputs/labels for masked language modeling: 80% MASK, 10% random, 10% original. |
|
|
""" |
|
|
labels = inputs.clone() |
|
|
|
|
|
probability_matrix = torch.full(labels.shape, mlm_probability) |
|
|
if special_tokens_mask is None: |
|
|
special_tokens_mask = [ |
|
|
tokenizer.get_special_tokens_mask(val, already_has_special_tokens=True) for val in labels.tolist() |
|
|
] |
|
|
special_tokens_mask = torch.tensor(special_tokens_mask, dtype=torch.bool) |
|
|
else: |
|
|
special_tokens_mask = special_tokens_mask.bool() |
|
|
|
|
|
probability_matrix.masked_fill_(special_tokens_mask, value=0.0) |
|
|
probability_matrix[torch.where(inputs==0)] = 0.0 |
|
|
masked_indices = torch.bernoulli(probability_matrix).bool() |
|
|
labels[~masked_indices] = -100 |
|
|
|
|
|
indices_replaced = torch.bernoulli(torch.full(labels.shape, 0.8)).bool() & masked_indices |
|
|
inputs[indices_replaced] = tokenizer.convert_tokens_to_ids(tokenizer.mask_token) |
|
|
|
|
|
indices_random = torch.bernoulli(torch.full(labels.shape, 0.5)).bool() & masked_indices & ~indices_replaced |
|
|
random_words = torch.randint(len(tokenizer), labels.shape, dtype=torch.long) |
|
|
inputs[indices_random] = random_words[indices_random] |
|
|
|
|
|
return inputs, labels |
|
|
|
|
|
class MemoryBank(object): |
|
|
def __init__(self, n, dim, num_classes, temperature): |
|
|
self.n = n |
|
|
self.dim = dim |
|
|
self.features = torch.FloatTensor(self.n, self.dim) |
|
|
self.targets = torch.LongTensor(self.n) |
|
|
self.ptr = 0 |
|
|
self.device = 'cpu' |
|
|
self.K = 100 |
|
|
self.temperature = temperature |
|
|
self.C = num_classes |
|
|
|
|
|
def weighted_knn(self, predictions): |
|
|
|
|
|
retrieval_one_hot = torch.zeros(self.K, self.C).to(self.device) |
|
|
batchSize = predictions.shape[0] |
|
|
correlation = torch.matmul(predictions, self.features.t()) |
|
|
yd, yi = correlation.topk(self.K, dim=1, largest=True, sorted=True) |
|
|
candidates = self.targets.view(1,-1).expand(batchSize, -1) |
|
|
retrieval = torch.gather(candidates, 1, yi) |
|
|
retrieval_one_hot.resize_(batchSize * self.K, self.C).zero_() |
|
|
retrieval_one_hot.scatter_(1, retrieval.view(-1, 1), 1) |
|
|
yd_transform = yd.clone().div_(self.temperature).exp_() |
|
|
probs = torch.sum(torch.mul(retrieval_one_hot.view(batchSize, -1 , self.C), |
|
|
yd_transform.view(batchSize, -1, 1)), 1) |
|
|
_, class_preds = probs.sort(1, True) |
|
|
class_pred = class_preds[:, 0] |
|
|
|
|
|
return class_pred |
|
|
|
|
|
def knn(self, predictions): |
|
|
|
|
|
correlation = torch.matmul(predictions, self.features.t()) |
|
|
sample_pred = torch.argmax(correlation, dim=1) |
|
|
class_pred = torch.index_select(self.targets, 0, sample_pred) |
|
|
return class_pred |
|
|
|
|
|
def mine_nearest_neighbors(self, topk, gpu_id, calculate_accuracy=True): |
|
|
|
|
|
import faiss |
|
|
features = self.features.cpu().numpy() |
|
|
n, dim = features.shape[0], features.shape[1] |
|
|
index = faiss.IndexFlatIP(dim) |
|
|
|
|
|
index = faiss.index_cpu_to_all_gpus(index) |
|
|
index.add(features) |
|
|
distances, indices = index.search(features, topk+1) |
|
|
|
|
|
|
|
|
if calculate_accuracy: |
|
|
targets = self.targets.cpu().numpy() |
|
|
neighbor_targets = np.take(targets, indices[:,1:], axis=0) |
|
|
anchor_targets = np.repeat(targets.reshape(-1,1), topk, axis=1) |
|
|
accuracy = np.mean(neighbor_targets == anchor_targets) |
|
|
return indices, accuracy |
|
|
|
|
|
else: |
|
|
return indices |
|
|
|
|
|
def reset(self): |
|
|
self.ptr = 0 |
|
|
|
|
|
def update(self, features, targets): |
|
|
b = features.size(0) |
|
|
|
|
|
assert(b + self.ptr <= self.n) |
|
|
|
|
|
self.features[self.ptr:self.ptr+b].copy_(features.detach()) |
|
|
self.targets[self.ptr:self.ptr+b].copy_(targets.detach()) |
|
|
self.ptr += b |
|
|
|
|
|
def to(self, device): |
|
|
self.features = self.features.to(device) |
|
|
self.targets = self.targets.to(device) |
|
|
self.device = device |
|
|
|
|
|
def cpu(self): |
|
|
self.to('cpu') |
|
|
|
|
|
def cuda(self): |
|
|
self.to('cuda:0') |
|
|
|
|
|
@torch.no_grad() |
|
|
def fill_memory_bank(self, loader, model, memory_bank): |
|
|
model.eval() |
|
|
memory_bank.reset() |
|
|
|
|
|
for i, batch in enumerate(loader): |
|
|
|
|
|
batch = tuple(t.to(self.device) for t in batch) |
|
|
input_ids, input_mask, segment_ids, label_ids = batch |
|
|
X = {"input_ids":input_ids, "attention_mask": input_mask, "token_type_ids": segment_ids} |
|
|
feature = model(X)["hidden_states"] |
|
|
|
|
|
memory_bank.update(feature, label_ids) |
|
|
if i % 100 == 0: |
|
|
print('Fill Memory Bank [%d/%d]' %(i, len(loader))) |
|
|
|
|
|
class view_generator: |
|
|
|
|
|
def __init__(self, tokenizer, args): |
|
|
self.tokenizer = tokenizer |
|
|
self.args = args |
|
|
|
|
|
def random_token_replace(self, ids): |
|
|
mask_id = self.tokenizer.convert_tokens_to_ids(self.tokenizer.mask_token) |
|
|
ids, _ = mask_tokens(ids, self.tokenizer, mlm_probability=self.args.rtr_prob) |
|
|
random_words = torch.randint(len(self.tokenizer), ids.shape, dtype=torch.long) |
|
|
indices_replaced = torch.where(ids == mask_id) |
|
|
ids[indices_replaced] = random_words[indices_replaced] |
|
|
return ids |
|
|
|
|
|
def shuffle_tokens(self, ids): |
|
|
view_pos = [] |
|
|
for inp in torch.unbind(ids): |
|
|
new_ids = copy.deepcopy(inp) |
|
|
special_tokens_mask = self.tokenizer.get_special_tokens_mask(inp, already_has_special_tokens=True) |
|
|
sent_tokens_inds = np.where(np.array(special_tokens_mask) == 0)[0] |
|
|
inds = np.arange(len(sent_tokens_inds)) |
|
|
np.random.shuffle(inds) |
|
|
shuffled_inds = sent_tokens_inds[inds] |
|
|
inp[sent_tokens_inds] = new_ids[shuffled_inds] |
|
|
view_pos.append(new_ids) |
|
|
view_pos = torch.stack(view_pos, dim=0) |
|
|
return view_pos |
|
|
|
|
|
def random_token_erase(self, input_ids, input_mask): |
|
|
|
|
|
aug_input_ids = [] |
|
|
aug_input_mask = [] |
|
|
|
|
|
for inp_i, inp_m in zip(input_ids, input_mask): |
|
|
|
|
|
special_tokens_mask = self.tokenizer.get_special_tokens_mask(inp_i, already_has_special_tokens=True) |
|
|
sent_tokens_inds = np.where(np.array(special_tokens_mask) == 0)[0] |
|
|
inds = np.arange(len(sent_tokens_inds)) |
|
|
masked_inds = np.random.choice(inds, size = int(len(inds) * self.args.re_prob), replace = False) |
|
|
sent_masked_inds = sent_tokens_inds[masked_inds] |
|
|
|
|
|
inp_i = np.delete(inp_i, sent_masked_inds) |
|
|
inp_i = F.pad(inp_i, (0, self.args.max_seq_length - len(inp_i)), 'constant', 0) |
|
|
|
|
|
inp_m = np.delete(inp_m, sent_masked_inds) |
|
|
inp_m = F.pad(inp_m, (0, self.args.max_seq_length - len(inp_m)), 'constant', 0) |
|
|
|
|
|
aug_input_ids.append(inp_i) |
|
|
aug_input_mask.append(inp_m) |
|
|
|
|
|
aug_input_ids = torch.stack(aug_input_ids, dim=0) |
|
|
aug_input_mask = torch.stack(aug_input_mask, dim=0) |
|
|
|
|
|
return aug_input_ids, aug_input_mask |
