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)
