manifold-smoothing.py

# Modified from https://github.com/Lingkai-Kong/Calibrated-BERT-Fine-Tuning/blob/main/manifold-smoothing.py

import os
# os.environ['CUDA_VISIBLE_DEVICES'] = '1'

from torch import nn
from torch.nn import functional as F
from torch.utils.data import TensorDataset, RandomSampler
from keras.preprocessing.sequence import pad_sequences
from transformers  import BertTokenizer, BertForSequenceClassification
from tqdm import trange
from torch.autograd.gradcheck import zero_gradients
import random
import numpy as np
import torch
from torch.utils.data import DataLoader, SequentialSampler
import argparse
from utils import set_seed, load_dataset, cos_dist

import warnings
# warnings.filterwarnings('ignore')


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=64, 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', 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='y', help='whether save the preprocessed pt file of the dataset')  #'n'

    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 in ['20news', '20news-15', 'trec', 'sst']:
            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.
                                truncation=True
                        )
        # 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,
                                truncation=True
                        )
            val_input_ids.append(encoded_sent)

        for sent in test_sentences:
            encoded_sent = tokenizer.encode(
                                sent,
                                add_special_tokens = True,
                                max_length = MAX_LEN,
                                truncation=True
                        )
            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 == 'sst':
        num_labels = 2
    elif args.dataset == 'trec':
        num_labels = 50
    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)
            dirname = '{}/BERT-manifold-smoothing-{}'.format(args.dataset, args.seed)
            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()