import torch from torch import nn, optim from torch.nn import functional as F from torch.utils.data import TensorDataset, DataLoader, RandomSampler, SequentialSampler import numpy as np from keras.preprocessing.sequence import pad_sequences from transformers import BertTokenizer from transformers import BertForSequenceClassification import random from sklearn.metrics import f1_score from utils import * import os import argparse import warnings warnings.filterwarnings("ignore") class ModelWithTemperature(nn.Module): """ A thin decorator, which wraps a model with temperature scaling model (nn.Module): A classification neural network NB: Output of the neural network should be the classification logits, NOT the softmax (or log softmax)! """ def __init__(self, model): super(ModelWithTemperature, self).__init__() self.model = model self.temperature = nn.Parameter(torch.ones(1) * 1.5) def forward(self, input_ids, token_type_ids, attention_mask): logits = self.model(input_ids, token_type_ids=token_type_ids, attention_mask=attention_mask)[0] return self.temperature_scale(logits) def temperature_scale(self, logits): """ Perform temperature scaling on logits """ # Expand temperature to match the size of logits temperature = self.temperature.unsqueeze(1).expand(logits.size(0), logits.size(1)) return logits / temperature # This function probably should live outside of this class, but whatever def set_temperature(self, valid_loader, args): """ Tune the tempearature of the model (using the validation set). We're going to set it to optimize NLL. valid_loader (DataLoader): validation set loader """ nll_criterion = nn.CrossEntropyLoss() ece_criterion = ECE().to(args.device) # First: collect all the logits and labels for the validation set logits_list = [] labels_list = [] with torch.no_grad(): for step, batch in enumerate(valid_loader): batch = tuple(t.to(args.device) for t in batch) b_input_ids, b_input_mask, b_labels = batch logits = self.model(b_input_ids, token_type_ids=None, attention_mask=b_input_mask)[0] logits_list.append(logits) labels_list.append(b_labels) logits = torch.cat(logits_list) labels = torch.cat(labels_list) # Calculate NLL and ECE before temperature scaling before_temperature_nll = nll_criterion(logits, labels).item() before_temperature_ece = ece_criterion(logits, labels).item() print('Before temperature - NLL: %.3f, ECE: %.3f' % (before_temperature_nll, before_temperature_ece)) # Next: optimize the temperature w.r.t. NLL optimizer = optim.LBFGS([self.temperature], lr=0.01, max_iter=50) def eval(): loss = nll_criterion(self.temperature_scale(logits), labels) loss.backward() return loss optimizer.step(eval) # Calculate NLL and ECE after temperature scaling after_temperature_nll = nll_criterion(self.temperature_scale(logits), labels).item() after_temperature_ece = ece_criterion(self.temperature_scale(logits), labels).item() print('Optimal temperature: %.3f' % self.temperature.item()) print('After temperature - NLL: %.3f, ECE: %.3f' % (after_temperature_nll, after_temperature_ece)) return self class ECE(nn.Module): def __init__(self, n_bins=15): """ n_bins (int): number of confidence interval bins """ super(ECE, self).__init__() bin_boundaries = torch.linspace(0, 1, n_bins + 1) self.bin_lowers = bin_boundaries[:-1] self.bin_uppers = bin_boundaries[1:] def forward(self, logits, labels): softmaxes = F.softmax(logits, dim=1) confidences, predictions = torch.max(softmaxes, 1) accuracies = predictions.eq(labels) ece = torch.zeros(1, device=logits.device) for bin_lower, bin_upper in zip(self.bin_lowers, self.bin_uppers): # Calculated |confidence - accuracy| in each bin in_bin = confidences.gt(bin_lower.item()) * confidences.le(bin_upper.item()) prop_in_bin = in_bin.float().mean() if prop_in_bin.item() > 0: accuracy_in_bin = accuracies[in_bin].float().mean() avg_confidence_in_bin = confidences[in_bin].mean() ece += torch.abs(avg_confidence_in_bin - accuracy_in_bin) * prop_in_bin return ece class ECE_v2(nn.Module): def __init__(self, n_bins=15): """ n_bins (int): number of confidence interval bins """ super(ECE_v2, self).__init__() bin_boundaries = torch.linspace(0, 1, n_bins + 1) self.bin_lowers = bin_boundaries[:-1] self.bin_uppers = bin_boundaries[1:] def forward(self, softmaxes, labels): confidences, predictions = torch.max(softmaxes, 1) accuracies = predictions.eq(labels) ece = torch.zeros(1, device=softmaxes.device) for bin_lower, bin_upper in zip(self.bin_lowers, self.bin_uppers): # Calculated |confidence - accuracy| in each bin in_bin = confidences.gt(bin_lower.item()) * confidences.le(bin_upper.item()) prop_in_bin = in_bin.float().mean() if prop_in_bin.item() > 0: accuracy_in_bin = accuracies[in_bin].float().mean() avg_confidence_in_bin = confidences[in_bin].mean() ece += torch.abs(avg_confidence_in_bin - accuracy_in_bin) * prop_in_bin return ece def accurate_nb(preds, labels): pred_flat = np.argmax(preds, axis=1).flatten() labels_flat = labels.flatten() return np.sum(pred_flat == labels_flat) def set_seed(args): random.seed(args.seed) np.random.seed(args.seed) torch.manual_seed(args.seed) def apply_dropout(m): if type(m) == nn.Dropout: m.train() def main(): parser = argparse.ArgumentParser(description='Test code - measure the detection peformance') parser.add_argument('--eva_iter', default=1, type=int, help='number of passes for mc-dropout when evaluation') parser.add_argument('--model', type=str, choices=['base', 'manifold-smoothing', 'mc-dropout','temperature'], default='base') parser.add_argument('--seed', type=int, default=0, help='random seed for test') parser.add_argument("--epochs", default=10, type=int, help="Number of epochs for training.") parser.add_argument('--index', type=int, default=0, help='random seed you used during training') parser.add_argument('--in_dataset', required=True, help='target dataset: 20news') parser.add_argument('--out_dataset', required=True, help='out-of-dist dataset') parser.add_argument('--eval_batch_size', type=int, default=32) parser.add_argument('--saved_dataset', type=str, default='n') parser.add_argument('--eps_out', default=0.001, type=float, help="Perturbation size of out-of-domain adversarial training") parser.add_argument("--eps_y", default=0.1, type=float, help="Perturbation size of label") parser.add_argument('--eps_in', default=0.0001, type=float, help="Perturbation size of in-domain adversarial training") args = parser.parse_args() device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') args.device = device set_seed(args) outf = 'test/'+args.model+'-'+str(args.index) if not os.path.isdir(outf): os.makedirs(outf) if args.model == 'base': dirname = '{}/BERT-base-{}'.format(args.in_dataset, args.index) pretrained_dir = './model_save/{}'.format(dirname) # Load a trained model and vocabulary that you have fine-tuned model = BertForSequenceClassification.from_pretrained(pretrained_dir) model.to(args.device) print('Load Tekenizer') elif args.model == 'mc-dropout': dirname = '{}/BERT-base-{}'.format(args.in_dataset, args.index) pretrained_dir = './model_save/{}'.format(dirname) # Load a trained model and vocabulary that you have fine-tuned model = BertForSequenceClassification.from_pretrained(pretrained_dir) model.to(args.device) elif args.model == 'temperature': dirname = '{}/BERT-base-{}'.format(args.in_dataset, args.index) pretrained_dir = './model_save/{}'.format(dirname) orig_model = BertForSequenceClassification.from_pretrained(pretrained_dir) orig_model.to(args.device) model = ModelWithTemperature(orig_model) model.to(args.device) elif args.model == 'manifold-smoothing': dirname = '{}/BERT-mf-{}-{}-{}-{}'.format(args.in_dataset, args.index, args.eps_in, args.eps_y, args.eps_out) print(dirname) pretrained_dir = './model_save/{}'.format(dirname) model = BertForSequenceClassification.from_pretrained(pretrained_dir) model.to(args.device) if args.saved_dataset == 'n': tokenizer = BertTokenizer.from_pretrained('bert-base-uncased', do_lower_case=True) train_sentences, val_sentences, test_sentences, train_labels, val_labels, test_labels = load_dataset(args.in_dataset) _, _, nt_test_sentences, _, _, nt_test_labels = load_dataset(args.out_dataset) val_input_ids = [] test_input_ids = [] nt_test_input_ids = [] if args.in_dataset == '20news' or args.in_dataset == '20news-15': MAX_LEN = 150 else: MAX_LEN = 256 for sent in val_sentences: encoded_sent = tokenizer.encode( sent, # Sentence to encode. add_special_tokens = True, # Add '[CLS]' and '[SEP]' truncation= True, max_length = MAX_LEN, # Truncate all sentences. #return_tensors = 'pt', # Return pytorch tensors. ) # Add the encoded sentence to the list. val_input_ids.append(encoded_sent) for sent in test_sentences: encoded_sent = tokenizer.encode( sent, # Sentence to encode. add_special_tokens = True, # Add '[CLS]' and '[SEP]' truncation= True, max_length = MAX_LEN, # Truncate all sentences. #return_tensors = 'pt', # Return pytorch tensors. ) # Add the encoded sentence to the list. test_input_ids.append(encoded_sent) for sent in nt_test_sentences: encoded_sent = tokenizer.encode( sent, add_special_tokens = True, truncation= True, max_length = MAX_LEN, ) nt_test_input_ids.append(encoded_sent) # Pad our input tokens val_input_ids = pad_sequences(val_input_ids, maxlen=MAX_LEN, dtype="long", truncating="post", padding="post") test_input_ids = pad_sequences(test_input_ids, maxlen=MAX_LEN, dtype="long", truncating="post", padding="post") nt_test_input_ids = pad_sequences(nt_test_input_ids, maxlen=MAX_LEN, dtype="long", truncating="post", padding="post") val_attention_masks = [] test_attention_masks = [] nt_test_attention_masks = [] for seq in val_input_ids: seq_mask = [float(i>0) for i in seq] val_attention_masks.append(seq_mask) for seq in test_input_ids: seq_mask = [float(i>0) for i in seq] test_attention_masks.append(seq_mask) for seq in nt_test_input_ids: seq_mask = [float(i>0) for i in seq] nt_test_attention_masks.append(seq_mask) val_inputs = torch.tensor(val_input_ids) val_labels = torch.tensor(val_labels) val_masks = torch.tensor(val_attention_masks) test_inputs = torch.tensor(test_input_ids) test_labels = torch.tensor(test_labels) test_masks = torch.tensor(test_attention_masks) nt_test_inputs = torch.tensor(nt_test_input_ids) nt_test_labels = torch.tensor(nt_test_labels) nt_test_masks = torch.tensor(nt_test_attention_masks) val_data = TensorDataset(val_inputs, val_masks, val_labels) test_data = TensorDataset(test_inputs, test_masks, test_labels) nt_test_data = TensorDataset(nt_test_inputs, nt_test_masks, nt_test_labels) dataset_dir = 'dataset/test' if not os.path.exists(dataset_dir): os.makedirs(dataset_dir) torch.save(val_data, dataset_dir+'/{}_val_in_domain.pt'.format(args.in_dataset)) torch.save(test_data, dataset_dir+'/{}_test_in_domain.pt'.format(args.in_dataset)) torch.save(nt_test_data, dataset_dir+'/{}_test_out_of_domain.pt'.format(args.out_dataset)) else: dataset_dir = 'dataset/test' val_data = torch.load(dataset_dir+'/{}_val_in_domain.pt'.format(args.in_dataset)) test_data = torch.load(dataset_dir+'/{}_test_in_domain.pt'.format(args.in_dataset)) nt_test_data = torch.load(dataset_dir+'/{}_test_out_of_domain.pt'.format(args.out_dataset)) ######## saved dataset test_sampler = SequentialSampler(test_data) test_dataloader = DataLoader(test_data, sampler=test_sampler, batch_size=args.eval_batch_size) nt_test_sampler = SequentialSampler(nt_test_data) nt_test_dataloader = DataLoader(nt_test_data, sampler=nt_test_sampler, batch_size=args.eval_batch_size) val_sampler = SequentialSampler(val_data) val_dataloader = DataLoader(val_data, sampler=val_sampler, batch_size=args.eval_batch_size) if args.model == 'temperature': model.set_temperature(val_dataloader, args) model.eval() if args.model == 'mc-dropout': model.apply(apply_dropout) correct = 0 total = 0 output_list = [] labels_list = [] ##### validation dat with torch.no_grad(): for step, batch in enumerate(val_dataloader): batch = tuple(t.to(args.device) for t in batch) b_input_ids, b_input_mask, b_labels = batch total += b_labels.shape[0] batch_output = 0 for j in range(args.eva_iter): if args.model == 'temperature': current_batch = model(input_ids=b_input_ids, token_type_ids=None, attention_mask=b_input_mask) #logits else: current_batch = model(input_ids=b_input_ids, token_type_ids=None, attention_mask=b_input_mask)[0] #logits batch_output = batch_output + F.softmax(current_batch, dim=1) batch_output = batch_output/args.eva_iter output_list.append(batch_output) labels_list.append(b_labels) score, predicted = batch_output.max(1) correct += predicted.eq(b_labels).sum().item() ###calculate accuracy and ECE val_eval_accuracy = correct/total print("Val Accuracy: {}".format(val_eval_accuracy)) ece_criterion = ECE_v2().to(args.device) softmaxes_ece = torch.cat(output_list) labels_ece = torch.cat(labels_list) val_ece = ece_criterion(softmaxes_ece, labels_ece).item() print('ECE on Val data: {}'.format(val_ece)) #### Test data correct = 0 total = 0 output_list = [] labels_list = [] predict_list = [] true_list = [] true_list_ood = [] predict_mis = [] predict_in = [] score_list = [] correct_index_all = [] ## test on in-distribution test set with torch.no_grad(): for step, batch in enumerate(test_dataloader): batch = tuple(t.to(args.device) for t in batch) b_input_ids, b_input_mask, b_labels = batch total += b_labels.shape[0] batch_output = 0 for j in range(args.eva_iter): if args.model == 'temperature': current_batch = model(input_ids=b_input_ids, token_type_ids=None, attention_mask=b_input_mask) #logits else: current_batch = model(input_ids=b_input_ids, token_type_ids=None, attention_mask=b_input_mask)[0] #logits batch_output = batch_output + F.softmax(current_batch, dim=1) batch_output = batch_output/args.eva_iter output_list.append(batch_output) labels_list.append(b_labels) score, predicted = batch_output.max(1) correct += predicted.eq(b_labels).sum().item() correct_index = (predicted == b_labels) correct_index_all.append(correct_index) score_list.append(score) ###calcutae accuracy eval_accuracy = correct/total print("Test Accuracy: {}".format(eval_accuracy)) ##calculate ece ece_criterion = ECE_v2().to(args.device) softmaxes_ece = torch.cat(output_list) labels_ece = torch.cat(labels_list) ece = ece_criterion(softmaxes_ece, labels_ece).item() print('ECE on Test data: {}'.format(ece)) #confidence for in-distribution data score_in_array = torch.cat(score_list) #indices of data that are classified correctly correct_array = torch.cat(correct_index_all) label_array = torch.cat(labels_list) ### test on out-of-distribution data predict_ood = [] score_ood_list = [] true_list_ood = [] with torch.no_grad(): for step, batch in enumerate(nt_test_dataloader): batch = tuple(t.to(args.device) for t in batch) b_input_ids, b_input_mask, b_labels = batch batch_output = 0 for j in range(args.eva_iter): if args.model == 'temperature': current_batch = model(b_input_ids, token_type_ids=None, attention_mask=b_input_mask) else: current_batch = model(b_input_ids, token_type_ids=None, attention_mask=b_input_mask)[0] batch_output = batch_output + F.softmax(current_batch, dim=1) batch_output = batch_output/args.eva_iter score_out, _ = batch_output.max(1) score_ood_list.append(score_out) score_ood_array = torch.cat(score_ood_list) label_array = label_array.cpu().numpy() score_ood_array = score_ood_array.cpu().numpy() score_in_array = score_in_array.cpu().numpy() correct_array = correct_array.cpu().numpy() ####### calculate NBAUCC for detection task predict_o = np.zeros(len(score_in_array)+len(score_ood_array)) true_o = np.ones(len(score_in_array)+len(score_ood_array)) true_o[:len(score_in_array)] = 0 ## in-distribution data as false, ood data as positive true_mis = np.ones(len(score_in_array)) true_mis[correct_array] = 0 ##true instances as false, misclassified instances as positive predict_mis = np.zeros(len(score_in_array)) ood_sum = 0 mis_sum = 0 ood_sum_list = [] mis_sum_list = [] #### upper bound of the threshold tau for NBAUCC stop_points = [0.50, 1.] for threshold in np.arange(0., 1.01, 0.02): predict_ood_index1 = (score_in_array < threshold) predict_ood_index2 = (score_ood_array < threshold) predict_ood_index = np.concatenate((predict_ood_index1, predict_ood_index2), axis=0) predict_o[predict_ood_index] = 1 predict_mis[score_in_array