def save_best_model(encoder_model, decoder_model, encoder_optimizer,
decoder_optimizer, hidden_size, layers, epoch, file_name):
"""
Save the best model
:param encoder_model: A trained encoder model
:param decoder_model: A trained decoder model
:param encoder_optimizer: Encoder optimizer
:param decoder_optimizer: Decoder optimizer
:param file_name: Output file name
:return:
"""
if os.path.isfile(file_name):
os.remove(file_name)
ModelHandler.save_checkpoint(
{
'encoder_state_dict': encoder_model.state_dict(),
'decoder_state_dict': decoder_model.state_dict(),
'encoder_optimizer': encoder_optimizer.state_dict(),
'decoder_optimizer': decoder_optimizer.state_dict(),
'hidden_size': hidden_size,
'gru_layers': layers,
'epochs': epoch,
}, file_name)
sys.stderr.write(TextColor.RED + "\nMODEL SAVED SUCCESSFULLY.\n" +
TextColor.END)
def do_test(test_file, batch_size, gpu_mode, num_workers, model_path,
output_directory, print_details):
"""
Train a model and save
:param test_file: A CSV file containing test image information
:param batch_size: Batch size for training
:param gpu_mode: If true the model will be trained on GPU
:param num_workers: Number of workers for data loading
:param model_path: Path to a saved model
:param num_classes: Number of output classes
:return:
"""
sys.stderr.write(TextColor.PURPLE + 'Loading data\n' + TextColor.END)
if os.path.isfile(model_path) is False:
sys.stderr.write(TextColor.RED + "ERROR: INVALID PATH TO MODEL\n")
exit(1)
sys.stderr.write(TextColor.GREEN + "INFO: MODEL LOADING\n" + TextColor.END)
transducer_model, hidden_size, gru_layers, prev_ite = \
ModelHandler.load_simple_model(model_path,
input_channels=ImageSizeOptions.IMAGE_CHANNELS,
image_features=ImageSizeOptions.IMAGE_HEIGHT,
seq_len=ImageSizeOptions.SEQ_LENGTH,
num_base_classes=ImageSizeOptions.TOTAL_BASE_LABELS,
num_rle_classes=ImageSizeOptions.TOTAL_RLE_LABELS)
sys.stderr.write(TextColor.GREEN + "INFO: MODEL LOADED\n" + TextColor.END)
if gpu_mode:
transducer_model = torch.nn.DataParallel(transducer_model).cuda()
stats_dictionary = test(
test_file,
batch_size,
gpu_mode,
transducer_model,
num_workers,
gru_layers,
hidden_size,
num_base_classes=ImageSizeOptions.TOTAL_BASE_LABELS,
num_rle_classes=ImageSizeOptions.TOTAL_RLE_LABELS,
output_directory=output_directory,
print_details=print_details)
save_rle_confusion_matrix(stats_dictionary)
save_base_confusion_matrix(stats_dictionary)
sys.stderr.write(TextColor.PURPLE + 'DONE\n' + TextColor.END)
def save_best_model(transducer_model, model_optimizer, hidden_size, layers,
epoch, file_name):
"""
Save the best model
:param transducer_model: A trained model
:param model_optimizer: Model optimizer
:param hidden_size: Number of hidden layers
:param layers: Number of GRU layers to use
:param epoch: Epoch/iteration number
:param file_name: Output file name
:return:
"""
if os.path.isfile(file_name):
os.remove(file_name)
ModelHandler.save_checkpoint(
{
'model_state_dict': transducer_model.state_dict(),
'model_optimizer': model_optimizer.state_dict(),
'hidden_size': hidden_size,
'gru_layers': layers,
'epochs': epoch,
}, file_name)
sys.stderr.write(TextColor.RED + "\nMODEL SAVED SUCCESSFULLY.\n" +
TextColor.END)
def train(train_file, test_file, batch_size, epoch_limit, gpu_mode,
num_workers, retrain_model, retrain_model_path, gru_layers,
hidden_size, lr, decay, model_dir, stats_dir, train_mode):
if train_mode is True:
train_loss_logger = open(stats_dir + "train_loss.csv", 'w')
test_loss_logger = open(stats_dir + "test_loss.csv", 'w')
confusion_matrix_logger = open(stats_dir + "confusion_matrix.txt", 'w')
else:
train_loss_logger = None
test_loss_logger = None
confusion_matrix_logger = None
sys.stderr.write(TextColor.PURPLE + 'Loading data\n' + TextColor.END)
train_data_set = SequenceDataset(train_file)
train_loader = DataLoader(train_data_set,
batch_size=batch_size,
shuffle=True,
num_workers=num_workers,
pin_memory=gpu_mode)
num_classes = ImageSizeOptions.TOTAL_LABELS
if retrain_model is True:
if os.path.isfile(retrain_model_path) is False:
sys.stderr.write(
TextColor.RED +
"ERROR: INVALID PATH TO RETRAIN PATH MODEL --retrain_model_path\n"
)
exit(1)
sys.stderr.write(TextColor.GREEN + "INFO: RETRAIN MODEL LOADING\n" +
TextColor.END)
transducer_model, hidden_size, gru_layers, prev_ite = \
ModelHandler.load_simple_model_for_training(retrain_model_path,
input_channels=ImageSizeOptions.IMAGE_CHANNELS,
image_features=ImageSizeOptions.IMAGE_HEIGHT,
seq_len=ImageSizeOptions.SEQ_LENGTH,
num_classes=num_classes)
if train_mode is True:
epoch_limit = prev_ite + epoch_limit
sys.stderr.write(TextColor.GREEN + "INFO: RETRAIN MODEL LOADED\n" +
TextColor.END)
else:
transducer_model = ModelHandler.get_new_gru_model(
input_channels=ImageSizeOptions.IMAGE_CHANNELS,
image_features=ImageSizeOptions.IMAGE_HEIGHT,
gru_layers=gru_layers,
hidden_size=hidden_size,
num_classes=num_classes)
prev_ite = 0
param_count = sum(p.numel() for p in transducer_model.parameters()
if p.requires_grad)
sys.stderr.write(TextColor.RED + "INFO: TOTAL TRAINABLE PARAMETERS:\t" +
str(param_count) + "\n" + TextColor.END)
if gpu_mode:
transducer_model = torch.nn.DataParallel(transducer_model).cuda()
class_weights = torch.Tensor(CLASS_WEIGHTS)
# Loss
criterion = nn.CrossEntropyLoss(class_weights)
if gpu_mode is True:
criterion = criterion.cuda()
model_optimizer = torch.optim.Adam(transducer_model.parameters(),
lr=lr,
weight_decay=decay)
lr_scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
model_optimizer, 'min')
if retrain_model is True:
sys.stderr.write(TextColor.GREEN + "INFO: OPTIMIZER LOADING\n" +
TextColor.END)
model_optimizer = ModelHandler.load_simple_optimizer(
model_optimizer, retrain_model_path, gpu_mode)
sys.stderr.write(TextColor.GREEN + "INFO: OPTIMIZER LOADED\n" +
TextColor.END)
start_epoch = prev_ite
# Train the Model
sys.stderr.write(TextColor.PURPLE + 'Training starting\n' + TextColor.END)
stats = dict()
stats['loss_epoch'] = []
stats['accuracy_epoch'] = []
sys.stderr.write(TextColor.BLUE + 'Start: ' + str(start_epoch + 1) +
' End: ' + str(epoch_limit) + "\n")
for epoch in range(start_epoch, epoch_limit, 1):
total_loss = 0
total_images = 0
sys.stderr.write(TextColor.BLUE + 'Train epoch: ' + str(epoch + 1) +
"\n")
# make sure the model is in train mode. BN is different in train and eval.
batch_no = 1
with tqdm(total=len(train_loader), desc='Loss', leave=True,
ncols=100) as progress_bar:
transducer_model.train()
for images, labels in train_loader:
labels = labels.type(torch.LongTensor)
images = images.type(torch.FloatTensor)
if gpu_mode:
# encoder_hidden = encoder_hidden.cuda()
images = images.cuda()
labels = labels.cuda()
hidden = torch.zeros(images.size(0),
2 * TrainOptions.GRU_LAYERS,
TrainOptions.HIDDEN_SIZE)
if gpu_mode:
hidden = hidden.cuda()
for i in range(0, ImageSizeOptions.SEQ_LENGTH,
TrainOptions.WINDOW_JUMP):
model_optimizer.zero_grad()
if i + TrainOptions.TRAIN_WINDOW > ImageSizeOptions.SEQ_LENGTH:
break
image_chunk = images[:, i:i + TrainOptions.TRAIN_WINDOW]
label_chunk = labels[:, i:i + TrainOptions.TRAIN_WINDOW]
output_, hidden = transducer_model(image_chunk, hidden)
loss = criterion(
output_.contiguous().view(-1, num_classes),
label_chunk.contiguous().view(-1))
loss.backward()
model_optimizer.step()
total_loss += loss.item()
total_images += image_chunk.size(0)
hidden = hidden.detach()
# update the progress bar
avg_loss = (total_loss / total_images) if total_images else 0
progress_bar.set_description("Loss: " + str(avg_loss))
if train_mode is True:
train_loss_logger.write(
str(epoch + 1) + "," + str(batch_no) + "," +
str(avg_loss) + "\n")
progress_bar.refresh()
progress_bar.update(1)
batch_no += 1
progress_bar.close()
stats_dictioanry = test(test_file,
batch_size,
gpu_mode,
transducer_model,
num_workers,
gru_layers,
hidden_size,
num_classes=ImageSizeOptions.TOTAL_LABELS)
stats['loss'] = stats_dictioanry['loss']
stats['accuracy'] = stats_dictioanry['accuracy']
stats['loss_epoch'].append((epoch, stats_dictioanry['loss']))
stats['accuracy_epoch'].append((epoch, stats_dictioanry['accuracy']))
lr_scheduler.step(stats['loss'])
# update the loggers
if train_mode is True:
# save the model after each epoch
# encoder_model, decoder_model, encoder_optimizer, decoder_optimizer, hidden_size, layers, epoch,
# file_name
save_best_model(
transducer_model, model_optimizer, hidden_size, gru_layers,
epoch,
model_dir + "_epoch_" + str(epoch + 1) + '_checkpoint.pkl')
test_loss_logger.write(
str(epoch + 1) + "," + str(stats['loss']) + "," +
str(stats['accuracy']) + "\n")
confusion_matrix_logger.write(
str(epoch + 1) + "\n" +
str(stats_dictioanry['confusion_matrix']) + "\n")
train_loss_logger.flush()
test_loss_logger.flush()
confusion_matrix_logger.flush()
else:
# this setup is for hyperband
if epoch + 1 >= 10 and stats['accuracy'] < 98:
sys.stderr.write(
TextColor.PURPLE +
'EARLY STOPPING AS THE MODEL NOT DOING WELL\n' +
TextColor.END)
return transducer_model, model_optimizer, stats
sys.stderr.write(TextColor.PURPLE + 'Finished training\n' + TextColor.END)
return transducer_model, model_optimizer, stats
def predict(test_file, output_filename, model_path, batch_size, threads, num_workers, gpu_mode):
"""
Create a prediction table/dictionary of an images set using a trained model.
:param test_file: File to predict on
:param batch_size: Batch size used for prediction
:param model_path: Path to a trained model
:param gpu_mode: If true, predictions will be done over GPU
:param threads: Number of threads to set for pytorch
:param num_workers: Number of workers to be used by the dataloader
:return: Prediction dictionary
"""
prediction_data_file = DataStore(output_filename, mode='w')
sys.stderr.write(TextColor.PURPLE + 'Loading data\n' + TextColor.END)
torch.set_num_threads(threads)
sys.stderr.write(TextColor.GREEN + 'INFO: TORCH THREADS SET TO: ' + str(torch.get_num_threads()) + ".\n"
+ TextColor.END)
sys.stderr.flush()
# data loader
test_data = SequenceDataset(test_file, load_labels=True)
test_loader = DataLoader(test_data,
batch_size=batch_size,
shuffle=False,
num_workers=num_workers)
transducer_model, hidden_size, gru_layers, prev_ite = \
ModelHandler.load_simple_model_for_training(model_path,
input_channels=ImageSizeOptions.IMAGE_CHANNELS,
image_features=ImageSizeOptions.IMAGE_HEIGHT,
seq_len=ImageSizeOptions.SEQ_LENGTH,
num_classes=ImageSizeOptions.TOTAL_LABELS)
transducer_model.eval()
if gpu_mode:
transducer_model = torch.nn.DataParallel(transducer_model).cuda()
sys.stderr.write(TextColor.CYAN + 'MODEL LOADED\n')
with torch.no_grad():
for contig, contig_start, contig_end, chunk_id, images, position, index, ref_seq, coverage, labels in tqdm(test_loader, ncols=50):
sys.stderr.flush()
images = images.type(torch.FloatTensor)
if gpu_mode:
images = images.cuda()
prediction_base_counter_h1 = np.zeros((images.size(0), ImageSizeOptions.SEQ_LENGTH, ImageSizeOptions.TOTAL_LABELS))
prediction_base_counter_h2 = np.zeros((images.size(0), ImageSizeOptions.SEQ_LENGTH, ImageSizeOptions.TOTAL_LABELS))
prediction_base_probs_h1 = np.zeros((images.size(0), ImageSizeOptions.SEQ_LENGTH, ImageSizeOptions.TOTAL_LABELS))
prediction_base_probs_h2 = np.zeros((images.size(0), ImageSizeOptions.SEQ_LENGTH, ImageSizeOptions.TOTAL_LABELS))
for i in range(0, ImageSizeOptions.SEQ_LENGTH, TrainOptions.WINDOW_JUMP):
if i + TrainOptions.TRAIN_WINDOW > ImageSizeOptions.SEQ_LENGTH:
break
chunk_start = i
chunk_end = i + TrainOptions.TRAIN_WINDOW
# chunk all the data
label_chunk_h1 = labels[:, 0, i:i+TrainOptions.TRAIN_WINDOW]
label_chunk_h2 = labels[:, 1, i:i+TrainOptions.TRAIN_WINDOW]
predicted_base_label_h1 = label_chunk_h1.numpy().tolist()
predicted_base_label_h2 = label_chunk_h2.numpy().tolist()
for ii in range(0, len(predicted_base_label_h1)):
chunk_pos = chunk_start
for base in predicted_base_label_h1[ii]:
prediction_base_counter_h1[ii][chunk_pos][base] += 1
chunk_pos += 1
for ii in range(0, len(predicted_base_label_h2)):
chunk_pos = chunk_start
for base in predicted_base_label_h2[ii]:
prediction_base_counter_h2[ii][chunk_pos][base] += 1
chunk_pos += 1
predicted_base_labels_h1 = np.argmax(np.array(prediction_base_counter_h1), axis=2)
predicted_base_labels_h2 = np.argmax(np.array(prediction_base_counter_h2), axis=2)
for i in range(images.size(0)):
prediction_data_file.write_prediction(contig[i], contig_start[i], contig_end[i], chunk_id[i],
position[i], index[i],
ref_seq[i],
coverage[i],
prediction_base_probs_h1[i],
prediction_base_probs_h2[i],
predicted_base_labels_h1[i],
predicted_base_labels_h2[i])
def predict(test_file, output_filename, model_path, batch_size, num_workers,
gpu_mode):
"""
Create a prediction table/dictionary of an images set using a trained model.
:param test_file: File to predict on
:param batch_size: Batch size used for prediction
:param model_path: Path to a trained model
:param gpu_mode: If true, predictions will be done over GPU
:param num_workers: Number of workers to be used by the dataloader
:return: Prediction dictionary
"""
prediction_data_file = DataStore(output_filename, mode='w')
sys.stderr.write(TextColor.PURPLE + 'Loading data\n' + TextColor.END)
# data loader
test_data = SequenceDataset(test_file)
test_loader = DataLoader(test_data,
batch_size=batch_size,
shuffle=False,
num_workers=num_workers)
transducer_model, hidden_size, gru_layers, prev_ite = \
ModelHandler.load_simple_model_for_training(model_path,
input_channels=ImageSizeOptions.IMAGE_CHANNELS,
image_features=ImageSizeOptions.IMAGE_HEIGHT,
seq_len=ImageSizeOptions.SEQ_LENGTH,
num_classes=ImageSizeOptions.TOTAL_LABELS)
transducer_model.eval()
if gpu_mode:
transducer_model = torch.nn.DataParallel(transducer_model).cuda()
sys.stderr.write(TextColor.CYAN + 'MODEL LOADED\n')
with torch.no_grad():
for contig, contig_start, contig_end, chunk_id, images, position, index in tqdm(
test_loader, ncols=50):
images = images.type(torch.FloatTensor)
if gpu_mode:
# encoder_hidden = encoder_hidden.cuda()
images = images.cuda()
hidden = torch.zeros(images.size(0), 2 * TrainOptions.GRU_LAYERS,
TrainOptions.HIDDEN_SIZE)
if gpu_mode:
hidden = hidden.cuda()
prediction_base_dict = np.zeros((images.size(0), images.size(1),
ImageSizeOptions.TOTAL_LABELS))
for i in range(0, ImageSizeOptions.SEQ_LENGTH,
TrainOptions.WINDOW_JUMP):
if i + TrainOptions.TRAIN_WINDOW > ImageSizeOptions.SEQ_LENGTH:
break
chunk_start = i
chunk_end = i + TrainOptions.TRAIN_WINDOW
# chunk all the data
image_chunk = images[:, chunk_start:chunk_end]
# run inference
output_base, hidden = transducer_model(image_chunk, hidden)
# do softmax and get prediction
m = nn.Softmax(dim=2)
soft_probs = m(output_base)
output_preds = soft_probs.cpu()
base_max_value, predicted_base_label = torch.max(output_preds,
dim=2)
# convert everything to list
base_max_value = base_max_value.numpy().tolist()
predicted_base_label = predicted_base_label.numpy().tolist()
assert (len(base_max_value) == len(predicted_base_label))
for ii in range(0, len(predicted_base_label)):
chunk_pos = chunk_start
for p_base, base in zip(base_max_value[ii],
predicted_base_label[ii]):
prediction_base_dict[ii][chunk_pos][base] += p_base
chunk_pos += 1
predicted_base_labels = np.argmax(np.array(prediction_base_dict),
axis=2)
for i in range(images.size(0)):
prediction_data_file.write_prediction(
contig[i], contig_start[i], contig_end[i], chunk_id[i],
position[i], index[i], predicted_base_labels[i])
def train(train_file, test_file, batch_size, epoch_limit, gpu_mode,
num_workers, retrain_model, retrain_model_path, gru_layers,
hidden_size, encoder_lr, encoder_decay, decoder_lr, decoder_decay,
model_dir, stats_dir, train_mode):
if train_mode is True:
train_loss_logger = open(stats_dir + "train_loss.csv", 'w')
test_loss_logger = open(stats_dir + "test_loss.csv", 'w')
confusion_matrix_logger = open(stats_dir + "confusion_matrix.txt", 'w')
else:
train_loss_logger = None
test_loss_logger = None
confusion_matrix_logger = None
sys.stderr.write(TextColor.PURPLE + 'Loading data\n' + TextColor.END)
train_data_set = SequenceDataset(train_file)
train_loader = DataLoader(train_data_set,
batch_size=batch_size,
shuffle=True,
num_workers=num_workers,
pin_memory=gpu_mode)
if retrain_model is True:
if os.path.isfile(retrain_model_path) is False:
sys.stderr.write(
TextColor.RED +
"ERROR: INVALID PATH TO RETRAIN PATH MODEL --retrain_model_path\n"
)
exit(1)
sys.stderr.write(TextColor.GREEN + "INFO: RETRAIN MODEL LOADING\n" +
TextColor.END)
encoder_model, decoder_model, hidden_size, gru_layers, prev_ite = \
ModelHandler.load_model_for_training(retrain_model_path,
input_channels=5,
seq_len=ImageSizeOptions.SEQ_LENGTH,
num_classes=3)
if train_mode is True:
epoch_limit = prev_ite + epoch_limit
sys.stderr.write(TextColor.GREEN + "INFO: RETRAIN MODEL LOADED\n" +
TextColor.END)
else:
encoder_model, decoder_model = ModelHandler.get_new_model(
input_channels=5,
gru_layers=gru_layers,
hidden_size=hidden_size,
seq_len=ImageSizeOptions.SEQ_LENGTH,
num_classes=3)
prev_ite = 0
encoder_optimizer = torch.optim.Adam(encoder_model.parameters(),
lr=encoder_lr,
weight_decay=encoder_decay)
decoder_optimizer = torch.optim.Adam(decoder_model.parameters(),
lr=decoder_lr,
weight_decay=decoder_decay)
if retrain_model is True:
sys.stderr.write(TextColor.GREEN + "INFO: OPTIMIZER LOADING\n" +
TextColor.END)
encoder_optimizer, decoder_optimizer = ModelHandler.load_optimizer(
encoder_optimizer, decoder_optimizer, retrain_model_path, gpu_mode)
sys.stderr.write(TextColor.GREEN + "INFO: OPTIMIZER LOADED\n" +
TextColor.END)
if gpu_mode:
encoder_model = torch.nn.DataParallel(encoder_model).cuda()
decoder_model = torch.nn.DataParallel(decoder_model).cuda()
class_weights = torch.FloatTensor(CLASS_WEIGHTS)
# Loss
criterion = nn.CrossEntropyLoss(weight=class_weights)
if gpu_mode is True:
criterion = criterion.cuda()
start_epoch = prev_ite
# Train the Model
sys.stderr.write(TextColor.PURPLE + 'Training starting\n' + TextColor.END)
stats = dict()
stats['loss_epoch'] = []
stats['accuracy_epoch'] = []
sys.stderr.write(TextColor.PURPLE + 'Start: ' + str(start_epoch + 1) +
' End: ' + str(epoch_limit + 1) + "\n")
for epoch in range(start_epoch, epoch_limit, 1):
total_loss = 0
total_images = 0
sys.stderr.write(TextColor.BLUE + 'Train epoch: ' + str(epoch + 1) +
"\n")
# make sure the model is in train mode. BN is different in train and eval.
encoder_model.train()
decoder_model.train()
batch_no = 1
with tqdm(total=len(train_loader), desc='Loss', leave=True,
ncols=100) as progress_bar:
for images, labels in train_loader:
# print(images.size(), labels.size())
# from modules.python.helper.tensor_analyzer import analyze_tensor
# for label in labels[0].data:
# print(label.item(), end='')
# print()
# analyze_tensor(images[0])
# exit()
if gpu_mode:
# encoder_hidden = encoder_hidden.cuda()
images = images.cuda()
labels = labels.cuda()
encoder_hidden = torch.FloatTensor(images.size(0),
gru_layers * 2,
hidden_size).zero_()
if gpu_mode:
encoder_hidden = encoder_hidden.cuda()
encoder_optimizer.zero_grad()
decoder_optimizer.zero_grad()
loss = 0
total_seq_length = images.size(2)
start_index = 0
end_index = images.size(2)
# from analysis.analyze_png_img import analyze_tensor
# print(labels[0, :].data.numpy())
# analyze_tensor(images[0, :, :, :])
# exit()
context_vector, hidden_encoder = encoder_model(
images, encoder_hidden)
for seq_index in range(start_index, end_index):
current_batch_size = images.size(0)
y = labels[:, seq_index - start_index]
attention_index = torch.from_numpy(
np.asarray([seq_index] * current_batch_size)).view(
-1, 1)
attention_index_onehot = torch.FloatTensor(
current_batch_size, total_seq_length)
attention_index_onehot.zero_()
attention_index_onehot.scatter_(1, attention_index, 1)
# print("\n", seq_index, attention_index_onehot)
# exit()
output_dec, decoder_hidden, attn = decoder_model(
attention_index_onehot,
context_vector=context_vector,
encoder_hidden=hidden_encoder)
# loss
loss += criterion(output_dec, y)
loss.backward()
encoder_optimizer.step()
decoder_optimizer.step()
total_loss += loss.item()
total_images += labels.size(0)
# update the progress bar
avg_loss = (total_loss / total_images) if total_images else 0
progress_bar.set_description("Loss: " + str(avg_loss))
if train_mode is True:
train_loss_logger.write(
str(epoch + 1) + "," + str(batch_no) + "," +
str(avg_loss) + "\n")
progress_bar.refresh()
progress_bar.update(1)
batch_no += 1
progress_bar.close()
stats_dictioanry = test(test_file,
batch_size,
gpu_mode,
encoder_model,
decoder_model,
num_workers,
gru_layers,
hidden_size,
num_classes=3)
stats['loss'] = stats_dictioanry['loss']
stats['accuracy'] = stats_dictioanry['accuracy']
stats['loss_epoch'].append((epoch, stats_dictioanry['loss']))
stats['accuracy_epoch'].append((epoch, stats_dictioanry['accuracy']))
# update the loggers
if train_mode is True:
# save the model after each epoch
# encoder_model, decoder_model, encoder_optimizer, decoder_optimizer, hidden_size, layers, epoch,
# file_name
save_best_model(
encoder_model, decoder_model, encoder_optimizer,
decoder_optimizer, hidden_size, gru_layers, epoch,
model_dir + "_epoch_" + str(epoch + 1) + '_checkpoint.pkl')
test_loss_logger.write(
str(epoch + 1) + "," + str(stats['loss']) + "," +
str(stats['accuracy']) + "\n")
confusion_matrix_logger.write(
str(epoch + 1) + "\n" +
str(stats_dictioanry['confusion_matrix']) + "\n")
train_loss_logger.flush()
test_loss_logger.flush()
confusion_matrix_logger.flush()
else:
# this setup is for hyperband
if epoch + 1 >= 2 and stats['accuracy'] < 90:
sys.stderr.write(
TextColor.PURPLE +
'EARLY STOPPING AS THE MODEL NOT DOING WELL\n' +
TextColor.END)
return encoder_model, decoder_model, encoder_optimizer, decoder_optimizer, stats
sys.stderr.write(TextColor.PURPLE + 'Finished training\n' + TextColor.END)
return encoder_model, decoder_model, encoder_optimizer, decoder_optimizer, stats
def predict(test_file, output_filename, model_path, batch_size, num_workers,
threads, gpu_mode):
"""
The predict method loads images generated by MarginPolish and produces base predictions using a
sequence transduction model based deep neural network. This method loads the model and iterates over
minibatch images to generate the predictions and saves the predictions to a hdf5 file.
:param test_file: File to predict on
:param output_filename: Name and path to the output file
:param batch_size: Batch size used for minibatch prediction
:param model_path: Path to a trained model
:param gpu_mode: If true, predictions will be done over GPU
:param num_workers: Number of workers to be used by the dataloader
:param threads: Number of threads to use with pytorch
:return: Prediction dictionary
"""
# create the output hdf5 file where all the predictions will be saved
prediction_data_file = DataStore(output_filename, mode='w')
torch.set_num_threads(threads)
sys.stderr.write(TextColor.GREEN + 'INFO: TORCH THREADS SET TO: ' +
str(torch.get_num_threads()) + ".\n" + TextColor.END)
# notify that the process has started and loading data
sys.stderr.write(TextColor.PURPLE + 'Loading data\n' + TextColor.END)
# create a pytorch dataset and dataloader that loads the data in mini_batches
test_data = SequenceDataset(test_file)
test_loader = DataLoader(test_data,
batch_size=batch_size,
shuffle=False,
num_workers=num_workers)
# load the model using the model path
transducer_model, hidden_size, gru_layers, prev_ite = \
ModelHandler.load_simple_model(model_path,
input_channels=ImageSizeOptions.IMAGE_CHANNELS,
image_features=ImageSizeOptions.IMAGE_HEIGHT,
seq_len=ImageSizeOptions.SEQ_LENGTH,
num_base_classes=ImageSizeOptions.TOTAL_BASE_LABELS,
num_rle_classes=ImageSizeOptions.TOTAL_RLE_LABELS)
# set the model to evaluation mode.
transducer_model.eval()
# if gpu mode is True, then load the model in the GPUs
if gpu_mode:
transducer_model = torch.nn.DataParallel(transducer_model).cuda()
# notify that the model has loaded successfully
sys.stderr.write(TextColor.CYAN + 'MODEL LOADED\n')
# iterate over the data in minibatches
with torch.no_grad():
# the dataloader loop, iterates in minibatches. tqdm is the progress logger.
for contig, contig_start, contig_end, chunk_id, images, position, filename in tqdm(
test_loader, ncols=50):
# the images are usually in uint8, convert them to FloatTensor
images = images.type(torch.FloatTensor)
# initialize the first hidden input as all zeros
hidden = torch.zeros(images.size(0), 2 * TrainOptions.GRU_LAYERS,
TrainOptions.HIDDEN_SIZE)
# if gpu_mode is True, transfer the image and hidden tensors to the GPU
if gpu_mode:
images = images.cuda()
hidden = hidden.cuda()
# this is a multi-task neural network where we predict a base and a run-length. We use two dictionaries
# to keep track of predictions.
# these two dictionaries save predictions for each of the chunks and later we aggregate all the predictions
# over the entire sequence to get a sequence prediction for the whole sequence.
prediction_base_tensor = torch.zeros(
(images.size(0), images.size(1),
ImageSizeOptions.TOTAL_BASE_LABELS))
prediction_rle_tensor = torch.zeros(
(images.size(0), images.size(1),
ImageSizeOptions.TOTAL_RLE_LABELS))
if gpu_mode:
prediction_base_tensor = prediction_base_tensor.cuda()
prediction_rle_tensor = prediction_rle_tensor.cuda()
# now the images usually contain 1000 bases, we iterate on a sliding window basis where we process
# the window size then jump to the next window
for i in range(0, ImageSizeOptions.SEQ_LENGTH,
TrainOptions.WINDOW_JUMP):
# if current position + window size goes beyond the size of the window, that means we've reached the end
if i + TrainOptions.TRAIN_WINDOW > ImageSizeOptions.SEQ_LENGTH:
break
chunk_start = i
chunk_end = i + TrainOptions.TRAIN_WINDOW
# get the image chunk
image_chunk = images[:, chunk_start:chunk_end]
# run inference
output_base, output_rle, hidden = transducer_model(
image_chunk, hidden)
# now calculate how much padding is on the top and bottom of this chunk so we can do a simple
# add operation
top_zeros = chunk_start
bottom_zeros = ImageSizeOptions.SEQ_LENGTH - chunk_end
# we run a softmax a padding to make the output tensor compatible for adding
inference_layers = nn.Sequential(
nn.Softmax(dim=2),
nn.ZeroPad2d((0, 0, top_zeros, bottom_zeros)))
if gpu_mode:
inference_layers = inference_layers.cuda()
# run the softmax and padding layers
base_prediction = inference_layers(output_base)
rle_prediction = inference_layers(output_rle)
# now simply add the tensor to the global counter
prediction_base_tensor = torch.add(prediction_base_tensor,
base_prediction)
prediction_rle_tensor = torch.add(prediction_rle_tensor,
rle_prediction)
# all done now create a SEQ_LENGTH long prediction list
prediction_base_tensor = prediction_base_tensor.cpu()
prediction_rle_tensor = prediction_rle_tensor.cpu()
base_values, base_labels = torch.max(prediction_base_tensor, 2)
rle_values, rle_labels = torch.max(prediction_rle_tensor, 2)
predicted_base_labels = base_labels.cpu().numpy()
predicted_rle_labels = rle_labels.cpu().numpy()
# go to each of the images and save the predictions to the file
for i in range(images.size(0)):
prediction_data_file.write_prediction(
contig[i], contig_start[i], contig_end[i], chunk_id[i],
position[i], predicted_base_labels[i],
predicted_rle_labels[i], filename[i])
def predict(test_file, output_filename, model_path, batch_size, threads,
num_workers, gpu_mode):
"""
Create a prediction table/dictionary of an images set using a trained model.
:param test_file: File to predict on
:param batch_size: Batch size used for prediction
:param model_path: Path to a trained model
:param gpu_mode: If true, predictions will be done over GPU
:param threads: Number of threads to set for pytorch
:param num_workers: Number of workers to be used by the dataloader
:return: Prediction dictionary
"""
prediction_data_file = DataStore(output_filename, mode='w')
sys.stderr.write(TextColor.PURPLE + 'Loading data\n' + TextColor.END)
torch.set_num_threads(threads)
sys.stderr.write(TextColor.GREEN + 'INFO: TORCH THREADS SET TO: ' +
str(torch.get_num_threads()) + ".\n" + TextColor.END)
sys.stderr.flush()
# data loader
test_data = SequenceDataset(test_file)
test_loader = DataLoader(test_data,
batch_size=batch_size,
shuffle=False,
num_workers=num_workers)
transducer_model, hidden_size, gru_layers, prev_ite = \
ModelHandler.load_simple_model_for_training(model_path,
input_channels=ImageSizeOptions.IMAGE_CHANNELS,
image_features=ImageSizeOptions.IMAGE_HEIGHT,
seq_len=ImageSizeOptions.SEQ_LENGTH,
num_classes=ImageSizeOptions.TOTAL_LABELS)
transducer_model.eval()
if gpu_mode:
transducer_model = torch.nn.DataParallel(transducer_model).cuda()
sys.stderr.write(TextColor.CYAN + 'MODEL LOADED\n')
with torch.no_grad():
for contig, contig_start, contig_end, chunk_id, images, position, index, ref_seq, labels in tqdm(
test_loader, ncols=50):
sys.stderr.flush()
images = images.type(torch.FloatTensor)
hidden = torch.zeros(images.size(0), 2 * TrainOptions.LSTM_LAYERS,
TrainOptions.HIDDEN_SIZE)
prediction_base_counter = torch.zeros(
(images.size(0), ImageSizeOptions.SEQ_LENGTH,
ImageSizeOptions.TOTAL_LABELS))
if gpu_mode:
images = images.cuda()
hidden = hidden.cuda()
prediction_base_counter = prediction_base_counter.cuda()
for i in range(0, ImageSizeOptions.SEQ_LENGTH,
TrainOptions.WINDOW_JUMP):
if i + TrainOptions.TRAIN_WINDOW > ImageSizeOptions.SEQ_LENGTH:
break
chunk_start = i
chunk_end = i + TrainOptions.TRAIN_WINDOW
# chunk all the data
image_chunk = images[:, i:i + TrainOptions.TRAIN_WINDOW]
# run inference
out, hidden = transducer_model(image_chunk, hidden)
# now calculate how much padding is on the top and bottom of this chunk so we can do a simple
# add operation
top_zeros = chunk_start
bottom_zeros = ImageSizeOptions.SEQ_LENGTH - chunk_end
# do softmax and get prediction
# we run a softmax a padding to make the output tensor compatible for adding
inference_layers = nn.Sequential(
nn.Softmax(dim=2),
nn.ZeroPad2d((0, 0, top_zeros, bottom_zeros)))
if gpu_mode:
inference_layers = inference_layers.cuda()
# run the softmax and padding layers
if gpu_mode:
base_prediction = (inference_layers(out) * 10).type(
torch.IntTensor).cuda()
else:
base_prediction = (inference_layers(out) * 10).type(
torch.IntTensor)
# now simply add the tensor to the global counter
prediction_base_counter = torch.add(prediction_base_counter,
base_prediction)
base_values, base_labels = torch.max(prediction_base_counter, 2)
predicted_base_labels = base_labels.cpu().numpy()
for i in range(images.size(0)):
prediction_data_file.write_prediction(contig[i],
contig_start[i],
contig_end[i],
chunk_id[i], position[i],
index[i], ref_seq[i],
predicted_base_labels[i])