def test_loop_fn(loader,
model,
criterion,
device,
epoch,
decoder,
alphabet):
running_loss = 0.0
total_words = 0
cumulative_wer = 0
iteration = 0
model.eval()
with torch.no_grad():
for inputs, input_lengths, labels, label_lengths in loader:
out = model(inputs)
loss = criterion(out, labels, input_lengths, label_lengths)
iteration += 1
running_loss += loss.item()
wers, n_words = compute_wer(out, labels, decoder, alphabet, print_output=True)
cumulative_wer += wers
total_words += n_words
avg_loss = running_loss / iteration
avg_wer = cumulative_wer / total_words
logging.info('[Val][%s] Loss=%.5f WER=%.3f Time=%s',
epoch, avg_loss, avg_wer, time.asctime())
return avg_loss
def _shared_eval(self, batch, batch_idx):
# prepare decode
images, labels, image_padding_mask, lengths = batch
# decode
pred = self.decode(images, self.beam_width, image_padding_mask)
tgt, tgt_len = labels[:, 1:], lengths - 2
# convert to strings
predicts = self.ctc_string_tf(pred)
groundtruth = self.string_tf(tgt, tgt_len)
cer_distances, num_chars = compute_cer(predicts,
groundtruth,
indistinguish=False)
wer_distances, num_words = compute_wer(predicts,
groundtruth,
indistinguish=False)
return {
'cer_distances': cer_distances,
'num_chars': num_chars,
'wer_distances': wer_distances,
'num_words': num_words
}
def train_loop_fn(loader,
optimizer,
model,
criterion,
device,
epoch,
decoder,
alphabet,
log_steps):
running_loss = 0.0
iteration = 0
model.train()
for step, (inputs, input_lengths, labels, label_lengths) in enumerate(loader, 1):
# zero the parameter gradients
optimizer.zero_grad()
out = model(inputs)
loss = criterion(out, labels, input_lengths, label_lengths)
loss_value = loss.item()
# Check to ensure valid loss was calculated
valid_loss, error = check_loss(loss, loss_value)
if valid_loss:
loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), 400)
optimizer.step()
else:
logging.error(error)
logging.info('Skipping grad update')
loss_value = 0
iteration += 1
running_loss += loss_value
if step % log_steps == 0:
wers, n_words = compute_wer(out, labels, decoder, alphabet, print_output=True)
batch_wer = wers / n_words
logging.info('Batch WER: %.3f', batch_wer)
avg_loss = running_loss / iteration
logging.info('[Train][%s] Loss=%.5f Time=%s',
epoch, avg_loss, time.asctime())