import torch from torch import nn from torch.nn import functional as F from torch.utils.data import TensorDataset, DataLoader, RandomSampler, SequentialSampler from keras.preprocessing.sequence import pad_sequences from sklearn.model_selection import train_test_split from transformers import BertTokenizer, BertConfig from transformers import AdamW, BertForSequenceClassification, get_linear_schedule_with_warmup from tqdm import tqdm, trange import pandas as pd import io import numpy as np import matplotlib.pyplot as plt from torch.autograd.gradcheck import zero_gradients import argparse import random from utils import * import os class softCrossEntropy(nn.Module): def __init__(self, reduce=True): super(softCrossEntropy, self).__init__() self.reduce = reduce return def forward(self, inputs, target): """ :param inputs: predictions :param target: target labels in vector form :return: loss """ log_likelihood = -F.log_softmax(inputs, dim=1) sample_num, class_num = target.shape if self.reduce: loss = torch.sum(torch.mul(log_likelihood, target)) / sample_num else: loss = torch.sum(torch.mul(log_likelihood, target), 1) return loss def one_hot_tensor(y_batch_tensor, num_classes, device): y_tensor = torch.FloatTensor(y_batch_tensor.size(0), num_classes).fill_(0).to(device) y_tensor[np.arange(len(y_batch_tensor)), y_batch_tensor] = 1.0 return y_tensor class on_manifold_samples(object): def __init__(self, epsilon_x=1e-4, epsilon_y=0.1): super(on_manifold_samples, self).__init__() self.epsilon_x = epsilon_x self.epsilon_y = epsilon_y def generate(self, input_ids, input_mask, y, model): model.eval() with torch.no_grad(): if torch.cuda.device_count() > 1: embedding = model.module.get_input_embeddings()(input_ids) else: embedding = model.get_input_embeddings()(input_ids) x = embedding.detach() inv_index = torch.arange(x.size(0) - 1, -1, -1).long() x_tilde = x[inv_index, :].detach() y_tilde = y[inv_index, :] x_init = x.detach() + torch.zeros_like(x).uniform_(-self.epsilon_x, self.epsilon_x) x_init.requires_grad_() zero_gradients(x_init) if x_init.grad is not None: x_init.grad.data.fill_(0) fea_b = model(inputs_embeds=x_init, token_type_ids=None, attention_mask=input_mask)[1][-1] fea_b = torch.mean(fea_b, 1) with torch.no_grad(): fea_t = model(inputs_embeds=x_tilde, token_type_ids=None, attention_mask=input_mask)[1][-1] fea_t = torch.mean(fea_t, 1) Dx = cos_dist(fea_b, fea_t) model.zero_grad() if torch.cuda.device_count() > 1: Dx = Dx.mean() Dx.backward() x_prime = x_init.data - self.epsilon_x * torch.sign(x_init.grad.data) x_prime = torch.min(torch.max(x_prime, embedding - self.epsilon_x), embedding + self.epsilon_x) y_prime = (1 - self.epsilon_y) * y + self.epsilon_y * y_tilde model.train() return x_prime.detach(), y_prime.detach() class off_manifold_samples(object): def __init__(self, eps=0.001, rand_init='n'): super(off_manifold_samples, self).__init__() self.eps = eps self.rand_init = rand_init def generate(self, model, input_ids, input_mask, labels): model.eval() ny = labels with torch.no_grad(): if torch.cuda.device_count() > 1: embedding = model.module.get_input_embeddings()(input_ids) else: embedding = model.get_input_embeddings()(input_ids) input_embedding = embedding.detach() #random init the adv samples if self.rand_init == 'y': input_embedding = input_embedding + torch.zeros_like(input_embedding).uniform_(-self.eps, self.eps) input_embedding.requires_grad = True zero_gradients(input_embedding) if input_embedding.grad is not None: input_embedding.grad.data.fill_(0) cost = model(inputs_embeds=input_embedding, token_type_ids=None, attention_mask=input_mask, labels=ny)[0] if torch.cuda.device_count() > 1: cost = cost.mean() model.zero_grad() cost.backward() off_samples = input_embedding + self.eps*torch.sign(input_embedding.grad.data) off_samples = torch.min(torch.max(off_samples, embedding - self.eps), embedding + self.eps) model.train() return off_samples.detach() 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 # Function to calculate the accuracy of our predictions vs labels 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 main(): parser = argparse.ArgumentParser() parser.add_argument("--lr", default=5e-5, type=float, help="The initial learning rate for Adam.") parser.add_argument("--train_batch_size", default=32, type=int, help="Batch size for training.") parser.add_argument("--eval_batch_size", default=128, type=int, help="Batch size for training.") parser.add_argument("--epochs", default=10, type=int, help="Number of epochs for training.") parser.add_argument("--seed", default=0, type=int, help="Number of epochs for training.") parser.add_argument("--dataset", default='20news-15', type=str, help="dataset") parser.add_argument("--weight_decay", default=0.0, type=float, help="Weight decay if we apply some.") parser.add_argument("--beta_on", default=1., type=float, help="Weight of on manifold reg") parser.add_argument("--beta_off", default=1., type=float, help="Weight of off manifold reg") parser.add_argument("--eps_in", default=1e-4, type=float, help="Perturbation size of on-manifold regularizer") parser.add_argument("--eps_y", default=0.1, type=float, help="Perturbation size of label") parser.add_argument('--eps_out', default=0.001, type=float, help="Perturbation size of out-of-domain adversarial training") parser.add_argument('--saved_dataset', type=str, default='n', help='whether save the preprocessed pt file of the dataset') args = parser.parse_args() print(args) device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') args.device = device set_seed(args) ece_criterion = ECE().to(args.device) soft_ce = softCrossEntropy() on_manifold = on_manifold_samples(epsilon_x=args.eps_in, epsilon_y=args.eps_y) off_manifold = off_manifold_samples(eps=args.eps_out) # load dataset if args.saved_dataset == 'n': train_sentences, val_sentences, test_sentences, train_labels, val_labels, test_labels = load_dataset(args.dataset) tokenizer = BertTokenizer.from_pretrained('bert-base-uncased', do_lower_case=True) train_input_ids = [] val_input_ids = [] test_input_ids = [] if args.dataset == '20news' or args.dataset == '20news-15': MAX_LEN = 150 else: MAX_LEN = 256 for sent in train_sentences: # `encode` will: # (1) Tokenize the sentence. # (2) Prepend the `[CLS]` token to the start. # (3) Append the `[SEP]` token to the end. # (4) Map tokens to their IDs. encoded_sent = tokenizer.encode( sent, # Sentence to encode. add_special_tokens = True, # Add '[CLS]' and '[SEP]' # This function also supports truncation and conversion # to pytorch tensors, but we need to do padding, so we # can't use these features :( . max_length = MAX_LEN, # Truncate all sentences. #return_tensors = 'pt', # Return pytorch tensors. ) # Add the encoded sentence to the list. train_input_ids.append(encoded_sent) for sent in val_sentences: encoded_sent = tokenizer.encode( sent, add_special_tokens = True, max_length = MAX_LEN, ) val_input_ids.append(encoded_sent) for sent in test_sentences: encoded_sent = tokenizer.encode( sent, add_special_tokens = True, max_length = MAX_LEN, ) test_input_ids.append(encoded_sent) # Pad our input tokens train_input_ids = pad_sequences(train_input_ids, maxlen=MAX_LEN, dtype="long", truncating="post", padding="post") 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") # Create attention masks train_attention_masks = [] val_attention_masks = [] test_attention_masks = [] # Create a mask of 1s for each token followed by 0s for padding for seq in train_input_ids: seq_mask = [float(i>0) for i in seq] train_attention_masks.append(seq_mask) 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) # Convert all of our data into torch tensors, the required datatype for our model train_inputs = torch.tensor(train_input_ids) validation_inputs = torch.tensor(val_input_ids) train_labels = torch.tensor(train_labels) validation_labels = torch.tensor(val_labels) train_masks = torch.tensor(train_attention_masks) validation_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) # Create an iterator of our data with torch DataLoader. train_data = TensorDataset(train_inputs, train_masks, train_labels) validation_data = TensorDataset(validation_inputs, validation_masks, validation_labels) prediction_data = TensorDataset(test_inputs, test_masks, test_labels) dataset_dir = 'dataset/{}'.format(args.dataset) if not os.path.exists(dataset_dir): os.makedirs(dataset_dir) torch.save(train_data, dataset_dir+'/train.pt') torch.save(validation_data, dataset_dir+'/val.pt') torch.save(prediction_data, dataset_dir+'/test.pt') else: dataset_dir = 'dataset/{}'.format(args.dataset) train_data = torch.load(dataset_dir+'/train.pt') validation_data = torch.load(dataset_dir+'/val.pt') prediction_data = torch.load(dataset_dir+'/test.pt') train_sampler = RandomSampler(train_data) train_dataloader = DataLoader(train_data, sampler=train_sampler, batch_size=args.train_batch_size) validation_sampler = SequentialSampler(validation_data) validation_dataloader = DataLoader(validation_data, sampler=validation_sampler, batch_size=args.eval_batch_size) prediction_sampler = SequentialSampler(prediction_data) prediction_dataloader = DataLoader(prediction_data, sampler=prediction_sampler, batch_size=args.eval_batch_size) if args.dataset == '20news': num_labels = 20 elif args.dataset == '20news-15': num_labels = 15 elif args.dataset == 'wos-in': num_labels = 100 elif args.dataset == 'wos': num_labels = 134 print(num_labels) model = BertForSequenceClassification.from_pretrained("bert-base-uncased", num_labels= num_labels, output_hidden_states=True) if torch.cuda.device_count() > 1: print("Let's use", torch.cuda.device_count(), "GPUs!") # dim = 0 [30, xxx] -> [10, ...], [10, ...], [10, ...] on 3 GPUs model = nn.DataParallel(model) model.to(args.device) #######train model param_optimizer = list(model.named_parameters()) no_decay = ['bias', 'gamma', 'beta'] optimizer_grouped_parameters = [ {'params': [p for n, p in param_optimizer if not any(nd in n for nd in no_decay)], 'weight_decay_rate': args.weight_decay}, {'params': [p for n, p in param_optimizer if any(nd in n for nd in no_decay)], 'weight_decay_rate': 0.0} ] optimizer = torch.optim.Adam(optimizer_grouped_parameters, lr=args.lr, eps=1e-9) scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, 'min', patience=2, factor=0.1) t_total = len(train_dataloader) * args.epochs # Store our loss and accuracy for plotting best_val = -np.inf # trange is a tqdm wrapper around the normal python range for epoch in trange(args.epochs, desc="Epoch"): # Training # Set our model to training mode (as opposed to evaluation mode) # Tracking variables tr_loss1, tr_loss2 = 0, 0 nb_tr_examples, nb_tr_steps = 0, 0 model.train() # Train the data for one epoch for step, batch in enumerate(train_dataloader): # Add batch to GPU batch = tuple(t.to(args.device) for t in batch) # Unpack the inputs from our dataloader b_input_ids, b_input_mask, b_labels = batch # generate on manifold samples targets_onehot = one_hot_tensor(b_labels, num_labels, args.device) on_manifold_x, on_manifold_y = on_manifold.generate(b_input_ids, b_input_mask, targets_onehot, model) model.train() # train with on manifold samples on_manifold_logits = model(token_type_ids=None, attention_mask=b_input_mask, inputs_embeds=on_manifold_x)[0] loss_on = soft_ce(on_manifold_logits, on_manifold_y) #generate off manifold samples off_manifold_x = off_manifold.generate(model, b_input_ids, b_input_mask, b_labels) model.train() # train with off manifold samples off_manifold_logits = model(token_type_ids=None, attention_mask=b_input_mask, inputs_embeds=off_manifold_x)[0] off_manifold_prob = F.softmax(off_manifold_logits, dim=1) loss_off = -torch.mean(-torch.sum(off_manifold_prob*torch.log(off_manifold_prob), dim=1)) loss_reg = args.beta_on*loss_on + args.beta_off*loss_off if torch.cuda.device_count() > 1: loss_reg = loss_reg.mean() # Clear out the gradients (by default they accumulate) optimizer.zero_grad() loss_reg.backward() loss_ce = model(b_input_ids, token_type_ids=None, attention_mask=b_input_mask, labels=b_labels)[0] if torch.cuda.device_count() > 1: loss_ce = loss_ce.mean() loss_ce.backward() torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) # Update parameters and take a step using the computed gradient optimizer.step() # Update tracking variables tr_loss1 += loss_ce.item() tr_loss2 += loss_reg.item() nb_tr_examples += b_input_ids.size(0) nb_tr_steps += 1 print("Train cross entropy loss: {} | reg loss: {}".format(tr_loss1/nb_tr_steps, tr_loss2/nb_tr_steps)) # Validation # Put model in evaluation mode to evaluate loss on the validation set model.eval() # Tracking variables eval_accurate_nb = 0 nb_eval_examples = 0 # Evaluate data for one epoch for batch in validation_dataloader: # Add batch to GPU batch = tuple(t.to(args.device) for t in batch) # Unpack the inputs from our dataloader b_input_ids, b_input_mask, b_labels = batch # Telling the model not to compute or store gradients, saving memory and speeding up validation with torch.no_grad(): # Forward pass, calculate logit predictions logits = model(b_input_ids, token_type_ids=None, attention_mask=b_input_mask)[0] # Move logits and labels to CPU logits = logits.detach().cpu().numpy() label_ids = b_labels.to('cpu').numpy() tmp_eval_nb = accurate_nb(logits, label_ids) eval_accurate_nb += tmp_eval_nb nb_eval_examples += label_ids.shape[0] eval_accuracy = eval_accurate_nb/nb_eval_examples print("Validation Accuracy: {}".format(eval_accuracy)) scheduler.step(eval_accuracy) if eval_accuracy > best_val: dirname = '{}/BERT-mf-{}-{}-{}-{}'.format(args.dataset, args.seed, args.eps_in, args.eps_y, args.eps_out) output_dir = './model_save/{}'.format(dirname) if not os.path.exists(output_dir): os.makedirs(output_dir) print("Saving model to %s" % output_dir) model_to_save = model.module if hasattr(model, 'module') else model model_to_save.save_pretrained(output_dir) #tokenizer.save_pretrained(output_dir) best_val = eval_accuracy # ##### test model on test data # Put model in evaluation mode model.eval() # Tracking variables predictions , true_labels = [], [] eval_accurate_nb = 0 nb_test_examples = 0 logits_list = [] labels_list = [] # Predict for batch in prediction_dataloader: # Add batch to GPU batch = tuple(t.to(args.device) for t in batch) # Unpack the inputs from our dataloader b_input_ids, b_input_mask, b_labels = batch # Telling the model not to compute or store gradients, saving memory and speeding up prediction with torch.no_grad(): # Forward pass, calculate logit predictions logits = model(b_input_ids, token_type_ids=None, attention_mask=b_input_mask)[0] logits_list.append(logits) labels_list.append(b_labels) # Move logits and labels to CPU logits = logits.detach().cpu().numpy() label_ids = b_labels.to('cpu').numpy() tmp_eval_nb = accurate_nb(logits, label_ids) eval_accurate_nb += tmp_eval_nb nb_test_examples += label_ids.shape[0] # Store predictions and true labels predictions.append(logits) true_labels.append(label_ids) print("Test Accuracy: {}".format(eval_accurate_nb/nb_test_examples)) logits_ece = torch.cat(logits_list) labels_ece = torch.cat(labels_list) ece = ece_criterion(logits_ece, labels_ece).item() print('ECE on test data: {}'.format(ece)) if __name__ == "__main__": main()