train_eval.py

# -*- coding: utf-8 -*-
import numpy as np
import torch
import matplotlib.pyplot as plt
from transformers import BertTokenizer, BertForMaskedLM
from transformers import AdamW, get_linear_schedule_with_warmup
from load_data import traindataloader, valdataloader
from sklearn import metrics
from tqdm import tqdm

device = "cuda" if torch.cuda.is_available() else 'cpu'
N_EPOCHS = 5
LR = 1e-6
WARMUP_PROPORTION = 0.1
MAX_GRAD_NORM = 1.0

def run():
    best_acc_score = 0
    tokenizer = BertTokenizer.from_pretrained('./bert-base-chinese')
    pos_id = tokenizer.convert_tokens_to_ids('很')
    neg_id = tokenizer.convert_tokens_to_ids('不')
    mask_idx = 1
    
    model = BertForMaskedLM.from_pretrained('../bert-base-chinese')
    for param in model.parameters():
        param.requires_grad = True    
    
    model.to(device)
    
    total_steps = len(traindataloader) * N_EPOCHS
    optimizer = AdamW(model.parameters(), lr=LR)
    scheduler = get_linear_schedule_with_warmup(optimizer, num_warmup_steps=int(WARMUP_PROPORTION * total_steps), num_training_steps=total_steps)
    
    loss_vals_train = []
    acc_vals_eval = []
    for epoch in range(N_EPOCHS):
        model.train()
        epoch_loss= []
        pbar = tqdm(traindataloader)
        pbar.set_description("[Train Epoch {}]".format(epoch)) 
    
        for batch_idx, batch_data in enumerate(pbar):
            input_ids = batch_data["input_ids"].to(device)
            attention_mask = batch_data["attention_mask"].to(device)
            labels = batch_data["labels"].to(device)
            
            model.zero_grad()
            outputs = model(input_ids, attention_mask, labels=labels)
            loss = outputs[0]#averaged loss
            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), MAX_GRAD_NORM)
            epoch_loss.append(loss.item())
            optimizer.step()
            scheduler.step()
        loss_vals_train.append(np.mean(epoch_loss))
        
        model.eval()
        pbar = tqdm(valdataloader)
        pbar.set_description("[Eval Epoch {}]".format(epoch))
        
        predict_all = np.array([], dtype=int)
        labels_all = np.array([], dtype=int)
        with torch.no_grad():
            for batch_idx, batch_data in enumerate(pbar):
                input_ids = batch_data["input_ids"].to(device)
                attention_mask = batch_data["attention_mask"].to(device)
                outputs = model(input_ids, attention_mask)
                prediction_scores = outputs[0]
                
                y_pred = prediction_scores[:, mask_idx, [neg_id, pos_id]].argmax(axis=1)
                predict_all = np.append(predict_all, y_pred.cpu().numpy())
                
                y_true = (batch_data["labels"][:, mask_idx] == pos_id).long()
                labels_all = np.append(labels_all, y_true.cpu().numpy()) 
                
        acc = metrics.accuracy_score(labels_all, predict_all)
        acc_vals_eval.append(acc)
        print(f'Epoch:{epoch}, ACC:{acc}')
    
        if acc > best_acc_score:
            best_acc_score = acc
            model_to_save = model.module if hasattr(model, "module") else model
            model_to_save.save_pretrained(f'model.pt')
        torch.cuda.empty_cache()
        
    l1, = plt.plot(np.linspace(1, N_EPOCHS, N_EPOCHS).astype(int), loss_vals_train)
    l2, = plt.plot(np.linspace(1, N_EPOCHS, N_EPOCHS).astype(int), acc_vals_eval)
    plt.legend(handles=[l1,l2],labels=['Train loss','Eval acc'],loc='best')
    
if __name__ == '__main__':
    run()