import os import argparse import time import torch import torch.nn as nn import numpy as np import pandas as pd import utils from .models.bert_labeler import bert_labeler from .datasets.impressions_dataset import ImpressionsDataset from .constants import * def collate_fn_labels(sample_list): """Custom collate function to pad reports in each batch to the max len @param sample_list (List): A list of samples. Each sample is a dictionary with keys 'imp', 'label', 'len' as returned by the __getitem__ function of ImpressionsDataset @returns batch (dictionary): A dictionary with keys 'imp', 'label', 'len' but now 'imp' is a tensor with padding and batch size as the first dimension. 'label' is a stacked tensor of labels for the whole batch with batch size as first dim. And 'len' is a list of the length of each sequence in batch """ tensor_list = [s['imp'] for s in sample_list] batched_imp = torch.nn.utils.rnn.pad_sequence(tensor_list, batch_first=True, padding_value=PAD_IDX) label_list = [s['label'] for s in sample_list] batched_label = torch.stack(label_list, dim=0) len_list = [s['len'] for s in sample_list] batch = {'imp': batched_imp, 'label': batched_label, 'len': len_list} return batch def load_data(train_csv_path, train_list_path, dev_csv_path, dev_list_path, train_weights=None, batch_size=BATCH_SIZE, shuffle=True, num_workers=NUM_WORKERS): """ Create ImpressionsDataset objects for train and test data @param train_csv_path (string): path to training csv file containing labels @param train_list_path (string): path to list of encoded impressions for train set @param dev_csv_path (string): same as train_csv_path but for dev set @param dev_list_path (string): same as train_list_path but for dev set @param train_weights (torch.Tensor): Tensor of shape (train_set_size) containing weights for each training example, for the purposes of batch sampling with replacement @param batch_size (int): the batch size. As per the BERT repository, the max batch size that can fit on a TITAN XP is 6 if the max sequence length is 512, which is our case. We have 3 TITAN XP's @param shuffle (bool): Whether to shuffle data before each epoch, ignored if train_weights is not None @param num_workers (int): How many worker processes to use to load data @returns dataloaders (tuple): tuple of two ImpressionsDataset objects, for train and dev sets """ collate_fn = collate_fn_labels train_dset = ImpressionsDataset(train_csv_path, train_list_path) dev_dset = ImpressionsDataset(dev_csv_path, dev_list_path) if train_weights is None: train_loader = torch.utils.data.DataLoader(train_dset, batch_size=batch_size, shuffle=shuffle, num_workers=num_workers, collate_fn=collate_fn) else: sampler = torch.utils.data.WeightedRandomSampler(weights=train_weights, num_samples=len(train_weights), replacement=True) train_loader = torch.utils.data.DataLoader(train_dset, batch_size=batch_size, num_workers=num_workers, collate_fn=collate_fn, sampler=sampler) dev_loader = torch.utils.data.DataLoader(dev_dset, batch_size=batch_size, shuffle=shuffle, num_workers=num_workers, collate_fn=collate_fn) dataloaders = (train_loader, dev_loader) return dataloaders def load_test_data(test_csv_path, test_list_path, batch_size=BATCH_SIZE, num_workers=NUM_WORKERS, shuffle=False): """ Create ImpressionsDataset object for the test set @param test_csv_path (string): path to test csv file containing labels @param test_list_path (string): path to list of encoded impressions @param batch_size (int): the batch size. As per the BERT repository, the max batch size that can fit on a TITAN XP is 6 if the max sequence length is 512, which is our case. We have 3 TITAN XP's @param num_workers (int): how many worker processes to use to load data @param shuffle (bool): whether to shuffle the data or not @returns test_loader (dataloader): dataloader object for test set """ collate_fn = collate_fn_labels test_dset = ImpressionsDataset(test_csv_path, test_list_path) test_loader = torch.utils.data.DataLoader(test_dset, batch_size=batch_size, shuffle=shuffle, num_workers=num_workers, collate_fn=collate_fn) return test_loader def train(save_path, dataloaders, f1_weights, model=None, device=None, optimizer=None, lr=LEARNING_RATE, log_every=LOG_EVERY, valid_niter=VALID_NITER, best_metric=0.0): """ Main training loop for the labeler @param save_path (string): Directory in which model weights are stored @param model (nn.Module): the labeler model to train, if applicable @param device (torch.device): device for the model. If model is not None, this parameter is required @param dataloaders (tuple): tuple of dataloader objects as returned by load_data @param f1_weights (dictionary): maps conditions to weights for blank, negation, uncertain and positxive f1 task averaging @param optimizer (torch.optim.Optimizer): the optimizer to use, if applicable @param lr (float): learning rate to use in the optimizer, ignored if optimizer is not None @param log_every (int): number of iterations to log after @param valid_niter (int): number of iterations after which to evaluate the model and save it if it is better than old best model @param best_metric (float): save checkpoints only if dev set performance is higher than best_metric """ if model and not device: print("train function error: Model specified but not device") return if model is None: model = bert_labeler(pretrain_path=PRETRAIN_PATH) model.train() #put the model into train mode device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") if torch.cuda.device_count() > 1: print("Using", torch.cuda.device_count(), "GPUs!") model = nn.DataParallel(model) #to utilize multiple GPU's model = model.to(device) else: model.train() if optimizer is None: optimizer = torch.optim.Adam(model.parameters(), lr=lr) begin_time = time.time() report_examples = 0 report_loss = 0.0 train_ld = dataloaders[0] dev_ld = dataloaders[1] loss_func = nn.CrossEntropyLoss(reduction='sum') print('begin labeler training') for epoch in range(NUM_EPOCHS): for i, data in enumerate(train_ld, 0): batch = data['imp'] #(batch_size, max_len) batch = batch.to(device) label = data['label'] #(batch_size, 14) label = label.permute(1, 0).to(device) src_len = data['len'] batch_size = batch.shape[0] attn_mask = utils.generate_attention_masks(batch, src_len, device) optimizer.zero_grad() out = model(batch, attn_mask) #list of 14 tensors batch_loss = 0.0 for j in range(len(out)): batch_loss += loss_func(out[j], label[j]) report_loss += batch_loss report_examples += batch_size loss = batch_loss / batch_size loss.backward() optimizer.step() if (i+1) % log_every == 0: print('epoch %d, iter %d, avg_loss %.3f, time_elapsed %.3f sec' % (epoch+1, i+1, report_loss/report_examples, time.time() - begin_time)) report_loss = 0.0 report_examples = 0 if (i+1) % valid_niter == 0: print('\n begin validation') metrics = utils.evaluate(model, dev_ld, device, f1_weights) weighted = metrics['weighted'] kappas = metrics['kappa'] for j in range(len(CONDITIONS)): print('%s kappa: %.3f' % (CONDITIONS[j], kappas[j])) print('average: %.3f' % (np.mean(kappas))) #for j in range(len(CONDITIONS)): # print('%s weighted_f1: %.3f' % (CONDITIONS[j], weighted[j])) #print('average of weighted_f1: %.3f' % (np.mean(weighted))) for j in range(len(CONDITIONS)): print('%s blank_f1: %.3f, negation_f1: %.3f, uncertain_f1: %.3f, positive: %.3f' % (CONDITIONS[j], metrics['blank'][j], metrics['negation'][j], metrics['uncertain'][j], metrics['positive'][j])) metric_avg = np.mean(kappas) if metric_avg > best_metric: #new best network print("saving new best network!\n") best_metric = metric_avg path = os.path.join(save_path, "model_epoch%d_iter%d" % (epoch+1, i+1)) torch.save({'epoch': epoch+1, 'model_state_dict': model.state_dict(), 'optimizer_state_dict': optimizer.state_dict()}, path) def model_from_ckpt(model, ckpt_path): """Load up model checkpoint @param model (nn.Module): the module to be loaded @param ckpt_path (string): path to a checkpoint. If this is None, then model is trained from scratch @return (tuple): tuple containing the model, optimizer and device """ device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") if torch.cuda.device_count() > 1: print("Using", torch.cuda.device_count(), "GPUs!") model = nn.DataParallel(model) #to utilize multiple GPU's model = model.to(device) optimizer = torch.optim.Adam(model.parameters()) checkpoint = torch.load(ckpt_path) model.load_state_dict(checkpoint['model_state_dict']) optimizer.load_state_dict(checkpoint['optimizer_state_dict']) return (model, optimizer, device) if __name__ == '__main__': parser = argparse.ArgumentParser(description='Train BERT-base model on task of labeling 14 medical conditions.') parser.add_argument('--train_csv', type=str, nargs='?', required=True, help='path to csv containing train reports.') parser.add_argument('--dev_csv', type=str, nargs='?', required=True, help='path to csv containing dev reports.') parser.add_argument('--train_imp_list', type=str, nargs='?', required=True, help='path to list of tokenized train set report impressions') parser.add_argument('--dev_imp_list', type=str, nargs='?', required=True, help='path to list of tokenized dev set report impressions') parser.add_argument('--output_dir', type=str, nargs='?', required=True, help='path to output directory where checkpoints will be saved') parser.add_argument('--checkpoint', type=str, nargs='?', required=False, help='path to existing checkpoint to initialize weights from') args = parser.parse_args() train_csv_path = args.train_csv dev_csv_path = args.dev_csv train_imp_path = args.train_imp_list dev_imp_path = args.dev_imp_list out_path = args.output_dir checkpoint_path = args.checkpoint if checkpoint_path: model, optimizer, device = model_from_ckpt(bert_labeler(), checkpoint_path) else: model, optimizer, device = None, None, None f1_weights = utils.get_weighted_f1_weights(dev_csv_path) dataloaders = load_data(train_csv_path, train_imp_path, dev_csv_path, dev_imp_path) train(save_path=out_path, dataloaders=dataloaders, model=model, optimizer=optimizer, device=device, f1_weights=f1_weights)