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("[" + str(datetime.now().strftime('%m-%d-%Y %H:%M:%S')) + "] INFO: MODEL" + file_name + " SAVED SUCCESSFULLY.\n")
def do_test(test_file, batch_size, gpu_mode, num_workers, model_path):
"""
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("[" + str(datetime.now().strftime('%m-%d-%Y %H:%M:%S')) +
"] INFO: LOADING DATA\n")
if os.path.isfile(model_path) is False:
sys.stderr.write("[" +
str(datetime.now().strftime('%m-%d-%Y %H:%M:%S')) +
"] ERROR: INVALID PATH TO MODEL\n")
exit(1)
sys.stderr.write("[" + str(datetime.now().strftime('%m-%d-%Y %H:%M:%S')) +
"] INFO: MODEL LOADING\n")
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)
sys.stderr.write("[" + str(datetime.now().strftime('%m-%d-%Y %H:%M:%S')) +
"] INFO: MODEL LOADED\n")
if gpu_mode:
transducer_model = torch.nn.DataParallel(transducer_model).cuda()
stats_dictioanry = test(test_file,
batch_size,
gpu_mode,
transducer_model,
num_workers,
gru_layers,
hidden_size,
num_classes=ImageSizeOptions.TOTAL_LABELS)
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, world_size,
rank, device_id):
if train_mode is True and rank == 0:
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
torch.cuda.set_device(device_id)
if rank == 0:
sys.stderr.write("[" +
str(datetime.now().strftime('%m-%d-%Y %H:%M:%S')) +
"] INFO: LOADING DATA\n")
train_data_set = SequenceDataset(train_file)
train_sampler = torch.utils.data.distributed.DistributedSampler(
train_data_set, num_replicas=world_size, rank=rank)
train_loader = torch.utils.data.DataLoader(dataset=train_data_set,
batch_size=batch_size,
shuffle=False,
num_workers=0,
pin_memory=True,
sampler=train_sampler)
num_classes = ImageSizeOptions.TOTAL_LABELS
if retrain_model is True:
if os.path.isfile(retrain_model_path) is False:
sys.stderr.write(
"[" + str(datetime.now().strftime('%m-%d-%Y %H:%M:%S')) +
"] ERROR: INVALID PATH TO RETRAIN PATH MODEL --retrain_model_path\n"
)
exit(1)
sys.stderr.write("[" +
str(datetime.now().strftime('%m-%d-%Y %H:%M:%S')) +
"] INFO: RETRAIN MODEL LOADING\n")
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("[" +
str(datetime.now().strftime('%m-%d-%Y %H:%M:%S')) +
"] INFO: RETRAIN MODEL LOADED\n")
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)
if rank == 0:
sys.stderr.write("[" +
str(datetime.now().strftime('%m-%d-%Y %H:%M:%S')) +
"] INFO: TOTAL TRAINABLE PARAMETERS:\t" +
str(param_count) + "\n")
sys.stderr.flush()
model_optimizer = torch.optim.Adam(transducer_model.parameters(),
lr=lr,
weight_decay=decay)
if retrain_model is True:
sys.stderr.write("[" +
str(datetime.now().strftime('%m-%d-%Y %H:%M:%S')) +
"] INFO: OPTIMIZER LOADING\n")
model_optimizer = ModelHandler.load_simple_optimizer(
model_optimizer, retrain_model_path, gpu_mode)
sys.stderr.write("[" +
str(datetime.now().strftime('%m-%d-%Y %H:%M:%S')) +
"] INFO: OPTIMIZER LOADED\n")
sys.stderr.flush()
if gpu_mode:
transducer_model = transducer_model.to(device_id)
transducer_model = nn.parallel.DistributedDataParallel(
transducer_model, device_ids=[device_id])
class_weights = torch.Tensor(ImageSizeOptions.class_weights)
# Loss
criterion = nn.CrossEntropyLoss(class_weights)
if gpu_mode is True:
criterion = criterion.to(device_id)
start_epoch = prev_ite
# Train the Model
if rank == 0:
sys.stderr.write("[" +
str(datetime.now().strftime('%m-%d-%Y %H:%M:%S')) +
"] INFO: TRAINING STARTING\n")
sys.stderr.write("[" +
str(datetime.now().strftime('%m-%d-%Y %H:%M:%S')) +
"] INFO: START: " + str(start_epoch + 1) + " END: " +
str(epoch_limit) + "\n")
stats = dict()
stats['loss_epoch'] = []
stats['accuracy_epoch'] = []
for epoch in range(start_epoch, epoch_limit, 1):
start_time = time.time()
total_loss = 0
total_images = 0
if rank == 0:
sys.stderr.write(
"[" + str(datetime.now().strftime('%m-%d-%Y %H:%M:%S')) +
"] INFO: TRAIN EPOCH: " + str(epoch + 1) + "\n")
# make sure the model is in train mode. BN is different in train and eval.
batch_no = 1
transducer_model.train()
for images, labels in train_loader:
labels = labels.type(torch.LongTensor)
images = images.type(torch.FloatTensor)
if gpu_mode:
images = images.to(device_id)
labels = labels.to(device_id)
hidden = torch.zeros(images.size(0), 2 * TrainOptions.GRU_LAYERS,
TrainOptions.HIDDEN_SIZE)
if gpu_mode:
hidden = hidden.to(device_id)
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]
#############################################################################################
########################
#######################
#Just a Wild Guess that this too needs a hidden cell_state for lstm
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
if train_mode is True and rank == 0:
train_loss_logger.write(
str(epoch + 1) + "," + str(batch_no) + "," +
str(avg_loss) + "\n")
if rank == 0 and batch_no % 10 == 0:
percent_complete = int((100 * batch_no) / len(train_loader))
time_now = time.time()
mins = int((time_now - start_time) / 60)
secs = int((time_now - start_time)) % 60
sys.stderr.write(
"[" + str(datetime.now().strftime('%m-%d-%Y %H:%M:%S')) +
"] INFO: " + "EPOCH: " + str(epoch + 1) + " BATCH: " +
str(batch_no) + "/" + str(len(train_loader)) + " LOSS: " +
str(avg_loss) + " COMPLETE (" + str(percent_complete) +
"%)" + " [ELAPSED TIME: " + str(mins) + " Min " +
str(secs) + " Sec]\n")
sys.stderr.flush()
batch_no += 1
dist.barrier()
if rank == 0:
stats_dictioanry = test(test_file,
batch_size * world_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']))
dist.barrier()
if rank == 0:
sys.stderr.write(
"[" + str(datetime.now().strftime('%m-%d-%Y %H:%M:%S')) +
"] INFO: TEST COMPLETED.\n")
# update the loggers
if train_mode is True and rank == 0:
# 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()
elif train_mode is False:
# this setup is for hyperband
if epoch + 1 >= 10 and stats['accuracy'] < 98:
sys.stderr.write(
"[" + str(datetime.now().strftime('%m-%d-%Y %H:%M:%S')) +
"] INFO: EARLY STOPPING AS THE MODEL NOT DOING WELL\n")
return transducer_model, model_optimizer, stats
if rank == 0:
time_now = time.time()
mins = int((time_now - start_time) / 60)
secs = int((time_now - start_time)) % 60
sys.stderr.write(
"[" + str(datetime.now().strftime('%m-%d-%Y %H:%M:%S')) +
"] INFO: ELAPSED TIME FOR ONE EPOCH: " + str(mins) + " Min " +
str(secs) + " Sec\n")
if rank == 0:
sys.stderr.write("[" +
str(datetime.now().strftime('%m-%d-%Y %H:%M:%S')) +
"] INFO: FINISHED TRAINING\n")
return transducer_model, model_optimizer, stats
def predict(input_filepath, file_chunks, output_filepath, model_path,
batch_size, num_workers, theads_per_caller, device_id, rank):
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()
transducer_model = transducer_model.eval()
# create output file
output_filename = output_filepath + "pepper_prediction_" + str(
rank) + ".hdf"
prediction_data_file = DataStore(output_filename, mode='w')
# data loader
input_data = SequenceDataset(input_filepath, file_chunks)
data_loader = DataLoader(input_data,
batch_size=batch_size,
shuffle=False,
num_workers=num_workers)
torch.set_num_threads(theads_per_caller)
torch.cuda.set_device(device_id)
transducer_model.to(device_id)
transducer_model.eval()
transducer_model = DistributedDataParallel(transducer_model,
device_ids=[device_id])
batch_completed = 0
total_batches = len(data_loader)
with torch.no_grad():
for contig, contig_start, contig_end, chunk_id, images, position, index, ref_seq in data_loader:
sys.stderr.flush()
images = images.type(torch.FloatTensor)
hidden = torch.zeros(images.size(0), 2 * TrainOptions.GRU_LAYERS,
TrainOptions.HIDDEN_SIZE)
prediction_base_tensor = torch.zeros(
(images.size(0), images.size(1),
ImageSizeOptions.TOTAL_LABELS))
images = images.to(device_id)
hidden = hidden.to(device_id)
prediction_base_tensor = prediction_base_tensor.to(device_id)
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]
#######################################################################################################
##############
##############
##############
#This is the lstm model that generates, calls for this, make LSTM
# run inference
output_base, 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)))
inference_layers = inference_layers.to(device_id)
# run the softmax and padding layers
base_prediction = (inference_layers(output_base) * 10).type(
torch.IntTensor).to(device_id)
# now simply add the tensor to the global counter
prediction_base_tensor = torch.add(prediction_base_tensor,
base_prediction)
del inference_layers
torch.cuda.empty_cache()
prediction_base_tensor = prediction_base_tensor.cpu().numpy(
).astype(int)
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], prediction_base_tensor[i],
ref_seq[i])
batch_completed += 1
if rank == 0:
sys.stderr.write(
"[" + str(datetime.now().strftime('%m-%d-%Y %H:%M:%S')) +
"] " + "INFO: BATCHES PROCESSED " + str(batch_completed) +
"/" + str(total_batches) + ".\n")
def predict_distributed_cpu(filepath, file_chunks, output_filepath, model_path,
batch_size, callers, threads, num_workers):
"""
Create a prediction table/dictionary of an images set using a trained model.
:param filepath: Path to image files to predict on
:param file_chunks: Path to chunked files
:param batch_size: Batch size used for prediction
:param model_path: Path to a trained model
:param output_filepath: Path to output directory
:param callers: Number of callers to start
:param threads: Number of threads per caller.
:param num_workers: Number of workers to be used by the dataloader
:return: Prediction dictionary
"""
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()
sys.stderr.write("[" + str(datetime.now().strftime('%m-%d-%Y %H:%M:%S')) +
"] INFO: MODEL LOADING TO ONNX\n")
x = torch.zeros(1, TrainOptions.TRAIN_WINDOW,
ImageSizeOptions.IMAGE_HEIGHT)
h = torch.zeros(1, 2 * TrainOptions.GRU_LAYERS, TrainOptions.HIDDEN_SIZE)
if not os.path.isfile(model_path + ".onnx"):
sys.stderr.write("[" +
str(datetime.now().strftime('%m-%d-%Y %H:%M:%S')) +
"] INFO: SAVING MODEL TO ONNX\n")
torch.onnx.export(transducer_model, (x, h),
model_path + ".onnx",
training=False,
opset_version=10,
do_constant_folding=True,
input_names=['input_image', 'input_hidden'],
output_names=['output_pred', 'output_hidden'],
dynamic_axes={
'input_image': {
0: 'batch_size'
},
'input_hidden': {
0: 'batch_size'
},
'output_pred': {
0: 'batch_size'
},
'output_hidden': {
0: 'batch_size'
}
})
transducer_model.eval()
args = (filepath, output_filepath, model_path, batch_size, threads,
num_workers)
mp.spawn(setup,
args=(callers, args, file_chunks),
nprocs=callers,
join=True)
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("[" + str(datetime.now().strftime('%m-%d-%Y %H:%M:%S')) +
"] INFO: LOADING DATA\n")
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(
"[" + str(datetime.now().strftime('%m-%d-%Y %H:%M:%S')) +
"] ERROR: INVALID PATH TO RETRAIN PATH MODEL --retrain_model_path\n"
)
exit(1)
sys.stderr.write("[" +
str(datetime.now().strftime('%m-%d-%Y %H:%M:%S')) +
"] INFO: RETRAIN MODEL LOADING\n")
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("[" +
str(datetime.now().strftime('%m-%d-%Y %H:%M:%S')) +
"] INFO: RETRAIN MODEL LOADED\n")
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("[" + str(datetime.now().strftime('%m-%d-%Y %H:%M:%S')) +
"] INFO: TOTAL TRAINABLE PARAMETERS:\t" +
str(param_count) + "\n")
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("[" +
str(datetime.now().strftime('%m-%d-%Y %H:%M:%S')) +
"] INFO: OPTIMIZER LOADING\n")
model_optimizer = ModelHandler.load_simple_optimizer(
model_optimizer, retrain_model_path, gpu_mode)
sys.stderr.write("[" +
str(datetime.now().strftime('%m-%d-%Y %H:%M:%S')) +
"] INFO: OPTIMIZER LOADED\n")
if gpu_mode:
transducer_model = torch.nn.DataParallel(transducer_model).cuda()
class_weights = torch.Tensor(ImageSizeOptions.class_weights)
# Loss
criterion = nn.CrossEntropyLoss(class_weights)
if gpu_mode is True:
criterion = criterion.cuda()
start_epoch = prev_ite
# Train the Model
sys.stderr.write("[" + str(datetime.now().strftime('%m-%d-%Y %H:%M:%S')) +
"] INFO: TRAINING STARTING\n")
stats = dict()
stats['loss_epoch'] = []
stats['accuracy_epoch'] = []
sys.stderr.write("[" + str(datetime.now().strftime('%m-%d-%Y %H:%M:%S')) +
"] INFO: 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("[" +
str(datetime.now().strftime('%m-%d-%Y %H:%M:%S')) +
"] INFO: 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 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
sys.stderr.write(
"[" + str(datetime.now().strftime('%m-%d-%Y %H:%M:%S')) +
"] INFO: " + "EPOCH: " + str(epoch + 1) + " BATCH: " +
str(batch_no) + "/" + str(len(train_loader)) + " LOSS: " +
str(avg_loss) + "\n")
if train_mode is True:
train_loss_logger.write(
str(epoch + 1) + "," + str(batch_no) + "," +
str(avg_loss) + "\n")
batch_no += 1
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 + "PEPPER_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(
"[" + str(datetime.now().strftime('%m-%d-%Y %H:%M:%S')) +
"] INFO: EARLY STOPPING AS THE MODEL NOT DOING WELL\n")
return transducer_model, model_optimizer, stats
sys.stderr.write("[" + str(datetime.now().strftime('%m-%d-%Y %H:%M:%S')) +
"] INFO: FINISHED TRAINING.\n")
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('Loading data\n')
torch.set_num_threads(threads)
sys.stderr.write('INFO: TORCH THREADS SET TO: ' +
str(torch.get_num_threads()) + ".\n")
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('MODEL LOADED\n')
with torch.no_grad():
for contig, contig_start, contig_end, chunk_id, images, position, index, ref_seq in tqdm(
test_loader, ncols=50):
sys.stderr.flush()
images = images.type(torch.FloatTensor)
hidden = torch.zeros(images.size(0), 2 * TrainOptions.GRU_LAYERS,
TrainOptions.HIDDEN_SIZE)
prediction_base_tensor = torch.zeros(
(images.size(0), images.size(1),
ImageSizeOptions.TOTAL_LABELS))
if gpu_mode:
images = images.cuda()
hidden = hidden.cuda()
prediction_base_tensor = prediction_base_tensor.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[:, chunk_start:chunk_end]
# run inference
output_base, 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(output_base) *
10).type(torch.IntTensor).cuda()
else:
base_prediction = (inference_layers(output_base) *
10).type(torch.IntTensor)
# now simply add the tensor to the global counter
prediction_base_tensor = torch.add(prediction_base_tensor,
base_prediction)
prediction_base_tensor = prediction_base_tensor.cpu().numpy(
).astype(int)
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], prediction_base_tensor[i],
ref_seq[i])