def att_plot(test_set, load_dir, plot_path, use_gpu, max_seq_len, beam_width):
"""
generate attention alignment plots
Args:
test_set: test dataset
load_dir: model dir
use_gpu: on gpu/cpu
max_seq_len
Returns:
"""
# check devide
print('cuda available: {}'.format(torch.cuda.is_available()))
use_gpu = use_gpu and torch.cuda.is_available()
# load model
# latest_checkpoint_path = Checkpoint.get_latest_checkpoint(load_dir)
# latest_checkpoint_path = Checkpoint.get_thirdlast_checkpoint(load_dir)
latest_checkpoint_path = load_dir
resume_checkpoint = Checkpoint.load(latest_checkpoint_path)
model = resume_checkpoint.model.to(device)
print('Model dir: {}'.format(latest_checkpoint_path))
print('Model laoded')
# reset batch_size:
model.reset_max_seq_len(max_seq_len)
model.reset_use_gpu(use_gpu)
model.reset_batch_size(test_set.batch_size)
# in plotting mode always turn off beam search
model.set_beam_width(beam_width=0)
model.check_var('ptr_net')
print('max seq len {}'.format(model.max_seq_len))
print('ptr_net {}'.format(model.ptr_net))
# load test
if type(test_set.attkey_path) == type(None):
test_batches, vocab_size = test_set.construct_batches(is_train=False)
else:
test_batches, vocab_size = test_set.construct_batches_with_ddfd_prob(is_train=False)
# start eval
model.eval()
match = 0
total = 0
count = 0
with torch.no_grad():
for batch in test_batches:
src_ids = batch['src_word_ids']
src_lengths = batch['src_sentence_lengths']
tgt_ids = batch['tgt_word_ids']
tgt_lengths = batch['tgt_sentence_lengths']
src_probs = None
if 'src_ddfd_probs' in batch:
src_probs = batch['src_ddfd_probs']
src_probs = _convert_to_tensor(src_probs, use_gpu).unsqueeze(2)
src_ids = _convert_to_tensor(src_ids, use_gpu)
tgt_ids = _convert_to_tensor(tgt_ids, use_gpu)
decoder_outputs, decoder_hidden, other = model(src_ids, tgt_ids,
is_training=False,
att_key_feats=src_probs,
beam_width=0)
# Evaluation
# default batch_size = 1
# attention: 31 * [1 x 1 x 32] ( tgt_len(query_len) * [ batch_size x 1 x src_len(key_len)] )
attention = other['attention_score']
seqlist = other['sequence'] # traverse over time not batch
bsize = test_set.batch_size
max_seq = test_set.max_seq_len
vocab_size = len(test_set.tgt_word2id)
for idx in range(len(decoder_outputs)): # loop over max_seq
step = idx
step_output = decoder_outputs[idx] # 64 x vocab_size
# count correct
target = tgt_ids[:, step+1]
non_padding = target.ne(PAD)
correct = seqlist[step].view(-1).eq(target).masked_select(non_padding).sum().item()
match += correct
total += non_padding.sum().item()
# Print sentence by sentence
srcwords = _convert_to_words_batchfirst(src_ids, test_set.src_id2word)
refwords = _convert_to_words_batchfirst(tgt_ids[:,1:], test_set.tgt_id2word)
seqwords = _convert_to_words(seqlist, test_set.tgt_id2word)
# print(type(attention))
# print(len(attention))
# print(type(attention[0]))
# print(attention[0].size())
# input('...')
n_q = len(attention)
n_k = attention[0].size(2)
b_size = attention[0].size(0)
att_score = torch.empty(n_q, n_k, dtype=torch.float)
# att_score = np.empty([n_q, n_k])
for i in range(len(seqwords)): # loop over sentences
outline_src = ' '.join(srcwords[i])
outline_ref = ' '.join(refwords[i])
outline_gen = ' '.join(seqwords[i])
print('SRC: {}'.format(outline_src))
print('REF: {}'.format(outline_ref))
print('GEN: {}'.format(outline_gen))
for j in range(len(attention)):
# i: idx of batch
# j: idx of query
gen = seqwords[i][j]
ref = refwords[i][j]
att = attention[j][i]
# record att scores
att_score[j] = att
# print('REF:GEN - {}:{}'.format(ref,gen))
# print('{}th ATT size: {}'.format(j, attention[j][i].size()))
# print(att)
# print(torch.argmax(att))
# print(sum(sum(att)))
# input('Press enter to continue ...')
# plotting
# print(att_score)
loc_eos_k = srcwords[i].index('</s>') + 1
loc_eos_q = seqwords[i].index('</s>') + 1
loc_eos_ref = refwords[i].index('</s>') + 1
print('eos_k: {}, eos_q: {}'.format(loc_eos_k, loc_eos_q))
att_score_trim = att_score[:loc_eos_q, :loc_eos_k] # each row (each query) sum up to 1
print(att_score_trim)
print('\n')
# import pdb; pdb.set_trace()
choice = input('Plot or not ? - y/n\n')
if choice:
if choice.lower()[0] == 'y':
print('plotting ...')
plot_dir = os.path.join(plot_path, '{}.png'.format(count))
src = srcwords[i][:loc_eos_k]
hyp = seqwords[i][:loc_eos_q]
ref = refwords[i][:loc_eos_ref]
# x-axis: src; y-axis: hyp
# plot_alignment(att_score_trim.numpy(), plot_dir, src=src, hyp=hyp, ref=ref)
plot_alignment(att_score_trim.numpy(), plot_dir, src=src, hyp=hyp, ref=None) # no ref
count += 1
input('Press enter to continue ...')
if total == 0:
accuracy = float('nan')
else:
accuracy = match / total
print(saccuracy)
def translate(test_set, load_dir, test_path_out, use_gpu, max_seq_len, beam_width, seqrev=False):
"""
no reference tgt given - Run translation.
Args:
test_set: test dataset
src, tgt using the same dir
test_path_out: output dir
load_dir: model dir
use_gpu: on gpu/cpu
"""
# load model
# latest_checkpoint_path = Checkpoint.get_latest_checkpoint(load_dir)
# latest_checkpoint_path = Checkpoint.get_thirdlast_checkpoint(load_dir)
latest_checkpoint_path = load_dir
resume_checkpoint = Checkpoint.load(latest_checkpoint_path)
model = resume_checkpoint.model.to(device)
print('Model dir: {}'.format(latest_checkpoint_path))
print('Model laoded')
# reset batch_size:
model.reset_max_seq_len(max_seq_len)
model.reset_use_gpu(use_gpu)
model.reset_batch_size(test_set.batch_size)
model.check_var('ptr_net')
print('max seq len {}'.format(model.max_seq_len))
sys.stdout.flush()
# load test
if type(test_set.attkey_path) == type(None):
test_batches, vocab_size = test_set.construct_batches(is_train=False)
else:
test_batches, vocab_size = test_set.construct_batches_with_ddfd_prob(is_train=False)
# f = open(os.path.join(test_path_out, 'translate.txt'), 'w') -> use proper encoding
with open(os.path.join(test_path_out, 'translate.txt'), 'w', encoding="utf8") as f:
model.eval()
match = 0
total = 0
with torch.no_grad():
for batch in test_batches:
src_ids = batch['src_word_ids']
src_lengths = batch['src_sentence_lengths']
src_probs = None
if 'src_ddfd_probs' in batch:
src_probs = batch['src_ddfd_probs']
src_probs = _convert_to_tensor(src_probs, use_gpu).unsqueeze(2)
src_ids = _convert_to_tensor(src_ids, use_gpu)
decoder_outputs, decoder_hidden, other = model(src=src_ids,
is_training=False,
att_key_feats=src_probs,
beam_width=beam_width)
# memory usage
mem_kb, mem_mb, mem_gb = get_memory_alloc()
mem_mb = round(mem_mb, 2)
print('Memory used: {0:.2f} MB'.format(mem_mb))
# write to file
seqlist = other['sequence']
seqwords = _convert_to_words(seqlist, test_set.src_id2word)
for i in range(len(seqwords)):
# skip padding sentences in batch (num_sent % batch_size != 0)
if src_lengths[i] == 0:
continue
words = []
for word in seqwords[i]:
if word == '<pad>':
continue
elif word == '</s>':
break
else:
words.append(word)
if len(words) == 0:
outline = ''
else:
if seqrev:
words = words[::-1]
outline = ' '.join(words)
f.write('{}\n'.format(outline))
# if i == 0:
# print(outline)
sys.stdout.flush()
def evaluate(test_set, load_dir, test_path_out, use_gpu, max_seq_len, beam_width, seqrev=False):
"""
with reference tgt given - Run translation.
Args:
test_set: test dataset
test_path_out: output dir
load_dir: model dir
use_gpu: on gpu/cpu
Returns:
accuracy (excluding PAD tokens)
"""
# load model
# latest_checkpoint_path = Checkpoint.get_latest_checkpoint(load_dir)
# latest_checkpoint_path = Checkpoint.get_thirdlast_checkpoint(load_dir)
latest_checkpoint_path = load_dir
resume_checkpoint = Checkpoint.load(latest_checkpoint_path)
model = resume_checkpoint.model.to(device)
print('Model dir: {}'.format(latest_checkpoint_path))
print('Model laoded')
# reset batch_size:
model.reset_max_seq_len(max_seq_len)
model.reset_use_gpu(use_gpu)
model.reset_batch_size(test_set.batch_size)
model.set_beam_width(beam_width)
model.check_var('ptr_net')
print('max seq len {}'.format(model.max_seq_len))
sys.stdout.flush()
# load test
if type(test_set.attkey_path) == type(None):
test_batches, vocab_size = test_set.construct_batches(is_train=False)
else:
test_batches, vocab_size = test_set.construct_batches_with_ddfd_prob(is_train=False)
# f = open(os.path.join(test_path_out, 'test.txt'), 'w')
with open(os.path.join(test_path_out, 'translate.txt'), 'w', encoding="utf8") as f:
model.eval()
match = 0
total = 0
with torch.no_grad():
for batch in test_batches:
src_ids = batch['src_word_ids']
src_lengths = batch['src_sentence_lengths']
tgt_ids = batch['tgt_word_ids']
tgt_lengths = batch['tgt_sentence_lengths']
src_probs = None
if 'src_ddfd_probs' in batch:
src_probs = batch['src_ddfd_probs']
src_probs = _convert_to_tensor(src_probs, use_gpu).unsqueeze(2)
src_ids = _convert_to_tensor(src_ids, use_gpu)
tgt_ids = _convert_to_tensor(tgt_ids, use_gpu)
decoder_outputs, decoder_hidden, other = model(src_ids, tgt_ids,
is_training=False,
att_key_feats=src_probs,
beam_width=beam_width)
# Evaluation
seqlist = other['sequence'] # traverse over time not batch
if beam_width > 1:
full_seqlist = other['topk_sequence']
decoder_outputs = decoder_outputs[:-1]
for step, step_output in enumerate(decoder_outputs):
target = tgt_ids[:, step+1]
non_padding = target.ne(PAD)
correct = seqlist[step].view(-1).eq(target)
.masked_select(non_padding).sum().item()
match += correct
total += non_padding.sum().item()
# write to file
seqwords = _convert_to_words(seqlist, test_set.tgt_id2word)
for i in range(len(seqwords)):
# skip padding sentences in batch (num_sent % batch_size != 0)
if src_lengths[i] == 0:
continue
words = []
for word in seqwords[i]:
if word == '<pad>':
continue
elif word == '</s>':
break
else:
words.append(word)
if len(words) == 0:
outline = ''
else:
if seqrev:
words = words[::-1]
outline = ' '.join(words)
f.write('{}\n'.format(outline))
if i == 0:
print(outline)
sys.stdout.flush()
if total == 0:
accuracy = float('nan')
else:
accuracy = match / total
def translate(test_set, load_dir, test_path_out,
use_gpu, max_seq_len, beam_width, mode='gec', seqrev=False):
"""
no reference tgt given - Run translation.
Args:
test_set: test dataset
src, tgt using the same dir
test_path_out: output dir
load_dir: model dir
use_gpu: on gpu/cpu
"""
# load model
# latest_checkpoint_path = Checkpoint.get_latest_checkpoint(load_dir)
# latest_checkpoint_path = Checkpoint.get_thirdlast_checkpoint(load_dir)
latest_checkpoint_path = load_dir
resume_checkpoint = Checkpoint.load(latest_checkpoint_path)
model = resume_checkpoint.model
print('Model dir: {}'.format(latest_checkpoint_path))
print('Model laoded')
# reset batch_size:
model.reset_max_seq_len(max_seq_len)
model.reset_use_gpu(use_gpu)
model.reset_batch_size(test_set.batch_size)
# fix compatibility
model.check_classvar('shared_embed')
model.check_classvar('additional_key_size')
model.check_classvar('gec_num_bilstm_dec')
model.check_classvar('num_unilstm_enc')
model.check_classvar('residual')
model.check_classvar('add_discriminator')
if model.ptr_net == 'none': model.ptr_net = 'null'
model.to(device)
print('max seq len {}'.format(model.max_seq_len))
sys.stdout.flush()
# load test
print('--- constrcut gec test set ---')
if type(test_set.tsv_path) == type(None):
test_batches, vocab_size = test_set.construct_batches(is_train=False)
else:
test_batches, vocab_size = test_set.construct_batches_with_ddfd_prob(is_train=False)
with open(os.path.join(test_path_out, 'translate.txt'), 'w', encoding="utf8") as f:
model.eval()
match = 0
total = 0
with torch.no_grad():
for batch in test_batches:
src_ids = batch['src_word_ids']
src_lengths = batch['src_sentence_lengths']
src_probs = None
if 'src_ddfd_probs' in batch and model.dd_additional_key_size > 0:
src_probs = batch['src_ddfd_probs']
src_probs = _convert_to_tensor(src_probs, use_gpu).unsqueeze(2)
src_ids = _convert_to_tensor(src_ids, use_gpu)
tgt_ids = None
if mode.lower() == 'gec' and beam_width <= 1:
gec_dd_decoder_outputs, gec_dd_dec_hidden, gec_dd_ret_dict, \
decoder_outputs, dec_hidden, ret_dict = \
model.gec_eval(src_ids, tgt_ids, is_training=False,
gec_dd_att_key_feats=src_probs, beam_width=beam_width)
elif mode.lower() == 'gec' and beam_width > 1:
decoder_outputs, dec_hidden, ret_dict = \
model.gec_eval(src_ids, tgt_ids, is_training=False,
gec_dd_att_key_feats=src_probs, beam_width=beam_width)
elif mode.lower() == 'dd':
decoder_outputs, dec_hidden, ret_dict = \
model.dd_eval(src_ids, tgt_ids, is_training=False,
dd_att_key_feats=src_probs, beam_width=beam_width)
else:
assert False, 'Unrecognised eval mode - choose from gec/dd'
# memory usage
mem_kb, mem_mb, mem_gb = get_memory_alloc()
mem_mb = round(mem_mb, 2)
print('Memory used: {0:.2f} MB'.format(mem_mb))
# gec output write to file
# import pdb; pdb.set_trace()
seqlist = ret_dict['sequence']
seqwords = _convert_to_words(seqlist, test_set.src_id2word)
for i in range(len(seqwords)):
# skip padding sentences in batch (num_sent % batch_size != 0)
if src_lengths[i] == 0:
continue
words = []
for word in seqwords[i]:
if word == '<pad>':
continue
elif word == '</s>':
break
else:
words.append(word)
if len(words) == 0:
outline = ''
else:
if seqrev:
words = words[::-1]
outline = ' '.join(words)
f.write('{}\n'.format(outline))
# if i == 0:
# print(outline)
sys.stdout.flush()
def debug_beam_search(test_set, load_dir, use_gpu, max_seq_len, beam_width):
"""
with reference tgt given - debug beam search.
Args:
test_set: test dataset
load_dir: model dir
use_gpu: on gpu/cpu
Returns:
accuracy (excluding PAD tokens)
"""
# load model
# latest_checkpoint_path = Checkpoint.get_latest_checkpoint(load_dir)
# latest_checkpoint_path = Checkpoint.get_thirdlast_checkpoint(load_dir)
latest_checkpoint_path = load_dir
resume_checkpoint = Checkpoint.load(latest_checkpoint_path)
model = resume_checkpoint.model
print('Model dir: {}'.format(latest_checkpoint_path))
print('Model laoded')
# reset batch_size:
model.reset_max_seq_len(max_seq_len)
model.reset_use_gpu(use_gpu)
model.reset_batch_size(test_set.batch_size)
print('max seq len {}'.format(model.max_seq_len))
sys.stdout.flush()
# load test
if type(test_set.attkey_path) == type(None):
test_batches, vocab_size = test_set.construct_batches(is_train=False)
else:
test_batches, vocab_size = test_set.construct_batches_with_ddfd_prob(
is_train=False)
model.eval()
match = 0
total = 0
with torch.no_grad():
for batch in test_batches:
src_ids = batch['src_word_ids']
src_lengths = batch['src_sentence_lengths']
tgt_ids = batch['tgt_word_ids']
tgt_lengths = batch['tgt_sentence_lengths']
src_probs = None
if 'src_ddfd_probs' in batch:
src_probs = batch['src_ddfd_probs']
src_probs = _convert_to_tensor(src_probs, use_gpu).unsqueeze(2)
src_ids = _convert_to_tensor(src_ids, use_gpu)
tgt_ids = _convert_to_tensor(tgt_ids, use_gpu)
decoder_outputs, decoder_hidden, other = model(
src_ids,
tgt_ids,
is_training=False,
att_key_feats=src_probs,
beam_width=beam_width)
# Evaluation
seqlist = other['sequence'] # traverse over time not batch
if beam_width > 1:
# print('dict:sequence')
# print(len(seqlist))
# print(seqlist[0].size())
full_seqlist = other['topk_sequence']
# print('dict:topk_sequence')
# print(len(full_seqlist))
# print((full_seqlist[0]).size())
# input('...')
seqlists = []
for i in range(beam_width):
seqlists.append([seq[:, i] for seq in full_seqlist])
# print(decoder_outputs[0].size())
# print('tgt id size {}'.format(tgt_ids.size()))
# input('...')
decoder_outputs = decoder_outputs[:-1]
# print(len(decoder_outputs))
for step, step_output in enumerate(
decoder_outputs): # loop over time steps
target = tgt_ids[:, step + 1]
non_padding = target.ne(PAD)
# print('step', step)
# print('target', target)
# print('hyp', seqlist[step])
# if beam_width > 1:
# print('full_seqlist', full_seqlist[step])
# input('...')
correct = seqlist[step].view(-1).eq(target).masked_select(
non_padding).sum().item()
match += correct
total += non_padding.sum().item()
# write to file
refwords = _convert_to_words_batchfirst(tgt_ids[:, 1:],
test_set.tgt_id2word)
seqwords = _convert_to_words(seqlist, test_set.tgt_id2word)
seqwords_list = []
for i in range(beam_width):
seqwords_list.append(
_convert_to_words(seqlists[i], test_set.tgt_id2word))
for i in range(len(seqwords)):
outline_ref = ' '.join(refwords[i])
print('REF', outline_ref)
outline_hyp = ' '.join(seqwords[i])
# print(outline_hyp)
outline_hyps = []
for j in range(beam_width):
outline_hyps.append(' '.join(seqwords_list[j][i]))
print('{}th'.format(j), outline_hyps[-1])
# skip padding sentences in batch (num_sent % batch_size != 0)
# if src_lengths[i] == 0:
# continue
# words = []
# for word in seqwords[i]:
# if word == '<pad>':
# continue
# elif word == '</s>':
# break
# else:
# words.append(word)
# if len(words) == 0:
# outline = ''
# else:
# outline = ' '.join(words)
input('...')
sys.stdout.flush()
if total == 0:
accuracy = float('nan')
else:
accuracy = match / total
return accuracy
def _train_epoches(self,
train_set,
model,
n_epochs,
start_epoch,
start_step,
dev_set=None):
log = self.logger
print_loss_total = 0 # Reset every print_every
epoch_loss_total = 0 # Reset every epoch
att_print_loss_total = 0 # Reset every print_every
att_epoch_loss_total = 0 # Reset every epoch
attcls_print_loss_total = 0 # Reset every print_every
attcls_epoch_loss_total = 0 # Reset every epoch
step = start_step
step_elapsed = 0
prev_acc = 0.0
count_no_improve = 0
count_num_rollback = 0
ckpt = None
# ******************** [loop over epochs] ********************
for epoch in range(start_epoch, n_epochs + 1):
for param_group in self.optimizer.optimizer.param_groups:
print('epoch:{} lr: {}'.format(epoch, param_group['lr']))
lr_curr = param_group['lr']
# ----------construct batches-----------
# allow re-shuffling of data
if type(train_set.attkey_path) == type(None):
print('--- construct train set ---')
train_batches, vocab_size = train_set.construct_batches(
is_train=True)
if dev_set is not None:
print('--- construct dev set ---')
dev_batches, vocab_size = dev_set.construct_batches(
is_train=False)
else:
print('--- construct train set ---')
train_batches, vocab_size = train_set.construct_batches_with_ddfd_prob(
is_train=True)
if dev_set is not None:
print('--- construct dev set ---')
assert type(dev_set.attkey_path) != type(
None), 'Dev set missing ddfd probabilities'
dev_batches, vocab_size = dev_set.construct_batches_with_ddfd_prob(
is_train=False)
# --------print info for each epoch----------
steps_per_epoch = len(train_batches)
total_steps = steps_per_epoch * n_epochs
log.info("steps_per_epoch {}".format(steps_per_epoch))
log.info("total_steps {}".format(total_steps))
log.debug(
" ----------------- Epoch: %d, Step: %d -----------------" %
(epoch, step))
mem_kb, mem_mb, mem_gb = get_memory_alloc()
mem_mb = round(mem_mb, 2)
print('Memory used: {0:.2f} MB'.format(mem_mb))
self.writer.add_scalar('Memory_MB', mem_mb, global_step=step)
sys.stdout.flush()
# ******************** [loop over batches] ********************
model.train(True)
for batch in train_batches:
# update macro count
step += 1
step_elapsed += 1
# load data
src_ids = batch['src_word_ids']
src_lengths = batch['src_sentence_lengths']
tgt_ids = batch['tgt_word_ids']
tgt_lengths = batch['tgt_sentence_lengths']
src_probs = None
src_labs = None
if 'src_ddfd_probs' in batch and model.additional_key_size > 0:
src_probs = batch['src_ddfd_probs']
src_probs = _convert_to_tensor(src_probs,
self.use_gpu).unsqueeze(2)
if 'src_ddfd_labs' in batch:
src_labs = batch['src_ddfd_labs']
src_labs = _convert_to_tensor(src_labs,
self.use_gpu).unsqueeze(2)
# sanity check src-tgt pair
if step == 1:
print('--- Check src tgt pair ---')
log_msgs = check_srctgt(src_ids, tgt_ids,
train_set.src_id2word,
train_set.tgt_id2word)
for log_msg in log_msgs:
sys.stdout.buffer.write(log_msg)
# convert variable to tensor
src_ids = _convert_to_tensor(src_ids, self.use_gpu)
tgt_ids = _convert_to_tensor(tgt_ids, self.use_gpu)
# Get loss
loss, att_loss, attcls_loss = self._train_batch(
src_ids,
tgt_ids,
model,
step,
total_steps,
src_probs=src_probs,
src_labs=src_labs)
print_loss_total += loss
epoch_loss_total += loss
att_print_loss_total += att_loss
att_epoch_loss_total += att_loss
attcls_print_loss_total += attcls_loss
attcls_epoch_loss_total += attcls_loss
if step % self.print_every == 0 and step_elapsed > self.print_every:
print_loss_avg = print_loss_total / self.print_every
att_print_loss_avg = att_print_loss_total / self.print_every
attcls_print_loss_avg = attcls_print_loss_total / self.print_every
print_loss_total = 0
att_print_loss_total = 0
attcls_print_loss_total = 0
log_msg = 'Progress: %d%%, Train nlll: %.4f, att: %.4f, attcls: %.4f' % (
step / total_steps * 100, print_loss_avg,
att_print_loss_avg, attcls_print_loss_avg)
# print(log_msg)
log.info(log_msg)
self.writer.add_scalar('train_loss',
print_loss_avg,
global_step=step)
self.writer.add_scalar('att_train_loss',
att_print_loss_avg,
global_step=step)
self.writer.add_scalar('attcls_train_loss',
attcls_print_loss_avg,
global_step=step)
# Checkpoint
if step % self.checkpoint_every == 0 or step == total_steps:
# save criteria
if dev_set is not None:
dev_loss, accuracy, dev_attlosses = \
self._evaluate_batches(model, dev_batches, dev_set)
dev_attloss = dev_attlosses['att_loss']
dev_attclsloss = dev_attlosses['attcls_loss']
log_msg = 'Progress: %d%%, Dev loss: %.4f, accuracy: %.4f, att: %.4f, attcls: %.4f' % (
step / total_steps * 100, dev_loss, accuracy,
dev_attloss, dev_attclsloss)
log.info(log_msg)
self.writer.add_scalar('dev_loss',
dev_loss,
global_step=step)
self.writer.add_scalar('dev_acc',
accuracy,
global_step=step)
self.writer.add_scalar('att_dev_loss',
dev_attloss,
global_step=step)
self.writer.add_scalar('attcls_dev_loss',
dev_attclsloss,
global_step=step)
# save
if prev_acc < accuracy:
# save best model
ckpt = Checkpoint(model=model,
optimizer=self.optimizer,
epoch=epoch,
step=step,
input_vocab=train_set.vocab_src,
output_vocab=train_set.vocab_tgt)
saved_path = ckpt.save(self.expt_dir)
print('saving at {} ... '.format(saved_path))
# reset
prev_acc = accuracy
count_no_improve = 0
count_num_rollback = 0
else:
count_no_improve += 1
# roll back
if count_no_improve > MAX_COUNT_NO_IMPROVE:
# resuming
latest_checkpoint_path = Checkpoint.get_latest_checkpoint(
self.expt_dir)
if type(latest_checkpoint_path) != type(None):
resume_checkpoint = Checkpoint.load(
latest_checkpoint_path)
print(
'epoch:{} step: {} - rolling back {} ...'.
format(epoch, step,
latest_checkpoint_path))
model = resume_checkpoint.model
self.optimizer = resume_checkpoint.optimizer
# A walk around to set optimizing parameters properly
resume_optim = self.optimizer.optimizer
defaults = resume_optim.param_groups[0]
defaults.pop('params', None)
defaults.pop('initial_lr', None)
self.optimizer.optimizer = resume_optim\
.__class__(model.parameters(), **defaults)
# start_epoch = resume_checkpoint.epoch
# step = resume_checkpoint.step
# reset
count_no_improve = 0
count_num_rollback += 1
# update learning rate
if count_num_rollback > MAX_COUNT_NUM_ROLLBACK:
# roll back
latest_checkpoint_path = Checkpoint.get_latest_checkpoint(
self.expt_dir)
if type(latest_checkpoint_path) != type(None):
resume_checkpoint = Checkpoint.load(
latest_checkpoint_path)
print(
'epoch:{} step: {} - rolling back {} ...'.
format(epoch, step,
latest_checkpoint_path))
model = resume_checkpoint.model
self.optimizer = resume_checkpoint.optimizer
# A walk around to set optimizing parameters properly
resume_optim = self.optimizer.optimizer
defaults = resume_optim.param_groups[0]
defaults.pop('params', None)
defaults.pop('initial_lr', None)
self.optimizer.optimizer = resume_optim\
.__class__(model.parameters(), **defaults)
start_epoch = resume_checkpoint.epoch
step = resume_checkpoint.step
# decrease lr
for param_group in self.optimizer.optimizer.param_groups:
param_group['lr'] *= 0.5
lr_curr = param_group['lr']
print('reducing lr ...')
print('step:{} - lr: {}'.format(
step, param_group['lr']))
# check early stop
if lr_curr < 0.000125:
print('early stop ...')
break
# reset
count_no_improve = 0
count_num_rollback = 0
model.train(mode=True)
if ckpt is None:
ckpt = Checkpoint(model=model,
optimizer=self.optimizer,
epoch=epoch,
step=step,
input_vocab=train_set.vocab_src,
output_vocab=train_set.vocab_tgt)
saved_path = ckpt.save(self.expt_dir)
ckpt.rm_old(self.expt_dir, keep_num=KEEP_NUM)
print('n_no_improve {}, num_rollback {}'.format(
count_no_improve, count_num_rollback))
sys.stdout.flush()
else:
if dev_set is None:
# save every epoch if no dev_set
ckpt = Checkpoint(model=model,
optimizer=self.optimizer,
epoch=epoch,
step=step,
input_vocab=train_set.vocab_src,
output_vocab=train_set.vocab_tgt)
# saved_path = ckpt.save(self.expt_dir)
saved_path = ckpt.save_epoch(self.expt_dir, epoch)
print('saving at {} ... '.format(saved_path))
continue
else:
continue
# break nested for loop
break
if step_elapsed == 0: continue
epoch_loss_avg = epoch_loss_total / min(steps_per_epoch,
step - start_step)
epoch_loss_total = 0
log_msg = "Finished epoch %d: Train %s: %.4f" % (
epoch, self.loss.name, epoch_loss_avg)
log.info('\n')
log.info(log_msg)
def _evaluate_batches(self, model, batches, dataset):
model.eval()
loss = NLLLoss()
loss.reset()
match = 0
total = 0
out_count = 0
with torch.no_grad():
for batch in batches:
src_ids = batch['src_word_ids']
src_lengths = batch['src_sentence_lengths']
tgt_ids = batch['tgt_word_ids']
tgt_lengths = batch['tgt_sentence_lengths']
src_probs = None
if 'src_ddfd_probs' in batch and model.additional_key_size > 0:
src_probs = batch['src_ddfd_probs']
src_probs = _convert_to_tensor(src_probs,
self.use_gpu).unsqueeze(2)
src_labs = None
if 'src_ddfd_labs' in batch:
src_labs = batch['src_ddfd_labs']
src_labs = _convert_to_tensor(src_labs,
self.use_gpu).unsqueeze(2)
src_ids = _convert_to_tensor(src_ids, self.use_gpu)
tgt_ids = _convert_to_tensor(tgt_ids, self.use_gpu)
non_padding_mask_tgt = tgt_ids.data.ne(PAD)
non_padding_mask_src = src_ids.data.ne(PAD)
decoder_outputs, decoder_hidden, other = model(
src_ids,
tgt_ids,
is_training=False,
att_key_feats=src_probs)
# Evaluation
logps = torch.stack(decoder_outputs, dim=1).to(device=device)
if not self.eval_with_mask:
loss.eval_batch(logps.reshape(-1, logps.size(-1)),
tgt_ids[:, 1:].reshape(-1))
else:
loss.eval_batch_with_mask(
logps.reshape(-1, logps.size(-1)),
tgt_ids[:, 1:].reshape(-1),
non_padding_mask_tgt[:, 1:].reshape(-1))
seqlist = other['sequence']
seqres = torch.stack(seqlist, dim=1).to(device=device)
correct = seqres.view(-1).eq(tgt_ids[:,1:].reshape(-1))\
.masked_select(non_padding_mask_tgt[:,1:].reshape(-1)).sum().item()
match += correct
total += non_padding_mask_tgt[:, 1:].sum().item()
if not self.eval_with_mask:
loss.norm_term = 1.0 * tgt_ids.size(
0) * tgt_ids[:, 1:].size(1)
else:
loss.norm_term = 1.0 * torch.sum(non_padding_mask_tgt[:,
1:])
loss.normalise()
if out_count < 3:
srcwords = _convert_to_words_batchfirst(
src_ids, dataset.tgt_id2word)
refwords = _convert_to_words_batchfirst(
tgt_ids[:, 1:], dataset.tgt_id2word)
seqwords = _convert_to_words(seqlist, dataset.tgt_id2word)
outsrc = 'SRC: {}\n'.format(' '.join(
srcwords[0])).encode('utf-8')
outref = 'REF: {}\n'.format(' '.join(
refwords[0])).encode('utf-8')
outline = 'GEN: {}\n'.format(' '.join(
seqwords[0])).encode('utf-8')
sys.stdout.buffer.write(outsrc)
sys.stdout.buffer.write(outref)
sys.stdout.buffer.write(outline)
out_count += 1
att_resloss = 0
attcls_resloss = 0
resloss = loss.get_loss()
if total == 0:
accuracy = float('nan')
else:
accuracy = match / total
torch.cuda.empty_cache()
losses = {}
losses['att_loss'] = att_resloss
losses['attcls_loss'] = attcls_resloss
return resloss, accuracy, losses
