def LogLoss(distribution, **unused_kwargs): # pylint: disable=invalid-name
"""Builds a log loss layer for a Distribution."""
return tl.Serial(
distribution.LogProb(),
tl.Negate(),
tl.WeightedSum()
)
'_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.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,
])
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,
'model_state', # Auxilliary state of the model.
])
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.CrossEntropyLoss(),
'accuracy': tl.Accuracy(),
'sequence_accuracy': tl.SequenceAccuracy(),
'neg_log_perplexity': tl.Serial(tl.CrossEntropyLoss(), tl.Negate()),
'weights_per_batch_per_core': tl.SumOfWeights(),
}
class Trainer(object):
"""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,
])
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.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,
output_dir=None, random_seed=None, n_devices=None,
checkpoints_at=None, should_save_checkpoints=True,
should_write_summaries=True,
