def LogLoss(distribution, has_weights, **unused_kwargs): # pylint: disable=invalid-name
"""Builds a log loss layer for a Distribution."""
return tl.Serial([
distribution.LogProb(),
tl.Negate(),
tl.WeightedSum() if has_weights else tl.Sum(),
])
'weights', # Model weights.
'slots', # Per-parameter optimizer state, e.g. gradient moments.
'opt_params', # Optimizer (hyper)parameters, e.g. learning rate, momentum.
])
_DEFAULT_METRICS = {
'loss':
tl.Serial(tl.LogSoftmax(), tl.CrossEntropyLoss()),
'accuracy':
tl.Accuracy(),
'sequence_accuracy':
tl.SequenceAccuracy(),
'neg_log_perplexity':
tl.Serial(tl.LogSoftmax(), tl.CrossEntropyLoss(), tl.Negate()),
'weights_per_batch_per_core':
tl.Serial(tl.Drop(), tl.Drop(), tl.Sum()),
}
class Trainer:
"""Trax trainer.
A trainer allows to make training steps, train for full epochs,
save the training state and access evaluation data.
"""
def __init__(self,
model,
loss_fn,
optimizer,
lr_schedule,
inputs,
OptState = collections.namedtuple(
'_OptState',
[
'weights', # Model weights.
'slots', # Per-parameter optimizer state, e.g. gradient moments.
'opt_params', # Optimizer (hyper)parameters, e.g. learning rate, momentum.
])
_DEFAULT_METRICS = {
'loss': tl.WeightedCategoryCrossEntropy(),
'accuracy': tl.WeightedCategoryAccuracy(),
'sequence_accuracy': tl.MaskedSequenceAccuracy(),
'neg_log_perplexity': tl.Serial(tl.WeightedCategoryCrossEntropy(),
tl.Negate()),
'weights_per_batch_per_core': tl.Serial(tl.Drop(), tl.Drop(), tl.Sum()),
}
class Trainer:
"""Trax trainer.
A trainer allows to make training steps, train for full epochs,
save the training state and access evaluation data.
"""
def __init__(self,
model,
loss_fn,
optimizer,
lr_schedule,
inputs,