def test(dataloader, loss_criterian, model): """ Function to test :param dataloader: Pytorch dataloader :param loss_criterian: Loss function with OHNM using MSE Loss :param model: Pytorch model of UNet-ResNet :return: all iteration loss values """ with torch.no_grad(): # For no gradient calculation model.eval() iterator = tqdm(dataloader) all_loss = [] all_accuracy = [] for no, (image, weight, weight_affinity) in enumerate(iterator): if config.use_cuda: image, weight, weight_affinity = image.cuda(), weight.cuda(), weight_affinity.cuda() output = model(image) loss = loss_criterian(output, weight, weight_affinity).mean() all_loss.append(loss.item()) if type(output) == list: output = torch.cat(output, dim=0) predicted_bbox = generate_word_bbox_batch( output[:, 0, :, :].data.cpu().numpy(), output[:, 1, :, :].data.cpu().numpy(), character_threshold=config.threshold_character, affinity_threshold=config.threshold_affinity) target_bbox = generate_word_bbox_batch( weight.data.cpu().numpy(), weight_affinity.data.cpu().numpy(), character_threshold=config.threshold_character, affinity_threshold=config.threshold_affinity) all_accuracy.append(calculate_batch_fscore(predicted_bbox, target_bbox, threshold=config.threshold_fscore)) iterator.set_description( 'Loss:' + str(int(loss.item() * 100000000) / 100000000) + ' Iterations:[' + str(no) + '/' + str( len(iterator)) + '] Average Loss:' + str(int(np.array(all_loss)[-min(1000, len(all_loss)):].mean()*100000000)/100000000) + '| Average F-Score: ' + str(int(np.array(all_accuracy)[-min(1000, len(all_accuracy)):].mean()*100000000)/100000000) ) if no % config.periodic_output == 0 and no != 0: if type(output) == list: output = torch.cat(output, dim=0) save(image, output, weight, weight_affinity, no) return all_loss
def train(dataloader, loss_criterian, model, optimizer, starting_no, all_loss, all_accuracy):
"""
Function to test
:param dataloader: Pytorch dataloader
:param loss_criterian: Loss function with OHNM using MSE Loss
:param model: Pytorch model of UNet-ResNet
:param optimizer: Adam Optimizer
:param starting_no: how many items to skip in the dataloader
:param all_loss: list of all loss values
:param all_accuracy: list of all f-scores
:return: all iteration loss values
"""
model.train()
optimizer.zero_grad()
iterator = tqdm(dataloader)
def change_lr(no_i):
for i in config.lr:
if i == no_i:
print('Learning Rate Changed to ', config.lr[i])
for param_group in optimizer.param_groups:
param_group['lr'] = config.lr[i]
for no, (image, weight, weight_affinity) in enumerate(iterator):
change_lr(no)
if config.pretrained:
if no == starting_no:
dataloader.start = True
continue
elif no < starting_no:
continue
if config.use_cuda:
image, weight, weight_affinity = image.cuda(), weight.cuda(), weight_affinity.cuda()
output = model(image)
loss = loss_criterian(output, weight, weight_affinity).mean()/4
all_loss.append(loss.item()*config.optimizer_iteration)
loss.backward()
if (no + 1) % config.optimizer_iteration == 0:
optimizer.step()
optimizer.zero_grad()
if no >= 0:
# Calculating the f-score after some iterations because initially there are a lot of stray contours
if no % config.periodic_fscore == 0:
if type(output) == list:
output = torch.cat(output, dim=0)
predicted_bbox = generate_word_bbox_batch(
output[:, 0, :, :].data.cpu().numpy(),
output[:, 1, :, :].data.cpu().numpy(),
character_threshold=config.threshold_character,
affinity_threshold=config.threshold_affinity,
word_threshold=config.threshold_word,
)
target_bbox = generate_word_bbox_batch(
weight.data.cpu().numpy(),
weight_affinity.data.cpu().numpy(),
character_threshold=config.threshold_character,
affinity_threshold=config.threshold_affinity,
word_threshold=config.threshold_word,
)
all_accuracy.append(
calculate_batch_fscore(
predicted_bbox, target_bbox, threshold=config.threshold_fscore, text_target=None
)
)
if len(all_accuracy) == 0:
iterator.set_description(
'Loss:' + str(
int(loss.item() * config.optimizer_iteration * 100000000) / 100000000) +
' Iterations:[' + str(no) + '/' + str(
len(iterator)) +
'] Average Loss:' + str(int(np.array(all_loss)[-min(1000, len(all_loss)):].mean() * 100000000) / 100000000))
else:
iterator.set_description(
'Loss:' + str(int(loss.item() * config.optimizer_iteration * 100000000) / 100000000) + ' Iterations:[' +
str(no) + '/' + str(len(iterator)) +
'] Average Loss:' + str(int(np.array(all_loss)[-min(1000, len(all_loss)):].mean()*100000000)/100000000) +
'| Average F-Score: ' + str(int(np.array(all_accuracy)[-min(1000, len(all_accuracy)):].mean()*100000000)/100000000)
)
if no % config.periodic_output == 0:
if type(output) == list:
output = torch.cat(output, dim=0)
save(image, output, weight, weight_affinity, no)
if no % config.periodic_save == 0:
torch.save(
{
'state_dict': model.state_dict(),
'optimizer': optimizer.state_dict()
}, config.save_path + '/' + str(no) + '_model.pkl')
np.save(config.save_path + '/loss_plot_training.npy', all_loss)
plt.plot(all_loss)
plt.savefig(config.save_path + '/loss_plot_training.png')
plt.clf()
return all_loss
def train(model, optimizer, iteration):
"""
Train the weak-supervised model iteratively
:param model: Pre-trained model on SynthText
:param optimizer: Pre-trained model's optimizer
:param iteration: current iteration of weak-supervision
:return: model, optimizer
"""
dataloader = DataLoader(DataLoaderMIX('train', iteration),
batch_size=config.batch_size['train'],
num_workers=8)
loss_criterian = DataParallelCriterion(Criterian())
model.train()
optimizer.zero_grad()
iterator = tqdm(dataloader)
"""
Currently not changing the learning rate while weak supervision
def change_lr(no):
# Change learning rate while training
for i in config.lr:
if i == no:
print('Learning Rate Changed to ', config.lr[i])
for param_group in optimizer.param_groups:
param_group['lr'] = config.lr[i]
change_lr(1)
"""
all_loss = []
all_accuracy = []
for no, (image, character_map, affinity_map, character_weight,
affinity_weight) in enumerate(iterator):
if config.use_cuda:
image, character_map, affinity_map = image.cuda(
), character_map.cuda(), affinity_map.cuda()
character_weight, affinity_weight = character_weight.cuda(
), affinity_weight.cuda()
output = model(image)
loss = loss_criterian(output, character_map, affinity_map,
character_weight, affinity_weight).mean()
all_loss.append(loss.item())
loss.backward()
optimizer.step()
optimizer.zero_grad()
if len(all_accuracy) == 0:
iterator.set_description(
'Loss:' + str(int(loss.item() * 100000000) / 100000000) +
' Iterations:[' + str(no) + '/' + str(len(iterator)) +
'] Average Loss:' + str(
int(
np.array(all_loss)[-min(1000, len(all_loss)):].mean() *
100000000) / 100000000))
else:
iterator.set_description(
'Loss:' + str(int(loss.item() * 100000000) / 100000000) +
' Iterations:[' + str(no) + '/' + str(len(iterator)) +
'] Average Loss:' + str(
int(
np.array(all_loss)[-min(1000, len(all_loss)):].mean() *
100000000) / 100000000) + '| Average F-Score: ' +
str(
int(
np.array(all_accuracy)[-min(1000, len(all_accuracy)):].
mean() * 100000000) / 100000000))
# ---------- Calculating the F-score ------------ #
if type(output) == list:
output = torch.cat(output, dim=0)
predicted_bbox = generate_word_bbox_batch(
output[:, 0, :, :].data.cpu().numpy(),
output[:, 1, :, :].data.cpu().numpy(),
character_threshold=config.threshold_character,
affinity_threshold=config.threshold_affinity)
target_bbox = generate_word_bbox_batch(
character_map.data.cpu().numpy(),
affinity_map.data.cpu().numpy(),
character_threshold=config.threshold_character,
affinity_threshold=config.threshold_affinity)
all_accuracy.append(
calculate_batch_fscore(predicted_bbox,
target_bbox,
threshold=config.threshold_fscore))
torch.cuda.empty_cache()
return model, optimizer