import argparse import os import torch import torch.nn as nn from torch.optim import Adam from torch.utils.data import DataLoader import pickle print("here1",os.getcwd()) from src.dataset import TokenizerDataset, TokenizerDatasetForCalibration from src.vocab import Vocab print("here3",os.getcwd()) from src.bert import BERT from src.seq_model import BERTSM from src.classifier_model import BERTForClassification, BERTForClassificationWithFeats # from src.new_finetuning.optim_schedule import ScheduledOptim import metrics, recalibration, visualization from recalibration import ModelWithTemperature import tqdm import sys import time import numpy as np from sklearn.metrics import precision_score, recall_score, f1_score, confusion_matrix, roc_curve, roc_auc_score import matplotlib.pyplot as plt import seaborn as sns import pandas as pd from collections import defaultdict print("here3",os.getcwd()) class BERTFineTuneTrainer: def __init__(self, bertFinetunedClassifierwithFeats: BERT, #BERTForClassificationWithFeats vocab_size: int, test_dataloader: DataLoader = None, lr: float = 1e-4, betas=(0.9, 0.999), weight_decay: float = 0.01, warmup_steps=10000, with_cuda: bool = True, cuda_devices=None, log_freq: int = 10, workspace_name=None, num_labels=2, log_folder_path: str = None): """ :param bert: BERT model which you want to train :param vocab_size: total word vocab size :param test_dataloader: test dataset data loader [can be None] :param lr: learning rate of optimizer :param betas: Adam optimizer betas :param weight_decay: Adam optimizer weight decay param :param with_cuda: traning with cuda :param log_freq: logging frequency of the batch iteration """ # Setup cuda device for BERT training, argument -c, --cuda should be true # cuda_condition = torch.cuda.is_available() and with_cuda # self.device = torch.device("cuda:0" if cuda_condition else "cpu") self.device = torch.device("cpu") #torch.device("cuda:0" if cuda_condition else "cpu") # print(cuda_condition, " Device used = ", self.device) print(" Device used = ", self.device) # available_gpus = list(range(torch.cuda.device_count())) # This BERT model will be saved every epoch self.model = bertFinetunedClassifierwithFeats.to("cpu") print(self.model.parameters()) for param in self.model.parameters(): param.requires_grad = False # Initialize the BERT Language Model, with BERT model # self.model = BERTForClassification(self.bert, vocab_size, num_labels).to(self.device) # self.model = BERTForClassificationWithFeats(self.bert, num_labels, 8).to(self.device) # self.model = bertFinetunedClassifierwithFeats # print(self.model.bert.parameters()) # for param in self.model.bert.parameters(): # param.requires_grad = False # BERTForClassificationWithFeats(self.bert, num_labels, 18).to(self.device) # self.model = BERTForClassificationWithFeats(self.bert, num_labels, 1).to(self.device) # Distributed GPU training if CUDA can detect more than 1 GPU # if with_cuda and torch.cuda.device_count() > 1: # print("Using %d GPUS for BERT" % torch.cuda.device_count()) # self.model = nn.DataParallel(self.model, device_ids=available_gpus) # Setting the train, validation and test data loader # self.train_data = train_dataloader # self.val_data = val_dataloader self.test_data = test_dataloader # self.optim = Adam(self.model.parameters(), lr=lr, weight_decay=weight_decay) #, eps=1e-9 self.optim = Adam(self.model.parameters(), lr=lr, betas=betas, weight_decay=weight_decay) # self.optim_schedule = ScheduledOptim(self.optim, self.model.bert.hidden, n_warmup_steps=warmup_steps) # self.scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, 'min', patience=2, factor=0.1) self.criterion = nn.CrossEntropyLoss() # if num_labels == 1: # self.criterion = nn.MSELoss() # elif num_labels == 2: # self.criterion = nn.BCEWithLogitsLoss() # # self.criterion = nn.CrossEntropyLoss() # elif num_labels > 2: # self.criterion = nn.CrossEntropyLoss() # self.criterion = nn.BCEWithLogitsLoss() self.log_freq = log_freq self.log_folder_path = log_folder_path # self.workspace_name = workspace_name # self.finetune_task = finetune_task # self.save_model = False # self.avg_loss = 10000 self.start_time = time.time() # self.probability_list = [] for fi in ['test']: #'val', f = open(self.log_folder_path+f"/log_{fi}_finetuned.txt", 'w') f.close() print("Total Parameters:", sum([p.nelement() for p in self.model.parameters()])) # def train(self, epoch): # self.iteration(epoch, self.train_data) # def val(self, epoch): # self.iteration(epoch, self.val_data, phase="val") def test(self, epoch): # if epoch == 0: # self.avg_loss = 10000 self.iteration(epoch, self.test_data, phase="test") def iteration(self, epoch, data_loader, phase="train"): """ loop over the data_loader for training or testing if on train status, backward operation is activated and also auto save the model every peoch :param epoch: current epoch index :param data_loader: torch.utils.data.DataLoader for iteration :param train: boolean value of is train or test :return: None """ # Setting the tqdm progress bar data_iter = tqdm.tqdm(enumerate(data_loader), desc="EP_%s:%d" % (phase, epoch), total=len(data_loader), bar_format="{l_bar}{r_bar}") avg_loss = 0.0 total_correct = 0 total_element = 0 plabels = [] tlabels = [] probabs = [] if phase == "train": self.model.train() else: self.model.eval() # self.probability_list = [] with open(self.log_folder_path+f"/log_{phase}_finetuned.txt", 'a') as f: sys.stdout = f for i, data in data_iter: # 0. batch_data will be sent into the device(GPU or cpu) data = {key: value.to(self.device) for key, value in data.items()} if phase == "train": logits = self.model.forward(data["input"], data["segment_label"], data["feat"]) else: with torch.no_grad(): logits = self.model.forward(data["input"].cpu(), data["segment_label"].cpu(), data["feat"].cpu()) logits = logits.cpu() loss = self.criterion(logits, data["label"]) # if torch.cuda.device_count() > 1: # loss = loss.mean() # 3. backward and optimization only in train # if phase == "train": # self.optim_schedule.zero_grad() # loss.backward() # self.optim_schedule.step_and_update_lr() # prediction accuracy probs = nn.Softmax(dim=-1)(logits) # Probabilities probabs.extend(probs.detach().cpu().numpy().tolist()) predicted_labels = torch.argmax(probs, dim=-1) #correct # self.probability_list.append(probs) # true_labels = torch.argmax(data["label"], dim=-1) plabels.extend(predicted_labels.cpu().numpy()) tlabels.extend(data['label'].cpu().numpy()) # Compare predicted labels to true labels and calculate accuracy correct = (data['label'] == predicted_labels).sum().item() avg_loss += loss.item() total_correct += correct # total_element += true_labels.nelement() total_element += data["label"].nelement() # print(">>>>>>>>>>>>>>", predicted_labels, true_labels, correct, total_correct, total_element) post_fix = { "epoch": epoch, "iter": i, "avg_loss": avg_loss / (i + 1), "avg_acc": total_correct / total_element * 100 if total_element != 0 else 0, "loss": loss.item() } if i % self.log_freq == 0: data_iter.write(str(post_fix)) precisions = precision_score(tlabels, plabels, average="weighted", zero_division=0) recalls = recall_score(tlabels, plabels, average="weighted") f1_scores = f1_score(tlabels, plabels, average="weighted") cmatrix = confusion_matrix(tlabels, plabels) end_time = time.time() auc_score = roc_auc_score(tlabels, plabels) final_msg = { "epoch": f"EP{epoch}_{phase}", "avg_loss": avg_loss / len(data_iter), "total_acc": total_correct * 100.0 / total_element, "precisions": precisions, "recalls": recalls, "f1_scores": f1_scores, # "confusion_matrix": f"{cmatrix}", # "true_labels": f"{tlabels}", # "predicted_labels": f"{plabels}", "time_taken_from_start": end_time - self.start_time, "auc_score":auc_score } with open("result.txt", 'w') as file: for key, value in final_msg.items(): file.write(f"{key}: {value}\n") print(final_msg) fpr, tpr, thresholds = roc_curve(tlabels, plabels) with open("roc_data.pkl", "wb") as f: pickle.dump((fpr, tpr, thresholds), f) print(final_msg) f.close() with open(self.log_folder_path+f"/log_{phase}_finetuned_info.txt", 'a') as f1: sys.stdout = f1 final_msg = { "epoch": f"EP{epoch}_{phase}", "confusion_matrix": f"{cmatrix}", "true_labels": f"{tlabels if epoch == 0 else ''}", "predicted_labels": f"{plabels}", "probabilities": f"{probabs}", "time_taken_from_start": end_time - self.start_time } print(final_msg) f1.close() sys.stdout = sys.__stdout__ sys.stdout = sys.__stdout__ class BERTFineTuneCalibratedTrainer: def __init__(self, bertFinetunedClassifierwithFeats: BERT, #BERTForClassificationWithFeats vocab_size: int, test_dataloader: DataLoader = None, lr: float = 1e-4, betas=(0.9, 0.999), weight_decay: float = 0.01, warmup_steps=10000, with_cuda: bool = True, cuda_devices=None, log_freq: int = 10, workspace_name=None, num_labels=2, log_folder_path: str = None): """ :param bert: BERT model which you want to train :param vocab_size: total word vocab size :param test_dataloader: test dataset data loader [can be None] :param lr: learning rate of optimizer :param betas: Adam optimizer betas :param weight_decay: Adam optimizer weight decay param :param with_cuda: traning with cuda :param log_freq: logging frequency of the batch iteration """ # Setup cuda device for BERT training, argument -c, --cuda should be true cuda_condition = torch.cuda.is_available() and with_cuda self.device = torch.device("cuda:0" if cuda_condition else "cpu") print(cuda_condition, " Device used = ", self.device) # available_gpus = list(range(torch.cuda.device_count())) # This BERT model will be saved every epoch self.model = bertFinetunedClassifierwithFeats print(self.model.parameters()) for param in self.model.parameters(): param.requires_grad = False # Initialize the BERT Language Model, with BERT model # self.model = BERTForClassification(self.bert, vocab_size, num_labels).to(self.device) # self.model = BERTForClassificationWithFeats(self.bert, num_labels, 8).to(self.device) # self.model = bertFinetunedClassifierwithFeats # print(self.model.bert.parameters()) # for param in self.model.bert.parameters(): # param.requires_grad = False # BERTForClassificationWithFeats(self.bert, num_labels, 18).to(self.device) # self.model = BERTForClassificationWithFeats(self.bert, num_labels, 1).to(self.device) # Distributed GPU training if CUDA can detect more than 1 GPU # if with_cuda and torch.cuda.device_count() > 1: # print("Using %d GPUS for BERT" % torch.cuda.device_count()) # self.model = nn.DataParallel(self.model, device_ids=available_gpus) # Setting the train, validation and test data loader # self.train_data = train_dataloader # self.val_data = val_dataloader self.test_data = test_dataloader # self.optim = Adam(self.model.parameters(), lr=lr, weight_decay=weight_decay) #, eps=1e-9 self.optim = Adam(self.model.parameters(), lr=lr, betas=betas, weight_decay=weight_decay) # self.optim_schedule = ScheduledOptim(self.optim, self.model.bert.hidden, n_warmup_steps=warmup_steps) # self.scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, 'min', patience=2, factor=0.1) self.criterion = nn.CrossEntropyLoss() # if num_labels == 1: # self.criterion = nn.MSELoss() # elif num_labels == 2: # self.criterion = nn.BCEWithLogitsLoss() # # self.criterion = nn.CrossEntropyLoss() # elif num_labels > 2: # self.criterion = nn.CrossEntropyLoss() # self.criterion = nn.BCEWithLogitsLoss() self.log_freq = log_freq self.log_folder_path = log_folder_path # self.workspace_name = workspace_name # self.finetune_task = finetune_task # self.save_model = False # self.avg_loss = 10000 self.start_time = time.time() # self.probability_list = [] for fi in ['test']: #'val', f = open(self.log_folder_path+f"/log_{fi}_finetuned.txt", 'w') f.close() print("Total Parameters:", sum([p.nelement() for p in self.model.parameters()])) # def train(self, epoch): # self.iteration(epoch, self.train_data) # def val(self, epoch): # self.iteration(epoch, self.val_data, phase="val") def test(self, epoch): # if epoch == 0: # self.avg_loss = 10000 self.iteration(epoch, self.test_data, phase="test") def iteration(self, epoch, data_loader, phase="train"): """ loop over the data_loader for training or testing if on train status, backward operation is activated and also auto save the model every peoch :param epoch: current epoch index :param data_loader: torch.utils.data.DataLoader for iteration :param train: boolean value of is train or test :return: None """ # Setting the tqdm progress bar data_iter = tqdm.tqdm(enumerate(data_loader), desc="EP_%s:%d" % (phase, epoch), total=len(data_loader), bar_format="{l_bar}{r_bar}") avg_loss = 0.0 total_correct = 0 total_element = 0 plabels = [] tlabels = [] probabs = [] if phase == "train": self.model.train() else: self.model.eval() # self.probability_list = [] with open(self.log_folder_path+f"/log_{phase}_finetuned.txt", 'a') as f: sys.stdout = f for i, data in data_iter: # 0. batch_data will be sent into the device(GPU or cpu) # print(data_pair[0]) data = {key: value.to(self.device) for key, value in data[0].items()} # print(f"data : {data}") # data = {key: value.to(self.device) for key, value in data.items()} # if phase == "train": # logits = self.model.forward(data["input"], data["segment_label"], data["feat"]) # else: with torch.no_grad(): # logits = self.model.forward(data["input"], data["segment_label"], data["feat"]) logits = self.model.forward(data) loss = self.criterion(logits, data["label"]) if torch.cuda.device_count() > 1: loss = loss.mean() # 3. backward and optimization only in train # if phase == "train": # self.optim_schedule.zero_grad() # loss.backward() # self.optim_schedule.step_and_update_lr() # prediction accuracy probs = nn.Softmax(dim=-1)(logits) # Probabilities probabs.extend(probs.detach().cpu().numpy().tolist()) predicted_labels = torch.argmax(probs, dim=-1) #correct # self.probability_list.append(probs) # true_labels = torch.argmax(data["label"], dim=-1) plabels.extend(predicted_labels.cpu().numpy()) tlabels.extend(data['label'].cpu().numpy()) positive_class_probs = [prob[1] for prob in probabs] # Compare predicted labels to true labels and calculate accuracy correct = (data['label'] == predicted_labels).sum().item() avg_loss += loss.item() total_correct += correct # total_element += true_labels.nelement() total_element += data["label"].nelement() # print(">>>>>>>>>>>>>>", predicted_labels, true_labels, correct, total_correct, total_element) post_fix = { "epoch": epoch, "iter": i, "avg_loss": avg_loss / (i + 1), "avg_acc": total_correct / total_element * 100 if total_element != 0 else 0, "loss": loss.item() } if i % self.log_freq == 0: data_iter.write(str(post_fix)) precisions = precision_score(tlabels, plabels, average="weighted", zero_division=0) recalls = recall_score(tlabels, plabels, average="weighted") f1_scores = f1_score(tlabels, plabels, average="weighted") cmatrix = confusion_matrix(tlabels, plabels) auc_score = roc_auc_score(tlabels, positive_class_probs) end_time = time.time() final_msg = { "epoch": f"EP{epoch}_{phase}", "avg_loss": avg_loss / len(data_iter), "total_acc": total_correct * 100.0 / total_element, "precisions": precisions, "recalls": recalls, "f1_scores": f1_scores, "auc_score":auc_score, # "confusion_matrix": f"{cmatrix}", # "true_labels": f"{tlabels}", # "predicted_labels": f"{plabels}", "time_taken_from_start": end_time - self.start_time } with open("result.txt", 'w') as file: for key, value in final_msg.items(): file.write(f"{key}: {value}\n") print(final_msg) fpr, tpr, thresholds = roc_curve(tlabels, positive_class_probs) f.close() with open(self.log_folder_path+f"/log_{phase}_finetuned_info.txt", 'a') as f1: sys.stdout = f1 final_msg = { "epoch": f"EP{epoch}_{phase}", "confusion_matrix": f"{cmatrix}", "true_labels": f"{tlabels if epoch == 0 else ''}", "predicted_labels": f"{plabels}", "probabilities": f"{probabs}", "time_taken_from_start": end_time - self.start_time } print(final_msg) f1.close() sys.stdout = sys.__stdout__ sys.stdout = sys.__stdout__ def train(): parser = argparse.ArgumentParser() parser.add_argument('-workspace_name', type=str, default=None) parser.add_argument('-code', type=str, default=None, help="folder for pretraining outputs and logs") parser.add_argument('-finetune_task', type=str, default=None, help="folder inside finetuning") parser.add_argument("-attention", type=bool, default=False, help="analyse attention scores") parser.add_argument("-diff_test_folder", type=bool, default=False, help="use for different test folder") parser.add_argument("-embeddings", type=bool, default=False, help="get and analyse embeddings") parser.add_argument('-embeddings_file_name', type=str, default=None, help="file name of embeddings") parser.add_argument("-pretrain", type=bool, default=False, help="pretraining: true, or false") # parser.add_argument('-opts', nargs='+', type=str, default=None, help='List of optional steps') parser.add_argument("-max_mask", type=int, default=0.15, help="% of input tokens selected for masking") # parser.add_argument("-p", "--pretrain_dataset", type=str, default="pretraining/pretrain.txt", help="pretraining dataset for bert") # parser.add_argument("-pv", "--pretrain_val_dataset", type=str, default="pretraining/test.txt", help="pretraining validation dataset for bert") # default="finetuning/test.txt", parser.add_argument("-vocab_path", type=str, default="pretraining/vocab.txt", help="built vocab model path with bert-vocab") parser.add_argument("-train_dataset_path", type=str, default="train.txt", help="fine tune train dataset for progress classifier") parser.add_argument("-val_dataset_path", type=str, default="val.txt", help="test set for evaluate fine tune train set") parser.add_argument("-test_dataset_path", type=str, default="test.txt", help="test set for evaluate fine tune train set") parser.add_argument("-num_labels", type=int, default=2, help="Number of labels") parser.add_argument("-train_label_path", type=str, default="train_label.txt", help="fine tune train dataset for progress classifier") parser.add_argument("-val_label_path", type=str, default="val_label.txt", help="test set for evaluate fine tune train set") parser.add_argument("-test_label_path", type=str, default="test_label.txt", help="test set for evaluate fine tune train set") ##### change Checkpoint for finetuning parser.add_argument("-pretrained_bert_checkpoint", type=str, default=None, help="checkpoint of saved pretrained bert model") parser.add_argument("-finetuned_bert_classifier_checkpoint", type=str, default=None, help="checkpoint of saved finetuned bert model") #."output_feb09/bert_trained.model.ep40" #."output_feb09/bert_trained.model.ep40" parser.add_argument('-check_epoch', type=int, default=None) parser.add_argument("-hs", "--hidden", type=int, default=64, help="hidden size of transformer model") #64 parser.add_argument("-l", "--layers", type=int, default=4, help="number of layers") #4 parser.add_argument("-a", "--attn_heads", type=int, default=4, help="number of attention heads") #8 parser.add_argument("-s", "--seq_len", type=int, default=5, help="maximum sequence length") parser.add_argument("-b", "--batch_size", type=int, default=500, help="number of batch_size") #64 parser.add_argument("-e", "--epochs", type=int, default=1)#1501, help="number of epochs") #501 # Use 50 for pretrain, and 10 for fine tune parser.add_argument("-w", "--num_workers", type=int, default=0, help="dataloader worker size") # Later run with cuda parser.add_argument("--with_cuda", type=bool, default=False, help="training with CUDA: true, or false") parser.add_argument("--log_freq", type=int, default=10, help="printing loss every n iter: setting n") # parser.add_argument("--corpus_lines", type=int, default=None, help="total number of lines in corpus") parser.add_argument("--cuda_devices", type=int, nargs='+', default=None, help="CUDA device ids") # parser.add_argument("--on_memory", type=bool, default=False, help="Loading on memory: true or false") parser.add_argument("--dropout", type=float, default=0.1, help="dropout of network") parser.add_argument("--lr", type=float, default=1e-05, help="learning rate of adam") #1e-3 parser.add_argument("--adam_weight_decay", type=float, default=0.01, help="weight_decay of adam") parser.add_argument("--adam_beta1", type=float, default=0.9, help="adam first beta value") parser.add_argument("--adam_beta2", type=float, default=0.98, help="adam first beta value") #0.999 parser.add_argument("-o", "--output_path", type=str, default="bert_trained.seq_encoder.model", help="ex)output/bert.model") # parser.add_argument("-o", "--output_path", type=str, default="output/bert_fine_tuned.model", help="ex)output/bert.model") args = parser.parse_args() for k,v in vars(args).items(): if 'path' in k: if v: if k == "output_path": if args.code: setattr(args, f"{k}", args.workspace_name+f"/output/{args.code}/"+v) elif args.finetune_task: setattr(args, f"{k}", args.workspace_name+f"/output/{args.finetune_task}/"+v) else: setattr(args, f"{k}", args.workspace_name+"/output/"+v) elif k != "vocab_path": if args.pretrain: setattr(args, f"{k}", args.workspace_name+"/pretraining/"+v) else: if args.code: setattr(args, f"{k}", args.workspace_name+f"/{args.code}/"+v) elif args.finetune_task: if args.diff_test_folder and "test" in k: setattr(args, f"{k}", args.workspace_name+f"/finetuning/"+v) else: setattr(args, f"{k}", args.workspace_name+f"/finetuning/{args.finetune_task}/"+v) else: setattr(args, f"{k}", args.workspace_name+"/finetuning/"+v) else: setattr(args, f"{k}", args.workspace_name+"/"+v) print(f"args.{k} : {getattr(args, f'{k}')}") print("Loading Vocab", args.vocab_path) vocab_obj = Vocab(args.vocab_path) vocab_obj.load_vocab() print("Vocab Size: ", len(vocab_obj.vocab)) print("Testing using finetuned model......") print("Loading Test Dataset", args.test_dataset_path) test_dataset = TokenizerDataset(args.test_dataset_path, args.test_label_path, vocab_obj, seq_len=args.seq_len) # test_dataset = TokenizerDatasetForCalibration(args.test_dataset_path, args.test_label_path, vocab_obj, seq_len=args.seq_len) print("Creating Dataloader...") test_data_loader = DataLoader(test_dataset, batch_size=args.batch_size, num_workers=args.num_workers) print("Load fine-tuned BERT classifier model with feats") # cuda_condition = torch.cuda.is_available() and args.with_cuda device = torch.device("cpu") #torch.device("cuda:0" if cuda_condition else "cpu") finetunedBERTclassifier = torch.load(args.finetuned_bert_classifier_checkpoint, map_location=device) if isinstance(finetunedBERTclassifier, torch.nn.DataParallel): finetunedBERTclassifier = finetunedBERTclassifier.module new_log_folder = f"{args.workspace_name}/logs" new_output_folder = f"{args.workspace_name}/output" if args.finetune_task: # is sent almost all the time new_log_folder = f"{args.workspace_name}/logs/{args.finetune_task}" new_output_folder = f"{args.workspace_name}/output/{args.finetune_task}" if not os.path.exists(new_log_folder): os.makedirs(new_log_folder) if not os.path.exists(new_output_folder): os.makedirs(new_output_folder) print("Creating BERT Fine Tuned Test Trainer") trainer = BERTFineTuneTrainer(finetunedBERTclassifier, len(vocab_obj.vocab), test_dataloader=test_data_loader, lr=args.lr, betas=(args.adam_beta1, args.adam_beta2), weight_decay=args.adam_weight_decay, with_cuda=args.with_cuda, cuda_devices = args.cuda_devices, log_freq=args.log_freq, workspace_name = args.workspace_name, num_labels=args.num_labels, log_folder_path=new_log_folder) # trainer = BERTFineTuneCalibratedTrainer(finetunedBERTclassifier, # len(vocab_obj.vocab), test_dataloader=test_data_loader, # lr=args.lr, betas=(args.adam_beta1, args.adam_beta2), weight_decay=args.adam_weight_decay, # with_cuda=args.with_cuda, cuda_devices = args.cuda_devices, log_freq=args.log_freq, # workspace_name = args.workspace_name, num_labels=args.num_labels, log_folder_path=new_log_folder) print("Testing fine-tuned model Start....") start_time = time.time() repoch = range(args.check_epoch, args.epochs) if args.check_epoch else range(args.epochs) counter = 0 # patience = 10 for epoch in repoch: print(f'Test Epoch {epoch} Starts, Time: {time.strftime("%D %T", time.localtime(time.time()))}') trainer.test(epoch) # pickle.dump(trainer.probability_list, open(f"{args.workspace_name}/output/aaai/change4_mid_prob_{epoch}.pkl","wb")) print(f'Test Epoch {epoch} Ends, Time: {time.strftime("%D %T", time.localtime(time.time()))} \n') end_time = time.time() print("Time Taken to fine-tune model = ", end_time - start_time) print(f'Pretraining Ends, Time: {time.strftime("%D %T", time.localtime(end_time))}') if __name__ == "__main__": train()