def query_eval(self, query_bsz: int, model: IModel, task_train_iter: Iterator,
criterion: BaseLoss, ):
meta_val_batch = next(task_train_iter)
if model.config['common']['device'] is not None:
meta_val_batch = batch_to_cuda(meta_val_batch)
query_loss = model.train_sl(meta_val_batch, criterion)
# query_loss = 0.0
with torch.no_grad():
loss_acc = 0.0
for _ in range(query_bsz - 1):
meta_val_batch = next(task_train_iter)
if model.config['common']['device'] is not None:
meta_val_batch = batch_to_cuda(meta_val_batch)
meta_val_loss = model.train_sl(meta_val_batch, criterion)
# query_loss = query_loss + meta_val_loss
loss_acc = loss_acc + meta_val_loss.item()
query_loss.data.add_(float(loss_acc)).div_(query_bsz)
return query_loss
def support_train(self, support_bsz: int, model: IModel, task_train_iter: Iterator,
criterion: BaseLoss, meta_optimizer: Optimizer, ):
for _ in range(support_bsz):
meta_train_batch = next(task_train_iter)
if model.config['common']['device'] is not None:
meta_train_batch = batch_to_cuda(meta_train_batch)
comment_loss = model.train_sl(meta_train_batch, criterion)
meta_optimizer.zero_grad()
comment_loss.backward()
gradient_accumulation(meta_optimizer, support_bsz)
def summarization_eval(model: IModel,
data_loader: DataLoader,
token_dicts: TokenDicts,
criterion: BaseLoss,
collate_func=None,
model_filename=None,
metrics=None) -> Any:
with torch.no_grad():
model.eval()
data_iter = iter(data_loader) # init
total_loss = 0.0
src_comment_all, trg_comment_all, pred_comment_all, src_code_all = \
[], [], [], []
if model.config['training']['pointer']:
oov_vocab = []
else:
oov_vocab = None
for iteration in range(1, 1 +
len(data_loader)): # 1 + len(data_loader)
batch = data_iter.__next__()
if collate_func is not None:
batch = collate_func(batch)
if model.config['common']['device'] is not None:
batch = batch_to_cuda(batch)
comment_pred, comment_logprobs, = model.eval_pipeline(batch)
if model_filename is None:
src_comment_all.extend(
[None] * model.config['training']['batch_size'])
src_code_all.extend([None] *
model.config['training']['batch_size'])
else:
src_codes, src_comments, = zip(*batch['case_study'])
src_comment_all.extend(src_comments)
src_code_all.extend(src_codes)
# comment
trg_comment_all.extend(batch['comment'][4])
pred_comment_all.extend(
[p if len(p) > 0 else '0' for p in comment_pred])
# oovs
if model.config['training']['pointer']:
oov_vocab.extend(batch['pointer'][-1])
# print(comment_logprobs.size())
# print(comment_target_padded.size())
if model.config['training']['pointer']:
comment_target = batch['pointer'][
1][:, :model.config['training']['max_predict_length']]
else:
comment_target = batch['comment'][
2][:, :model.config['training']['max_predict_length']]
# print('comment_logprobs: ', comment_logprobs.size())
# print('comment_target: ', comment_target.size())
comment_loss = criterion(
comment_logprobs[:, :comment_target.size(1)],
comment_target)
total_loss += comment_loss.item()
total_loss /= len(data_loader)
LOGGER.debug('Summarization test loss: {:.4}'.format(total_loss))
if model_filename is None:
pred_filename = None
else:
pred_filename = model_filename.replace('.pt', '.pred')
bleu1, bleu2, bleu3, bleu4, pycoco_bleu, \
meteor, pycoco_meteor, \
rouge1, rouge2, rouge3, rouge4, rougel, pycoco_rouge, \
cider, \
srcs_return, trgs_return, preds_return, = \
eval_metrics(src_comment_all, trg_comment_all, pred_comment_all, src_code_all,
oov_vocab, token_dicts, pred_filename,
metrics=model.config['testing']['metrics'] if metrics is None else metrics, )
bleu1, bleu2, bleu3, bleu4, pycoco_bleu, \
meteor, pycoco_meteor, \
rouge1, rouge2, rouge3, rouge4, rougel, pycoco_rouge, \
cider = \
map(lambda array: sum(array) / len(array), (bleu1, bleu2, bleu3, bleu4, pycoco_bleu, \
meteor, pycoco_meteor, \
rouge1, rouge2, rouge3, rouge4, rougel, pycoco_rouge, \
cider))
return bleu1, bleu2, bleu3, bleu4, pycoco_bleu, \
meteor, pycoco_meteor, \
rouge1, rouge2, rouge3, rouge4, rougel, pycoco_rouge, \
cider
def train(
self,
model: IModel,
dataset: UnilangDataloader,
criterion: BaseLoss,
optimizer: Optimizer,
SAVE_DIR=None,
start_time=None,
model_name_prefix=None,
scheduler=None,
) -> Dict:
super().train()
start_time = time.time() if start_time is None else start_time
if model_name_prefix is not None:
model_name_prefix += "-tt{}".format(self.__class__.__name__)
if scheduler is None:
scheduler = create_scheduler(optimizer,
self.config['sl']['warmup_epochs'],
self.config['sl']['warmup_factor'],
self.config['sl']['lr_milestones'],
self.config['sl']['lr_gamma'])
bleu_best, rougel_best, cider_best = 0.0, 0.0, 0.0
return_model = {'bleu': None, 'cider': None, 'rouge': None}
for epoch in range(1, 1 + self.config['training']['train_epoch']):
model.train()
train_data_iter = iter(dataset['train'])
total_loss = 0.0
for iteration in range(1, 1 + len(dataset['train'])):
batch = train_data_iter.__next__()
if self.config['common']['device'] is not None:
batch = batch_to_cuda(batch)
sl_loss = model.train_sl(batch, criterion)
LOGGER.debug('{} batch loss: {:.8f}'.format(
self.__class__.__name__, sl_loss.item()))
optimizer.zero_grad()
sl_loss.backward()
total_loss += sl_loss.item()
if model.config['sl']['max_grad_norm'] != -1:
nn.utils.clip_grad_norm_(
model.parameters(),
model.config['sl']['max_grad_norm'])
optimizer.step()
if iteration % self.config['training'][
'log_interval'] == 0 and iteration > 0:
LOGGER.info(
'Epoch: {:0>3d}/{:0>3d}, batches: {:0>3d}/{:0>3d}, avg_loss: {:.6f}; lr: {:.6f}, time: {}'. \
format(epoch, self.config['training']['train_epoch'], iteration, len(dataset['train']),
total_loss / iteration, scheduler.get_lr()[0],
str(datetime.timedelta(seconds=int(time.time() - start_time)))))
scheduler.step(epoch)
# Validation on each epoch
bleu1, bleu2, _, _, pycoco_bleu, meteor, pycoco_meteor, rouge1, rouge2, _, _, rougel, pycoco_rouge, \
cider = \
Evaluator.summarization_eval(model, dataset['valid'], dataset.token_dicts, criterion,
metrics=model.config['testing']['metrics'])
headers = [
'B1', 'B2', 'PycocoB4', 'Meteor', 'PycocoMeteor', 'R1', 'R2',
'RL', 'PycocoRL', 'Cider'
]
result_table = [[
round(i, 4) for i in [
bleu1, bleu2, pycoco_bleu, meteor, pycoco_meteor, rouge1,
rouge2, rougel, pycoco_rouge, cider
]
]]
LOGGER.info('Evaluation results:\n{}'.format(
tabulate(
result_table,
headers=headers,
tablefmt=model.config['common']['result_table_format'])))
# Dump the model if save_dir exists.
if SAVE_DIR is not None:
if model_name_prefix is None:
model_name_prefix = '{}-bs{}-lr{}-attn{}-pointer{}-tt{}'.format(
'8'.join(self.config['training']['code_modalities']),
self.config['training']['batch_size'],
self.config['sl']['lr'],
self.config['training']['attn_type'],
self.config['training']['pointer'],
self.__class__.__name__)
model_name = '{}-ep{}'.format(model_name_prefix, epoch)
model_path = os.path.join(
SAVE_DIR,
'{}.pt'.format(model_name),
)
torch.save(model.state_dict(), model_path)
LOGGER.info('Dumping model in {}'.format(model_path))
if 'bleu' in self.config['testing']['metrics']:
if pycoco_bleu > bleu_best:
bleu_best = pycoco_bleu
model_path = os.path.join(
SAVE_DIR,
'{}-best-bleu1.pt'.format(model_name_prefix),
)
return_model['bleu'] = model_path
torch.save(model.state_dict(), model_path)
LOGGER.info('Dumping best bleu1 model in {}'.format(
model_path))
if 'cider' in self.config['testing']['metrics']:
if cider > cider_best:
cider_best = cider
model_path = os.path.join(
SAVE_DIR,
'{}-best-cider.pt'.format(model_name_prefix),
)
return_model['cider'] = model_path
torch.save(model.state_dict(), model_path)
LOGGER.info('Dumping best cider model in {}'.format(
model_path))
if 'rouge' in self.config['testing']['metrics']:
if rougel > rougel_best:
rougel_best = rougel
model_path = os.path.join(
SAVE_DIR,
'{}-best-rougel.pt'.format(model_name_prefix),
)
return_model['rouge'] = model_path
torch.save(model.state_dict(), model_path)
LOGGER.info('Dumping best rouge model in {}'.format(
model_path))
LOGGER.info('{} train end'.format(self))
return return_model
def case_study_eval_code2seq(
model: IModel,
data_loader: DataLoader,
token_dicts: TokenDicts,
collate_func=None,
model_filename=None,
other_modal=None,
lng='ruby',
) -> Any:
# load ast size
# import glob, ujson
# ast_files = sorted([filename for filename in glob.glob('{}/*'.format(os.path.join(
# model.config['dataset']['dataset_dir'], 'ruby', 'ast'
# ))) if 'test' in filename])
#
# ast_len = []
# for fl in ast_files:
# with open(fl, 'r') as reader:
# line = reader.readline().strip()
# while len(line) > 0:
# line = ujson.loads(line)
# ast_len.append(len(line))
# line = reader.readline().strip()
len_info = {}
if other_modal is None:
modal_list = ['ast']
elif (set(other_modal) == set(['cfg', 'ast'])):
modal_list = ['cfg', 'ast']
for modal in modal_list:
len_info[modal] = load_data(model, datatype=modal, lng=lng)
# ast_len = load_data(model,datatype='ast',lng=lng)
tok_len = load_data(model, datatype='tok', lng=lng)
with torch.no_grad():
model.eval()
data_iter = iter(data_loader) # init
src_comment_all, trg_comment_all, pred_comment_all, src_code_all = \
[], [], [], []
# tok_len, comment_len = [], []
comment_len = []
if model.config['training']['pointer']:
oov_vocab = []
else:
oov_vocab = None
for iteration in range(1, 1 +
len(data_loader)): # 1 + len(data_loader)
batch = data_iter.__next__()
if collate_func is not None:
batch = collate_func(batch)
if model.config['common']['device'] is not None:
batch = batch_to_cuda(batch)
comment_pred, comment_logprobs, = model.eval_pipeline(batch)
if model_filename is None:
src_comment_all.extend(
[None] * model.config['training']['batch_size'])
src_code_all.extend([None] *
model.config['training']['batch_size'])
else:
src_codes, src_comments, = zip(*batch['case_study'])
src_comment_all.extend(src_comments)
src_code_all.extend(src_codes)
# comment
trg_comment_all.extend(batch['comment'][4])
pred_comment_all.extend(comment_pred)
# oovs
if model.config['training']['pointer']:
oov_vocab.extend(batch['pointer'][-1])
# tok_len.extend(batch['tok'][1].tolist())
comment_len.extend(batch['comment'][-2].tolist())
if model_filename is None:
pred_filename = None
else:
pred_filename = model_filename.replace('.pt', '.pred')
eval_per_metrics(
src_comment_all,
trg_comment_all,
pred_comment_all,
src_code_all,
oov_vocab,
tok_len,
comment_len,
len_info,
token_dicts,
pred_filename,
)
LOGGER.info(
'write test case-study info in {}'.format(pred_filename))