def init_meters(self, args):
self.meters = OrderedDict()
self.meters['train_loss'] = AverageMeter()
self.meters['train_nll_loss'] = AverageMeter()
self.meters['valid_loss'] = AverageMeter()
self.meters['valid_nll_loss'] = AverageMeter()
self.meters['wps'] = TimeMeter() # words per second
self.meters['ups'] = TimeMeter() # updates per second
self.meters['wpb'] = AverageMeter() # words per batch
self.meters['bsz'] = AverageMeter() # sentences per batch
self.meters['gnorm'] = AverageMeter() # gradient norm
self.meters['clip'] = AverageMeter() # % of updates clipped
self.meters['oom'] = AverageMeter() # out of memory
if args.fp16:
self.meters['loss_scale'] = AverageMeter() # dynamic loss scale
self.meters['wall'] = TimeMeter() # wall time in seconds
self.meters['train_wall'] = StopwatchMeter(
) # train wall time in seconds
def init_meters(self, args):
self.meters = OrderedDict()
self.meters["train_loss"] = AverageMeter()
self.meters["train_nll_loss"] = AverageMeter()
self.meters["valid_loss"] = AverageMeter()
self.meters["valid_nll_loss"] = AverageMeter()
self.meters["wps"] = TimeMeter() # words per second
self.meters["ups"] = TimeMeter() # updates per second
self.meters["wpb"] = AverageMeter() # words per batch
self.meters["bsz"] = AverageMeter() # sentences per batch
self.meters["gnorm"] = AverageMeter() # gradient norm
self.meters["clip"] = AverageMeter() # % of updates clipped
self.meters["oom"] = AverageMeter() # out of memory
if args.fp16:
self.meters["loss_scale"] = AverageMeter() # dynamic loss scale
self.meters["wall"] = TimeMeter() # wall time in seconds
self.meters["train_wall"] = StopwatchMeter(
) # train wall time in seconds
def __init__(self, args, task, model, criterion, dummy_batch):
if not torch.cuda.is_available():
raise NotImplementedError('Training on CPU is not supported')
self.args = args
self.task = task
# copy model and criterion to current device
self.criterion = criterion.cuda()
if args.fp16:
self._model = model.half().cuda()
else:
self._model = model.cuda()
# initialize meters
self.meters = OrderedDict()
self.meters['train_loss'] = AverageMeter()
self.meters['train_nll_loss'] = AverageMeter()
self.meters['valid_loss'] = AverageMeter()
self.meters['valid_nll_loss'] = AverageMeter()
self.meters['train_acc'] = AverageMeter()
self.meters['valid_acc'] = AverageMeter()
self.meters['wps'] = TimeMeter() # words per second
self.meters['ups'] = TimeMeter() # updates per second
self.meters['wpb'] = AverageMeter() # words per batch
self.meters['bsz'] = AverageMeter() # sentences per batch
self.meters['gnorm'] = AverageMeter() # gradient norm
self.meters['clip'] = AverageMeter() # % of updates clipped
self.meters['oom'] = AverageMeter() # out of memory
if args.fp16:
self.meters['loss_scale'] = AverageMeter() # dynamic loss scale
self.meters['wall'] = TimeMeter() # wall time in seconds
self.meters['train_wall'] = StopwatchMeter(
) # train wall time in seconds
self._dummy_batch = dummy_batch
self._num_updates = 0
self._optim_history = None
self._optimizer = None
self._wrapped_model = None
def init_meters(self, args):
self.meters = OrderedDict()
self.meters['train_loss'] = AverageMeter()
self.meters['train_nll_loss'] = AverageMeter()
self.meters['valid_loss'] = AverageMeter()
self.meters['valid_nll_loss'] = AverageMeter()
self.meters['wps'] = TimeMeter() # words per second
self.meters['ups'] = TimeMeter() # updates per second
self.meters['wpb'] = AverageMeter() # words per batch
self.meters['bsz'] = AverageMeter() # sentences per batch
self.meters['gnorm'] = AverageMeter() # gradient norm
self.meters['clip'] = AverageMeter() # % of updates clipped
self.meters['oom'] = AverageMeter() # out of memory
if args.fp16:
self.meters['loss_scale'] = AverageMeter() # dynamic loss scale
self.meters['wall'] = TimeMeter() # wall time in seconds
self.meters['train_wall'] = StopwatchMeter(
) # train wall time in seconds
# add_norm
for i in range(args.encoder_layers):
self.meters['enc{}_attn_input_gnorm'.format(
i)] = AverageMeter() # gradient norm
self.meters['enc{}_ffn_fc1_gnorm'.format(i)] = AverageMeter()
def main(args, init_distributed=False):
utils.import_user_module(args)
assert args.max_tokens is not None or args.max_sentences is not None, \
'Must specify batch size either with --max-tokens or --max-sentences'
# Initialize CUDA and distributed training
if torch.cuda.is_available() and not args.cpu:
torch.cuda.set_device(args.device_id)
np.random.seed(args.seed)
torch.manual_seed(args.seed)
if init_distributed:
args.distributed_rank = distributed_utils.distributed_init(args)
if distributed_utils.is_master(args):
checkpoint_utils.verify_checkpoint_directory(args.save_dir)
# Print args
print(args)
# Setup task, e.g., translation, language modeling, etc.
task = tasks.setup_task(args)
# Load valid dataset (we load training data below, based on the latest checkpoint)
for valid_sub_split in args.valid_subset.split(','):
task.load_dataset(valid_sub_split, combine=False, epoch=0)
# Build model and criterion
model = task.build_model(args)
criterion = task.build_criterion(args)
print(model)
print('| model {}, criterion {}'.format(args.arch,
criterion.__class__.__name__))
print('| num. model params: {} (num. trained: {})'.format(
sum(p.numel() for p in model.parameters()),
sum(p.numel() for p in model.parameters() if p.requires_grad),
))
# Build trainer
trainer = Trainer(args, task, model, criterion)
print('| training on {} GPUs'.format(args.distributed_world_size))
print('| max tokens per GPU = {} and max sentences per GPU = {}'.format(
args.max_tokens,
args.max_sentences,
))
# Load the latest checkpoint if one is available and restore the
# corresponding train iterator
extra_state, epoch_itr = checkpoint_utils.load_checkpoint(args, trainer)
# Train until the learning rate gets too small
max_epoch = args.max_epoch or math.inf
max_update = args.max_update or math.inf
lr = trainer.get_lr()
train_meter = StopwatchMeter()
train_meter.start()
valid_subsets = args.valid_subset.split(',')
while lr > args.min_lr and epoch_itr.epoch < max_epoch and trainer.get_num_updates(
) < max_update:
# train for one epoch
train(args, trainer, task, epoch_itr)
if not args.disable_validation and epoch_itr.epoch % args.validate_interval == 0:
valid_losses = validate(args, trainer, task, epoch_itr,
valid_subsets)
else:
valid_losses = [None]
# only use first validation loss to update the learning rate
lr = trainer.lr_step(epoch_itr.epoch, valid_losses[0])
# save checkpoint
if epoch_itr.epoch % args.save_interval == 0:
checkpoint_utils.save_checkpoint(args, trainer, epoch_itr,
valid_losses[0])
reload_dataset = ':' in getattr(args, 'data', '')
# sharded data: get train iterator for next epoch
epoch_itr = trainer.get_train_iterator(epoch_itr.epoch,
load_dataset=reload_dataset)
train_meter.stop()
print('| done training in {:.1f} seconds'.format(train_meter.sum))
def main(args):
assert args.path is not None, '--path required for generation!'
assert not args.sampling or args.nbest == args.beam, \
'--sampling requires --nbest to be equal to --beam'
utils.import_user_module(args)
if args.max_tokens is None and args.max_sentences is None:
args.max_tokens = 12000
print(args)
use_cuda = torch.cuda.is_available() and not args.cpu
# Load dataset split
task = tasks.setup_task(args)
task.load_dataset(args.gen_subset)
# Load ensemble
print('| loading model(s) from {}'.format(args.path))
model_paths = args.path.split(':')
models, _model_args = checkpoint_utils.load_model_ensemble(
model_paths,
arg_overrides=eval(args.model_overrides),
task=task,
)
# Optimize ensemble for generation
for model in models:
# model.make_generation_fast_()
if args.fp16:
model.half()
if use_cuda:
model.cuda()
# Load dataset (possibly sharded)
itr = task.get_batch_iterator(
dataset=task.dataset(args.gen_subset),
max_tokens=args.max_tokens,
max_sentences=args.max_sentences,
max_positions=utils.resolve_max_positions(
task.max_positions(),
*[model.max_positions() for model in models]),
ignore_invalid_inputs=args.skip_invalid_size_inputs_valid_test,
required_batch_size_multiple=args.required_batch_size_multiple,
num_shards=args.num_shards,
shard_id=args.shard_id,
num_workers=args.num_workers,
).next_epoch_itr(shuffle=False)
# Initialize generator
gen_timer = StopwatchMeter()
generator = task.build_generator(args)
# Generate and compute score
coco = task.dataset(args.gen_subset).coco
iou_types = ['bbox']
scorer = CocoEvaluator(coco, iou_types)
num_images = 0
with progress_bar.build_progress_bar(
args,
itr,
prefix='inference on \'{}\' subset'.format(args.gen_subset),
no_progress_bar='simple',
) as progress:
wps_meter = TimeMeter()
for sample in progress:
sample = utils.move_to_cuda(sample) if use_cuda else sample
gen_timer.start()
hypos = task.inference_step(generator, models, sample)
num_generated_boxes = sum(len(h['scores']) for h in hypos)
gen_timer.stop(num_generated_boxes)
result = {}
for i, sample_id in enumerate(sample['id'].tolist()):
result[sample_id] = hypos[i]
scorer.update(result)
wps_meter.update(num_generated_boxes)
progress.log({'wps': round(wps_meter.avg)})
num_images += sample['nsentences']
print(
'| Detected {} images ({} tokens) in {:.1f}s ({:.2f} images/s, {:.2f} tokens/s)'
.format(num_images, gen_timer.n, gen_timer.sum,
num_images / gen_timer.sum, 1. / gen_timer.avg))
# gather the stats from all processes
scorer.synchronize_between_processes()
# accumulate predictions from all images
scorer.accumulate()
scorer.summarize()
return scorer
def _generate_score(models, args, task, dataset, optimize=True):
use_cuda = torch.cuda.is_available() and not args.cpu
# Load ensemble
if not args.quiet:
print("| loading model(s) from {}".format(", ".join(args.path.split(":"))))
# Optimize ensemble for generation
if optimize:
for model in models:
model.make_generation_fast_(
beamable_mm_beam_size=None if args.no_beamable_mm else args.beam,
need_attn=True,
)
translator = build_sequence_generator(args, task, models)
# Load alignment dictionary for unknown word replacement
# (None if no unknown word replacement, empty if no path to align dictionary)
align_dict = utils.load_align_dict(args.replace_unk)
# Keep track of translations
# Initialize with empty translations
# and zero probs scores
translated_sentences = [""] * len(dataset)
translated_scores = [0.0] * len(dataset)
# Generate and compute BLEU score
dst_dict = task.target_dictionary
scorer = bleu.Scorer(dst_dict.pad(), dst_dict.eos(), dst_dict.unk())
itr = get_eval_itr(args, models, task, dataset)
num_sentences = 0
translation_samples = []
with progress_bar.build_progress_bar(args, itr) as t:
wps_meter = TimeMeter()
gen_timer = StopwatchMeter()
translations = translator.generate_batched_itr(
t,
maxlen_a=args.max_len_a,
maxlen_b=args.max_len_b,
cuda=use_cuda,
timer=gen_timer,
prefix_size=1 if pytorch_translate_data.is_multilingual(args) else 0,
)
if pytorch_translate_data.is_multilingual(args):
first_best_translations = _iter_first_best_multilingual
else:
first_best_translations = _iter_first_best_bilingual
for trans_info in first_best_translations(
args, task, dataset, translations, align_dict
):
scorer.add(trans_info.target_tokens, trans_info.hypo_tokens)
translated_sentences[trans_info.sample_id] = trans_info.hypo_str
translated_scores[trans_info.sample_id] = trans_info.hypo_score
translation_samples.append(
collections.OrderedDict(
{
"sample_id": trans_info.sample_id.item(),
"src_str": trans_info.src_str,
"target_str": trans_info.target_str,
"hypo_str": trans_info.hypo_str,
}
)
)
wps_meter.update(trans_info.src_tokens.size(0))
t.log({"wps": round(wps_meter.avg)})
num_sentences += 1
# If applicable, save the translations to the output file
# For eg. external evaluation
if getattr(args, "translation_output_file", False):
with open(args.translation_output_file, "w") as out_file:
for hypo_str in translated_sentences:
print(hypo_str, file=out_file)
if getattr(args, "translation_probs_file", False):
with open(args.translation_probs_file, "w") as out_file:
for hypo_score in translated_scores:
print(np.exp(hypo_score), file=out_file)
return scorer, num_sentences, gen_timer, translation_samples
def main(args):
check_args(args)
import_user_module(args)
if args.max_tokens is None and args.max_sentences is None:
args.max_tokens = 30000
logger.info(args)
use_cuda = torch.cuda.is_available() and not args.cpu
# Load dataset splits
task = tasks.setup_task(args)
task.load_dataset(args.gen_subset)
logger.info("| {} {} {} examples".format(
args.data, args.gen_subset, len(task.dataset(args.gen_subset))))
# Set dictionary
tgt_dict = task.target_dictionary
logger.info("| decoding with criterion {}".format(args.criterion))
# Load ensemble
logger.info("| loading model(s) from {}".format(args.path))
models, criterions, _model_args = load_models_and_criterions(
args.path.split(":"),
arg_overrides=eval(args.model_overrides), # noqa
task=task,
)
optimize_models(args, use_cuda, models)
# hack to pass transitions to W2lDecoder
if args.criterion == "asg_loss":
trans = criterions[0].asg.trans.data
args.asg_transitions = torch.flatten(trans).tolist()
# Load dataset (possibly sharded)
itr = get_dataset_itr(args, task)
# Initialize generator
gen_timer = StopwatchMeter()
generator = task.build_generator(args)
num_sentences = 0
if not os.path.exists(args.results_path):
os.makedirs(args.results_path)
sp = spm.SentencePieceProcessor()
sp.Load(os.path.join(args.data, "spm.model"))
res_files = prepare_result_files(args)
with progress_bar.build_progress_bar(args, itr) as t:
wps_meter = TimeMeter()
for sample in t:
sample = utils.move_to_cuda(sample) if use_cuda else sample
if "net_input" not in sample:
continue
prefix_tokens = None
if args.prefix_size > 0:
prefix_tokens = sample["target"][:, :args.prefix_size]
gen_timer.start()
hypos = task.inference_step(generator, models, sample,
prefix_tokens)
num_generated_tokens = sum(len(h[0]["tokens"]) for h in hypos)
gen_timer.stop(num_generated_tokens)
for i, sample_id in enumerate(sample["id"].tolist()):
speaker = task.dataset(
args.gen_subset).speakers[int(sample_id)]
id = task.dataset(args.gen_subset).ids[int(sample_id)]
target_tokens = (utils.strip_pad(sample["target"][i, :],
tgt_dict.pad()).int().cpu())
# Process top predictions
process_predictions(args, hypos[i], sp, tgt_dict,
target_tokens, res_files, speaker, id)
wps_meter.update(num_generated_tokens)
t.log({"wps": round(wps_meter.avg)})
num_sentences += sample["nsentences"]
logger.info("| Processed {} sentences ({} tokens) in {:.1f}s ({:.2f}"
"sentences/s, {:.2f} tokens/s)".format(
num_sentences,
gen_timer.n,
gen_timer.sum,
num_sentences / gen_timer.sum,
1.0 / gen_timer.avg,
))
logger.info("| Generate {} with beam={}".format(args.gen_subset,
args.beam))
def main(args):
import_user_module(args)
if args.max_tokens is None:
args.max_tokens = 6000
print(args)
if torch.cuda.is_available() and not args.cpu:
torch.cuda.set_device(args.device_id)
torch.manual_seed(args.seed)
# Setup task, e.g., translation, language modeling, etc.
task = tasks.setup_task(args)
# Load dataset splits
load_dataset_splits(task, ['train', 'valid'])
# Build model and criterion
model = task.build_model(args)
criterion = task.build_criterion(args)
print(model)
print('| model {}, criterion {}'.format(args.arch,
criterion.__class__.__name__))
print('| num. model params: {} (num. trained: {})'.format(
sum(p.numel() for p in model.parameters()),
sum(p.numel() for p in model.parameters() if p.requires_grad),
))
# Make a dummy batch to (i) warm the caching allocator and (ii) as a
# placeholder DistributedDataParallel when there's an uneven number of
# batches per worker.
max_positions = utils.resolve_max_positions(
task.max_positions(),
model.max_positions(),
)
dummy_batch = task.dataset('train').get_dummy_batch(
args.max_tokens, max_positions)
oom_batch = task.dataset('train').get_dummy_batch(1, max_positions)
# Build trainer
trainer = Trainer(args, task, model, criterion, dummy_batch, oom_batch)
print('| training on {} GPUs'.format(args.distributed_world_size))
print('| max tokens per GPU = {} and max sentences per GPU = {}'.format(
args.max_tokens,
args.max_sentences,
))
# Initialize dataloader
epoch_itr = task.get_batch_iterator(
dataset=task.dataset(args.train_subset),
max_tokens=args.max_tokens,
max_sentences=args.max_sentences,
max_positions=max_positions,
ignore_invalid_inputs=True,
required_batch_size_multiple=8,
seed=args.seed,
num_shards=args.distributed_world_size,
shard_id=args.distributed_rank,
num_workers=args.num_workers,
)
# Load the latest checkpoint if one is available
if not load_checkpoint(args, trainer, epoch_itr):
trainer.dummy_train_step([dummy_batch])
# Train until the learning rate gets too small
max_epoch = args.max_epoch or math.inf
max_update = args.max_update or math.inf
lr = trainer.get_lr()
train_meter = StopwatchMeter()
train_meter.start()
valid_losses = [None]
valid_subsets = args.valid_subset.split(',')
while lr > args.min_lr and epoch_itr.epoch < max_epoch and trainer.get_num_updates(
) < max_update:
# train for one epoch
train(args, trainer, task, epoch_itr)
if epoch_itr.epoch % args.validate_interval == 0:
valid_losses = validate(args, trainer, task, epoch_itr,
valid_subsets)
# only use first validation loss to update the learning rate
lr = trainer.lr_step(epoch_itr.epoch, valid_losses[0])
# save checkpoint
if epoch_itr.epoch % args.save_interval == 0:
save_checkpoint(args, trainer, epoch_itr, valid_losses[0])
train_meter.stop()
print('| done training in {:.1f} seconds'.format(train_meter.sum))
def main(parsed_args):
assert parsed_args.path is not None, '--path required for evaluation!'
print(parsed_args)
use_cuda = torch.cuda.is_available() and not parsed_args.cpu
task = tasks.setup_task(parsed_args)
# Load ensemble
print('| loading model(s) from {}'.format(parsed_args.path))
models, args = utils.load_ensemble_for_inference(
parsed_args.path.split(':'),
task,
model_arg_overrides=eval(parsed_args.model_overrides))
for arg in vars(parsed_args).keys():
if arg not in {
'self_target', 'future_target', 'past_target',
'tokens_per_sample', 'output_size_dictionary'
}:
setattr(args, arg, getattr(parsed_args, arg))
task = tasks.setup_task(args)
# Load dataset splits
task.load_dataset(args.gen_subset)
print('| {} {} {} examples'.format(args.data, args.gen_subset,
len(task.dataset(args.gen_subset))))
# Optimize ensemble for generation and set the source and dest dicts on the model (required by scorer)
for model in models:
model.make_generation_fast_()
if args.fp16:
model.half()
assert len(models) > 0
print('num. model params: {}'.format(
sum(p.numel() for p in models[0].parameters())))
itr = task.get_batch_iterator(
dataset=task.dataset(args.gen_subset),
max_tokens=args.max_tokens or 36000,
max_sentences=args.max_sentences,
max_positions=utils.resolve_max_positions(
*[model.max_positions() for model in models]),
num_shards=args.num_shards,
shard_id=args.shard_id,
ignore_invalid_inputs=True,
).next_epoch_itr(shuffle=False)
gen_timer = StopwatchMeter()
scorer = SequenceScorer(models, task.target_dictionary)
if use_cuda:
scorer.cuda()
score_sum = 0.
count = 0
if args.remove_bpe is not None:
bpe_cont = args.remove_bpe.rstrip()
bpe_toks = set(i for i in range(len(task.dictionary))
if task.dictionary[i].endswith(bpe_cont))
bpe_len = len(bpe_cont)
else:
bpe_toks = None
bpe_len = 0
word_stats = dict()
with progress_bar.build_progress_bar(args, itr) as t:
results = scorer.score_batched_itr(t, cuda=use_cuda, timer=gen_timer)
wps_meter = TimeMeter()
for _, src_tokens, __, hypos in results:
for hypo in hypos:
pos_scores = hypo['positional_scores']
skipped_toks = 0
if bpe_toks is not None:
for i in range(len(hypo['tokens']) - 1):
if hypo['tokens'][i].item() in bpe_toks:
skipped_toks += 1
pos_scores[i + 1] += pos_scores[i]
pos_scores[i] = 0
inf_scores = pos_scores.eq(float('inf')) | pos_scores.eq(
float('-inf'))
if inf_scores.any():
print(
'| Skipping tokens with inf scores:',
task.target_dictionary.string(
hypo['tokens'][inf_scores.nonzero()]))
pos_scores = pos_scores[(~inf_scores).nonzero()]
score_sum += pos_scores.sum().cpu()
count += pos_scores.numel() - skipped_toks
if args.output_word_probs or args.output_word_stats:
w = ''
word_prob = []
is_bpe = False
for i in range(len(hypo['tokens'])):
w_ind = hypo['tokens'][i].item()
w += task.dictionary[w_ind]
if bpe_toks is not None and w_ind in bpe_toks:
w = w[:-bpe_len]
is_bpe = True
else:
word_prob.append((w, pos_scores[i].item()))
next_prob = None
ind = i + 1
while ind < len(hypo['tokens']):
if pos_scores[ind].item() != 0:
next_prob = pos_scores[ind]
break
ind += 1
word_stats.setdefault(w, WordStat(w, is_bpe)).add(
pos_scores[i].item(), next_prob)
is_bpe = False
w = ''
if args.output_word_probs:
print('\t'.join('{} [{:2f}]'.format(x[0], x[1])
for x in word_prob))
wps_meter.update(src_tokens.size(0))
t.log({'wps': round(wps_meter.avg)})
avg_nll_loss = -score_sum / count
print('| Evaluated {} tokens in {:.1f}s ({:.2f} tokens/s)'.format(
gen_timer.n, gen_timer.sum, 1. / gen_timer.avg))
print('| Loss: {:.4f}, Perplexity: {:.2f}'.format(avg_nll_loss,
np.exp(avg_nll_loss)))
if args.output_word_stats:
for ws in sorted(word_stats.values(),
key=lambda x: x.count,
reverse=True):
print(ws)
def main(args):
if args.max_tokens is None:
args.max_tokens = 6000
print(args)
if not torch.cuda.is_available():
raise NotImplementedError('Training on CPU is not supported')
torch.cuda.set_device(args.device_id)
torch.manual_seed(args.seed)
# Setup task, e.g., translation, language modeling, etc.
task = tasks.setup_task(args)
# Load dataset splits
load_dataset_splits(task, ['train', 'valid'])
# Build model and criterion
model = task.build_model(args)
criterion = task.build_criterion(args)
print('| model {}, criterion {}'.format(args.arch,
criterion.__class__.__name__))
print('| num. model params: {}'.format(
sum(p.numel() for p in model.parameters())))
# Make a dummy batch to (i) warm the caching allocator and (ii) as a
# placeholder DistributedDataParallel when there's an uneven number of
# batches per worker.
max_positions = utils.resolve_max_positions(
task.max_positions(),
model.max_positions(),
)
dummy_batch = task.dataset('train').get_dummy_batch(
args.max_tokens, max_positions)
# Build trainer
trainer = Trainer(args, task, model, criterion, dummy_batch)
print('| training on {} GPUs'.format(args.distributed_world_size))
print('| max tokens per GPU = {} and max sentences per GPU = {}'.format(
args.max_tokens,
args.max_sentences,
))
## "Prune" heads (actually mask but shh...)
#if len(args.transformer_mask_heads) > 0:
# # Determine which head to prune
# to_prune = parse_head_pruning_descriptors(
# args.transformer_mask_heads,
# reverse_descriptors=args.transformer_mask_all_but_one_head,
# n_heads=model.encoder.layers[0].self_attn.num_heads
# )
# print(to_prune)
# # Apply pruning
# mask_heads(model, to_prune, args.transformer_mask_rescale)
# Save initial model
initial = os.path.join(args.save_dir, "checkpoint_initial.pt")
trainer.save_checkpoint(initial, {})
# Initialize dataloader
epoch_itr = task.get_batch_iterator(
dataset=task.dataset(args.train_subset),
max_tokens=args.max_tokens,
max_sentences=args.max_sentences,
max_positions=max_positions,
ignore_invalid_inputs=True,
required_batch_size_multiple=8,
seed=args.seed,
num_shards=args.distributed_world_size,
shard_id=args.distributed_rank,
)
# Load the latest checkpoint if one is available
if not load_checkpoint(args, trainer, epoch_itr):
trainer.dummy_train_step([dummy_batch])
# Train until the learning rate gets too small
max_epoch = args.max_epoch or math.inf
max_update = args.max_update or math.inf
lr = trainer.get_lr()
train_meter = StopwatchMeter()
train_meter.start()
valid_losses = [None]
valid_subsets = args.valid_subset.split(',')
while lr > args.min_lr and epoch_itr.epoch < max_epoch and trainer.get_num_updates(
) < max_update:
# train for one epoch
train(args, trainer, task, epoch_itr)
if epoch_itr.epoch % args.validate_interval == 0:
valid_losses = validate(args, trainer, task, epoch_itr,
valid_subsets)
# only use first validation loss to update the learning rate
lr = trainer.lr_step(epoch_itr.epoch, valid_losses[0])
#if epoch_itr.epoch % args.save_interval == 0:
# save_checkpoint(args, trainer, epoch_itr, valid_losses[0])
# ***********************************Below Changed******************************
# save checkpoint
#if epoch_itr.epoch % args.save_interval == 0:
save_interval = 5 # prune and save checkpoint for every five epoch
if epoch_itr.epoch % save_interval == 0: #****** changed
# do prunning before saving
prune2(args, task, model, trainer, epoch_itr) #****** changed2
# save checkpoint
save_checkpoint(args, trainer, epoch_itr, valid_losses[0])
prune2(args, task, model, trainer, epoch_itr
) #****** changed2 do last prunning on the last chekcpoint saved
# ***********************************Above Changed******************************
train_meter.stop()
print('| done training in {:.1f} seconds'.format(train_meter.sum))
def prune(args):
if args.max_tokens is None:
args.max_tokens = 6000
print(args)
if not torch.cuda.is_available():
raise NotImplementedError('Training on CPU is not supported')
torch.cuda.set_device(args.device_id)
torch.manual_seed(args.seed)
# Setup task, e.g., translation, language modeling, etc.
task = tasks.setup_task(args)
# Load dataset splits
load_dataset_splits(task, ['train', "valid"])
# Build model and criterion
model = task.build_model(args)
criterion = task.build_criterion(args)
print('| model {}, criterion {},'.format(args.arch,
criterion.__class__.__name__))
print('| num. model params: {}'.format(
sum(p.numel() for p in model.parameters())))
# Make a dummy batch to (i) warm the caching allocator and (ii) as a
# placeholder DistributedDataParallel when there's an uneven number of
# batches per worker.
max_positions = utils.resolve_max_positions(
task.max_positions(),
model.max_positions(),
)
dummy_batch = task.dataset('train').get_dummy_batch(
args.max_tokens, max_positions)
# Build trainer
trainer = Trainer(args, task, model, criterion, dummy_batch)
print('| training on {} GPUs'.format(args.distributed_world_size))
print('| max tokens per GPU = {} and max sentences per GPU = {}'.format(
args.max_tokens,
args.max_sentences,
))
print('| Optimizer {}'.format(trainer.optimizer.__class__.__name__))
# Initialize dataloader
epoch_itr = task.get_batch_iterator(
dataset=task.dataset(args.train_subset),
max_tokens=args.max_tokens,
max_sentences=args.max_sentences,
max_positions=max_positions,
ignore_invalid_inputs=True,
required_batch_size_multiple=8,
seed=args.seed,
num_shards=args.distributed_world_size,
shard_id=args.distributed_rank,
)
# Load the latest checkpoint if one is available
if not load_checkpoint(args, trainer, epoch_itr):
trainer.dummy_train_step([dummy_batch])
# Train until the learning rate gets too small
prune_meter = StopwatchMeter()
prune_meter.start()
# Estimate head importance scores
head_importance, head_stats = estimate_head_importance(
args, trainer, task, epoch_itr)
prune_meter.stop()
print('| done estimating head importance in {:.1f} seconds'.format(
prune_meter.sum))
torch.save(head_stats,
f"{os.path.dirname(args.restore_file)}/heads_stats.bin")
# Print
print("Head importances")
print("Encoder self attention")
for layer in range(head_importance["encoder_self"].size(0)):
print("\t".join(f"{x:.5f}"
for x in head_importance["encoder_self"][layer]))
print("Encoder decoder attention")
for layer in range(head_importance["encoder_decoder"].size(0)):
print("\t".join(f"{x:.5f}"
for x in head_importance["encoder_decoder"][layer]))
print("Decoder self attention")
for layer in range(head_importance["decoder_self"].size(0)):
print("\t".join(f"{x:.5f}"
for x in head_importance["decoder_self"][layer]))
# Print sorted pruning profile
encoder_self_profile = get_profile(head_importance["encoder_self"],
prefix="E")
encoder_decoder_profile = get_profile(head_importance["encoder_decoder"],
prefix="A")
decoder_self_profile = get_profile(head_importance["decoder_self"],
prefix="D")
# Join all
all_profiles = {}
if not (args.decoder_self_only or args.encoder_decoder_only):
all_profiles.update(encoder_self_profile)
if not (args.encoder_self_only or args.decoder_self_only):
all_profiles.update(encoder_decoder_profile)
if not (args.encoder_self_only or args.encoder_decoder_only):
all_profiles.update(decoder_self_profile)
sorted_profiles = sorted(all_profiles.items(),
key=lambda x: x[1],
reverse=args.one_minus)
print("Heads sorted by importance:")
print(" ".join(p for p, _ in sorted_profiles))
print("Sorted head importance scores:")
print(" ".join(f"{v.data:.5f}" for _, v in sorted_profiles))
if args.only_importance:
return
tot_n_heads = len(sorted_profiles)
# Eval pruning
if args.one_head:
kept_layers = set()
to_prune_profile = []
for p, _ in reversed(sorted_profiles):
layer_name = ":".join(p.split(":")[:-1])
if layer_name not in kept_layers:
kept_layers.add(layer_name)
continue
else:
to_prune_profile.insert(0, p)
to_prune = parse_head_pruning_descriptors(to_prune_profile,
reverse_descriptors=False)
print(f"Evaluating following profile: \t{' '.join(to_prune_profile)}")
# Apply pruning
mask_heads(model, to_prune, args.transformer_mask_rescale)
bleu = eval_bleu_score(
model,
task,
task.dataset(args.valid_subset),
beam=args.beam,
replace_unk=args.replace_unk,
lenpen=args.lenpen,
buffer_size=100,
use_cuda=torch.cuda.is_available() and not args.cpu,
remove_bpe=args.remove_bpe,
max_sentences=args.max_sentences,
max_tokens=args.max_tokens,
stop_early=not args.no_early_stop,
normalize_scores=not args.unnormalized,
min_len=args.min_len,
)
print(f"BLEU score: \t{bleu.score:.2f}")
sys.stdout.flush()
return
for i in range(0, 10):
n_to_prune = int(ceil(tot_n_heads * i / 10))
to_prune_profile = [p for p, _ in sorted_profiles[:n_to_prune]]
to_prune = parse_head_pruning_descriptors(to_prune_profile,
reverse_descriptors=False)
print(f"Evaluating following profile: \t{' '.join(to_prune_profile)}")
# Apply pruning
mask_heads(model, to_prune, args.transformer_mask_rescale)
bleu = eval_bleu_score(
model,
task,
task.dataset(args.valid_subset),
beam=args.beam,
replace_unk=args.replace_unk,
lenpen=args.lenpen,
buffer_size=100,
use_cuda=torch.cuda.is_available() and not args.cpu,
remove_bpe=args.remove_bpe,
max_sentences=args.max_sentences,
max_tokens=args.max_tokens,
stop_early=not args.no_early_stop,
normalize_scores=not args.unnormalized,
min_len=args.min_len,
)
print(f"BLEU score: \t{bleu.score:.2f}")
sys.stdout.flush()
def prune2(args, task, model, trainer, epoch_itr): # changed2
if args.max_tokens is None:
args.max_tokens = 6000
print(args)
# avoid aliasing
task = copy.deepcopy(task)
model = copy.deepcopy(model)
trainer = copy.deepcopy(trainer)
epoch_itr = copy.deepcopy(trainer)
if not torch.cuda.is_available():
raise NotImplementedError('Training on CPU is not supported')
torch.cuda.set_device(args.device_id)
torch.manual_seed(args.seed)
# Make a dummy batch to (i) warm the caching allocator and (ii) as a
# placeholder DistributedDataParallel when there's an uneven number of
# batches per worker.
max_positions = utils.resolve_max_positions(
task.max_positions(),
model.max_positions(),
)
dummy_batch = task.dataset('train').get_dummy_batch(
args.max_tokens, max_positions)
print('| training on {} GPUs'.format(args.distributed_world_size))
print('| max tokens per GPU = {} and max sentences per GPU = {}'.format(
args.max_tokens,
args.max_sentences,
))
print('| Optimizer {}'.format(trainer.optimizer.__class__.__name__))
# Train until the learning rate gets too small
prune_meter = StopwatchMeter()
prune_meter.start()
# Estimate head importance scores
head_importance, head_stats = estimate_head_importance(
args, trainer, task, epoch_itr)
prune_meter.stop()
print('| done estimating head importance in {:.1f} seconds'.format(
prune_meter.sum))
torch.save(head_stats,
f"{os.path.dirname(args.restore_file)}/heads_stats.bin")
# Print
print("Head importances")
print("Encoder self attention")
for layer in range(head_importance["encoder_self"].size(0)):
print("\t".join(f"{x:.5f}"
for x in head_importance["encoder_self"][layer]))
print("Encoder decoder attention")
for layer in range(head_importance["encoder_decoder"].size(0)):
print("\t".join(f"{x:.5f}"
for x in head_importance["encoder_decoder"][layer]))
print("Decoder self attention")
for layer in range(head_importance["decoder_self"].size(0)):
print("\t".join(f"{x:.5f}"
for x in head_importance["decoder_self"][layer]))
# Print sorted pruning profile
encoder_self_profile = get_profile(head_importance["encoder_self"],
prefix="E")
encoder_decoder_profile = get_profile(head_importance["encoder_decoder"],
prefix="A")
decoder_self_profile = get_profile(head_importance["decoder_self"],
prefix="D")
# Join all
all_profiles = {}
if not (args.decoder_self_only or args.encoder_decoder_only):
all_profiles.update(encoder_self_profile)
if not (args.encoder_self_only or args.decoder_self_only):
all_profiles.update(encoder_decoder_profile)
if not (args.encoder_self_only or args.encoder_decoder_only):
all_profiles.update(decoder_self_profile)
sorted_profiles = sorted(all_profiles.items(),
key=lambda x: x[1],
reverse=args.one_minus)
print("Heads sorted by importance:")
print(" ".join(p for p, _ in sorted_profiles))
print("Sorted head importance scores:")
print(" ".join(f"{v.data:.5f}" for _, v in sorted_profiles))
tot_n_heads = len(sorted_profiles)
for i in range(0, 10):
n_to_prune = int(ceil(tot_n_heads * i / 10))
to_prune_profile = [p for p, _ in sorted_profiles[:n_to_prune]]
to_prune = parse_head_pruning_descriptors(to_prune_profile,
reverse_descriptors=False)
print(f"Evaluating following profile: \t{' '.join(to_prune_profile)}")
# Apply pruning
mask_heads(model, to_prune, args.transformer_mask_rescale)
bleu = eval_bleu_score(
model,
task,
task.dataset(args.valid_subset),
beam=args.beam,
replace_unk=args.replace_unk,
lenpen=args.lenpen,
buffer_size=100,
use_cuda=torch.cuda.is_available() and not args.cpu,
remove_bpe=args.remove_bpe,
max_sentences=args.max_sentences,
max_tokens=args.max_tokens,
stop_early=not args.no_early_stop,
normalize_scores=not args.unnormalized,
min_len=args.min_len,
)
print(f"BLEU score: \t{bleu.score:.2f}")
sys.stdout.flush()
def main(args, init_distributed=False):
utils.import_user_module(args)
assert args.max_tokens is not None or args.max_sentences is not None, \
'Must specify batch size either with --max-tokens or --max-sentences'
# Initialize CUDA and distributed training
if torch.cuda.is_available() and not args.cpu:
torch.cuda.set_device(args.device_id)
np.random.seed(args.seed)
torch.manual_seed(args.seed)
if init_distributed:
args.distributed_rank = distributed_utils.distributed_init(args)
if distributed_utils.is_master(args):
checkpoint_utils.verify_checkpoint_directory(args.save_dir)
# Print args
logger.info(args)
# Setup task, e.g., translation, language modeling, etc.
task = tasks.setup_task(args)
# Load valid dataset (we load training data below, based on the latest checkpoint)
for valid_sub_split in args.valid_subset.split(','):
task.load_dataset(valid_sub_split, combine=False, epoch=0)
# Build model and criterion
model = task.build_model(args)
criterion = task.build_criterion(args)
logger.info(model)
logger.info('model {}, criterion {}'.format(args.arch,
criterion.__class__.__name__))
logger.info('num. model params: {} (num. trained: {})'.format(
sum(p.numel() for p in model.parameters()),
sum(p.numel() for p in model.parameters() if p.requires_grad),
))
# Build trainer
trainer = Trainer(args, task, model, criterion)
logger.info('training on {} GPUs'.format(args.distributed_world_size))
logger.info(
'max tokens per GPU = {} and max sentences per GPU = {}'.format(
args.max_tokens,
args.max_sentences,
))
# Load the latest checkpoint if one is available and restore the
# corresponding train iterator
extra_state, epoch_itr = checkpoint_utils.load_checkpoint(args, trainer)
# Train until the learning rate gets too small
train_meter = StopwatchMeter()
train_meter.start()
valid_subsets = args.valid_subset.split(',')
tokenize = sacrebleu.DEFAULT_TOKENIZER if not args.eval_tokenized_bleu else 'none'
hyps, refs = validate(args, trainer, task, epoch_itr, valid_subsets)
for h, r, split in zip(hyps, refs, args.valid_subset.split(',')):
assert len(h) == len(r)
sacrebleu_score, _, _ = sacrebleu.corpus_bleu(
h, [r], tokenize=tokenize), hyps, refs
bleu = compute_cvpr_bleu(h, r)
rouge_score = rouge.rouge(h, r)
print('{} set has {} samples,\n'
'sacrebleu: {},\n'
'CVPR BLEU scripts: {}\n'
'CVPR ROUGE: {}'.format(split, len(h), sacrebleu_score, bleu,
rouge_score))
print('performance: {:.2f} {}'.format(
rouge_score['rouge_l/f_score'] * 100,
' '.join([str(b) for b in bleu])))
def _generate_score(models, args, dataset, dataset_split):
use_cuda = torch.cuda.is_available() and not args.cpu
# Load ensemble
if not args.quiet:
print("| loading model(s) from {}".format(", ".join(args.path)))
# Optimize ensemble for generation
for model in models:
model.make_generation_fast_(
beamable_mm_beam_size=None if args.no_beamable_mm else args.beam)
# Initialize generator
model_weights = None
if args.model_weights:
model_weights = [
float(w.strip()) for w in args.model_weights.split(",")
]
translator = beam_decode.SequenceGenerator(
models,
beam_size=args.beam,
stop_early=(not args.no_early_stop),
normalize_scores=(not args.unnormalized),
len_penalty=args.lenpen,
unk_penalty=args.unkpen,
word_reward=args.word_reward,
model_weights=model_weights,
)
if use_cuda:
translator.cuda()
# Load alignment dictionary for unknown word replacement
# (None if no unknown word replacement, empty if no path to align dictionary)
align_dict = utils.load_align_dict(args.replace_unk)
# Generate and compute BLEU score
scorer = bleu.Scorer(dataset.dst_dict.pad(), dataset.dst_dict.eos(),
dataset.dst_dict.unk())
max_positions = min(model.max_encoder_positions() for model in models)
itr = dataset.eval_dataloader(
dataset_split,
max_sentences=args.max_sentences,
max_positions=max_positions,
skip_invalid_size_inputs_valid_test=(
args.skip_invalid_size_inputs_valid_test),
)
if args.num_shards > 1:
if args.shard_id < 0 or args.shard_id >= args.num_shards:
raise ValueError("--shard-id must be between 0 and num_shards")
itr = data.sharded_iterator(itr, args.num_shards, args.shard_id)
num_sentences = 0
with progress_bar.build_progress_bar(args, itr) as t:
wps_meter = TimeMeter()
gen_timer = StopwatchMeter()
translations = translator.generate_batched_itr(
t,
maxlen_a=args.max_len_a,
maxlen_b=args.max_len_b,
cuda=use_cuda,
timer=gen_timer,
)
for sample_id, src_tokens, target_tokens, hypos in translations:
# Process input and ground truth
target_tokens = target_tokens.int().cpu()
# Either retrieve the original sentences or regenerate them from tokens.
if align_dict is not None:
src_str = dataset.splits[dataset_split].src.get_original_text(
sample_id)
target_str = dataset.splits[
dataset_split].dst.get_original_text(sample_id)
else:
src_str = dataset.src_dict.string(src_tokens, args.remove_bpe)
target_str = dataset.dst_dict.string(target_tokens,
args.remove_bpe,
escape_unk=True)
if not args.quiet:
print(f"S-{sample_id}\t{src_str}")
print(f"T-{sample_id}\t{target_str}")
# Process top predictions
for i, hypo in enumerate(hypos[:min(len(hypos), args.nbest)]):
hypo_tokens, hypo_str, alignment = utils.post_process_prediction(
hypo_tokens=hypo["tokens"].int().cpu(),
src_str=src_str,
alignment=hypo["alignment"].int().cpu(),
align_dict=align_dict,
dst_dict=dataset.dst_dict,
remove_bpe=args.remove_bpe,
)
if not args.quiet:
print(f"H-{sample_id}\t{hypo['score']}\t{hypo_str}")
print("A-{}\t{}".format(
sample_id,
" ".join(map(lambda x: str(utils.item(x)), alignment)),
))
# Score only the top hypothesis
if i == 0:
if align_dict is not None or args.remove_bpe is not None:
# Convert back to tokens for evaluation with unk replacement
# and/or without BPE
target_tokens = tokenizer.Tokenizer.tokenize(
target_str,
dataset.dst_dict,
add_if_not_exist=True)
scorer.add(target_tokens, hypo_tokens)
wps_meter.update(src_tokens.size(0))
t.log({"wps": round(wps_meter.avg)})
num_sentences += 1
return scorer, num_sentences, gen_timer
def main(args, init_distributed=False):
utils.import_user_module(args)
assert args.max_tokens is not None or args.max_sentences is not None, \
'Must specify batch size either with --max-tokens or --max-sentences'
# Initialize CUDA and distributed training
if torch.cuda.is_available() and not args.cpu:
torch.cuda.set_device(args.device_id)
np.random.seed(args.seed)
torch.manual_seed(args.seed)
if init_distributed:
args.distributed_rank = distributed_utils.distributed_init(args)
if distributed_utils.is_master(args):
checkpoint_utils.verify_checkpoint_directory(args.save_dir)
# Print args
logger.info(args)
# Setup task, e.g., translation, language modeling, etc.
task = tasks.setup_task(args)
# Load valid dataset (we load training data below, based on the latest checkpoint)
for valid_sub_split in args.valid_subset.split(','):
task.load_dataset(valid_sub_split, combine=False, epoch=0)
# Build model and criterion
model = task.build_model(args)
criterion = task.build_criterion(args)
logger.info(model)
logger.info('model {}, criterion {}'.format(args.arch, criterion.__class__.__name__))
logger.info('num. model params: {} (num. trained: {})'.format(
sum(p.numel() for p in model.parameters()),
sum(p.numel() for p in model.parameters() if p.requires_grad),
))
# Build trainer
trainer = Trainer(args, task, model, criterion)
logger.info('training on {} GPUs'.format(args.distributed_world_size))
logger.info('max tokens per GPU = {} and max sentences per GPU = {}'.format(
args.max_tokens,
args.max_sentences,
))
# Load the latest checkpoint if one is available and restore the
# corresponding train iterator
extra_state, epoch_itr = checkpoint_utils.load_checkpoint(args, trainer)
# Train until the learning rate gets too small
max_epoch = args.max_epoch or math.inf
max_update = args.max_update or math.inf
lr = trainer.get_lr()
train_meter = StopwatchMeter()
train_meter.start()
valid_subsets = args.valid_subset.split(',')
print(args.multi_views)
while (
lr > args.min_lr
and (
epoch_itr.epoch < max_epoch
# allow resuming training from the final checkpoint
or epoch_itr._next_epoch_itr is not None
)
and trainer.get_num_updates() < max_update
):
# train for one epoch
train(args, trainer, task, epoch_itr)
if not args.disable_validation and epoch_itr.epoch % args.validate_interval == 0:
valid_losses = validate(args, trainer, task, epoch_itr, valid_subsets)
else:
valid_losses = [None]
# only use first validation loss to update the learning rate
lr = trainer.lr_step(epoch_itr.epoch, valid_losses[0])
bart = BARTHubInterface(args, task, trainer.model).cuda()
#print(bart.device)
bart.eval()
count = 1
bsz = 8
print("Test on val set: ")
with open('../data/val_sent_trans_cons_label.source') as source, open('../data/val_sent_c99_label.source') as source2, open('./val_best_multi_attn_'+str(args.lr_weight)+'_.hypo', 'wt', encoding='utf-8') as fout:
s1 = source.readlines()
s2 = source2.readlines()
slines = [s1[0].strip()]
slines2 = [s2[0].strip()]
for i in tqdm(range(1, len(s1))):
if count % bsz == 0:
with torch.no_grad():
if args.multi_views:
hypotheses_batch = bart.sample(slines, sentences2 = slines2, balance = True, beam=4, lenpen=2.0, max_len_b=100, min_len=5, no_repeat_ngram_size=3)
else:
hypotheses_batch = bart.sample(slines, beam=4, lenpen=2.0, max_len_b=100, min_len=5, no_repeat_ngram_size=3)
for hypothesis in hypotheses_batch:
fout.write(hypothesis + '\n')
fout.flush()
slines = []
slines2 = []
slines.append(s1[i].strip())
slines2.append(s2[i].strip())
count += 1
if slines != []:
if args.multi_views:
hypotheses_batch = bart.sample(slines, sentences2 = slines2, balance = True, beam=4, lenpen=2.0, max_len_b=100, min_len=5, no_repeat_ngram_size=3)
else:
hypotheses_batch = bart.sample(slines, beam=4, lenpen=2.0, max_len_b=100, min_len=5, no_repeat_ngram_size=3)
#hypotheses_batch = bart.sample(slines, sentences2 = slines2, balance = True, beam=4, lenpen=2.0, max_len_b=100, min_len=5, no_repeat_ngram_size=3)
for hypothesis in hypotheses_batch:
fout.write(hypothesis + '\n')
fout.flush()
hyp_path = './val_best_multi_attn_'+str(args.lr_weight)+'_.hypo'
ref_path = '../data/val_sent_trans_cons_label.target'
hypothesis = []
with open(hyp_path, 'r') as f:
lines = f.readlines()
for l in lines:
hypothesis.append(l[:-1])
reference = []
with open(ref_path, 'r') as f:
lines = f.readlines()
for l in lines:
reference.append(l[:-1])
rouge = Rouge()
print("Val", rouge.get_scores(hypothesis, reference, avg = True))
# save checkpoint
if epoch_itr.epoch % args.save_interval == 0:
checkpoint_utils.save_checkpoint(args, trainer, epoch_itr, valid_losses[0])
print("Test on testing set: ")
count = 1
bsz = 8
with open('../data/test_sent_trans_cons_label.source') as source, open('../data/test_sent_c99_label.source') as source2, open('./test_best_multi_attn_'+str(args.lr_weight)+'_.hypo', 'wt', encoding='utf-8') as fout:
s1 = source.readlines()
s2 = source2.readlines()
slines = [s1[0].strip()]
slines2 = [s2[0].strip()]
for i in tqdm(range(1, len(s1))):
if count % bsz == 0:
with torch.no_grad():
if args.multi_views:
hypotheses_batch = bart.sample(slines, sentences2 = slines2, balance = True, beam=4, lenpen=2.0, max_len_b=100, min_len=5, no_repeat_ngram_size=3)
else:
hypotheses_batch = bart.sample(slines, beam=4, lenpen=2.0, max_len_b=100, min_len=5, no_repeat_ngram_size=3)
for hypothesis in hypotheses_batch:
fout.write(hypothesis + '\n')
fout.flush()
slines = []
slines2 = []
slines.append(s1[i].strip())
slines2.append(s2[i].strip())
count += 1
if slines != []:
if args.multi_views:
hypotheses_batch = bart.sample(slines, sentences2 = slines2, balance = True, beam=4, lenpen=2.0, max_len_b=100, min_len=5, no_repeat_ngram_size=3)
else:
hypotheses_batch = bart.sample(slines, beam=4, lenpen=2.0, max_len_b=100, min_len=5, no_repeat_ngram_size=3)
for hypothesis in hypotheses_batch:
fout.write(hypothesis + '\n')
fout.flush()
hyp_path = './test_best_multi_attn_'+str(args.lr_weight)+'_.hypo'
ref_path = '../data/test_sent_trans_cons_label.target'
hypothesis = []
with open(hyp_path, 'r') as f:
lines = f.readlines()
for l in lines:
hypothesis.append(l[:-1])
reference = []
with open(ref_path, 'r') as f:
lines = f.readlines()
for l in lines:
reference.append(l[:-1])
rouge = Rouge()
print('Test', rouge.get_scores(hypothesis, reference, avg = True))
# early stop
if should_stop_early(args, valid_losses[0]):
logger.info('early stop since valid performance hasn\'t improved for last {} runs'.format(args.patience))
break
epoch_itr = trainer.get_train_iterator(
epoch_itr.epoch,
# sharded data: get train iterator for next epoch
load_dataset=(os.pathsep in getattr(args, 'data', '')),
)
train_meter.stop()
logger.info('done training in {:.1f} seconds'.format(train_meter.sum))
def _generate_score(models, args, task, dataset, modify_target_dict):
use_cuda = torch.cuda.is_available() and not args.cpu
# Load ensemble
if not args.quiet:
print("| loading model(s) from {}".format(", ".join(args.path.split(":"))))
# Optimize ensemble for generation
for model in models:
model.make_generation_fast_(
beamable_mm_beam_size=None if args.no_beamable_mm else args.beam,
need_attn=True,
)
translator = build_sequence_generator(args, task, models)
# Load alignment dictionary for unknown word replacement
# (None if no unknown word replacement, empty if no path to align dictionary)
align_dict = utils.load_align_dict(args.replace_unk)
print("seed number is" + str(args.max_examples_to_evaluate_seed))
if args.max_examples_to_evaluate > 0:
pytorch_translate_data.subsample_pair_dataset(
dataset, args.max_examples_to_evaluate, args.max_examples_to_evaluate_seed
)
# Keep track of translations
# Initialize with empty translations
# and zero probs scores
translated_sentences = [""] * len(dataset)
translated_scores = [0.0] * len(dataset)
hypos_list = []
collect_output_hypos = getattr(args, "output_hypos_binary_path", False)
if collect_output_hypos:
output_hypos_token_arrays = [None] * len(dataset)
# Generate and compute BLEU score
dst_dict = task.target_dictionary
if args.sacrebleu:
scorer = bleu.SacrebleuScorer()
else:
scorer = bleu.Scorer(dst_dict.pad(), dst_dict.eos(), dst_dict.unk())
itr = task.get_batch_iterator(
dataset=dataset,
max_tokens=args.max_tokens,
max_sentences=args.max_sentences,
max_positions=utils.resolve_max_positions(
task.max_positions(), *[model.max_positions() for model in models]
),
ignore_invalid_inputs=args.skip_invalid_size_inputs_valid_test,
required_batch_size_multiple=8,
num_shards=args.num_shards,
shard_id=args.shard_id,
num_workers=args.num_workers,
).next_epoch_itr(shuffle=False)
oracle_scorer = None
if args.report_oracle_bleu:
oracle_scorer = bleu.Scorer(dst_dict.pad(), dst_dict.eos(), dst_dict.unk())
rescorer = None
num_sentences = 0
translation_samples = []
translation_info_list = []
with progress_bar.build_progress_bar(args, itr) as t:
wps_meter = TimeMeter()
gen_timer = StopwatchMeter()
translations = translator.generate_batched_itr(
t,
maxlen_a=args.max_len_a,
maxlen_b=args.max_len_b,
cuda=use_cuda,
timer=gen_timer,
prefix_size=1
if pytorch_translate_data.is_multilingual_many_to_one(args)
else 0,
)
for trans_info in _iter_translations(
args, task, dataset, translations, align_dict, rescorer, modify_target_dict
):
if hasattr(scorer, "add_string"):
scorer.add_string(trans_info.target_str, trans_info.hypo_str)
else:
scorer.add(trans_info.target_tokens, trans_info.hypo_tokens)
if oracle_scorer is not None:
oracle_scorer.add(trans_info.target_tokens, trans_info.best_hypo_tokens)
if getattr(args, "translation_output_file", False):
translated_sentences[trans_info.sample_id] = trans_info.hypo_str
if getattr(args, "translation_probs_file", False):
translated_scores[trans_info.sample_id] = trans_info.hypo_score
if getattr(args, "hypotheses_export_path", False):
hypos_list.append(trans_info.hypos)
if collect_output_hypos:
output_hypos_token_arrays[
trans_info.sample_id
] = trans_info.best_hypo_tokens
if args.translation_info_export_path is not None:
# Strip expensive data from hypotheses before saving
hypos = [
{k: v for k, v in hypo.items() if k in ["tokens", "score"]}
for hypo in trans_info.hypos
]
# Make sure everything is on cpu before exporting
hypos = [
{"score": hypo["score"], "tokens": hypo["tokens"].cpu()}
for hypo in hypos
]
translation_info_list.append(
{
"src_tokens": trans_info.src_tokens.cpu(),
"target_tokens": trans_info.target_tokens,
"hypos": hypos,
}
)
translation_samples.append(
collections.OrderedDict(
{
"sample_id": trans_info.sample_id.item(),
"src_str": trans_info.src_str,
"target_str": trans_info.target_str,
"hypo_str": trans_info.hypo_str,
}
)
)
wps_meter.update(trans_info.src_tokens.size(0))
t.log({"wps": round(wps_meter.avg)})
num_sentences += 1
# If applicable, save collected hypothesis tokens to binary output file
if collect_output_hypos:
output_dataset = pytorch_translate_data.InMemoryIndexedDataset()
output_dataset.load_from_sequences(output_hypos_token_arrays)
output_dataset.save(args.output_hypos_binary_path)
if args.output_source_binary_path:
dataset.src.save(args.output_source_binary_path)
if args.translation_info_export_path is not None:
f = open(args.translation_info_export_path, "wb")
pickle.dump(translation_info_list, f)
f.close()
# If applicable, save the translations and scores to the output files
# These two ouputs are used in dual learning for weighted backtranslation
if getattr(args, "translation_output_file", False) and getattr(
args, "translation_probs_file", False
):
with open(args.translation_output_file, "w") as translation_file, open(
args.translation_probs_file, "w"
) as score_file:
for hypo_str, hypo_score in zip(translated_sentences, translated_scores):
if len(hypo_str.strip()) > 0:
print(hypo_str, file=translation_file)
print(np.exp(hypo_score), file=score_file)
# For eg. external evaluation
if getattr(args, "hypotheses_export_path", False):
with open(args.hypotheses_export_path, "w") as out_file:
for hypos in hypos_list:
for hypo in hypos:
print(
task.tgt_dict.string(
hypo["tokens"], bpe_symbol=args.remove_bpe
),
file=out_file,
)
if oracle_scorer is not None:
print(f"| Oracle BLEU (best hypo in beam): {oracle_scorer.result_string()}")
return scorer, num_sentences, gen_timer, translation_samples
def main(args, init_distributed=False):
import_user_module(args)
if args.max_tokens is None:
args.max_tokens = 6000
print(args)
#args.distributed_world_size = 1
if torch.cuda.is_available() and not args.cpu:
torch.cuda.set_device(args.device_id)
torch.manual_seed(args.seed)
# Setup task, e.g., translation, language modeling, etc.
task = tasks.setup_task(args)
# Load dataset splits
load_dataset_splits(args, task)
# Build model and criterion
model = task.build_model(args)
criterion = task.build_criterion(args)
print(model)
print('| model {}, criterion {}'.format(args.arch,
criterion.__class__.__name__))
print('| num. model params: {} (num. trained: {})'.format(
sum(p.numel() for p in model.parameters()),
sum(p.numel() for p in model.parameters() if p.requires_grad),
))
# Make a dummy batch to (i) warm the caching allocator and (ii) as a
# placeholder DistributedDataParallel when there's an uneven number of
# batches per worker.
max_positions = utils.resolve_max_positions(
task.max_positions(),
model.max_positions(),
)
dummy_batch = task.dataset(args.train_subset).get_dummy_batch(
args.max_tokens, max_positions)
oom_batch = task.dataset(args.train_subset).get_dummy_batch(
1, max_positions)
# Build trainer
print("Building trainer...")
trainer = Trainer(args, task, model, criterion, dummy_batch, oom_batch)
print('| training on {} GPUs'.format(args.distributed_world_size))
print('| max tokens per GPU = {} and max sentences per GPU = {}'.format(
args.max_tokens,
args.max_sentences,
))
# Initialize dataloader
print("Initialize dataloader...")
epoch_itr = task.get_batch_iterator(
dataset=task.dataset(args.train_subset),
max_tokens=args.max_tokens,
max_sentences=args.max_sentences,
max_positions=max_positions,
ignore_invalid_inputs=True,
required_batch_size_multiple=args.required_batch_size_multiple,
seed=args.seed,
num_shards=args.distributed_world_size,
shard_id=args.distributed_rank,
num_workers=args.num_workers,
)
# Initialize distributed training (after data loading)
print("Initialize distributed training (after data loading)...")
if init_distributed:
import socket
args.distributed_rank = distributed_utils.distributed_init(args)
print('| initialized host {} as rank {}'.format(
socket.gethostname(), args.distributed_rank))
# Load the latest checkpoint if one is available
print("Load the latest checkpoint if one is available...")
if not load_checkpoint(args, trainer, epoch_itr):
trainer.dummy_train_step([dummy_batch])
if args.reset_target_embedding:
trainer.init_meters(args)
print("reset trainer.meters")
# Train until the learning rate gets too small
max_epoch = args.max_epoch or math.inf
max_update = args.max_update or math.inf
lr = trainer.get_lr()
train_meter = StopwatchMeter()
train_meter.start()
valid_losses = [None]
valid_subsets = args.valid_subset.split(',')
if args.distributed_rank == 0:
if os.path.basename(args.save_dir) != "":
log_file = os.path.join(
args.save_dir,
"({0})-params.log".format(os.path.basename(args.save_dir)))
else:
log_file = os.path.join(
args.save_dir,
"({0})-params.log".format(args.save_dir.split('/')[-2]))
# create log file
args.log_file = log_file
if os.path.exists(log_file):
w = open(log_file, "a+", encoding="utf-8")
else:
w = open(log_file, "w", encoding="utf-8")
w.write(str(args).replace(", ", ",\n") + "\n")
w.write(str(model) + "\n")
w.flush()
w.close()
print("saving params file into{}...".format(log_file))
while lr > args.min_lr and epoch_itr.epoch < max_epoch and trainer.get_num_updates(
) < max_update:
# train for one epoch
train(args, trainer, task, epoch_itr)
if epoch_itr.epoch % args.validate_interval == 0:
valid_losses = validate(args, trainer, task, epoch_itr,
valid_subsets)
# only use first validation loss to update the learning rate
lr = trainer.lr_step(epoch_itr.epoch, valid_losses[0])
# save checkpoint
if epoch_itr.epoch % args.save_interval == 0:
save_checkpoint(args, trainer, epoch_itr, valid_losses[0])
train_meter.stop()
print('| done training in {:.1f} seconds'.format(train_meter.sum))
def main(args):
assert args.path is not None, '--path required for generation!'
assert not args.sampling or args.nbest == args.beam, \
'--sampling requires --nbest to be equal to --beam'
assert args.replace_unk is None or args.raw_text, \
'--replace-unk requires a raw text dataset (--raw-text)'
import_user_module(args)
if args.max_tokens is None and args.max_sentences is None:
args.max_tokens = 12000
print(args)
use_cuda = torch.cuda.is_available() and not args.cpu
# ========== for bartsv task, rebuild dictionary after model args are loaded ==========
# assert not hasattr(args, 'node_freq_min'), 'node_freq_min should be read from model args'
# args.node_freq_min = 5 # temporarily set before model loading, as this is needed in tasks.setup_task(args)
# =====================================================================================
# Load dataset splits
task = tasks.setup_task(args)
# Note: states are not needed since they will be provided by the state
# machine
task.load_dataset(args.gen_subset, state_machine=False)
# Set dictionaries
try:
src_dict = getattr(task, 'source_dictionary', None)
except NotImplementedError:
src_dict = None
tgt_dict = task.target_dictionary
# Load ensemble
print('| loading model(s) from {}'.format(args.path))
try:
models, _model_args = checkpoint_utils.load_model_ensemble(
args.path.split(':'),
arg_overrides=eval(args.model_overrides),
task=task,
)
except:
# NOTE this is for "bartsv" models when default "args.node_freq_min" (5) is not equal to the model
# when loading model with the above task there will be an error when building the model with the task's
# target vocabulary, which would be of different size
# TODO better handle these cases (without sacrificing compatibility with other model archs)
models, _model_args = checkpoint_utils.load_model_ensemble(
args.path.split(':'),
arg_overrides=eval(args.model_overrides),
task=None,
)
# ========== for bartsv task, rebuild the dictionary based on model args ==========
if 'bartsv' in _model_args.arch and args.node_freq_min != _model_args.node_freq_min:
args.node_freq_min = _model_args.node_freq_min
# Load dataset splits
task = tasks.setup_task(args)
# Note: states are not needed since they will be provided by the state machine
task.load_dataset(args.gen_subset, state_machine=False)
# Set dictionaries
try:
src_dict = getattr(task, 'source_dictionary', None)
except NotImplementedError:
src_dict = None
tgt_dict = task.target_dictionary
# ==================================================================================
# import pdb; pdb.set_trace()
# print(_model_args)
# ========== for previous model trained when new arguments were not there ==========
if not hasattr(_model_args, 'shift_pointer_value'):
_model_args.shift_pointer_value = 1
# ==================================================================================
# Optimize ensemble for generation
for model in models:
model.make_generation_fast_(
beamable_mm_beam_size=None if args.no_beamable_mm else args.beam,
need_attn=args.print_alignment,
)
if args.fp16:
model.half()
if use_cuda:
model.cuda()
# Load alignment dictionary for unknown word replacement
# (None if no unknown word replacement, empty if no path to align
# dictionary)
align_dict = utils.load_align_dict(args.replace_unk)
# Load dataset (possibly sharded)
itr = task.get_batch_iterator(
dataset=task.dataset(args.gen_subset),
max_tokens=args.max_tokens,
max_sentences=args.max_sentences,
max_positions=None,
ignore_invalid_inputs=args.skip_invalid_size_inputs_valid_test,
required_batch_size_multiple=args.required_batch_size_multiple,
num_shards=args.num_shards,
shard_id=args.shard_id,
num_workers=args.num_workers,
# large_sent_first=False # not in fairseq
).next_epoch_itr(shuffle=False)
# Initialize generator
gen_timer = StopwatchMeter()
generator = task.build_generator(args, _model_args)
# Generate and compute BLEU score
if args.sacrebleu:
scorer = bleu.SacrebleuScorer()
else:
scorer = bleu.Scorer(tgt_dict.pad(), tgt_dict.eos(), tgt_dict.unk())
num_sentences = 0
has_target = True
examples = Examples(args.path, args.results_path, args.gen_subset,
args.nbest)
error_stats = {'num_sub_start': 0}
with progress_bar.build_progress_bar(args, itr) as t:
wps_meter = TimeMeter()
for sample in t:
sample = utils.move_to_cuda(sample) if use_cuda else sample
if 'net_input' not in sample:
raise Exception("Did not expect empty sample")
continue
prefix_tokens = None
if args.prefix_size > 0:
prefix_tokens = sample['target'][:, :args.prefix_size]
# breakpoint()
gen_timer.start()
hypos = task.inference_step(generator, models, sample, args,
prefix_tokens)
num_generated_tokens = sum(len(h[0]['tokens']) for h in hypos)
gen_timer.stop(num_generated_tokens)
# breakpoint()
for i, sample_id in enumerate(sample['id'].tolist()):
has_target = sample['target'] is not None
# Remove padding
src_tokens = utils.strip_pad(
sample['net_input']['src_tokens'][i, :], tgt_dict.pad())
target_tokens = None
if has_target:
target_tokens = utils.strip_pad(
sample['target'][i, :], tgt_dict.pad()).int().cpu()
# Either retrieve the original sentences or regenerate them from tokens.
if align_dict is not None:
src_str = task.dataset(
args.gen_subset).src.get_original_text(sample_id)
target_str = task.dataset(
args.gen_subset).tgt.get_original_text(sample_id)
else:
if src_dict is not None:
src_str = src_dict.string(src_tokens, args.remove_bpe)
else:
src_str = ""
if has_target:
target_str = tgt_dict.string(target_tokens,
args.remove_bpe,
escape_unk=True)
# debug: '<<unk>>' is added to the dictionary
# if 'unk' in target_str:
# breakpoint()
# ==========> NOTE we do not really have the ground truth target (with the same alignments)
# target_str might have <unk> as the target dictionary is only built on training data
# but it doesn't matter. It should not affect the target dictionary!
if not args.quiet:
if src_dict is not None:
print('S-{}\t{}'.format(sample_id, src_str))
if has_target:
print('T-{}\t{}'.format(sample_id, target_str))
# Process top predictions
for j, hypo in enumerate(hypos[i][:args.nbest]):
# hypo_tokens, hypo_str, alignment = utils.post_process_prediction(
# hypo_tokens=hypo['tokens'].int().cpu(),
# src_str=src_str,
# alignment=hypo['alignment'].int().cpu() if hypo['alignment'] is not None else None,
# align_dict=align_dict,
# tgt_dict=tgt_dict,
# remove_bpe=args.remove_bpe,
# # FIXME: AMR specific
# split_token="\t",
# line_tokenizer=task.tokenize,
# )
if 'bartsv' in _model_args.arch:
if not tgt_dict[hypo['tokens'][0]].startswith(
tgt_dict.bpe.INIT):
error_stats['num_sub_start'] += 1
try:
actions_nopos, actions_pos, actions = post_process_action_pointer_prediction_bartsv(
hypo, tgt_dict)
except:
breakpoint()
else:
actions_nopos, actions_pos, actions = post_process_action_pointer_prediction(
hypo, tgt_dict)
# breakpoint()
if args.clean_arcs:
actions_nopos, actions_pos, actions, invalid_idx = clean_pointer_arcs(
actions_nopos, actions_pos, actions)
# TODO these are just dummy for the reference below to run
hypo_tokens = hypo['tokens'].int().cpu()
hypo_str = '/t'.join(actions)
alignment = None
# update the list of examples
examples.append({
'actions_nopos': actions_nopos,
'actions_pos': actions_pos,
'actions': actions,
'reference': target_str,
'src_str': src_str,
'sample_id': sample_id
})
if not args.quiet:
print('H-{}\t{}\t{}'.format(sample_id, hypo_str,
hypo['score']))
print('P-{}\t{}'.format(
sample_id, ' '.join(
map(
lambda x: '{:.4f}'.format(x),
hypo['positional_scores'].tolist(),
))))
if args.print_alignment:
print('A-{}\t{}'.format(
sample_id, ' '.join(
map(lambda x: str(utils.item(x)),
alignment))))
# Score only the top hypothesis
if has_target and j == 0:
if align_dict is not None or args.remove_bpe is not None:
# Convert back to tokens for evaluation with unk replacement and/or without BPE
target_tokens = tgt_dict.encode_line(
target_str, add_if_not_exist=False)
# NOTE do not modify the tgt dictionary with 'add_if_not_exist=True'!
if hasattr(scorer, 'add_string'):
scorer.add_string(target_str, hypo_str)
else:
scorer.add(target_tokens, hypo_tokens)
wps_meter.update(num_generated_tokens)
t.log({'wps': round(wps_meter.avg)})
num_sentences += sample['nsentences']
# Save examples to files
examples.save()
print('| Error case (handled by manual fix) statistics:')
print(error_stats)
print(
'| Translated {} sentences ({} tokens) in {:.1f}s ({:.2f} sentences/s, {:.2f} tokens/s)'
.format(num_sentences, gen_timer.n, gen_timer.sum,
num_sentences / gen_timer.sum, 1. / gen_timer.avg))
if has_target:
print('| Generate {} with beam={}: {}'.format(args.gen_subset,
args.beam,
scorer.result_string()))
return scorer
def train(args, extra_state, trainer, dataset):
# offset for current epoch (may be different from checkpoint offset)
starting_offset = extra_state["batch_offset"]
# Train until the learning rate gets too small
max_epoch = args.max_epoch or math.inf
lr = trainer.get_lr()
train_meter = StopwatchMeter()
train_meter.start()
while lr > args.min_lr and extra_state["epoch"] <= max_epoch:
"""Train the model for one epoch."""
itr, progress, extra_meters = setup_epoch(
args=args,
epoch=extra_state["epoch"],
batch_offset=starting_offset,
trainer=trainer,
dataset=dataset,
)
for i, sample in enumerate(itr, start=starting_offset):
log_output = trainer.train_step(sample)
train_stats = log_mid_epoch_stats(
trainer=trainer,
progress=progress,
extra_meters=extra_meters,
log_output=log_output,
)
if (args.continuous_averaging_after_epochs >= 0
and extra_state["epoch"] >
args.continuous_averaging_after_epochs):
model_param_dict = trainer.model.state_dict()
if "param_totals" not in extra_state:
extra_state["param_totals"] = {}
for name, value in model_param_dict.items():
extra_state["param_totals"][name] = value.clone()
extra_state["param_accum_count"] = 1
else:
for name, value in model_param_dict.items():
extra_state["param_totals"][name] += value
extra_state["param_accum_count"] += 1
if i == starting_offset:
# ignore the first mini-batch in words-per-second calculation
trainer.get_meter("wps").reset()
num_updates = trainer.get_num_updates()
do_validate = (args.subepoch_validate_interval > 0 and
num_updates % args.subepoch_validate_interval == 0)
do_save = (not args.no_save and args.save_interval > 0
and num_updates % args.save_interval == 0)
do_eval_bleu = (
# We can only do BLEU eval when we have a new checkpoint to load.
do_save and args.generate_bleu_eval_interval > 0 and
num_updates - extra_state["last_bleu_eval"] >=
args.generate_bleu_eval_interval)
if do_eval_bleu:
extra_state["last_bleu_eval"] = num_updates
extra_state["batch_offset"] = i + 1
(
_,
val_ppl,
val_bleu,
stop_training_mid_epoch,
) = validate_save_and_evaluate_bleu(
args=args,
trainer=trainer,
dataset=dataset,
extra_state=extra_state,
do_validate=do_validate,
do_save=do_save,
do_eval_bleu=do_eval_bleu,
)
yield (
trainer.get_num_updates(),
{
"train_ppl": train_stats["ppl"],
"tune_ppl": val_ppl,
"tune_bleu": val_bleu,
},
)
if stop_training_mid_epoch:
break
# log end-of-epoch stats
train_stats = log_end_epoch_stats(trainer=trainer,
progress=progress,
extra_meters=extra_meters)
if stop_training_mid_epoch:
break
# batch_offset being None denotes the end of an epoch.
extra_state["batch_offset"] = None
(
val_loss,
val_ppl,
val_bleu,
stop_training_end_of_epoch,
) = validate_save_and_evaluate_bleu(
args=args,
trainer=trainer,
dataset=dataset,
extra_state=extra_state,
do_validate=True,
do_save=not args.no_save and not args.no_end_of_epoch_checkpoints,
do_eval_bleu=args.generate_bleu_eval_per_epoch,
)
extra_state["val_loss"] = val_loss
yield (
trainer.get_num_updates(),
{
"train_ppl": train_stats["ppl"],
"tune_ppl": val_ppl,
"tune_bleu": val_bleu,
},
)
if stop_training_end_of_epoch:
break
lr = trainer.lr_step(extra_state["epoch"], val_loss)
extra_state["epoch"] += 1
extra_state["batch_offset"] = 0
starting_offset = 0
if is_training_over_time_limit(extra_state["start_time"],
args.stop_time_hr):
break
train_meter.stop()
print(f"| done training in {train_meter.sum:.1f} seconds")
if "evaluate_bleu" in extra_state:
print(
f"| Best BLEU score of {extra_state['evaluate_bleu']['best']} was from "
f"epoch {extra_state['evaluate_bleu']['best_epoch']}")
def main(args):
print(args)
if not torch.cuda.is_available():
raise NotImplementedError('Training on CPU is not supported')
torch.cuda.set_device(args.device_id)
torch.manual_seed(args.seed)
# Load dataset
splits = ['train', 'valid']
if data.has_binary_files(args.data, splits):
dataset = data.load_dataset(
args.data, splits, args.source_lang, args.target_lang)
else:
dataset = data.load_raw_text_dataset(
args.data, splits, args.source_lang, args.target_lang)
if args.source_lang is None or args.target_lang is None:
# record inferred languages in args, so that it's saved in checkpoints
args.source_lang, args.target_lang = dataset.src, dataset.dst
print('| [{}] dictionary: {} types'.format(dataset.src, len(dataset.src_dict)))
print('| [{}] dictionary: {} types'.format(dataset.dst, len(dataset.dst_dict)))
for split in splits:
print('| {} {} {} examples'.format(args.data, split, len(dataset.splits[split])))
# Build model and criterion
model = models.build_model(args, dataset.src_dict, dataset.dst_dict)
criterion = criterions.build_criterion(args, dataset.src_dict, dataset.dst_dict)
print('| model {}, criterion {}'.format(args.arch, criterion.__class__.__name__))
print('| num. model params: {}'.format(sum(p.data.numel() for p in model.parameters())))
# Build trainer
trainer = Trainer(args, model, criterion)
print('| training on {} GPUs'.format(args.distributed_world_size))
print('| max tokens per GPU = {} and max sentences per GPU = {}'.format(
args.max_tokens,
args.max_sentences,
))
# Load the latest checkpoint if one is available
os.makedirs(args.save_dir, exist_ok=True)
checkpoint_path = os.path.join(args.save_dir, args.restore_file)
extra_state = trainer.load_checkpoint(checkpoint_path)
if extra_state is not None:
epoch = extra_state['epoch']
batch_offset = extra_state['batch_offset']
print('| loaded checkpoint {} (epoch {})'.format(checkpoint_path, epoch))
if batch_offset == 0:
trainer.lr_step(epoch)
epoch += 1
else:
epoch, batch_offset = 1, 0
# Train until the learning rate gets too small
max_epoch = args.max_epoch or math.inf
lr = trainer.get_lr()
train_meter = StopwatchMeter()
train_meter.start()
while lr > args.min_lr and epoch <= max_epoch:
# train for one epoch
train(args, trainer, dataset, epoch, batch_offset)
# evaluate on validate set
for k, subset in enumerate(args.valid_subset.split(',')):
val_loss = validate(args, trainer, dataset, subset, epoch)
if k == 0:
# only use first validation loss to update the learning schedule
lr = trainer.lr_step(epoch, val_loss)
# save checkpoint
if not args.no_save:
save_checkpoint(trainer, args, epoch, 0, val_loss)
epoch += 1
batch_offset = 0
train_meter.stop()
print('| done training in {:.1f} seconds'.format(train_meter.sum))
def main(args, init_distributed=False):
utils.import_user_module(args)
assert args.max_tokens is not None or args.max_sentences is not None, \
'Must specify batch size either with --max-tokens or --max-sentences'
# Initialize CUDA and distributed training
if torch.cuda.is_available() and not args.cpu:
torch.cuda.set_device(args.device_id)
torch.manual_seed(args.seed)
if init_distributed:
args.distributed_rank = distributed_utils.distributed_init(args)
# Print args
print(args)
# Setup task, e.g., translation, language modeling, etc.
task = tasks.setup_task(args)
# Load dataset splits
task.load_dataset(args.train_subset, combine=True, epoch=0)
for valid_sub_split in args.valid_subset.split(','):
task.load_dataset(valid_sub_split, combine=True, epoch=0)
# Build model and criterion
model = task.build_model(args)
criterion = task.build_criterion(args)
print(model)
print('| model {}, criterion {}'.format(args.arch,
criterion.__class__.__name__))
print('| num. model params: {} (num. trained: {})'.format(
sum(p.numel() for p in model.parameters()),
sum(p.numel() for p in model.parameters() if p.requires_grad),
))
# Build trainer
trainer = Trainer(args, task, model, criterion)
print('| training on {} GPUs'.format(args.distributed_world_size))
print('| max tokens per GPU = {} and max sentences per GPU = {}'.format(
args.max_tokens,
args.max_sentences,
))
max_positions = utils.resolve_max_positions(
task.max_positions(),
model.max_positions(),
)
# Initialize dataloader
epoch_itr = task.get_batch_iterator(
dataset=task.dataset(args.train_subset),
max_tokens=args.max_tokens,
max_sentences=args.max_sentences,
max_positions=max_positions,
ignore_invalid_inputs=True,
required_batch_size_multiple=args.required_batch_size_multiple,
seed=args.seed,
num_shards=args.distributed_world_size,
shard_id=args.distributed_rank,
num_workers=args.num_workers,
)
# Load the latest checkpoint if one is available
load_checkpoint(args, trainer, epoch_itr, max_positions, task)
# Train until the learning rate gets too small
max_epoch = args.max_epoch or math.inf
max_update = args.max_update or math.inf
lr = trainer.get_lr()
train_meter = StopwatchMeter()
train_meter.start()
valid_losses = [None]
valid_subsets = args.valid_subset.split(',')
while lr > args.min_lr and epoch_itr.epoch < max_epoch and trainer.get_num_updates(
) < max_update:
# train for one epoch
train(args, trainer, task, epoch_itr)
if epoch_itr.epoch % args.validate_interval == 0:
valid_losses = validate(args, trainer, task, epoch_itr,
valid_subsets)
# only use first validation loss to update the learning rate
lr = trainer.lr_step(epoch_itr.epoch, valid_losses[0])
# save checkpoint
if epoch_itr.epoch % args.save_interval == 0:
save_checkpoint(args, trainer, epoch_itr, valid_losses[0])
epoch_itr = reload_train(args, epoch_itr, max_positions, task)
train_meter.stop()
print('| done training in {:.1f} seconds'.format(train_meter.sum))
def main(args):
assert args.path is not None, '--path required for generation!'
assert not args.sampling or args.nbest == args.beam, \
'--sampling requires --nbest to be equal to --beam'
assert args.replace_unk is None or args.raw_text, \
'--replace-unk requires a raw text dataset (--raw-text)'
import_user_module(args)
if args.max_tokens is None and args.max_sentences is None:
args.max_tokens = 12000
print(args)
tgt_file = None
hypo_file = None
if args.output_dir:
os.makedirs(args.output_dir, exist_ok=True)
tgt_fn = os.path.join(args.output_dir, 'gold')
hypo_fn = os.path.join(args.output_dir, 'candidate')
tgt_file = open(tgt_fn, 'w', encoding='utf-8')
hypo_file = open(hypo_fn, 'w', encoding='utf-8')
use_cuda = torch.cuda.is_available() and not args.cpu
# Load dataset splits
task = tasks.setup_task(args)
task.load_dataset(args.gen_subset)
print('| {} {} {} examples'.format(args.data, args.gen_subset,
len(task.dataset(args.gen_subset))))
# Set dictionaries
try:
src_dict = getattr(task, 'source_dictionary', None)
except NotImplementedError:
src_dict = None
tgt_dict = task.target_dictionary
# Load ensemble
print('| loading model(s) from {}'.format(args.path))
models, _model_args = utils.load_ensemble_for_inference(
args.path.split(':'),
task,
model_arg_overrides=eval(args.model_overrides),
)
# Optimize ensemble for generation
for model in models:
model.make_generation_fast_(
beamable_mm_beam_size=None if args.no_beamable_mm else args.beam,
need_attn=args.print_alignment,
)
if args.fp16:
model.half()
if use_cuda:
model.cuda()
# Load alignment dictionary for unknown word replacement
# (None if no unknown word replacement, empty if no path to align dictionary)
align_dict = utils.load_align_dict(args.replace_unk)
# Load dataset (possibly sharded)
itr = task.get_batch_iterator(
dataset=task.dataset(args.gen_subset),
max_tokens=args.max_tokens,
max_sentences=args.max_sentences,
max_positions=utils.resolve_max_positions(
task.max_positions(),
*[model.max_positions() for model in models]),
ignore_invalid_inputs=args.skip_invalid_size_inputs_valid_test,
required_batch_size_multiple=args.required_batch_size_multiple,
num_shards=args.num_shards,
shard_id=args.shard_id,
num_workers=args.num_workers,
).next_epoch_itr(shuffle=False)
# Initialize generator
gen_timer = StopwatchMeter()
generator = task.build_generator(args)
# Generate and compute BLEU score
if args.sacrebleu:
scorer = bleu.SacrebleuScorer()
else:
scorer = bleu.Scorer(tgt_dict.pad(), tgt_dict.eos(), tgt_dict.unk())
num_sentences = 0
has_target = True
with progress_bar.build_progress_bar(args, itr) as t:
wps_meter = TimeMeter()
for sample in t:
sample = utils.move_to_cuda(sample) if use_cuda else sample
if 'net_input' not in sample:
continue
prefix_tokens = None
if args.prefix_size > 0:
prefix_tokens = sample['target'][:, :args.prefix_size]
gen_timer.start()
hypos = task.inference_step(generator, models, sample,
prefix_tokens)
num_generated_tokens = sum(len(h[0]['tokens']) for h in hypos)
gen_timer.stop(num_generated_tokens)
for i, sample_id in enumerate(sample['id'].tolist()):
has_target = sample['target'] is not None
# Remove padding
src_tokens = utils.strip_pad(
sample['net_input']['src_tokens'][i, :], tgt_dict.pad())
target_tokens = None
if has_target:
target_tokens = utils.strip_pad(
sample['target'][i, :], tgt_dict.pad()).int().cpu()
# Either retrieve the original sentences or regenerate them from tokens.
if align_dict is not None:
src_str = task.dataset(
args.gen_subset).src.get_original_text(sample_id)
target_str = task.dataset(
args.gen_subset).tgt.get_original_text(sample_id)
else:
if src_dict is not None:
src_str = src_dict.string(src_tokens, args.remove_bpe)
else:
src_str = ""
if has_target:
target_str = tgt_dict.string(target_tokens,
args.remove_bpe,
escape_unk=True)
if not args.quiet:
if src_dict is not None:
print('S-{}\t{}'.format(
sample_id, src_str.encode(encoding='utf-8')))
if has_target:
print('T-{}\t{}'.format(
sample_id, target_str.encode(encoding='utf-8')))
# Process top predictions
for j, hypo in enumerate(
hypos[i][:min(len(hypos[i]), args.nbest)]):
hypo_tokens, hypo_str, alignment = utils.post_process_prediction(
hypo_tokens=hypo['tokens'].int().cpu(),
src_str=src_str,
alignment=hypo['alignment'].int().cpu()
if hypo['alignment'] is not None else None,
align_dict=align_dict,
tgt_dict=tgt_dict,
remove_bpe=args.remove_bpe,
)
if not args.quiet:
print('H-{}\t{}\t{}'.format(
sample_id, hypo['score'],
hypo_str.encode(encoding='utf-8')))
print('P-{}\t{}'.format(
sample_id, ' '.join(
map(
lambda x: '{:.4f}'.format(x),
hypo['positional_scores'].tolist(),
))))
if args.print_alignment:
print('A-{}\t{}'.format(
sample_id, ' '.join(
map(lambda x: str(utils.item(x)),
alignment))))
# Score only the top hypothesis
if has_target and j == 0:
if align_dict is not None or args.remove_bpe is not None:
# Convert back to tokens for evaluation with unk replacement and/or without BPE
target_tokens = tgt_dict.encode_line(
target_str, add_if_not_exist=True)
if hasattr(scorer, 'add_string'):
scorer.add_string(target_str, hypo_str)
else:
scorer.add(target_tokens, hypo_tokens)
if args.output_dir:
tgt_file.writelines(target_str + '\n')
hypo_file.writelines(hypo_str + '\n')
wps_meter.update(num_generated_tokens)
t.log({'wps': round(wps_meter.avg)})
num_sentences += sample['nsentences']
tgt_file.close()
hypo_file.close()
print(
'| Translated {} sentences ({} tokens) in {:.1f}s ({:.2f} sentences/s, {:.2f} tokens/s)'
.format(num_sentences, gen_timer.n, gen_timer.sum,
num_sentences / gen_timer.sum, 1. / gen_timer.avg))
if has_target:
print('| Generate {} with beam={}: {}'.format(args.gen_subset,
args.beam,
scorer.result_string()))
return scorer
def save_checkpoint(args, trainer, epoch_itr, val_loss):
if args.no_save or not distributed_utils.is_master(args):
return
write_timer = StopwatchMeter()
write_timer.start()
epoch = epoch_itr.epoch
end_of_epoch = epoch_itr.end_of_epoch()
updates = trainer.get_num_updates()
checkpoint_conds = collections.OrderedDict()
checkpoint_conds['checkpoint{}.pt'.format(epoch)] = (
end_of_epoch and not args.no_epoch_checkpoints
and epoch % args.save_interval == 0)
checkpoint_conds['checkpoint_{}_{}.pt'.format(
epoch, updates)] = (not end_of_epoch and args.save_interval_updates > 0
and updates % args.save_interval_updates == 0)
checkpoint_conds['checkpoint_best.pt'] = (
val_loss is not None and (not hasattr(save_checkpoint, 'best')
or val_loss < save_checkpoint.best))
checkpoint_conds[
'checkpoint_last.pt'] = True # keep this last so that it's a symlink
prev_best = getattr(save_checkpoint, 'best', val_loss)
if val_loss is not None:
save_checkpoint.best = min(val_loss, prev_best)
extra_state = {
'train_iterator': epoch_itr.state_dict(),
'val_loss': val_loss,
}
if hasattr(save_checkpoint, 'best'):
extra_state.update({'best': save_checkpoint.best})
checkpoints = [
os.path.join(args.save_dir, fn)
for fn, cond in checkpoint_conds.items() if cond
]
if len(checkpoints) > 0:
trainer.save_checkpoint(checkpoints[0], extra_state)
for cp in checkpoints[1:]:
shutil.copyfile(checkpoints[0], cp)
write_timer.stop()
print(
'| saved checkpoint {} (epoch {} @ {} updates) (writing took {} seconds)'
.format(checkpoints[0], epoch, updates, write_timer.sum))
if not end_of_epoch and args.keep_interval_updates > 0:
# remove old checkpoints; checkpoints are sorted in descending order
checkpoints = checkpoint_utils.checkpoint_paths(
args.save_dir,
pattern=r'checkpoint_\d+_(\d+)\.pt',
)
for old_chk in checkpoints[args.keep_interval_updates:]:
if os.path.lexists(old_chk):
os.remove(old_chk)
if args.keep_last_epochs > 0:
# remove old epoch checkpoints; checkpoints are sorted in descending order
checkpoints = checkpoint_utils.checkpoint_paths(
args.save_dir,
pattern=r'checkpoint(\d+)\.pt',
)
for old_chk in checkpoints[args.keep_last_epochs:]:
if os.path.lexists(old_chk):
os.remove(old_chk)
def main(args):
if args.max_tokens is None:
args.max_tokens = 6000
print(args)
if not torch.cuda.is_available():
raise NotImplementedError('Training on CPU is not supported')
torch.cuda.set_device(args.device_id)
torch.manual_seed(args.seed)
# Setup task, e.g., translation, language modeling, etc.
task = tasks.setup_task(args)
# Load dataset splits
load_dataset_splits(args, task, ['train', 'valid'])
# Build model and criterion
model = task.build_model(args)
criterion = task.build_criterion(args)
print('| model {}, criterion {}'.format(args.arch,
criterion.__class__.__name__))
print('| num. model params: {}'.format(
sum(p.numel() for p in model.parameters())))
# Build trainer
if args.fp16:
trainer = FP16Trainer(args, task, model, criterion)
else:
if torch.cuda.get_device_capability(0)[0] >= 7:
print(
'| NOTICE: your device may support faster training with --fp16'
)
trainer = Trainer(args, task, model, criterion)
print('| training on {} GPUs'.format(args.distributed_world_size))
print('| max tokens per GPU = {} and max sentences per GPU = {}'.format(
args.max_tokens,
args.max_sentences,
))
# Initialize dataloader
max_positions = trainer.get_model().max_positions()
epoch_itr = data.EpochBatchIterator(
dataset=task.dataset(args.train_subset),
max_tokens=args.max_tokens,
max_sentences=args.max_sentences_valid,
max_positions=max_positions,
ignore_invalid_inputs=True,
required_batch_size_multiple=8,
seed=args.seed,
num_shards=args.distributed_world_size,
shard_id=args.distributed_rank,
)
# Load the latest checkpoint if one is available
load_checkpoint(args, trainer, epoch_itr)
# Send a dummy batch to warm the caching allocator
dummy_batch = task.dataset('train').get_dummy_batch(
args.max_tokens, max_positions)
trainer.dummy_train_step(dummy_batch)
# Train until the learning rate gets too small
max_epoch = args.max_epoch or math.inf
max_update = args.max_update or math.inf
lr = trainer.get_lr()
train_meter = StopwatchMeter()
train_meter.start()
valid_losses = [None]
valid_subsets = args.valid_subset.split(',')
while lr > args.min_lr and epoch_itr.epoch <= max_epoch and trainer.get_num_updates(
) < max_update:
# train for one epoch
train(args, trainer, task, epoch_itr)
if epoch_itr.epoch % args.validate_interval == 0:
valid_losses = validate(args, trainer, task, epoch_itr,
valid_subsets)
# only use first validation loss to update the learning rate
lr = trainer.lr_step(epoch_itr.epoch, valid_losses[0])
# save checkpoint
if epoch_itr.epoch % args.save_interval == 0:
save_checkpoint(args, trainer, epoch_itr, valid_losses[0])
train_meter.stop()
print('| done training in {:.1f} seconds'.format(train_meter.sum))
def main(args):
assert args.path is not None, '--path required for generation!'
assert not args.sampling or args.nbest == args.beam, \
'--sampling requires --nbest to be equal to --beam'
assert args.replace_unk is None or args.raw_text, \
'--replace-unk requires a raw text dataset (--raw-text)'
if args.max_tokens is None and args.max_sentences is None:
args.max_tokens = 12000
print(args)
use_cuda = torch.cuda.is_available() and not args.cpu
# Load dataset splits
task = tasks.setup_task(args)
task.load_dataset(args.gen_subset)
print('| {} {} {} examples'.format(args.data, args.gen_subset, len(task.dataset(args.gen_subset))))
# Set dictionaries
src_dict = task.source_dictionary
tgt_dict = task.target_dictionary
# Load ensemble
print('| loading model(s) from {}'.format(args.path))
models, _model_args = utils.load_ensemble_for_inference(
args.path.split(':'), task, model_arg_overrides=eval(args.model_overrides),
)
# Optimize ensemble for generation
for model in models:
model.make_generation_fast_(
beamable_mm_beam_size=None if args.no_beamable_mm else args.beam,
need_attn=args.print_alignment,
)
if args.fp16:
model.half()
# Load alignment dictionary for unknown word replacement
# (None if no unknown word replacement, empty if no path to align dictionary)
align_dict = utils.load_align_dict(args.replace_unk)
# Load dataset (possibly sharded)
itr = task.get_batch_iterator(
dataset=task.dataset(args.gen_subset),
max_tokens=args.max_tokens,
max_sentences=args.max_sentences,
max_positions=utils.resolve_max_positions(
task.max_positions(),
*[model.max_positions() for model in models]
),
ignore_invalid_inputs=args.skip_invalid_size_inputs_valid_test,
required_batch_size_multiple=8,
num_shards=args.num_shards,
shard_id=args.shard_id,
num_workers=args.num_workers,
).next_epoch_itr(shuffle=False)
# Initialize generator
gen_timer = StopwatchMeter()
if args.score_reference:
translator = SequenceScorer(models, task.target_dictionary)
else:
translator = SequenceGenerator(
models, task.target_dictionary, beam_size=args.beam, minlen=args.min_len,
stop_early=(not args.no_early_stop), normalize_scores=(not args.unnormalized),
len_penalty=args.lenpen, unk_penalty=args.unkpen,
sampling=args.sampling, sampling_topk=args.sampling_topk, sampling_temperature=args.sampling_temperature,
diverse_beam_groups=args.diverse_beam_groups, diverse_beam_strength=args.diverse_beam_strength,
match_source_len=args.match_source_len, no_repeat_ngram_size=args.no_repeat_ngram_size,
)
if use_cuda:
translator.cuda()
# Generate and compute BLEU score
scorer = bleu.Scorer(tgt_dict.pad(), tgt_dict.eos(), tgt_dict.unk())
num_sentences = 0
has_target = True
# output the result
result=['']*21678
with progress_bar.build_progress_bar(args, itr) as t:
if args.score_reference:
translations = translator.score_batched_itr(t, cuda=use_cuda, timer=gen_timer)
else:
translations = translator.generate_batched_itr(
t, maxlen_a=args.max_len_a, maxlen_b=args.max_len_b,
cuda=use_cuda, timer=gen_timer, prefix_size=args.prefix_size,
)
wps_meter = TimeMeter()
for sample_id, src_tokens, target_tokens, hypos in translations:
# Process input and ground truth
has_target = target_tokens is not None
target_tokens = target_tokens.int().cpu() if has_target else None
# Either retrieve the original sentences or regenerate them from tokens.
if align_dict is not None:
src_str = task.dataset(args.gen_subset).src.get_original_text(sample_id)
target_str = task.dataset(args.gen_subset).tgt.get_original_text(sample_id)
else:
src_str = src_dict.string(src_tokens, args.remove_bpe)
if has_target:
target_str = tgt_dict.string(target_tokens, args.remove_bpe, escape_unk=True)
if not args.quiet:
print('S-{}\t{}'.format(sample_id, src_str))
if has_target:
print('T-{}\t{}'.format(sample_id, target_str))
# Process top predictions
for i, hypo in enumerate(hypos[:min(len(hypos), args.nbest)]):
hypo_tokens, hypo_str, alignment = utils.post_process_prediction(
hypo_tokens=hypo['tokens'].int().cpu(),
src_str=src_str,
alignment=hypo['alignment'].int().cpu() if hypo['alignment'] is not None else None,
align_dict=align_dict,
tgt_dict=tgt_dict,
remove_bpe=args.remove_bpe,
)
result[sample_id]=hypo_str
if not args.quiet:
print('H-{}\t{}\t{}'.format(sample_id, hypo['score'], hypo_str))
print('P-{}\t{}'.format(
sample_id,
' '.join(map(
lambda x: '{:.4f}'.format(x),
hypo['positional_scores'].tolist(),
))
))
if args.print_alignment:
print('A-{}\t{}'.format(
sample_id,
' '.join(map(lambda x: str(utils.item(x)), alignment))
))
# Score only the top hypothesis
if has_target and i == 0:
if align_dict is not None or args.remove_bpe is not None:
# Convert back to tokens for evaluation with unk replacement and/or without BPE
target_tokens = tokenizer.Tokenizer.tokenize(
target_str, tgt_dict, add_if_not_exist=True)
scorer.add(target_tokens, hypo_tokens)
wps_meter.update(src_tokens.size(0))
t.log({'wps': round(wps_meter.avg)})
num_sentences += 1
print('| Translated {} sentences ({} tokens) in {:.1f}s ({:.2f} sentences/s, {:.2f} tokens/s)'.format(
num_sentences, gen_timer.n, gen_timer.sum, num_sentences / gen_timer.sum, 1. / gen_timer.avg))
if has_target:
print('| Generate {} with beam={}: {}'.format(args.gen_subset, args.beam, scorer.result_string()))
#output the result
return result
def main(args):
# we should not do this!
'''
if args.max_tokens is None:
args.max_tokens = 6000
'''
utils.xpprint(args)
if not torch.cuda.is_available():
raise NotImplementedError('Training on CPU is not supported')
torch.cuda.set_device(args.device_id)
torch.manual_seed(args.seed)
# Setup task, e.g., translation, language modeling, etc.
task = tasks.setup_task(args)
utils.xprintln('setup task done!')
# Load dataset splits
load_dataset_splits(args, task, ['train'])
valid_dataset = args.valid_subset.split(',')
load_dataset_splits(args, task, valid_dataset, shuffle=False)
utils.xprintln('load dataset done!')
if args.task.startswith('extractive_summarization'):
if distributed_utils.is_master(args):
from sum_eval import MultiProcSumEval
sum_eval_pool = MultiProcSumEval(args.ncpu_eval)
sum_valid_pool_params = dict(
article_file=args.raw_valid + '.article',
summary_file=args.raw_valid + '.summary',
entity_map_file=None,
length=-1,
eval_type='predict',
topk=args.topk_sent_eval,
rerank=False,
with_m=False,
cmd='-a -c 95 -m -n 4 -w 1.2',
trigram_block=args.trigram_block,
)
sum_test_pool_params = dict(
article_file=args.raw_test + '.article',
summary_file=args.raw_test + '.summary',
entity_map_file=None,
length=-1,
eval_type='predict',
topk=args.topk_sent_eval,
rerank=False,
with_m=False,
cmd='-a -c 95 -m -n 4 -w 1.2',
trigram_block=args.trigram_block,
)
sum_pool_params = dict(valid=sum_valid_pool_params,
test=sum_test_pool_params)
def make_params(default_dict,
result_file,
out_rouge_file,
rerank=False,
with_m=False):
para_dict = dict(default_dict)
para_dict['result_file'] = result_file
para_dict['out_rouge_file'] = out_rouge_file
para_dict['rerank'] = rerank
para_dict['with_m'] = with_m
return para_dict
# Build model and criterion
model = task.build_model(args)
criterion = task.build_criterion(args)
print('| model {}, criterion {}'.format(args.arch,
criterion.__class__.__name__))
print('| num. model params: {}'.format(
sum(p.numel() for p in model.parameters())))
# print(model)
import sys
sys.stdout.flush()
# if summarization try to load pretrained model
# if args.task.startswith('extractive_summarization') or args.task == 'pretrain_document_modeling':
# # assume this is a single GPU program
if args.init_from_pretrained_doc_model:
task.load_pretrained_model(model, args.pretrained_doc_model_path)
sys.stdout.flush()
# Build trainer
trainer = Trainer(args, task, model, criterion)
print('| training on {} GPUs'.format(args.distributed_world_size))
print('| max tokens per GPU = {} and max sentences per GPU = {}'.format(
args.max_tokens,
args.max_sentences,
))
# Initialize dataloader
max_positions = trainer.get_model().max_positions()
epoch_itr = trainer.get_train_iterator(epoch=0, load_dataset=False)
# Load the latest checkpoint if one is available
# load_checkpoint(args, trainer, epoch_itr)
# make sure training from a different checkpoint will use different random seed
cur_dataset = task.dataset('train')
if hasattr(cur_dataset, 'rng'):
print('epoch ', epoch_itr.epoch)
cur_dataset.rng = numpy.random.RandomState(args.seed + epoch_itr.epoch)
# Train until the learning rate gets too small
max_epoch = args.max_epoch or math.inf
max_update = args.max_update or math.inf
lr = trainer.get_lr()
train_meter = StopwatchMeter()
train_meter.start()
valid_losses = [None]
valid_subsets = args.valid_subset.split(',')
for alpha in range(10, 9, -1):
# train for one epoch
# train(args, trainer, task, epoch_itr)
epoch_itr.next_epoch_itr()
if epoch_itr.epoch % args.validate_interval == 0:
if args.task.startswith('extractive_summarization'):
if distributed_utils.is_master(args):
validate_metric(args, trainer, task, epoch_itr,
valid_subsets)
def infer_onebyone(args, models, task, input):
assert args.path is not None, '--path required for generation!'
assert not args.sampling or args.nbest == args.beam, \
'--sampling requires --nbest to be equal to --beam'
assert args.replace_unk is None or args.raw_text, \
'--replace-unk requires a raw text dataset (--raw-text)'
utils.import_user_module(args)
if args.max_tokens is None and args.max_sentences is None:
args.max_tokens = 12000
print(args)
use_cuda = torch.cuda.is_available() and not args.cpu
# Set dictionaries
try:
src_dict = getattr(task, 'source_dictionary', None)
except NotImplementedError:
src_dict = None
tgt_dict = task.target_dictionary
input = ' '.join([i for i in input])
src_tokens = tgt_dict.encode_line(input).type(torch.LongTensor)
input_sample = {
'id': torch.Tensor([0]),
'nsentences': 1,
'ntokens': len(src_tokens),
'net_input': {
'src_tokens': src_tokens.unsqueeze(0),
'src_lengths': torch.tensor([len(src_tokens)]),
'prev_output_tokens': torch.tensor([[tgt_dict.eos()]])
},
'target': src_tokens.unsqueeze(0),
}
# Load ensemble
print('| loading model(s) from {}'.format(args.path))
# Optimize ensemble for generation
# Load alignment dictionary for unknown word replacement
# (None if no unknown word replacement, empty if no path to align dictionary)
align_dict = utils.load_align_dict(args.replace_unk)
# Load dataset (possibly sharded)
# Initialize generator
gen_timer = StopwatchMeter()
generator = task.build_generator(args)
# Generate and compute BLEU score
if args.sacrebleu:
scorer = bleu.SacrebleuScorer()
else:
scorer = bleu.Scorer(tgt_dict.pad(), tgt_dict.eos(), tgt_dict.unk())
num_sentences = 0
has_target = True
wps_meter = TimeMeter()
sample = input_sample
sample = utils.move_to_cuda(sample) if use_cuda else sample
prefix_tokens = None
if args.prefix_size > 0:
prefix_tokens = sample['target'][:, :args.prefix_size]
gen_timer.start()
#pdb.set_trace()
hypos = task.inference_step(generator, models, sample, prefix_tokens)
num_generated_tokens = sum(len(h[0]['tokens']) for h in hypos)
gen_timer.stop(num_generated_tokens)
for i, sample_id in enumerate(sample['id'].tolist()):
has_target = sample['target'] is not None
# Remove padding
src_tokens = utils.strip_pad(sample['net_input']['src_tokens'][i, :],
tgt_dict.pad())
target_tokens = None
if has_target:
target_tokens = utils.strip_pad(sample['target'][i, :],
tgt_dict.pad()).int().cpu()
# Either retrieve the original sentences or regenerate them from tokens.
if align_dict is not None:
src_str = task.dataset(
args.gen_subset).src.get_original_text(sample_id)
target_str = task.dataset(
args.gen_subset).tgt.get_original_text(sample_id)
else:
if src_dict is not None:
src_str = src_dict.string(src_tokens, args.remove_bpe)
else:
src_str = ""
if has_target:
target_str = tgt_dict.string(target_tokens,
args.remove_bpe,
escape_unk=True)
if not args.quiet:
if src_dict is not None:
print('S-{}\t{}'.format(sample_id, src_str))
if has_target:
print('T-{}\t{}'.format(sample_id, target_str))
# Process top predictions
for j, hypo in enumerate(hypos[i][:args.nbest]):
hypo_tokens, hypo_str, alignment = utils.post_process_prediction(
hypo_tokens=hypo['tokens'].int().cpu(),
src_str=src_str,
alignment=hypo['alignment'],
align_dict=align_dict,
tgt_dict=tgt_dict,
remove_bpe=args.remove_bpe,
)
if not args.quiet:
output = ''.join(hypo_str.split(' '))
print('H-{}\t{}\t{}'.format(sample_id, hypo['score'],
hypo_str))
print('P-{}\t{}'.format(
sample_id, ' '.join(
map(
lambda x: '{:.4f}'.format(x),
hypo['positional_scores'].tolist(),
))))
if args.print_alignment:
print('A-{}\t{}'.format(
sample_id, ' '.join([
'{}-{}'.format(src_idx, tgt_idx)
for src_idx, tgt_idx in alignment
])))
if args.print_step:
print('I-{}\t{}'.format(sample_id, hypo['steps']))
# Score only the top hypothesis
if has_target and j == 0:
if align_dict is not None or args.remove_bpe is not None:
# Convert back to tokens for evaluation with unk replacement and/or without BPE
target_tokens = tgt_dict.encode_line(target_str,
add_if_not_exist=True)
if hasattr(scorer, 'add_string'):
scorer.add_string(target_str, hypo_str)
else:
scorer.add(target_tokens, hypo_tokens)
wps_meter.update(num_generated_tokens)
num_sentences += sample['nsentences']
print(
'| Translated {} sentences ({} tokens) in {:.1f}s ({:.2f} sentences/s, {:.2f} tokens/s)'
.format(num_sentences, gen_timer.n, gen_timer.sum,
num_sentences / gen_timer.sum, 1. / gen_timer.avg))
if has_target:
print('| Generate {} with beam={}: {}'.format(args.gen_subset,
args.beam,
scorer.result_string()))
return scorer, output
def main(args):
assert args.path is not None, '--path required for recognition!'
assert not args.sampling or args.nbest == args.beam, \
'--sampling requires --nbest to be equal to --beam'
utils.import_user_module(args)
if args.max_tokens is None and args.max_sentences is None:
args.max_tokens = 12000
print(args)
use_cuda = torch.cuda.is_available() and not args.cpu
# Load dataset split
task = tasks.setup_task(args)
task.load_dataset(args.gen_subset)
# Set dictionary
dict = task.target_dictionary
# Load ensemble
print('| loading model(s) from {}'.format(args.path))
models, _model_args = checkpoint_utils.load_model_ensemble(
args.path.split(':'),
arg_overrides=eval(args.model_overrides),
task=task,
)
for i, m in enumerate(models):
if hasattr(m, 'is_wordlm') and m.is_wordlm:
# assume subword LM comes before word LM
if isinstance(models[i - 1], FairseqLanguageModel):
models[i - 1] = MultiLevelLanguageModel(
m,
models[i - 1],
subwordlm_weight=args.subwordlm_weight,
oov_penalty=args.oov_penalty,
open_vocab=not args.disable_open_vocab,
)
del models[i]
print('| LM fusion with Multi-level LM')
else:
models[i] = TensorizedLookaheadLanguageModel(
m,
dict,
oov_penalty=args.oov_penalty,
open_vocab=not args.disable_open_vocab,
)
print('| LM fusion with Look-ahead Word LM')
# assume subword LM comes after E2E models
elif i == len(models) - 1 and isinstance(m, FairseqLanguageModel):
print('| LM fusion with Subword LM')
if args.lm_weight != 0.0:
print('| using LM fusion with lm-weight={:.2f}'.format(args.lm_weight))
# Optimize ensemble for generation
for model in models:
model.make_generation_fast_(
beamable_mm_beam_size=None if args.no_beamable_mm else args.beam,
need_attn=args.print_alignment,
)
if args.fp16:
model.half()
if use_cuda:
model.cuda()
# Load dataset (possibly sharded)
itr = task.get_batch_iterator(
dataset=task.dataset(args.gen_subset),
max_tokens=args.max_tokens,
max_sentences=args.max_sentences,
max_positions=utils.resolve_max_positions(
task.max_positions(), *[
model.max_positions() if hasattr(model, 'encoder') else
(None, model.max_positions()) for model in models
]),
ignore_invalid_inputs=args.skip_invalid_size_inputs_valid_test,
required_batch_size_multiple=8,
num_shards=args.num_shards,
shard_id=args.shard_id,
num_workers=args.num_workers,
).next_epoch_itr(shuffle=False)
# Initialize generator
if args.match_source_len:
print('| The option match_source_len is not applicable to '
'speech recognition. Ignoring it.')
gen_timer = StopwatchMeter()
generator = task.build_generator(args)
# Generate and compute WER
scorer = wer.Scorer(dict, wer_output_filter=args.wer_output_filter)
num_sentences = 0
has_target = True
with progress_bar.build_progress_bar(args, itr) as t:
wps_meter = TimeMeter()
for sample in t:
sample = utils.move_to_cuda(sample) if use_cuda else sample
if 'net_input' not in sample:
continue
prefix_tokens = None
if args.prefix_size > 0:
prefix_tokens = sample['target'][:, :args.prefix_size]
gen_timer.start()
hypos = task.inference_step(
generator,
models,
sample,
prefix_tokens,
lm_weight=args.lm_weight,
)
num_generated_tokens = sum(len(h[0]['tokens']) for h in hypos)
gen_timer.stop(num_generated_tokens)
# obtain nonpad mask of encoder output to plot attentions
if args.print_alignment:
net_input = sample['net_input']
src_tokens = net_input['src_tokens']
output_lengths = models[0].encoder.output_lengths(
net_input['src_lengths'])
nonpad_idxs = sequence_mask(
output_lengths,
models[0].encoder.output_lengths(src_tokens.size(1)))
for i, sample_id in enumerate(sample['id'].tolist()):
has_target = sample['target'] is not None
utt_id = sample['utt_id'][i]
# Retrieve the original sentences
if has_target:
target_str = task.dataset(
args.gen_subset).tgt.get_original_tokens(sample_id)
if not args.quiet:
target_sent = dict.tokens_to_sentence(
target_str,
use_unk_sym=False,
bpe_symbol=args.remove_bpe,
)
print('T-{}\t{}'.format(utt_id, target_sent))
# Process top predictions
for j, hypo in enumerate(hypos[i][:args.nbest]):
hypo_str = dict.string(hypo['tokens'].int().cpu()
) # not removing bpe at this point
if not args.quiet or i == 0:
hypo_sent = dict.tokens_to_sentence(
hypo_str, bpe_symbol=args.remove_bpe)
if not args.quiet:
print('H-{}\t{}\t{}'.format(utt_id, hypo_sent,
hypo['score']))
# Score and obtain attention only the top hypothesis
if j == 0:
# src_len x tgt_len
attention = hypo['attention'][nonpad_idxs[i]].float().cpu() \
if args.print_alignment and hypo['attention'] is not None else None
if args.print_alignment and attention is not None:
save_dir = os.path.join(args.results_path,
'attn_plots')
os.makedirs(save_dir, exist_ok=True)
plot_attention(attention, hypo_sent, utt_id,
save_dir)
scorer.add_prediction(utt_id,
hypo_str,
bpe_symbol=args.remove_bpe)
if has_target:
scorer.add_evaluation(utt_id,
target_str,
hypo_str,
bpe_symbol=args.remove_bpe)
wps_meter.update(num_generated_tokens)
t.log({'wps': round(wps_meter.avg)})
num_sentences += sample['nsentences']
print(
'| Recognized {} utterances ({} tokens) in {:.1f}s ({:.2f} sentences/s, {:.2f} tokens/s)'
.format(num_sentences, gen_timer.n, gen_timer.sum,
num_sentences / gen_timer.sum, 1. / gen_timer.avg))
if args.print_alignment:
print('| Saved attention plots in ' + save_dir)
if has_target:
assert args.test_text_files is not None
scorer.add_ordered_utt_list(*args.test_text_files)
os.makedirs(args.results_path, exist_ok=True)
fn = 'decoded_char_results.txt'
with open(os.path.join(args.results_path, fn), 'w', encoding='utf-8') as f:
f.write(scorer.print_char_results())
print('| Decoded char results saved as ' + f.name)
fn = 'decoded_results.txt'
with open(os.path.join(args.results_path, fn), 'w', encoding='utf-8') as f:
f.write(scorer.print_results())
print('| Decoded results saved as ' + f.name)
if has_target:
header = ' Recognize {} with beam={}: '.format(args.gen_subset,
args.beam)
fn = 'wer'
with open(os.path.join(args.results_path, fn), 'w',
encoding='utf-8') as f:
res = 'WER={:.2f}%, Sub={:.2f}%, Ins={:.2f}%, Del={:.2f}%'.format(
*(scorer.wer()))
print('|' + header + res)
f.write(res + '\n')
print('| WER saved in ' + f.name)
fn = 'cer'
with open(os.path.join(args.results_path, fn), 'w',
encoding='utf-8') as f:
res = 'CER={:.2f}%, Sub={:.2f}%, Ins={:.2f}%, Del={:.2f}%'.format(
*(scorer.cer()))
print('|' + ' ' * len(header) + res)
f.write(res + '\n')
print('| CER saved in ' + f.name)
fn = 'aligned_results.txt'
with open(os.path.join(args.results_path, fn), 'w',
encoding='utf-8') as f:
f.write(scorer.print_aligned_results())
print('| Aligned results saved as ' + f.name)
return scorer
def main():
parser = options.get_parser('Trainer')
dataset_args = options.add_dataset_args(parser)
dataset_args.add_argument('--max-tokens',
default=0,
type=int,
metavar='N',
help='maximum number of tokens in a batch')
dataset_args.add_argument('--batch-size',
default=32,
type=int,
metavar='N',
help='batch size')
dataset_args.add_argument('--test-batch-size',
default=32,
type=int,
metavar='N',
help='batch size for test set')
dataset_args.add_argument('--valid-batch-size',
default=32,
type=int,
metavar='N',
help='batch size for validation set')
dataset_args.add_argument(
'--train-subset',
default='train',
metavar='SPLIT',
choices=['train', 'valid', 'test'],
help='data subset to use for training (train, valid, test)')
dataset_args.add_argument(
'--valid-subset',
default='valid',
metavar='SPLIT',
help='comma separated list ofdata subsets '
' to use for validation (train, valid, valid1,test, test1)')
dataset_args.add_argument('--test-subset',
default='test',
metavar='SPLIT',
help='comma separated list ofdata subset '
'to use for testing (train, valid, test)')
dataset_args.add_argument(
'--valid-script',
nargs='+',
metavar='PATH',
help='path to external validation script (optional).')
options.add_optimization_args(parser)
options.add_checkpoint_args(parser)
options.add_model_args(parser)
args = utils.parse_args_and_arch(parser)
print(args)
if args.no_progress_bar:
progress_bar.enabled = False
progress_bar.print_interval = args.log_interval
if not os.path.exists(args.save_dir):
os.makedirs(args.save_dir)
torch.manual_seed(args.seed)
# Setting args.max_tokens to infinity(same as setting to None)
if args.max_tokens == 0:
args.max_tokens = None
# Load dataset
dataset = data.load_with_check(args.data, args.source_lang,
args.target_lang)
if args.source_lang is None or args.target_lang is None:
# record inferred languages in args, so that it's saved in checkpoints
args.source_lang, args.target_lang = dataset.src, dataset.dst
print('| [{}] dictionary: {} types'.format(dataset.src,
len(dataset.src_dict)))
print('| [{}] dictionary: {} types'.format(dataset.dst,
len(dataset.dst_dict)))
for split in dataset.splits:
print('| {} {} {} examples'.format(args.data, split,
len(dataset.splits[split])))
if not torch.cuda.is_available():
raise NotImplementedError('Training on CPU is not supported')
num_gpus = torch.cuda.device_count()
print('| using {} GPUs (with max tokens per GPU = {})'.format(
num_gpus, args.max_tokens))
# Build model
print('| model {}'.format(args.arch))
model = utils.build_model(args, dataset)
criterion = utils.build_criterion(args, dataset)
# Start multiprocessing
trainer = MultiprocessingTrainer(args, model)
# Load the latest checkpoint if one is available
epoch, batch_offset = trainer.load_checkpoint(
os.path.join(args.save_dir, args.restore_file))
# Train until the learning rate gets too small
val_loss = None
max_epoch = args.max_epoch or math.inf
lr = trainer.get_lr()
train_meter = StopwatchMeter()
train_meter.start()
while lr > args.min_lr and epoch <= max_epoch:
# train for one epoch
train(args, epoch, batch_offset, trainer, criterion, dataset, num_gpus)
# evaluate on validate set
for k, subset in enumerate(args.valid_subset.split(',')):
val_loss = validate(args, epoch, trainer, criterion, dataset,
subset, num_gpus)
if k == 0:
if not args.no_save:
# save checkpoint
trainer.save_checkpoint(
args,
epoch,
0,
val_loss,
validation_script=args.valid_script)
# only use first validation loss to update the learning schedule
lr = trainer.lr_step(val_loss, epoch)
epoch += 1
batch_offset = 0
train_meter.stop()
print('| done training in {:.1f} seconds'.format(train_meter.sum))
# Generate on test set and compute BLEU score
for beam in [1, 5, 10, 20]:
for subset in args.test_subset.split(','):
scorer = score_test(args,
trainer.get_model(),
dataset,
subset,
beam,
cuda_device=(0 if num_gpus > 0 else None))
print('| Test on {} with beam={}: {}'.format(
subset, beam, scorer.result_string()))
# Stop multiprocessing
trainer.stop()
