def train(model_path, _log, _run, window=2):
"""Train a CRF model."""
train_reader = read_train_corpus()
sents, tags = separate_tagged_sents(train_reader.tagged_sents())
sents = preprocess(sents)
_log.info('Extracting features from train corpus')
train_itemseq = ItemSequence(
[fs for sent in sents for fs in extract_crf_features(sent)])
train_labels = [tag for tags_ in tags for tag in tags_]
trainer = make_crf_trainer()
trainer.append(train_itemseq, train_labels, group=0)
dev_reader = read_dev_corpus()
if dev_reader is not None:
_log.info('Extracting features from dev corpus')
sents, tags = separate_tagged_sents(dev_reader.tagged_sents())
sents = preprocess(sents)
dev_itemseq = ItemSequence(
[fs for sent in sents for fs in extract_crf_features(sent)])
dev_labels = [tag for tags_ in tags for tag in tags_]
trainer.append(dev_itemseq, dev_labels, group=1)
_log.info('Begin training; saving model to %s', model_path)
holdout = -1 if dev_reader is None else 1
trainer.train(model_path, holdout=holdout)
if SACRED_OBSERVE_FILES:
_run.add_artifact(model_path)
def make_preds(tagger, sents, _log):
sents = preprocess(sents)
_log.info('Extracting features')
itemseq = ItemSequence(
[fs for sent in sents for fs in extract_crf_features(sent)])
_log.info('Making predictions with the model')
return tagger.tag(itemseq)
def make_preds(field_odict, model, sents, _log, test_batch_size=32, device=-1):
TAGS = field_odict['tags']
# We need to set tags field to None because we don't have tags at test time; setting it
# to None skips it when creating examples
field_odict = collections.OrderedDict(field_odict) # shallow copy
field_odict['tags'] = None
# We need index field to sort the examples back to its original order because
# the iterator will sort them according to length to minimize padding
field_odict['index'] = Field(sequential=False, use_vocab=False)
dataset = make_dataset(preprocess(sents), field_odict.items())
sort_key = lambda ex: len(ex.words) # noqa: E731
iterator = BucketIterator(dataset,
test_batch_size,
sort_key=sort_key,
device=device,
train=False)
_log.info('Making predictions with the model')
ix_preds = []
for minibatch in iterator:
# shape: (batch_size,)
ix = minibatch.index
# shape: (batch_size, seq_length)
inputs = [minibatch.words]
if hasattr(minibatch, 'prefs_2'):
# shape: (batch_size, seq_length)
inputs.append(minibatch.prefs_2)
if hasattr(minibatch, 'prefs_3'):
# shape: (batch_size, seq_length)
inputs.append(minibatch.prefs_3)
if hasattr(minibatch, 'suffs_2'):
# shape: (batch_size, seq_length)
inputs.append(minibatch.suffs_2)
if hasattr(minibatch, 'suffs_3'):
# shape: (batch_size, seq_length)
inputs.append(minibatch.suffs_3)
if hasattr(minibatch, 'chars'):
# shape: (batch_size, seq_length, num_chars)
inputs.append(minibatch.chars)
preds = model.decode(inputs)
for i, pred, in zip(ix, preds):
pred = pred[1:-1] # strip init and eos tokens
ix_preds.append((int(i), pred))
ix_preds.sort()
return [TAGS.vocab.itos[p] for _, ps in ix_preds for p in ps]
def train(model_path, _log, _run, window=2):
"""Train a memorization model."""
train_reader = read_train_corpus()
sents, tags = separate_tagged_sents(train_reader.tagged_sents())
sents = preprocess(sents)
_log.info('Start training model')
model = MemorizationTagger.train(sents, tags, window=window)
_log.info('Saving model to %s', model_path)
with open(model_path, 'w') as f:
print(dump(model), file=f)
if SACRED_OBSERVE_FILES:
_run.add_artifact(model_path)
def train(save_dir,
_log,
_run,
batch_size=16,
test_batch_size=32,
device=-1,
print_every=10,
max_epochs=20,
stopping_patience=5,
scheduler_patience=2,
tol=0.01,
scheduler_verbose=False,
resume_from=None,
overwrite=False,
comparing=Comparing.F1.value,
grad_norm_threshold=1.):
"""Train a neural tagger."""
set_random_seed()
_log.info('Creating save directory %s if it does not exist', save_dir)
os.makedirs(save_dir, exist_ok=overwrite)
# Create fields
field_odict = create_fields()
WORDS = field_odict['words']
# Create datasets and iterators
reader = read_train_corpus()
sents, tags = separate_tagged_sents(reader.tagged_sents())
train_dataset = make_dataset(preprocess(sents),
field_odict.items(),
tags=tags)
sort_key = lambda ex: len(ex.words) # noqa: E731
train_iter = BucketIterator(train_dataset,
batch_size,
sort_key=sort_key,
device=device,
repeat=False)
train_eval_iter = BucketIterator(train_dataset,
test_batch_size,
sort_key=sort_key,
device=device,
train=False)
dev_iter = None
reader = read_dev_corpus()
if reader is not None:
sents, tags = separate_tagged_sents(reader.tagged_sents())
dev_dataset = make_dataset(preprocess(sents),
field_odict.items(),
tags=tags)
dev_iter = BucketIterator(dev_dataset,
test_batch_size,
sort_key=sort_key,
device=device,
train=False)
# Build vocabularies and save fields
build_vocab(field_odict.items(), train_dataset)
save_fields(field_odict)
# Create model and restore from checkpoint if given
metadata = get_metadata(field_odict)
checkpoint = None if resume_from is None else load_checkpoint()
model = make_model(metadata,
checkpoint=checkpoint,
pretrained_embedding=WORDS.vocab.vectors)
model.train()
save_metadata(metadata)
# Create optimizer and learning rate scheduler
optimizer = make_optimizer(model, checkpoint=checkpoint)
comp = Comparing(comparing)
scheduler = optim.lr_scheduler.ReduceLROnPlateau(
optimizer,
mode='max' if comp is Comparing.F1 else 'min',
factor=0.5,
patience=scheduler_patience,
threshold=tol,
threshold_mode='abs',
verbose=scheduler_verbose)
# Create engine, meters, timers, etc
engine = tnt.engine.Engine()
loss_meter = tnt.meter.AverageValueMeter()
speed_meter = tnt.meter.AverageValueMeter()
references = []
hypotheses = []
train_timer = tnt.meter.TimeMeter(None)
epoch_timer = tnt.meter.TimeMeter(None)
batch_timer = tnt.meter.TimeMeter(None)
comp_op = operator.ge if comp is Comparing.F1 else operator.le
sign = 1 if comp is Comparing.F1 else -1
def net(minibatch):
# shape: (batch_size, seq_length)
inputs = [minibatch.words]
if hasattr(minibatch, 'prefs_2'):
# shape: (batch_size, seq_length)
inputs.append(minibatch.prefs_2)
if hasattr(minibatch, 'prefs_3'):
# shape: (batch_size, seq_length)
inputs.append(minibatch.prefs_3)
if hasattr(minibatch, 'suffs_2'):
# shape: (batch_size, seq_length)
inputs.append(minibatch.suffs_2)
if hasattr(minibatch, 'suffs_3'):
# shape: (batch_size, seq_length)
inputs.append(minibatch.suffs_3)
if hasattr(minibatch, 'chars'):
# shape: (batch_size, seq_length, num_chars)
inputs.append(minibatch.chars)
# shape: (1,)
loss = model(inputs, minibatch.tags)
return loss, model.decode(inputs)
def reset_meters():
loss_meter.reset()
speed_meter.reset()
nonlocal references, hypotheses
references = []
hypotheses = []
def make_checkpoint(state, is_best=False):
save_checkpoint({
'epoch': state['epoch'],
't': state['t'],
'best_score': state['best_score'],
'num_bad_epochs': state['num_bad_epochs'],
'model': model.state_dict(),
'optimizer': state['optimizer'].state_dict(),
}, is_best=is_best) # yapf: disable
def evaluate_on(name):
assert name in ('train', 'dev')
iterator = train_eval_iter if name == 'train' else dev_iter
_log.info('Evaluating on %s', name)
engine.test(net, iterator)
loss = loss_meter.mean
f1 = f1_score(references, hypotheses, average='weighted')
_log.info(
'** Result on %s (%.2fs): %.2f samples/s | loss %.4f | ppl %.4f | f1 %s',
name.upper(), epoch_timer.value(), speed_meter.mean, loss,
math.exp(loss), f'{f1:.2%}')
_run.log_scalar(f'loss({name})', loss)
_run.log_scalar(f'ppl({name})', math.exp(loss))
_run.log_scalar(f'f1({name})', f1)
return loss, f1
def on_start(state):
if state['train']:
state.update({
'best_score': -sign * float('inf'),
'num_bad_epochs': 0,
})
if checkpoint is not None:
_log.info('Resuming training from the checkpoint')
state.update({
'epoch': checkpoint['epoch'],
't': checkpoint['t'],
'best_score': checkpoint['best_score'],
'num_bad_epochs': checkpoint['num_bad_epochs'],
})
make_checkpoint(state, is_best=True)
_log.info('Start training')
train_timer.reset()
else:
reset_meters()
epoch_timer.reset()
model.eval()
def on_start_epoch(state):
_log.info('Starting epoch %s', state['epoch'] + 1)
reset_meters()
epoch_timer.reset()
model.train()
def on_sample(state):
batch_timer.reset()
def on_forward(state):
if state['train']:
torch.nn.utils.clip_grad_norm(
(p for p in model.parameters() if p.requires_grad),
grad_norm_threshold)
batch_loss = float(state['loss'])
loss_meter.add(batch_loss)
# shape: (batch_size, seq_length)
golds = state['sample'].tags
elapsed_time = batch_timer.value()
batch_speed = golds.size(0) / elapsed_time
speed_meter.add(batch_speed)
if not state['train']:
for gold, pred in zip(golds, state['output']):
gold = gold.data[:len(pred)]
assert gold[0] == WORDS.vocab.stoi[WORDS.init_token]
assert gold[-1] == WORDS.vocab.stoi[WORDS.eos_token]
gold, pred = gold[1:-1], pred[1:
-1] # strip init and eos tokens
references.extend(gold)
hypotheses.extend(pred)
elif (state['t'] + 1) % print_every == 0:
batch_ref, batch_hyp = [], []
for gold, pred in zip(golds, state['output']):
gold = gold.data[:len(pred)]
assert gold[0] == WORDS.vocab.stoi[WORDS.init_token]
assert gold[-1] == WORDS.vocab.stoi[WORDS.eos_token]
gold, pred = gold[1:-1], pred[1:
-1] # strip init and eos tokens
batch_ref.extend(gold)
batch_hyp.extend(pred)
batch_f1 = f1_score(batch_ref, batch_hyp, average='weighted')
epoch = (state['t'] + 1) / len(state['iterator'])
_log.info(
'Epoch %.2f (%5.2fms): %.2f samples/s | loss %.4f | ppl %.4f | f1 %s',
epoch, 1000 * elapsed_time, batch_speed, batch_loss,
math.exp(batch_loss), f'{batch_f1:.2%}')
_run.log_scalar('batch_loss(train)', batch_loss, step=state['t'])
_run.log_scalar('batch_ppl(train)',
math.exp(batch_loss),
step=state['t'])
_run.log_scalar('batch_f1(train)', batch_f1, step=state['t'])
def on_end_epoch(state):
_log.info(
'Epoch %d done (%.2fs): mean speed %.2f samples/s | mean loss %.4f | mean ppl %.4f',
state['epoch'], epoch_timer.value(), speed_meter.mean,
loss_meter.mean, math.exp(loss_meter.mean))
evaluate_on('train')
is_best = False
if dev_iter is not None:
dev_loss, dev_f1 = evaluate_on('dev')
score = dev_f1 if comp is Comparing.F1 else dev_loss
scheduler.step(score, epoch=state['epoch'])
if comp_op(score, state['best_score'] + sign * tol):
_log.info('** NEW best result on dev corpus')
state.update({'best_score': score, 'num_bad_epochs': 0})
is_best = True
else:
state['num_bad_epochs'] += 1
if state['num_bad_epochs'] >= stopping_patience * (
scheduler_patience + 1):
num_reduction = state['num_bad_epochs'] // (
scheduler_patience + 1)
_log.info(
f"No improvements after {num_reduction} LR reductions, stopping early"
)
state['maxepoch'] = -1 # force training loop to stop
make_checkpoint(state, is_best=is_best)
def on_end(state):
if state['train']:
_log.info('Training done in %.2fs', train_timer.value())
engine.hooks['on_start'] = on_start
engine.hooks['on_start_epoch'] = on_start_epoch
engine.hooks['on_sample'] = on_sample
engine.hooks['on_forward'] = on_forward
engine.hooks['on_end_epoch'] = on_end_epoch
engine.hooks['on_end'] = on_end
try:
engine.train(net, train_iter, max_epochs, optimizer)
except KeyboardInterrupt:
_log.info('Training interrupted, aborting')
def make_preds(model, sents, _log):
sents = preprocess(sents)
_log.info('Making predictions with the model')
return [tag for sent in sents for tag in model.predict(sent)]