class GradientDescentTrainer(Trainer):
"""
A trainer for doing supervised learning with gradient descent. It just takes a labeled dataset
and a `DataLoader`, and uses the supplied `Optimizer` to learn the weights for your model over
some fixed number of epochs. You can also pass in a validation dataloader and enable early
stopping. There are many other bells and whistles as well.
Registered as a `Trainer` with the name "gradient_descent" (and is also the default `Trainer`).
The constructor that is registered is `from_partial_objects` - see the arguments to that
function for the exact keys that should be used, if you are using a configuration file. They
largely match the arguments to `__init__`, and we don't repeat their docstrings in
`from_partial_objects`.
# Parameters
model : `Model`, required.
An AllenNLP model to be optimized. Pytorch Modules can also be optimized if
their `forward` method returns a dictionary with a "loss" key, containing a
scalar tensor representing the loss function to be optimized.
If you are training your model using GPUs, your model should already be
on the correct device. (If you are using our `train` command this will be
handled for you.)
optimizer : `torch.nn.Optimizer`, required.
An instance of a Pytorch Optimizer, instantiated with the parameters of the
model to be optimized.
data_loader : `DataLoader`, required.
A pytorch `DataLoader` containing your `Dataset`, yielding padded indexed batches.
patience : Optional[int] > 0, optional (default=None)
Number of epochs to be patient before early stopping: the training is stopped
after `patience` epochs with no improvement. If given, it must be `> 0`.
If None, early stopping is disabled.
validation_metric : str, optional (default="loss")
Validation metric to measure for whether to stop training using patience
and whether to serialize an `is_best` model each epoch. The metric name
must be prepended with either "+" or "-", which specifies whether the metric
is an increasing or decreasing function.
validation_dataloader : `DataLoader`, optional (default=None)
A `DataLoader` to use for the validation set. If `None`, then
use the training `DataLoader` with the validation data.
num_epochs : int, optional (default = 20)
Number of training epochs.
serialization_dir : str, optional (default=None)
Path to directory for saving and loading model files. Models will not be saved if
this parameter is not passed.
checkpointer : `Checkpointer`, optional (default=None)
A `Checkpointer` is responsible for periodically saving model weights. If none is given
here, we will construct one with default parameters.
cuda_device : `int`, optional (default = -1)
An integer specifying the CUDA device(s) to use for this process. If -1, the CPU is used.
Data parallelism is controlled at the allennlp train level, so each trainer will have a single
GPU.
grad_norm : `float`, optional, (default = None).
If provided, gradient norms will be rescaled to have a maximum of this value.
grad_clipping : `float`, optional (default = `None`).
If provided, gradients will be clipped `during the backward pass` to have an (absolute)
maximum of this value. If you are getting `NaNs` in your gradients during training
that are not solved by using `grad_norm`, you may need this.
learning_rate_scheduler : `LearningRateScheduler`, optional (default = None)
If specified, the learning rate will be decayed with respect to
this schedule at the end of each epoch (or batch, if the scheduler implements
the `step_batch` method). If you use `torch.optim.lr_scheduler.ReduceLROnPlateau`,
this will use the `validation_metric` provided to determine if learning has plateaued.
To support updating the learning rate on every batch, this can optionally implement
`step_batch(batch_num_total)` which updates the learning rate given the batch number.
momentum_scheduler : `MomentumScheduler`, optional (default = None)
If specified, the momentum will be updated at the end of each batch or epoch
according to the schedule.
tensorboard_writer : `TensorboardWriter`, optional
If this is not provided, we will construct a `TensorboardWriter` with default
parameters and use that.
moving_average : `MovingAverage`, optional, (default = None)
If provided, we will maintain moving averages for all parameters. During training, we
employ a shadow variable for each parameter, which maintains the moving average. During
evaluation, we backup the original parameters and assign the moving averages to corresponding
parameters. Be careful that when saving the checkpoint, we will save the moving averages of
parameters. This is necessary because we want the saved model to perform as well as the validated
model if we load it later. But this may cause problems if you restart the training from checkpoint.
batch_callbacks : `List[BatchCallback]`, optional (default = None)
A list of callbacks that will be called at the end of every batch, during both train and
validation.
epoch_callbacks : `List[EpochCallback]`, optional (default = None)
A list of callbacks that will be called at the end of every epoch, and at the start of
training (with epoch = -1).
distributed : `bool`, optional, (default = False)
If set, PyTorch's `DistributedDataParallel` is used to train the model in multiple GPUs. This also
requires `world_size` to be greater than 1.
local_rank : `int`, optional, (default = 0)
This is the unique identifier of the `Trainer` in a distributed process group. The GPU device id is
used as the rank.
world_size : `int`, (default = 1)
The number of `Trainer` workers participating in the distributed training.
num_gradient_accumulation_steps : `int`, optional, (default = 1)
Gradients are accumulated for the given number of steps before doing an optimizer step. This can
be useful to accommodate batches that are larger than the RAM size. Refer Thomas Wolf's
[post](https://tinyurl.com/y5mv44fw) for details on Gradient Accumulation.
opt_level : `str`, optional, (default = `None`)
Each opt_level establishes a set of properties that govern Amp’s implementation of pure or mixed
precision training. Must be a choice of `"O0"`, `"O1"`, `"O2"`, or `"O3"`.
See the Apex [documentation](https://nvidia.github.io/apex/amp.html#opt-levels-and-properties) for
more details. If `None`, Amp is not used. Defaults to `None`.
"""
def __init__(
self,
model: Model,
optimizer: torch.optim.Optimizer,
data_loader: torch.utils.data.DataLoader,
patience: Optional[int] = None,
validation_metric: str = "-loss",
validation_data_loader: torch.utils.data.DataLoader = None,
num_epochs: int = 20,
serialization_dir: Optional[str] = None,
checkpointer: Checkpointer = None,
cuda_device: int = -1,
grad_norm: Optional[float] = None,
grad_clipping: Optional[float] = None,
learning_rate_scheduler: Optional[LearningRateScheduler] = None,
momentum_scheduler: Optional[MomentumScheduler] = None,
tensorboard_writer: TensorboardWriter = None,
moving_average: Optional[MovingAverage] = None,
batch_callbacks: List[BatchCallback] = None,
epoch_callbacks: List[EpochCallback] = None,
distributed: bool = False,
local_rank: int = 0,
world_size: int = 1,
num_gradient_accumulation_steps: int = 1,
opt_level: Optional[str] = None,
) -> None:
super().__init__(serialization_dir, cuda_device, distributed,
local_rank, world_size)
# I am not calling move_to_gpu here, because if the model is
# not already on the GPU then the optimizer is going to be wrong.
self.model = model
self.data_loader = data_loader
self._validation_data_loader = validation_data_loader
self.optimizer = optimizer
if patience is None: # no early stopping
if validation_data_loader:
logger.warning(
"You provided a validation dataset but patience was set to None, "
"meaning that early stopping is disabled")
elif (not isinstance(patience, int)) or patience <= 0:
raise ConfigurationError(
'{} is an invalid value for "patience": it must be a positive integer '
"or None (if you want to disable early stopping)".format(
patience))
# For tracking is_best_so_far and should_stop_early
self._metric_tracker = MetricTracker(patience, validation_metric)
# Get rid of + or -
self._validation_metric = validation_metric[1:]
self._num_epochs = num_epochs
if checkpointer is not None:
self._checkpointer = checkpointer
else:
self._checkpointer = Checkpointer(serialization_dir)
self._grad_norm = grad_norm
self._grad_clipping = grad_clipping
self._learning_rate_scheduler = learning_rate_scheduler
self._momentum_scheduler = momentum_scheduler
self._moving_average = moving_average
self._batch_callbacks = batch_callbacks or []
self._epoch_callbacks = epoch_callbacks or []
# We keep the total batch number as an instance variable because it
# is used inside a closure for the hook which logs activations in
# `_enable_activation_logging`.
self._batch_num_total = 0
self._tensorboard = tensorboard_writer or TensorboardWriter(
serialization_dir)
self._tensorboard.get_batch_num_total = lambda: self._batch_num_total
self._tensorboard.enable_activation_logging(self.model)
self._last_log = 0.0 # time of last logging
self._num_gradient_accumulation_steps = num_gradient_accumulation_steps
# Enable automatic mixed precision training with NVIDIA Apex.
self._opt_level = opt_level
if self._opt_level is not None:
if amp is None:
raise ConfigurationError((
"Apex not installed but opt_level was provided. Please install NVIDIA's Apex to enable"
" automatic mixed precision (AMP) training. See: https://github.com/NVIDIA/apex."
))
self.model, self.optimizer = amp.initialize(
self.model, self.optimizer, opt_level=self._opt_level)
# Using `DistributedDataParallel`(ddp) brings in a quirk wrt AllenNLP's `Model` interface and its
# usage. A `Model` object is wrapped by `ddp`, but assigning the wrapped model to `self.model`
# will break the usages such as `Model.get_regularization_penalty`, `Model.get_metrics`, etc.
#
# Hence a reference to Pytorch's object is maintained in the case of distributed training and in the
# normal case, reference to `Model` is retained. This reference is only used in
# these places: `model.__call__`, `model.train` and `model.eval`.
if self._distributed:
self._pytorch_model = DistributedDataParallel(
self.model,
device_ids=[self.cuda_device],
find_unused_parameters=True)
else:
self._pytorch_model = self.model
def rescale_gradients(self) -> Optional[float]:
"""
Performs gradient rescaling. Is a no-op if gradient rescaling is not enabled.
"""
if self._grad_norm:
if self._opt_level is not None:
# See: https://nvidia.github.io/apex/advanced.html#gradient-clipping
parameters_to_clip = [
p for p in amp.master_params(self.optimizer)
if p.grad is not None
]
else:
parameters_to_clip = [
p for p in self.model.parameters() if p.grad is not None
]
return training_util.sparse_clip_norm(parameters_to_clip,
self._grad_norm)
else:
return None
def batch_outputs(self, batch: TensorDict,
for_training: bool) -> Dict[str, torch.Tensor]:
"""
Does a forward pass on the given batch and returns the output dictionary that the model
returns, after adding any specified regularization penalty to the loss (if training).
"""
batch = nn_util.move_to_device(batch, self.cuda_device)
output_dict = self._pytorch_model(**batch)
if for_training:
try:
regularization_penalty = self.model.get_regularization_penalty(
)
loss = output_dict["loss"]
# Handle model without regularization
if regularization_penalty == 0.0:
regularization_penalty = loss.new_full(size=[],
fill_value=0.0)
output_dict["reg_loss"] = regularization_penalty
output_dict["loss"] += regularization_penalty
except KeyError:
if for_training:
raise RuntimeError(
"The model you are trying to optimize does not contain a"
" 'loss' key in the output of model.forward(inputs).")
return output_dict
def _train_epoch(self, epoch: int) -> Dict[str, float]:
"""
Trains one epoch and returns metrics.
"""
logger.info("Epoch %d/%d", epoch, self._num_epochs - 1)
peak_cpu_usage = common_util.peak_memory_mb()
logger.info(f"Peak CPU memory usage MB: {peak_cpu_usage}")
gpu_usage = []
for gpu, memory in common_util.gpu_memory_mb().items():
gpu_usage.append((gpu, memory))
logger.info(f"GPU {gpu} memory usage MB: {memory}")
train_loss = 0.0
train_reg_loss = 0.0
# Set the model to "train" mode.
self._pytorch_model.train()
# Get tqdm for the training batches
batch_generator = iter(self.data_loader)
batch_group_generator = common_util.lazy_groups_of(
batch_generator, self._num_gradient_accumulation_steps)
logger.info("Training")
num_training_batches = math.ceil(
len(self.data_loader) / self._num_gradient_accumulation_steps)
# Having multiple tqdm bars in case of distributed training will be a mess. Hence only the master's
# progress is shown
if self._master:
batch_group_generator_tqdm = Tqdm.tqdm(batch_group_generator,
total=num_training_batches)
else:
batch_group_generator_tqdm = batch_group_generator
self._last_log = time.time()
batches_this_epoch = 0
if self._batch_num_total is None:
self._batch_num_total = 0
done_early = False
for batch_group in batch_group_generator_tqdm:
if self._distributed:
# Check whether the other workers have stopped already (due to differing amounts of
# data in each). If so, we can't proceed because we would hang when we hit the
# barrier implicit in Model.forward. We use a IntTensor instead a BoolTensor
# here because NCCL process groups apparently don't support BoolTensor.
done = torch.tensor(0, device=self.cuda_device)
torch.distributed.all_reduce(done,
torch.distributed.ReduceOp.SUM)
if done.item() > 0:
done_early = True
logger.warning(
f"Worker {torch.distributed.get_rank()} finishing training early! "
"This implies that there is an imbalance in your training "
"data across the workers and that some amount of it will be "
"ignored. A small amount of this is fine, but a major imbalance "
"should be avoided. Note: This warning will appear unless your "
"data is perfectly balanced.")
break
batches_this_epoch += 1
self._batch_num_total += 1
batch_num_total = self._batch_num_total
self.optimizer.zero_grad()
batch_group_outputs = []
for batch in batch_group:
batch_outputs = self.batch_outputs(batch, for_training=True)
batch_group_outputs.append(batch_outputs)
loss = batch_outputs["loss"]
reg_loss = batch_outputs["reg_loss"]
if torch.isnan(loss):
raise ValueError("nan loss encountered")
loss = loss / len(batch_group)
reg_loss = reg_loss / len(batch_group)
if self._opt_level is not None:
with amp.scale_loss(loss, self.optimizer) as scaled_loss:
scaled_loss.backward()
else:
loss.backward()
train_loss += loss.item()
train_reg_loss += reg_loss.item()
batch_grad_norm = self.rescale_gradients()
# This does nothing if batch_num_total is None or you are using a
# scheduler which doesn't update per batch.
if self._learning_rate_scheduler:
self._learning_rate_scheduler.step_batch(batch_num_total)
if self._momentum_scheduler:
self._momentum_scheduler.step_batch(batch_num_total)
param_updates = None
if self._tensorboard.should_log_histograms_this_batch(
) and self._master:
# Get the magnitude of parameter updates for logging. We need to do some
# computation before and after the optimizer step, and it's expensive because of
# GPU/CPU copies (necessary for large models, and for shipping to tensorboard), so
# we don't do this every batch, only when it's requested.
param_updates = {
name: param.detach().cpu().clone()
for name, param in self.model.named_parameters()
}
self.optimizer.step()
for name, param in self.model.named_parameters():
param_updates[name].sub_(param.detach().cpu())
else:
self.optimizer.step()
# Update moving averages
if self._moving_average is not None:
self._moving_average.apply(batch_num_total)
# Update the description with the latest metrics
metrics = training_util.get_metrics(
self.model,
train_loss,
train_reg_loss,
batches_this_epoch,
world_size=self._world_size,
cuda_device=[self.cuda_device],
)
# Updating tqdm only for the master as the trainers wouldn't have one
if self._master:
description = training_util.description_from_metrics(metrics)
batch_group_generator_tqdm.set_description(description,
refresh=False)
self._tensorboard.log_batch(self.model, self.optimizer,
batch_grad_norm, metrics,
batch_group, param_updates)
if self._master:
self._checkpointer.maybe_save_checkpoint(
self, epoch, batches_this_epoch)
for callback in self._batch_callbacks:
callback(
self,
batch_group,
batch_group_outputs,
epoch,
batches_this_epoch,
is_training=True,
)
if self._distributed and not done_early:
logger.warning(
f"Worker {torch.distributed.get_rank()} completed its entire epoch (training)."
)
# Indicate that we're done so that any workers that have remaining data stop the epoch early.
done = torch.tensor(1, device=self.cuda_device)
torch.distributed.all_reduce(done, torch.distributed.ReduceOp.SUM)
assert done.item()
# Let all workers finish their epoch before computing
# the final statistics for the epoch.
if self._distributed:
dist.barrier()
metrics = training_util.get_metrics(
self.model,
train_loss,
train_reg_loss,
batches_this_epoch,
reset=True,
world_size=self._world_size,
cuda_device=[self.cuda_device],
)
metrics["cpu_memory_MB"] = peak_cpu_usage
for (gpu_num, memory) in gpu_usage:
metrics["gpu_" + str(gpu_num) + "_memory_MB"] = memory
return metrics
def _validation_loss(self, epoch: int) -> Tuple[float, float, int]:
"""
Computes the validation loss. Returns it and the number of batches.
"""
logger.info("Validating")
self._pytorch_model.eval()
# Replace parameter values with the shadow values from the moving averages.
if self._moving_average is not None:
self._moving_average.assign_average_value()
if self._validation_data_loader is not None:
validation_data_loader = self._validation_data_loader
else:
raise ConfigurationError(
"Validation results cannot be calculated without a validation_data_loader"
)
val_generator_tqdm = Tqdm.tqdm(validation_data_loader)
batches_this_epoch = 0
val_loss = 0
val_reg_loss = 0
done_early = False
for batch in val_generator_tqdm:
if self._distributed:
# Check whether the other workers have stopped already (due to differing amounts of
# data in each). If so, we can't proceed because we would hang when we hit the
# barrier implicit in Model.forward. We use a IntTensor instead a BoolTensor
# here because NCCL process groups apparently don't support BoolTensor.
done = torch.tensor(0, device=self.cuda_device)
torch.distributed.all_reduce(done,
torch.distributed.ReduceOp.SUM)
if done.item() > 0:
done_early = True
logger.warning(
f"Worker {torch.distributed.get_rank()} finishing validation early! "
"This implies that there is an imbalance in your validation "
"data across the workers and that some amount of it will be "
"ignored. A small amount of this is fine, but a major imbalance "
"should be avoided. Note: This warning will appear unless your "
"data is perfectly balanced.")
break
batch_outputs = self.batch_outputs(batch, for_training=False)
loss = batch_outputs.get("loss")
reg_loss = batch_outputs.get("reg_loss")
if loss is not None:
# You shouldn't necessarily have to compute a loss for validation, so we allow for
# `loss` to be None. We need to be careful, though - `batches_this_epoch` is
# currently only used as the divisor for the loss function, so we can safely only
# count those batches for which we actually have a loss. If this variable ever
# gets used for something else, we might need to change things around a bit.
batches_this_epoch += 1
val_loss += loss.detach().cpu().numpy()
if reg_loss is not None:
val_reg_loss += reg_loss.detach().cpu().numpy()
# Update the description with the latest metrics
val_metrics = training_util.get_metrics(
self.model,
val_loss,
val_reg_loss,
batches_this_epoch,
world_size=self._world_size,
cuda_device=[self.cuda_device],
)
description = training_util.description_from_metrics(val_metrics)
val_generator_tqdm.set_description(description, refresh=False)
if self._master:
for callback in self._batch_callbacks:
callback(
self,
[batch],
[batch_outputs],
epoch,
batches_this_epoch,
is_training=False,
)
if self._distributed and not done_early:
logger.warning(
f"Worker {torch.distributed.get_rank()} completed its entire epoch (validation)."
)
# Indicate that we're done so that any workers that have remaining data stop validation early.
done = torch.tensor(1, device=self.cuda_device)
torch.distributed.all_reduce(done, torch.distributed.ReduceOp.SUM)
assert done.item()
# Now restore the original parameter values.
if self._moving_average is not None:
self._moving_average.restore()
return val_loss, val_reg_loss, batches_this_epoch
def train(self) -> Dict[str, Any]:
"""
Trains the supplied model with the supplied parameters.
"""
try:
epoch_counter = self._restore_checkpoint()
except RuntimeError:
traceback.print_exc()
raise ConfigurationError(
"Could not recover training from the checkpoint. Did you mean to output to "
"a different serialization directory or delete the existing serialization "
"directory?")
training_util.enable_gradient_clipping(self.model, self._grad_clipping)
logger.info("Beginning training.")
val_metrics: Dict[str, float] = {}
this_epoch_val_metric: float = None
metrics: Dict[str, Any] = {}
epochs_trained = 0
training_start_time = time.time()
metrics["best_epoch"] = self._metric_tracker.best_epoch
for key, value in self._metric_tracker.best_epoch_metrics.items():
metrics["best_validation_" + key] = value
for callback in self._epoch_callbacks:
callback(self, metrics={}, epoch=-1)
for epoch in range(epoch_counter, self._num_epochs):
epoch_start_time = time.time()
train_metrics = self._train_epoch(epoch)
# get peak of memory usage
if "cpu_memory_MB" in train_metrics:
metrics["peak_cpu_memory_MB"] = max(
metrics.get("peak_cpu_memory_MB", 0),
train_metrics["cpu_memory_MB"])
for key, value in train_metrics.items():
if key.startswith("gpu_"):
metrics["peak_" + key] = max(metrics.get("peak_" + key, 0),
value)
if self._validation_data_loader is not None:
with torch.no_grad():
# We have a validation set, so compute all the metrics on it.
val_loss, val_reg_loss, num_batches = self._validation_loss(
epoch)
# It is safe again to wait till the validation is done. This is
# important to get the metrics right.
if self._distributed:
dist.barrier()
val_metrics = training_util.get_metrics(
self.model,
val_loss,
val_reg_loss,
num_batches,
reset=True,
world_size=self._world_size,
cuda_device=[self.cuda_device],
)
# Check validation metric for early stopping
this_epoch_val_metric = val_metrics[
self._validation_metric]
self._metric_tracker.add_metric(this_epoch_val_metric)
if self._metric_tracker.should_stop_early():
logger.info("Ran out of patience. Stopping training.")
break
if self._master:
self._tensorboard.log_metrics(
train_metrics,
val_metrics=val_metrics,
log_to_console=True,
epoch=epoch + 1) # +1 because tensorboard doesn't like 0
# Create overall metrics dict
training_elapsed_time = time.time() - training_start_time
metrics["training_duration"] = str(
datetime.timedelta(seconds=training_elapsed_time))
metrics["training_start_epoch"] = epoch_counter
metrics["training_epochs"] = epochs_trained
metrics["epoch"] = epoch
for key, value in train_metrics.items():
metrics["training_" + key] = value
for key, value in val_metrics.items():
metrics["validation_" + key] = value
if self._metric_tracker.is_best_so_far():
# Update all the best_ metrics.
# (Otherwise they just stay the same as they were.)
metrics["best_epoch"] = epoch
for key, value in val_metrics.items():
metrics["best_validation_" + key] = value
self._metric_tracker.best_epoch_metrics = val_metrics
if self._serialization_dir and self._master:
common_util.dump_metrics(
os.path.join(self._serialization_dir,
f"metrics_epoch_{epoch}.json"), metrics)
# The Scheduler API is agnostic to whether your schedule requires a validation metric -
# if it doesn't, the validation metric passed here is ignored.
if self._learning_rate_scheduler:
self._learning_rate_scheduler.step(this_epoch_val_metric)
if self._momentum_scheduler:
self._momentum_scheduler.step(this_epoch_val_metric)
if self._master:
self._checkpointer.save_checkpoint(
epoch,
self,
is_best_so_far=self._metric_tracker.is_best_so_far())
# Wait for the master to finish saving the checkpoint
if self._distributed:
dist.barrier()
for callback in self._epoch_callbacks:
callback(self, metrics=metrics, epoch=epoch)
epoch_elapsed_time = time.time() - epoch_start_time
logger.info("Epoch duration: %s",
datetime.timedelta(seconds=epoch_elapsed_time))
if epoch < self._num_epochs - 1:
training_elapsed_time = time.time() - training_start_time
estimated_time_remaining = training_elapsed_time * (
(self._num_epochs - epoch_counter) /
float(epoch - epoch_counter + 1) - 1)
formatted_time = str(
datetime.timedelta(seconds=int(estimated_time_remaining)))
logger.info("Estimated training time remaining: %s",
formatted_time)
epochs_trained += 1
# make sure pending events are flushed to disk and files are closed properly
self._tensorboard.close()
# Load the best model state before returning
best_model_state = self._checkpointer.best_model_state()
if best_model_state:
self.model.load_state_dict(best_model_state)
return metrics
@contextmanager
def get_checkpoint_state(
self) -> Iterator[Tuple[Dict[str, Any], Dict[str, Any]]]:
if self._moving_average is not None:
# Assigning average value to model parameters. The checkpointer will call
# `restore_state_after_checkpointing` when it is done to put this back to what it was.
self._moving_average.assign_average_value()
model_state = self.model.state_dict()
# These are the training states we need to persist.
training_states = {
"metric_tracker": self._metric_tracker.state_dict(),
"optimizer": self.optimizer.state_dict(),
"batch_num_total": self._batch_num_total,
}
# If we have a learning rate or momentum scheduler, we should persist them too.
if self._learning_rate_scheduler is not None:
training_states[
"learning_rate_scheduler"] = self._learning_rate_scheduler.state_dict(
)
if self._momentum_scheduler is not None:
training_states[
"momentum_scheduler"] = self._momentum_scheduler.state_dict()
# If model was trained with amp, we should persist the amp state.
if self._opt_level is not None:
training_states["amp"] = amp.state_dict()
try:
yield model_state, training_states
finally:
if self._moving_average is not None:
self._moving_average.restore()
def _restore_checkpoint(self) -> int:
"""
Restores the model and training state from the last saved checkpoint.
This includes an epoch count and optimizer state, which is serialized separately
from model parameters. This function should only be used to continue training -
if you wish to load a model for inference/load parts of a model into a new
computation graph, you should use the native Pytorch functions:
` model.load_state_dict(torch.load("/path/to/model/weights.th"))`
If `self._serialization_dir` does not exist or does not contain any checkpointed weights,
this function will do nothing and return 0.
# Returns
epoch: int
The epoch at which to resume training, which should be one after the epoch
in the saved training state.
"""
model_state, training_state = self._checkpointer.restore_checkpoint()
if not training_state:
# No checkpoint to restore, start at 0
return 0
# The apex docs recommend calling amp.initialize before calling load_state_dict.
if self._opt_level is not None and "amp" in training_state:
self.model, self.optimizer = amp.initialize(
self.model, self.optimizer, opt_level=self._opt_level)
amp.load_state_dict(training_state["amp"])
self.model.load_state_dict(model_state)
self.optimizer.load_state_dict(training_state["optimizer"])
if (self._learning_rate_scheduler is not None
and "learning_rate_scheduler" in training_state):
self._learning_rate_scheduler.load_state_dict(
training_state["learning_rate_scheduler"])
if self._momentum_scheduler is not None and "momentum_scheduler" in training_state:
self._momentum_scheduler.load_state_dict(
training_state["momentum_scheduler"])
training_util.move_optimizer_to_cuda(self.optimizer)
# Currently the `training_state` contains a serialized `MetricTracker`.
if "metric_tracker" in training_state:
self._metric_tracker.load_state_dict(
training_state["metric_tracker"])
# It used to be the case that we tracked `val_metric_per_epoch`.
elif "val_metric_per_epoch" in training_state:
self._metric_tracker.clear()
self._metric_tracker.add_metrics(
training_state["val_metric_per_epoch"])
# And before that we didn't track anything.
else:
self._metric_tracker.clear()
if isinstance(training_state["epoch"], int):
epoch_to_return = training_state["epoch"] + 1
else:
epoch_to_return = int(training_state["epoch"].split(".")[0]) + 1
# For older checkpoints with batch_num_total missing, default to old behavior where
# it is unchanged.
batch_num_total = training_state.get("batch_num_total")
if batch_num_total is not None:
self._batch_num_total = batch_num_total
return epoch_to_return
@classmethod
def from_partial_objects(
cls,
model: Model,
serialization_dir: str,
data_loader: DataLoader,
validation_data_loader: DataLoader = None,
local_rank: int = 0,
patience: int = None,
validation_metric: str = "-loss",
num_epochs: int = 20,
cuda_device: int = -1,
grad_norm: float = None,
grad_clipping: float = None,
distributed: bool = None,
world_size: int = 1,
num_gradient_accumulation_steps: int = 1,
opt_level: Optional[str] = None,
no_grad: List[str] = None,
optimizer: Lazy[Optimizer] = None,
learning_rate_scheduler: Lazy[LearningRateScheduler] = None,
momentum_scheduler: Lazy[MomentumScheduler] = None,
tensorboard_writer: Lazy[TensorboardWriter] = None,
moving_average: Lazy[MovingAverage] = None,
checkpointer: Lazy[Checkpointer] = None,
batch_callbacks: List[BatchCallback] = None,
epoch_callbacks: List[EpochCallback] = None,
) -> "Trainer":
"""
This method exists so that we can have a documented method to construct this class using
`FromParams`. If you are not using `FromParams` or config files, you can safely ignore this
method.
The reason we can't just use `__init__` with `FromParams` here is because there are
sequential dependencies to this class's arguments. Anything that has a `Lazy[]` type
annotation needs something from one of the non-`Lazy` arguments. The `Optimizer` needs to
have the parameters from the `Model` before it's constructed, and the `Schedulers` need to
have the `Optimizer`. Because of this, the typical way we construct things `FromParams`
doesn't work, so we use `Lazy` to allow for constructing the objects sequentially.
If you're not using `FromParams`, you can just construct these arguments in the right order
yourself in your code and call the constructor directly.
"""
check_for_gpu(cuda_device)
if cuda_device >= 0:
# Moving model to GPU here so that the optimizer state gets constructed on
# the right device.
model = model.cuda(cuda_device)
if no_grad:
for name, parameter in model.named_parameters():
if any(re.search(regex, name) for regex in no_grad):
parameter.requires_grad_(False)
common_util.log_frozen_and_tunable_parameter_names(model)
parameters = [[n, p] for n, p in model.named_parameters()
if p.requires_grad]
optimizer_ = optimizer.construct(model_parameters=parameters)
if not optimizer_:
optimizer_ = Optimizer.default(parameters)
try:
batches_per_epoch = len(data_loader)
except TypeError:
# If the dataset is lazy, it won't have a length.
batches_per_epoch = None
moving_average_ = moving_average.construct(parameters=parameters)
learning_rate_scheduler_ = learning_rate_scheduler.construct(
optimizer=optimizer_,
num_epochs=num_epochs,
num_steps_per_epoch=batches_per_epoch)
momentum_scheduler_ = momentum_scheduler.construct(
optimizer=optimizer_)
checkpointer_ = checkpointer.construct() or Checkpointer(
serialization_dir)
tensorboard_writer_ = tensorboard_writer.construct(
) or TensorboardWriter(serialization_dir)
return cls(
model,
optimizer_,
data_loader,
patience=patience,
validation_metric=validation_metric,
validation_data_loader=validation_data_loader,
num_epochs=num_epochs,
serialization_dir=serialization_dir,
cuda_device=cuda_device,
grad_norm=grad_norm,
grad_clipping=grad_clipping,
learning_rate_scheduler=learning_rate_scheduler_,
momentum_scheduler=momentum_scheduler_,
tensorboard_writer=tensorboard_writer_,
checkpointer=checkpointer_,
moving_average=moving_average_,
batch_callbacks=batch_callbacks,
epoch_callbacks=epoch_callbacks,
distributed=distributed,
local_rank=local_rank,
world_size=world_size,
num_gradient_accumulation_steps=num_gradient_accumulation_steps,
opt_level=opt_level,
)
class MultiTaskTrainer(TrainerBase):
"""
A simple trainer that works in our multi-task setup.
Really the main thing that makes this task not fit into our
existing trainer is the multiple datasets.
"""
def __init__(self,
model: Model,
serialization_dir: str,
iterator: DataIterator,
mingler: DatasetMingler,
optimizer: torch.optim.Optimizer,
datasets: Dict[str, Iterable[Instance]],
num_epochs: int = 10,
num_serialized_models_to_keep: int = 10) -> None:
super().__init__(serialization_dir)
self.model = model
self.iterator = iterator
self.mingler = mingler
self.optimizer = optimizer
self.datasets = datasets
self.num_epochs = num_epochs
self.checkpointer = Checkpointer(
serialization_dir,
num_serialized_models_to_keep=num_serialized_models_to_keep)
def save_checkpoint(self, epoch: int) -> None:
training_state = {
"epoch": epoch,
"optimizer": self.optimizer.state_dict()
}
self.checkpointer.save_checkpoint(epoch, self.model.state_dict(),
training_state, True)
def restore_checkpoint(self) -> int:
model_state, trainer_state = self.checkpointer.restore_checkpoint()
if not model_state and not trainer_state:
return 0
else:
self.model.load_state_dict(model_state)
self.optimizer.load_state_dict(trainer_state["optimizer"])
return trainer_state["epoch"] + 1
def train(self) -> Dict:
start_epoch = self.restore_checkpoint()
self.model.train()
for epoch in range(start_epoch, self.num_epochs):
total_loss = 0.0
batches = tqdm.tqdm(
self.iterator(self.mingler.mingle(self.datasets),
num_epochs=1))
for i, batch in enumerate(batches):
self.optimizer.zero_grad()
loss = self.model.forward(**batch)['loss'] # type: ignore
loss.backward()
total_loss += loss.item()
self.optimizer.step()
batches.set_description(
f"epoch: {epoch} loss: {total_loss / (i + 1)}")
# Save checkpoint
self.save_checkpoint(epoch)
return {}
@classmethod
def from_params(
cls, # type: ignore
params: Params,
serialization_dir: str,
recover: bool = False,
cache_directory: str = None,
cache_prefix: str = None) -> 'MultiTaskTrainer':
readers = {
name: DatasetReader.from_params(reader_params)
for name, reader_params in params.pop(
"train_dataset_readers").items()
}
train_file_paths = params.pop("train_file_paths").as_dict()
datasets = {
name: reader.read(train_file_paths[name])
for name, reader in readers.items()
}
instances = (instance for dataset in datasets.values()
for instance in dataset)
vocab = Vocabulary.from_params(Params({}), instances)
model = Model.from_params(params.pop('model'), vocab=vocab)
iterator = DataIterator.from_params(params.pop('iterator'))
iterator.index_with(vocab)
mingler = DatasetMingler.from_params(params.pop('mingler'))
parameters = [[n, p] for n, p in model.named_parameters()
if p.requires_grad]
optimizer = Optimizer.from_params(parameters, params.pop('optimizer'))
num_epochs = params.pop_int("num_epochs", 10)
_ = params.pop("trainer", Params({}))
params.assert_empty(__name__)
return MultiTaskTrainer(model, serialization_dir, iterator, mingler,
optimizer, datasets, num_epochs)
class PtDistTrainer(TrainerBase):
def __init__(
self,
model: Model,
optimizer: torch.optim.Optimizer,
train_dataset: Iterable[Instance],
validation_dataset: Optional[Iterable[Instance]] = None,
batch_size: int = 1,
validation_metric: str = "-loss",
shuffle: bool = True,
num_epochs: int = 20,
serialization_dir: Optional[str] = None,
num_serialized_models_to_keep: int = 20,
checkpointer: Checkpointer = None,
cuda_device: Union[int, List] = -1,
grad_clipping: Optional[float] = None,
learning_rate_scheduler: Optional[LearningRateScheduler] = None
) -> None:
super().__init__(serialization_dir, cuda_device)
self.local_rank = dist.get_rank()
self.local_device = torch.device("cuda", self.local_rank)
self.model = DDP(model,
device_ids=[self.local_rank],
output_device=self.local_rank)
self.batch_size = batch_size
self.shuffle = shuffle
self.optimizer = optimizer
self.train_data = train_dataset
self._validation_data = validation_dataset
self._metric_tracker = MetricTracker(metric_name=validation_metric)
self._validation_metric = validation_metric[1:]
self._num_epochs = num_epochs
if checkpointer is not None:
self._checkpointer = checkpointer
else:
self._checkpointer = Checkpointer(serialization_dir, None,
num_serialized_models_to_keep)
self._grad_clipping = grad_clipping
self._learning_rate_scheduler = learning_rate_scheduler
self._batch_num_total = 0
self._last_log = 0.0 # time of last logging
def batch_loss(self, batch_group: List[TensorDict],
for_training: bool) -> torch.Tensor:
output_dict = self.model(**batch_group)
loss = output_dict["loss"]
return loss
def _train_epoch(self, epoch: int) -> Dict[str, float]:
train_loss = 0.0
self.model.train()
num_gpus = len(self._cuda_devices)
if getattr(self, "train_dataset", None) is None:
self.train_dataset = DMDataSet(data=self.train_data[0],
batch_size=self.batch_size,
num_gpus=num_gpus,
shuffle=True,
distributed=True,
data_slice=True)
self.train_dataset.set_epoch(epoch)
num_training_batches = math.ceil(
len(self.train_dataset) / self.batch_size / num_gpus)
self._last_log = time.time()
batches_this_epoch = 0
if self._batch_num_total is None:
self._batch_num_total = 0
logger.info("Training")
train_generator_tqdm = Tqdm.tqdm(self.train_dataset,
total=num_training_batches)
for batch_group in train_generator_tqdm:
# print('batch_size: ', len(batch_group["source_tokens"]["tokens"]))
# gpu_data = batch_group
# src_data = gpu_data["source_tokens"]["tokens"]
# tgt_data = gpu_data["target_tokens"]["tokens"]
# for sdata, tdata in zip(src_data, tgt_data):
# s = ''.join([self.get_model().vocab.get_token_from_index(x, "source_tokens") if x != 0 else '' for x in
# sdata.numpy()])
# t = ''.join([self.get_model().vocab.get_token_from_index(x, "target_tokens") if x != 0 else '' for x in
# tdata.numpy()])
# print(s)
# print(t)
batches_this_epoch += 1
self._batch_num_total += 1
batch_num_total = self._batch_num_total
self.optimizer.zero_grad()
loss = self.batch_loss(batch_group, for_training=True)
if torch.isnan(loss):
raise ValueError("nan loss encountered")
loss.backward()
train_loss += loss.item()
if self._learning_rate_scheduler:
self._learning_rate_scheduler.step_batch(batch_num_total)
self.optimizer.step()
metrics = training_util.get_metrics(self.get_model(), train_loss,
batches_this_epoch)
description = self.get_desc_from_metrics(metrics, epoch)
train_generator_tqdm.set_description(description, refresh=False)
metrics = training_util.get_metrics(self.get_model(),
train_loss,
batches_this_epoch,
reset=True)
return metrics
def get_desc_from_metrics(self, metrics, epoch=None):
description = training_util.description_from_metrics(metrics)
if epoch is None:
description = f'epoch: -- rank: {dist.get_rank()} || {description}'
else:
description = f'epoch: {epoch} rank: {dist.get_rank()} || {description}'
return description
def get_model(self):
return self.model.module
def _validation_loss(self) -> Tuple[float, int]:
logger.info("Validating")
self.model.eval()
num_gpus = len(self._cuda_devices)
if getattr(self, "val_dataset", None) is None:
self.val_dataset = DMDataSet(data=self._validation_data[0],
batch_size=self.batch_size,
num_gpus=num_gpus,
shuffle=False,
distributed=True,
data_slice=False)
num_validation_batches = math.ceil(
len(self.val_dataset) / self.batch_size / num_gpus)
val_generator_tqdm = Tqdm.tqdm(self.val_dataset,
total=num_validation_batches)
batches_this_epoch = 0
val_loss = 0
for batch_group in val_generator_tqdm:
loss = self.batch_loss(batch_group, for_training=False)
if loss is not None:
batches_this_epoch += 1
val_loss += loss.detach().cpu().numpy()
# Update the description with the latest metrics
val_metrics = training_util.get_metrics(self.get_model(), val_loss,
batches_this_epoch)
description = self.get_desc_from_metrics(val_metrics)
val_generator_tqdm.set_description(description, refresh=False)
return val_loss, batches_this_epoch
def train(self) -> Dict[str, Any]:
try:
epoch_counter = self._restore_checkpoint()
except RuntimeError:
traceback.print_exc()
raise ConfigurationError(
"Could not recover training from the checkpoint. Did you mean to output to "
"a different serialization directory or delete the existing serialization "
"directory?")
training_util.enable_gradient_clipping(self.model, self._grad_clipping)
logger.info("Beginning training.")
train_metrics: Dict[str, float] = {}
val_metrics: Dict[str, float] = {}
this_epoch_val_metric: float = None
metrics: Dict[str, Any] = {}
epochs_trained = 0
training_start_time = time.time()
metrics['best_epoch'] = self._metric_tracker.best_epoch
for key, value in self._metric_tracker.best_epoch_metrics.items():
metrics["best_validation_" + key] = value
for epoch in range(epoch_counter, self._num_epochs):
epoch_start_time = time.time()
train_metrics = self._train_epoch(epoch)
if self._validation_data is not None:
with torch.no_grad():
val_loss, num_batches = self._validation_loss()
val_metrics = training_util.get_metrics(self.get_model(),
val_loss,
num_batches,
reset=True)
this_epoch_val_metric = val_metrics[
self._validation_metric]
self._metric_tracker.add_metric(this_epoch_val_metric)
if self._metric_tracker.should_stop_early():
logger.info("Ran out of patience. Stopping training.")
break
# Create overall metrics dict
training_elapsed_time = time.time() - training_start_time
metrics["training_duration"] = str(
datetime.timedelta(seconds=training_elapsed_time))
metrics["training_start_epoch"] = epoch_counter
metrics["training_epochs"] = epochs_trained
metrics["epoch"] = epoch
for key, value in train_metrics.items():
metrics["training_" + key] = value
for key, value in val_metrics.items():
metrics["validation_" + key] = value
if self._metric_tracker.is_best_so_far():
metrics['best_epoch'] = epoch
for key, value in val_metrics.items():
metrics["best_validation_" + key] = value
self._metric_tracker.best_epoch_metrics = val_metrics
if self._serialization_dir and is_master_rank():
dump_metrics(
os.path.join(self._serialization_dir,
f'metrics_epoch_{epoch}.json'), metrics)
if self._learning_rate_scheduler:
self._learning_rate_scheduler.step(this_epoch_val_metric,
epoch)
if is_master_rank():
self._save_checkpoint(epoch)
epoch_elapsed_time = time.time() - epoch_start_time
logger.info("Epoch duration: %s",
datetime.timedelta(seconds=epoch_elapsed_time))
if epoch < self._num_epochs - 1:
training_elapsed_time = time.time() - training_start_time
estimated_time_remaining = training_elapsed_time * \
((self._num_epochs - epoch_counter) / float(epoch - epoch_counter + 1) - 1)
formatted_time = str(
datetime.timedelta(seconds=int(estimated_time_remaining)))
logger.info("Estimated training time remaining: %s",
formatted_time)
epochs_trained += 1
best_model_state = self._checkpointer.best_model_state()
if best_model_state:
self.model.load_state_dict(best_model_state)
return metrics
def _save_checkpoint(self, epoch: Union[int, str]) -> None:
training_states = {
"metric_tracker": self._metric_tracker.state_dict(),
"optimizer": self.optimizer.state_dict(),
"batch_num_total": self._batch_num_total
}
if self._learning_rate_scheduler is not None:
training_states[
"learning_rate_scheduler"] = self._learning_rate_scheduler.state_dict(
)
self._checkpointer.save_checkpoint(
model_state=self.model.state_dict(),
epoch=epoch,
training_states=training_states,
is_best_so_far=self._metric_tracker.is_best_so_far())
def _restore_checkpoint(self) -> int:
model_state, training_state = self._checkpointer.restore_checkpoint()
if not training_state:
# No checkpoint to restore, start at 0
return 0
self.model.load_state_dict(model_state)
self.optimizer.load_state_dict(training_state["optimizer"])
if self._learning_rate_scheduler is not None and "learning_rate_scheduler" in training_state:
self._learning_rate_scheduler.load_state_dict(
training_state["learning_rate_scheduler"])
training_util.move_optimizer_to_cuda(self.optimizer)
# Currently the ``training_state`` contains a serialized ``MetricTracker``.
if "metric_tracker" in training_state:
self._metric_tracker.load_state_dict(
training_state["metric_tracker"])
# It used to be the case that we tracked ``val_metric_per_epoch``.
elif "val_metric_per_epoch" in training_state:
self._metric_tracker.clear()
self._metric_tracker.add_metrics(
training_state["val_metric_per_epoch"])
# And before that we didn't track anything.
else:
self._metric_tracker.clear()
if isinstance(training_state["epoch"], int):
epoch_to_return = training_state["epoch"] + 1
else:
epoch_to_return = int(training_state["epoch"].split('.')[0]) + 1
# For older checkpoints with batch_num_total missing, default to old behavior where
# it is unchanged.
batch_num_total = training_state.get('batch_num_total')
if batch_num_total is not None:
self._batch_num_total = batch_num_total
return epoch_to_return
# Requires custom from_params.
@classmethod
def from_params(cls,
params: Params,
serialization_dir: str,
recover: bool = False,
cache_directory: str = None,
cache_prefix: str = None) -> 'PtDistTrainer':
all_datasets = training_util.datasets_from_params(
params, cache_directory, cache_prefix)
vocab = Vocabulary.from_files(params.vocabulary.directory_path)
model = Model.from_params(vocab=vocab, params=params.pop('model'))
model.extend_embedder_vocab()
if is_master_rank():
vocab.save_to_files(os.path.join(serialization_dir, "vocabulary"))
train_data = all_datasets['train']
validation_data = all_datasets.get('validation')
batch_size = params.iterator.batch_size
trainer_params = params.pop("trainer")
keys = [key for key in params]
for key in keys:
params.pop(key)
params = trainer_params
validation_metric = params.pop("validation_metric", "-loss")
shuffle = params.pop_bool("shuffle", True)
num_epochs = params.pop_int("num_epochs", 20)
cuda_device = parse_cuda_device(params.pop("cuda_device", -1))
grad_clipping = params.pop_float("grad_clipping", None)
lr_scheduler_params = params.pop("learning_rate_scheduler", None)
pretrain_file = params.pop("pretrain_file", None)
no_grad_regexes = params.pop("no_grad", ())
for name, parameter in model.named_parameters():
if any(re.search(regex, name) for regex in no_grad_regexes):
parameter.requires_grad_(False)
frozen_parameter_names, tunable_parameter_names = \
get_frozen_and_tunable_parameter_names(model)
logger.info("Following parameters are Frozen (without gradient):")
for name in frozen_parameter_names:
logger.info(name)
logger.info("Following parameters are Tunable (with gradient):")
for name in tunable_parameter_names:
logger.info(name)
model = model.cuda(dist.get_rank())
if pretrain_file:
model_state = torch.load(pretrain_file,
map_location=nn_util.device_mapping(
dist.get_rank()))
model.load_state_dict(model_state)
parameters = [[n, p] for n, p in model.named_parameters()
if p.requires_grad]
# print([n for n, p in model.named_parameters() if p.requires_grad])
optimizer = Optimizer.from_params(parameters, params.pop("optimizer"))
if lr_scheduler_params:
lr_scheduler = LearningRateScheduler.from_params(
optimizer, lr_scheduler_params)
else:
lr_scheduler = None
num_serialized_models_to_keep = params.pop_int(
"num_serialized_models_to_keep", 20)
checkpointer = Checkpointer(
serialization_dir=serialization_dir,
num_serialized_models_to_keep=num_serialized_models_to_keep,
keep_serialized_model_every_num_seconds=None)
return cls(model,
optimizer,
train_data,
validation_data,
batch_size=batch_size,
validation_metric=validation_metric,
shuffle=shuffle,
num_epochs=num_epochs,
serialization_dir=serialization_dir,
cuda_device=cuda_device,
grad_clipping=grad_clipping,
learning_rate_scheduler=lr_scheduler,
checkpointer=checkpointer)
class MyGradientDescentTrainer(Trainer):
"""
A trainer for doing supervised learning with gradient descent. It just takes a labeled dataset
and a `DataLoader`, and uses the supplied `Optimizer` to learn the weights for your model over
some fixed number of epochs. You can also pass in a validation data_loader and enable early
stopping. There are many other bells and whistles as well.
Registered as a `Trainer` with the name "gradient_descent" (and is also the default `Trainer`).
The constructor that is registered is [`from_partial_objects`](#from_partial_objects) -
see the arguments to that function for the exact keys that should be used, if you are using
a configuration file. They largely match the arguments to `__init__`, and we don't repeat their
docstrings in `from_partial_objects`.
[0]: https://tinyurl.com/y5mv44fw
# Parameters
model : `Model`, required.
An AllenNLP model to be optimized. Pytorch Modules can also be optimized if
their `forward` method returns a dictionary with a "loss" key, containing a
scalar tensor representing the loss function to be optimized.
If you are training your model using GPUs, your model should already be
on the correct device. (If you are using our `train` command this will be
handled for you.)
In a typical AllenNLP configuration file, this parameter does not get an entry under the
"trainer", it gets constructed separately.
optimizer : `torch.nn.Optimizer`, required.
An instance of a Pytorch Optimizer, instantiated with the parameters of the
model to be optimized.
data_loader : `DataLoader`, required.
A `DataLoader` containing your `Dataset`, yielding padded indexed batches.
In a typical AllenNLP configuration file, this parameter does not get an entry under the
"trainer", it gets constructed separately.
patience : `Optional[int] > 0`, optional (default=`None`)
Number of epochs to be patient before early stopping: the training is stopped
after `patience` epochs with no improvement. If given, it must be `> 0`.
If None, early stopping is disabled.
validation_metric : `Union[str, List[str]]`, optional (default=`"-loss"`)
Validation metric to measure for whether to stop training using patience
and whether to serialize an `is_best` model each epoch. The metric name
must be prepended with either "+" or "-", which specifies whether the metric
is an increasing or decreasing function. If you specify more than one metric,
the metrics will be summed to make the `is_best` decision.
validation_data_loader : `DataLoader`, optional (default=`None`)
A `DataLoader` to use for the validation set. If `None`, then
use the training `DataLoader` with the validation data.
In a typical AllenNLP configuration file, this parameter does not get an entry under the
"trainer", it gets constructed separately.
num_epochs : `int`, optional (default = `20`)
Number of training epochs.
serialization_dir : `str`, optional (default=`None`)
Path to directory for saving and loading model files. Models will not be saved if
this parameter is not passed.
In a typical AllenNLP configuration file, this parameter does not get an entry under the
"trainer", it gets constructed separately.
checkpointer : `Checkpointer`, optional (default=`None`)
A `Checkpointer` is responsible for periodically saving model weights. If none is given
here, we will construct one with default parameters.
cuda_device : `int`, optional (default = `-1`)
An integer specifying the CUDA device(s) to use for this process. If -1, the CPU is used.
Data parallelism is controlled at the allennlp train level, so each trainer will have a single
GPU.
grad_norm : `float`, optional, (default = `None`).
If provided, gradient norms will be rescaled to have a maximum of this value.
grad_clipping : `float`, optional (default = `None`).
If provided, gradients will be clipped `during the backward pass` to have an (absolute)
maximum of this value. If you are getting `NaNs` in your gradients during training
that are not solved by using `grad_norm`, you may need this.
learning_rate_scheduler : `LearningRateScheduler`, optional (default = `None`)
If specified, the learning rate will be decayed with respect to
this schedule at the end of each epoch (or batch, if the scheduler implements
the `step_batch` method). If you use `torch.optim.lr_scheduler.ReduceLROnPlateau`,
this will use the `validation_metric` provided to determine if learning has plateaued.
To support updating the learning rate on every batch, this can optionally implement
`step_batch(batch_num_total)` which updates the learning rate given the batch number.
momentum_scheduler : `MomentumScheduler`, optional (default = `None`)
If specified, the momentum will be updated at the end of each batch or epoch
according to the schedule.
moving_average : `MovingAverage`, optional, (default = `None`)
If provided, we will maintain moving averages for all parameters. During training, we
employ a shadow variable for each parameter, which maintains the moving average. During
evaluation, we backup the original parameters and assign the moving averages to corresponding
parameters. Be careful that when saving the checkpoint, we will save the moving averages of
parameters. This is necessary because we want the saved model to perform as well as the validated
model if we load it later. But this may cause problems if you restart the training from checkpoint.
callbacks : `List[Lazy[TrainerCallback]]`, optional (default = `None`)
A list of callbacks that can be called at certain events: e.g. each batch, epoch, and at the start
and end of training, etc.
distributed : `bool`, optional, (default = `False`)
If set, PyTorch's `DistributedDataParallel` is used to train the model in multiple GPUs. This also
requires `world_size` to be greater than 1.
In a typical AllenNLP configuration file, this parameter does not get an entry under the
"trainer", it gets constructed separately (you need a top-level "distributed" key, next to
the "trainer" entry, that specifies a list of "cuda_devices").
local_rank : `int`, optional, (default = `0`)
This is the unique identifier of the `Trainer` in a distributed process group. The GPU device id is
used as the rank.
In a typical AllenNLP configuration file, this parameter does not get an entry under the
"trainer", it gets constructed separately.
world_size : `int`, (default = `1`)
The number of `Trainer` workers participating in the distributed training.
In a typical AllenNLP configuration file, this parameter does not get an entry under the
"trainer", it gets constructed separately.
num_gradient_accumulation_steps : `int`, optional, (default = `1`)
Gradients are accumulated for the given number of steps before doing an optimizer step. This can
be useful to accommodate batches that are larger than the RAM size. Refer [Thomas Wolf's
post][0] for details on Gradient Accumulation.
use_amp : `bool`, optional, (default = `False`)
If `True`, we'll train using [Automatic Mixed Precision](https://pytorch.org/docs/stable/amp.html).
enable_default_callbacks : `bool`, optional (default = `True`)
When `True`, the [`DEFAULT_CALLBACKS`](#default_callbacks) will be used in
addition to any other callbacks listed in the `callbacks` parameter.
When set to `False`, `DEFAULT_CALLBACKS` are not used.
run_sanity_checks : `bool`, optional (default = `True`)
Determines whether model sanity checks, such as
[`NormalizationBiasVerification`](../../sanity_checks/normalization_bias_verification/),
are ran.
"""
def __init__(
self,
model: Model,
optimizer: torch.optim.Optimizer,
data_loader: DataLoader,
patience: Optional[int] = None,
validation_metric: Union[str, List[str]] = "-loss",
validation_data_loader: DataLoader = None,
num_epochs: int = 20,
serialization_dir: Optional[str] = None,
checkpointer: Checkpointer = None,
cuda_device: Optional[Union[int, torch.device]] = None,
grad_norm: Optional[float] = None,
grad_clipping: Optional[float] = None,
learning_rate_scheduler: Optional[LearningRateScheduler] = None,
momentum_scheduler: Optional[MomentumScheduler] = None,
moving_average: Optional[MovingAverage] = None,
callbacks: List[TrainerCallback] = None,
distributed: bool = False,
local_rank: int = 0,
world_size: int = 1,
num_gradient_accumulation_steps: int = 1,
use_amp: bool = False,
enable_default_callbacks: bool = True,
run_sanity_checks: bool = True,
val_loss_steps: int = 50
) -> None:
super().__init__(serialization_dir, cuda_device, distributed, local_rank, world_size)
# I am not calling move_to_gpu here, because if the model is
# not already on the GPU then the optimizer is going to be wrong.
self.model = model
self.data_loader = data_loader
self.data_loader.set_target_device(self.cuda_device)
self._validation_data_loader = validation_data_loader
if self._validation_data_loader is not None:
self._validation_data_loader.set_target_device(self.cuda_device)
self.optimizer = optimizer
if patience is None: # no early stopping
if validation_data_loader is not None:
logger.warning(
"You provided a validation dataset but patience was set to None, "
"meaning that early stopping is disabled"
)
elif (not isinstance(patience, int)) or patience <= 0:
raise ConfigurationError(
'{} is an invalid value for "patience": it must be a positive integer '
"or None (if you want to disable early stopping)".format(patience)
)
# For tracking is_best_so_far and should_stop_early
self._metric_tracker = MetricTracker(validation_metric, patience)
self._num_epochs = num_epochs
self._checkpointer: Optional[Checkpointer] = checkpointer
if checkpointer is None and serialization_dir is not None:
self._checkpointer = Checkpointer(serialization_dir)
self._grad_norm = grad_norm
self._grad_clipping = grad_clipping
self._learning_rate_scheduler = learning_rate_scheduler
self._momentum_scheduler = momentum_scheduler
self._moving_average = moving_average
self._callbacks = callbacks or []
default_callbacks = list(DEFAULT_CALLBACKS) if enable_default_callbacks else []
if run_sanity_checks:
default_callbacks.append(SanityChecksCallback)
for callback_cls in default_callbacks:
for callback in self._callbacks:
if callback.__class__ == callback_cls:
break
else:
self._callbacks.append(callback_cls(self._serialization_dir))
self._batch_num_total = 0
self._last_log = 0.0 # time of last logging
self._num_gradient_accumulation_steps = num_gradient_accumulation_steps
# Enable automatic mixed precision training.
self._scaler: Optional[amp.GradScaler] = None
self._use_amp = use_amp
if self._use_amp:
if self.cuda_device == torch.device("cpu"):
raise ValueError("Using AMP requires a cuda device")
self._scaler = amp.GradScaler()
# Using `DistributedDataParallel`(ddp) brings in a quirk wrt AllenNLP's `Model` interface and its
# usage. A `Model` object is wrapped by `ddp`, but assigning the wrapped model to `self.model`
# will break the usages such as `Model.get_regularization_penalty`, `Model.get_metrics`, etc.
#
# Hence a reference to Pytorch's object is maintained in the case of distributed training and in the
# normal case, reference to `Model` is retained. This reference is only used in
# these places: `model.__call__`, `model.train` and `model.eval`.
if self._distributed:
self._pytorch_model = DistributedDataParallel(
self.model,
device_ids=None if self.cuda_device == torch.device("cpu") else [self.cuda_device],
find_unused_parameters=True,
)
else:
self._pytorch_model = self.model
self.val_loss_steps = val_loss_steps
def batch_outputs(self, batch: TensorDict, for_training: bool) -> Dict[str, torch.Tensor]:
"""
Does a forward pass on the given batch and returns the output dictionary that the model
returns, after adding any specified regularization penalty to the loss (if training).
"""
output_dict = self._pytorch_model(**batch)
if for_training:
try:
assert "loss" in output_dict
regularization_penalty = self.model.get_regularization_penalty()
if regularization_penalty is not None:
output_dict["reg_loss"] = regularization_penalty
output_dict["loss"] += regularization_penalty
except AssertionError:
if for_training:
raise RuntimeError(
"The model you are trying to optimize does not contain a"
" 'loss' key in the output of model.forward(inputs)."
)
return output_dict
# @overrides
def _train_epoch(self, epoch: int) -> Dict[str, float]:
"""
Trains one epoch and returns metrics.
"""
logger.info("Epoch %d/%d", epoch, self._num_epochs - 1)
cpu_memory_usage = []
for worker, memory in common_util.peak_cpu_memory().items():
cpu_memory_usage.append((worker, memory))
logger.info(f"Worker {worker} memory usage: {common_util.format_size(memory)}")
gpu_memory_usage = []
for gpu, memory in common_util.peak_gpu_memory().items():
gpu_memory_usage.append((gpu, memory))
logger.info(f"GPU {gpu} memory usage: {common_util.format_size(memory)}")
regularization_penalty = self.model.get_regularization_penalty()
train_loss = 0.0
batch_loss = 0.0
train_reg_loss = None if regularization_penalty is None else 0.0
batch_reg_loss = None if regularization_penalty is None else 0.0
# Set the model to "train" mode.
self._pytorch_model.train()
# Get tqdm for the training batches
batch_generator = iter(self.data_loader)
batch_group_generator = common_util.lazy_groups_of(
batch_generator, self._num_gradient_accumulation_steps
)
logger.info("Training")
num_training_batches: Union[int, float]
try:
len_data_loader = len(self.data_loader)
num_training_batches = math.ceil(
len_data_loader / self._num_gradient_accumulation_steps
)
except TypeError:
num_training_batches = float("inf")
# Having multiple tqdm bars in case of distributed training will be a mess. Hence only the primary's
# progress is shown
if self._primary:
batch_group_generator_tqdm = Tqdm.tqdm(
batch_group_generator, total=num_training_batches
)
else:
batch_group_generator_tqdm = batch_group_generator
self._last_log = time.time()
batches_this_epoch = 0
if self._batch_num_total is None:
self._batch_num_total = 0
done_early = False
for batch_group in batch_group_generator_tqdm:
if done_early:
break
batches_this_epoch += 1
self._batch_num_total += 1
batch_num_total = self._batch_num_total
# Zero gradients.
# NOTE: this is actually more efficient than calling `self.optimizer.zero_grad()`
# because it avoids a read op when the gradients are first updated below.
for param_group in self.optimizer.param_groups:
for p in param_group["params"]:
p.grad = None
batch_loss = 0.0
batch_group_outputs = []
for batch in batch_group:
with amp.autocast(self._use_amp):
batch_outputs = self.batch_outputs(batch, for_training=True)
batch_group_outputs.append(batch_outputs)
loss = batch_outputs["loss"]
reg_loss = batch_outputs.get("reg_loss")
if torch.isnan(loss):
raise ValueError("nan loss encountered")
loss = loss / len(batch_group)
batch_loss += loss.item()
if reg_loss is not None:
reg_loss = reg_loss / len(batch_group)
batch_reg_loss = reg_loss.item()
train_reg_loss += batch_reg_loss # type: ignore
if self._scaler is not None:
self._scaler.scale(loss).backward()
else:
loss.backward()
if len(batch_group_outputs) <= 0:
continue
train_loss += batch_loss
batch_grad_norm = self.rescale_gradients()
# This does nothing if batch_num_total is None or you are using a
# scheduler which doesn't update per batch.
if self._learning_rate_scheduler:
self._learning_rate_scheduler.step_batch(batch_num_total)
if self._momentum_scheduler:
self._momentum_scheduler.step_batch(batch_num_total)
if self._scaler is not None:
self._scaler.step(self.optimizer)
self._scaler.update()
else:
self.optimizer.step()
# Update moving averages
if self._moving_average is not None:
self._moving_average.apply(batch_num_total)
# Update the description with the latest metrics
metrics = training_util.get_metrics(
self.model,
train_loss,
train_reg_loss,
batch_loss,
batch_reg_loss,
batches_this_epoch,
world_size=self._world_size,
cuda_device=self.cuda_device,
)
if batch_num_total % self.val_loss_steps == 0:
logger.info("%s: %.4f" % ('train_loss', train_loss / batches_this_epoch))
if self._validation_data_loader is not None:
with torch.no_grad():
# We have a validation set, so compute all the metrics on it.
val_loss, val_reg_loss, num_batches = self._validation_loss_n_step(batch_num_total)
val_metrics = training_util.get_metrics(
self.model,
val_loss,
val_reg_loss,
num_batches=num_batches,
batch_loss=None,
batch_reg_loss=None,
reset=True,
world_size=self._world_size,
cuda_device=self.cuda_device,
)
# description = training_util.description_from_metrics(val_metrics)
logger.info("%s: %.4f" % ('val_loss', val_loss / num_batches))
# batch_group_generator_tqdm.set_description(description, refresh=False)
self._pytorch_model.train()
if self._primary:
# Updating tqdm only for the primary as the trainers wouldn't have one
description = training_util.description_from_metrics(metrics)
batch_group_generator_tqdm.set_description(description, refresh=False)
if self._checkpointer is not None:
self._checkpointer.maybe_save_checkpoint(self, epoch, batches_this_epoch)
for callback in self._callbacks:
callback.on_batch(
self,
batch_group,
batch_group_outputs,
metrics,
epoch,
batches_this_epoch,
is_training=True,
is_primary=self._primary,
batch_grad_norm=batch_grad_norm,
)
metrics = training_util.get_metrics(
self.model,
train_loss,
train_reg_loss,
batch_loss=None,
batch_reg_loss=None,
num_batches=batches_this_epoch,
reset=True,
world_size=self._world_size,
cuda_device=self.cuda_device,
)
for (worker, memory) in cpu_memory_usage:
metrics["worker_" + str(worker) + "_memory_MB"] = memory / (1024 * 1024)
for (gpu_num, memory) in gpu_memory_usage:
metrics["gpu_" + str(gpu_num) + "_memory_MB"] = memory / (1024 * 1024)
return metrics
def _validation_loss(self, epoch: int) -> Tuple[float, Optional[float], int]:
"""
Computes the validation loss. Returns it and the number of batches.
"""
logger.info("Validating")
self._pytorch_model.eval()
# Replace parameter values with the shadow values from the moving averages.
if self._moving_average is not None:
self._moving_average.assign_average_value()
if self._validation_data_loader is not None:
validation_data_loader = self._validation_data_loader
else:
raise ConfigurationError(
"Validation results cannot be calculated without a validation_data_loader"
)
regularization_penalty = self.model.get_regularization_penalty()
# Having multiple tqdm bars in case of distributed training will be a mess. Hence only the primary's
# progress is shown
if self._primary:
val_generator_tqdm = Tqdm.tqdm(validation_data_loader)
else:
val_generator_tqdm = validation_data_loader
batches_this_epoch = 0
val_loss = 0.0
val_batch_loss = 0.0
val_reg_loss = None if regularization_penalty is None else 0.0
val_batch_reg_loss = None if regularization_penalty is None else 0.0
done_early = False
for batch in val_generator_tqdm:
if self._distributed:
# Check whether the other workers have stopped already (due to differing amounts of
# data in each). If so, we can't proceed because we would hang when we hit the
# barrier implicit in Model.forward. We use a IntTensor instead a BoolTensor
# here because NCCL process groups apparently don't support BoolTensor.
done = torch.tensor(0, device=self.cuda_device)
torch.distributed.all_reduce(done, torch.distributed.ReduceOp.SUM)
if done.item() > 0:
done_early = True
logger.warning(
f"Worker {torch.distributed.get_rank()} finishing validation early! "
"This implies that there is an imbalance in your validation "
"data across the workers and that some amount of it will be "
"ignored. A small amount of this is fine, but a major imbalance "
"should be avoided. Note: This warning will appear unless your "
"data is perfectly balanced."
)
break
with amp.autocast(self._use_amp):
batch_outputs = self.batch_outputs(batch, for_training=False)
loss = batch_outputs.get("loss")
reg_loss = batch_outputs.get("reg_loss")
if loss is not None:
# You shouldn't necessarily have to compute a loss for validation, so we allow for
# `loss` to be None. We need to be careful, though - `batches_this_epoch` is
# currently only used as the divisor for the loss function, so we can safely only
# count those batches for which we actually have a loss. If this variable ever
# gets used for something else, we might need to change things around a bit.
batches_this_epoch += 1
val_batch_loss = loss.item()
val_loss += val_batch_loss
if reg_loss is not None:
val_batch_reg_loss = reg_loss.item()
val_reg_loss += val_batch_reg_loss # type: ignore
# Update the description with the latest metrics
val_metrics = training_util.get_metrics(
self.model,
val_loss,
val_reg_loss,
val_batch_loss,
val_batch_reg_loss,
batches_this_epoch,
world_size=self._world_size,
cuda_device=self.cuda_device,
)
description = training_util.description_from_metrics(val_metrics)
if self._primary:
val_generator_tqdm.set_description(description, refresh=False)
for callback in self._callbacks:
callback.on_batch(
self,
[batch],
[batch_outputs],
val_metrics,
epoch,
batches_this_epoch,
is_training=False,
is_primary=self._primary,
)
if self._distributed and not done_early:
logger.warning(
f"Worker {torch.distributed.get_rank()} completed its entire epoch (validation)."
)
# Indicate that we're done so that any workers that have remaining data stop validation early.
done = torch.tensor(1, device=self.cuda_device)
torch.distributed.all_reduce(done, torch.distributed.ReduceOp.SUM)
assert done.item()
# Now restore the original parameter values.
if self._moving_average is not None:
self._moving_average.restore()
return val_loss, val_reg_loss, batches_this_epoch
def _validation_loss_n_step(self, step: int) -> Tuple[float, Optional[float], int]:
"""
Computes the validation loss. Returns it and the number of batches.
"""
logger.info("Validating on %d steps" % step)
self._pytorch_model.eval()
# Replace parameter values with the shadow values from the moving averages.
if self._moving_average is not None:
self._moving_average.assign_average_value()
if self._validation_data_loader is not None:
validation_data_loader = self._validation_data_loader
else:
raise ConfigurationError(
"Validation results cannot be calculated without a validation_data_loader"
)
regularization_penalty = self.model.get_regularization_penalty()
# Having multiple tqdm bars in case of distributed training will be a mess. Hence only the primary's
# progress is shown
val_batch_generator = iter(validation_data_loader)
val_batch_group_generator = common_util.lazy_groups_of(
val_batch_generator, 5
)
num_training_batches: Union[int, float]
try:
len_data_loader = len(validation_data_loader)
num_training_batches = math.ceil(
len_data_loader / 5
)
except TypeError:
num_training_batches = float("inf")
# Having multiple tqdm bars in case of distributed training will be a mess. Hence only the primary's
# progress is shown
if self._primary:
val_generator_tqdm = Tqdm.tqdm(
val_batch_group_generator, total=num_training_batches
)
else:
val_generator_tqdm = val_batch_group_generator
batches_this_epoch = 0
val_loss = 0.0
val_reg_loss = None if regularization_penalty is None else 0.0
for val_batch_group in val_generator_tqdm:
for val_batch in val_batch_group:
with amp.autocast(self._use_amp):
batches_this_epoch += 1
batch_outputs = self.batch_outputs(val_batch, for_training=False)
loss = batch_outputs.get("loss")
reg_loss = batch_outputs.get("reg_loss")
if loss is not None:
# You shouldn't necessarily have to compute a loss for validation, so we allow for
# `loss` to be None. We need to be careful, though - `batches_this_epoch` is
# currently only used as the divisor for the loss function, so we can safely only
# count those batches for which we actually have a loss. If this variable ever
# gets used for something else, we might need to change things around a bit.
val_batch_loss = loss.item()
val_loss += val_batch_loss
if reg_loss is not None:
val_batch_reg_loss = reg_loss.item()
val_reg_loss += val_batch_reg_loss # type: ignore
return val_loss, val_reg_loss, batches_this_epoch
def train(self) -> Dict[str, Any]:
"""
Trains the supplied model with the supplied parameters.
"""
for callback in self._callbacks:
callback.on_start(self, is_primary=self._primary)
# Set default values in case of failure
epoch = None
metrics = None
try:
metrics, epoch = self._try_train()
return metrics
finally:
for callback in self._callbacks:
callback.on_end(self, metrics=metrics, epoch=epoch, is_primary=self._primary)
# @overrides
def _try_train(self) -> Tuple[Dict[str, Any], int]:
try:
epoch_counter = self._restore_checkpoint()
except RuntimeError:
traceback.print_exc()
raise ConfigurationError(
"Could not recover training from the checkpoint. Did you mean to output to "
"a different serialization directory or delete the existing serialization "
"directory?"
)
training_util.enable_gradient_clipping(self.model, self._grad_clipping)
logger.info("Beginning training.")
val_metrics: Dict[str, float] = {}
metrics: Dict[str, Any] = {}
epochs_trained = 0
training_start_time = time.time()
metrics["best_epoch"] = self._metric_tracker.best_epoch
for key, value in self._metric_tracker.best_epoch_metrics.items():
metrics["best_validation_" + key] = value
for epoch in range(epoch_counter, self._num_epochs):
epoch_start_time = time.time()
train_metrics = self._train_epoch(epoch)
# Back up the model now, in case something goes wrong later with the evaluation
if self._primary and self._checkpointer is not None:
self._checkpointer.shelve_model(epoch, self)
# Wait for the primary process to finish saving the model checkpoint
if self._distributed:
dist.barrier()
# get peak of memory usage
for key, value in train_metrics.items():
if key.startswith("gpu_") and key.endswith("_memory_MB"):
metrics["peak_" + key] = max(metrics.get("peak_" + key, 0), value)
elif key.startswith("worker_") and key.endswith("_memory_MB"):
metrics["peak_" + key] = max(metrics.get("peak_" + key, 0), value)
this_epoch_val_metric: float = 0.0
if self._validation_data_loader is not None:
with torch.no_grad():
# We have a validation set, so compute all the metrics on it.
val_loss, val_reg_loss, num_batches = self._validation_loss(epoch)
# It is safe again to wait till the validation is done. This is
# important to get the metrics right.
if self._distributed:
dist.barrier()
val_metrics = training_util.get_metrics(
self.model,
val_loss,
val_reg_loss,
batch_loss=None,
batch_reg_loss=None,
num_batches=num_batches,
reset=True,
world_size=self._world_size,
cuda_device=self.cuda_device,
)
# Check validation metric for early stopping
this_epoch_val_metric = self._metric_tracker.combined_score(val_metrics)
self._metric_tracker.add_metrics(val_metrics)
# Create overall metrics dict
training_elapsed_time = time.time() - training_start_time
metrics["training_duration"] = str(datetime.timedelta(seconds=training_elapsed_time))
metrics["training_start_epoch"] = epoch_counter
metrics["training_epochs"] = epochs_trained
metrics["epoch"] = epoch
for key, value in train_metrics.items():
metrics["training_" + key] = value
for key, value in val_metrics.items():
metrics["validation_" + key] = value
if self._metric_tracker.is_best_so_far():
# Update all the best_ metrics.
# (Otherwise they just stay the same as they were.)
metrics["best_epoch"] = epoch
for key, value in val_metrics.items():
metrics["best_validation_" + key] = value
self._metric_tracker.best_epoch_metrics = val_metrics
if self._serialization_dir and self._primary:
common_util.dump_metrics(
os.path.join(self._serialization_dir, f"metrics_epoch_{epoch}.json"),
metrics,
)
# The Scheduler API is agnostic to whether your schedule requires a validation metric -
# if it doesn't, the validation metric passed here is ignored.
if self._learning_rate_scheduler:
self._learning_rate_scheduler.step(this_epoch_val_metric)
if self._momentum_scheduler:
self._momentum_scheduler.step(this_epoch_val_metric)
# The checkpointer saves state from the learning rate scheduler and the momentum
# scheduler, so we have to make sure those are updated before we save the checkpoint here.
if self._primary and self._checkpointer is not None:
self._checkpointer.save_checkpoint(
epoch, self, is_best_so_far=self._metric_tracker.is_best_so_far()
)
# Wait for the primary process to finish saving the checkpoint
if self._distributed:
dist.barrier()
for callback in self._callbacks:
callback.on_epoch(self, metrics=metrics, epoch=epoch, is_primary=self._primary)
epoch_elapsed_time = time.time() - epoch_start_time
logger.info("Epoch duration: %s", datetime.timedelta(seconds=epoch_elapsed_time))
if epoch < self._num_epochs - 1:
training_elapsed_time = time.time() - training_start_time
estimated_time_remaining = training_elapsed_time * (
(self._num_epochs - epoch_counter) / float(epoch - epoch_counter + 1) - 1
)
formatted_time = str(datetime.timedelta(seconds=int(estimated_time_remaining)))
logger.info("Estimated training time remaining: %s", formatted_time)
epochs_trained += 1
if self._metric_tracker.should_stop_early():
logger.info("Ran out of patience. Stopping training.")
break
else:
epoch = self._num_epochs - 1
# Load the best model state before returning
best_model_state = (
None if self._checkpointer is None else self._checkpointer.best_model_state()
)
if best_model_state:
self.model.load_state_dict(best_model_state)
return metrics, epoch
@contextmanager
def get_checkpoint_state(self) -> Iterator[Tuple[Dict[str, Any], Dict[str, Any]]]:
if self._moving_average is not None:
# Assigning average value to model parameters. The checkpointer will call
# `restore_state_after_checkpointing` when it is done to put this back to what it was.
self._moving_average.assign_average_value()
model_state = self.model.state_dict()
# These are the training states we need to persist.
training_states = {
"metric_tracker": self._metric_tracker.state_dict(),
"optimizer": self.optimizer.state_dict(),
"batch_num_total": self._batch_num_total,
}
# If we have a learning rate or momentum scheduler, we should persist them too.
if self._learning_rate_scheduler is not None:
training_states["learning_rate_scheduler"] = self._learning_rate_scheduler.state_dict()
if self._momentum_scheduler is not None:
training_states["momentum_scheduler"] = self._momentum_scheduler.state_dict()
try:
yield model_state, training_states
finally:
if self._moving_average is not None:
self._moving_average.restore()
def _restore_checkpoint(self) -> int:
"""
Restores the model and training state from the last saved checkpoint.
This includes an epoch count and optimizer state, which is serialized separately
from model parameters. This function should only be used to continue training -
if you wish to load a model for inference/load parts of a model into a new
computation graph, you should use the native Pytorch functions:
` model.load_state_dict(torch.load("/path/to/model/weights.th"))`
If `self._serialization_dir` does not exist or does not contain any checkpointed weights,
this function will do nothing and return 0.
# Returns
epoch: `int`
The epoch at which to resume training, which should be one after the epoch
in the saved training state.
"""
if self._checkpointer is None:
return 0
model_state, training_state = self._checkpointer.restore_checkpoint()
if not training_state:
# No checkpoint to restore, start at 0
return 0
self.model.load_state_dict(model_state)
self.optimizer.load_state_dict(training_state["optimizer"])
if (
self._learning_rate_scheduler is not None
and "learning_rate_scheduler" in training_state
):
self._learning_rate_scheduler.load_state_dict(training_state["learning_rate_scheduler"])
if self._momentum_scheduler is not None and "momentum_scheduler" in training_state:
self._momentum_scheduler.load_state_dict(training_state["momentum_scheduler"])
training_util.move_optimizer_to_cuda(self.optimizer)
# Currently the `training_state` contains a serialized `MetricTracker`.
if "metric_tracker" in training_state:
self._metric_tracker.load_state_dict(training_state["metric_tracker"])
else:
self._metric_tracker.clear()
if isinstance(training_state["epoch"], int):
epoch_to_return = training_state["epoch"] + 1
else:
epoch_to_return = int(training_state["epoch"].split(".")[0]) + 1
# For older checkpoints with batch_num_total missing, default to old behavior where
# it is unchanged.
batch_num_total = training_state.get("batch_num_total")
if batch_num_total is not None:
self._batch_num_total = batch_num_total
return epoch_to_return
def rescale_gradients(self) -> float:
"""
Performs gradient rescaling. Is a no-op if gradient rescaling is not enabled.
Returns the norm of the gradients.
"""
parameters_to_clip = [p for p in self.model.parameters() if p.grad is not None]
if self._grad_norm:
if self._scaler is not None:
# Need to first unscale gradients in order to clip as usual.
self._scaler.unscale_(self.optimizer)
return clip_grad_norm_(parameters_to_clip, self._grad_norm)
else:
return torch.norm(
torch.stack([torch.norm(p.grad.detach()) for p in parameters_to_clip])
)
@classmethod
def from_partial_objects(
cls,
model: Model,
serialization_dir: str,
data_loader: DataLoader,
validation_data_loader: DataLoader = None,
local_rank: int = 0,
patience: int = None,
validation_metric: Union[str, List[str]] = "-loss",
num_epochs: int = 20,
cuda_device: Optional[Union[int, torch.device]] = None,
grad_norm: float = None,
grad_clipping: float = None,
distributed: bool = False,
world_size: int = 1,
num_gradient_accumulation_steps: int = 1,
use_amp: bool = False,
no_grad: List[str] = None,
optimizer: Lazy[Optimizer] = Lazy(Optimizer.default),
learning_rate_scheduler: Lazy[LearningRateScheduler] = None,
momentum_scheduler: Lazy[MomentumScheduler] = None,
moving_average: Lazy[MovingAverage] = None,
checkpointer: Lazy[Checkpointer] = Lazy(Checkpointer),
callbacks: List[Lazy[TrainerCallback]] = None,
enable_default_callbacks: bool = True,
run_sanity_checks: bool = True,
) -> "Trainer":
"""
This method exists so that we can have a documented method to construct this class using
`FromParams`. If you are not using `FromParams` or config files, you can safely ignore this
method.
The reason we can't just use `__init__` with `FromParams` here is because there are
sequential dependencies to this class's arguments. Anything that has a `Lazy[]` type
annotation needs something from one of the non-`Lazy` arguments. The `Optimizer` needs to
have the parameters from the `Model` before it's constructed, and the `Schedulers` need to
have the `Optimizer`. Because of this, the typical way we construct things `FromParams`
doesn't work, so we use `Lazy` to allow for constructing the objects sequentially.
If you're not using `FromParams`, you can just construct these arguments in the right order
yourself in your code and call the constructor directly.
"""
if cuda_device is None:
from torch import cuda
if cuda.device_count() > 0:
cuda_device = 0
else:
cuda_device = -1
check_for_gpu(cuda_device)
if cuda_device >= 0:
# Moving model to GPU here so that the optimizer state gets constructed on
# the right device.
model = model.cuda(cuda_device)
if no_grad:
for name, parameter in model.named_parameters():
if any(re.search(regex, name) for regex in no_grad):
parameter.requires_grad_(False)
parameters = [[n, p] for n, p in model.named_parameters() if p.requires_grad]
optimizer_ = optimizer.construct(model_parameters=parameters)
common_util.log_frozen_and_tunable_parameter_names(model)
batches_per_epoch: Optional[int]
try:
batches_per_epoch = len(data_loader)
batches_per_epoch = math.ceil(batches_per_epoch / num_gradient_accumulation_steps)
except TypeError:
batches_per_epoch = None
moving_average_ = (
None if moving_average is None else moving_average.construct(parameters=parameters)
)
learning_rate_scheduler_ = (
None
if learning_rate_scheduler is None
else learning_rate_scheduler.construct(
optimizer=optimizer_, num_epochs=num_epochs, num_steps_per_epoch=batches_per_epoch
)
)
momentum_scheduler_ = (
None
if momentum_scheduler is None
else momentum_scheduler.construct(optimizer=optimizer_)
)
checkpointer_ = checkpointer.construct(serialization_dir=serialization_dir)
callbacks_: List[TrainerCallback] = []
for callback_ in callbacks or []:
callbacks_.append(callback_.construct(serialization_dir=serialization_dir))
return cls(
model,
optimizer_,
data_loader,
patience=patience,
validation_metric=validation_metric,
validation_data_loader=validation_data_loader,
num_epochs=num_epochs,
serialization_dir=serialization_dir,
cuda_device=cuda_device,
grad_norm=grad_norm,
grad_clipping=grad_clipping,
learning_rate_scheduler=learning_rate_scheduler_,
momentum_scheduler=momentum_scheduler_,
checkpointer=checkpointer_,
moving_average=moving_average_,
callbacks=callbacks_,
distributed=distributed,
local_rank=local_rank,
world_size=world_size,
num_gradient_accumulation_steps=num_gradient_accumulation_steps,
use_amp=use_amp,
enable_default_callbacks=enable_default_callbacks,
run_sanity_checks=run_sanity_checks,
)
class MyTrainer(TrainerBase):
def __init__(
self,
model: Model,
optimizer: torch.optim.Optimizer,
iterator: DataIterator,
train_dataset: Iterable[Instance],
validation_dataset: Optional[Iterable[Instance]] = None,
patience: Optional[int] = None,
validation_metric: str = "-loss",
validation_iterator: DataIterator = None,
shuffle: bool = True,
num_epochs: int = 20,
serialization_dir: Optional[str] = None,
num_serialized_models_to_keep: int = 20,
keep_serialized_model_every_num_seconds: int = None,
checkpointer: Checkpointer = None,
model_save_interval: float = None,
cuda_device: Union[int, List] = -1,
grad_norm: Optional[float] = None,
grad_clipping: Optional[float] = None,
learning_rate_scheduler: Optional[LearningRateScheduler] = None,
momentum_scheduler: Optional[MomentumScheduler] = None,
summary_interval: int = 100,
histogram_interval: int = None,
should_log_parameter_statistics: bool = True,
should_log_learning_rate: bool = False,
log_batch_size_period: Optional[int] = None,
moving_average: Optional[MovingAverage] = None,
) -> None:
super().__init__(serialization_dir, cuda_device)
# I am not calling move_to_gpu here, because if the model is
# not already on the GPU then the optimizer is going to be wrong.
self.model = model
self.iterator = iterator
self._validation_iterator = validation_iterator
self.shuffle = shuffle
self.optimizer = optimizer
self.train_data = train_dataset
self._validation_data = validation_dataset
if patience is None: # no early stopping
if validation_dataset:
logger.warning(
"You provided a validation dataset but patience was set to None, "
"meaning that early stopping is disabled")
elif (not isinstance(patience, int)) or patience <= 0:
raise ConfigurationError(
'{} is an invalid value for "patience": it must be a positive integer '
"or None (if you want to disable early stopping)".format(
patience))
# Frequency of metric checking (in batch)
self._summary_interval = summary_interval
# For tracking is_best_so_far and should_stop_early
self._metric_tracker = MetricTracker(patience, validation_metric)
# Get rid of + or -
self._validation_metric = validation_metric[1:]
self._num_epochs = num_epochs
if checkpointer is not None:
# We can't easily check if these parameters were passed in, so check against their default values.
# We don't check against serialization_dir since it is also used by the parent class.
if (num_serialized_models_to_keep != 20
or keep_serialized_model_every_num_seconds is not None):
raise ConfigurationError(
"When passing a custom Checkpointer, you may not also pass in separate checkpointer "
"args 'num_serialized_models_to_keep' or 'keep_serialized_model_every_num_seconds'."
)
self._checkpointer = checkpointer
else:
self._checkpointer = Checkpointer(
serialization_dir,
keep_serialized_model_every_num_seconds,
num_serialized_models_to_keep,
)
self._model_save_interval = model_save_interval
self._grad_norm = grad_norm
self._grad_clipping = grad_clipping
self._learning_rate_scheduler = learning_rate_scheduler
self._momentum_scheduler = momentum_scheduler
self._moving_average = moving_average
# We keep the total batch number as an instance variable because it
# is used inside a closure for the hook which logs activations in
# ``_enable_activation_logging``.
self._batch_num_total = 0
self._tensorboard = TensorboardWriter(
get_batch_num_total=lambda: self._batch_num_total,
serialization_dir=serialization_dir,
summary_interval=summary_interval,
histogram_interval=histogram_interval,
should_log_parameter_statistics=should_log_parameter_statistics,
should_log_learning_rate=should_log_learning_rate,
)
self._log_batch_size_period = log_batch_size_period
self._last_log = 0.0 # time of last logging
# Enable activation logging.
if histogram_interval is not None:
self._tensorboard.enable_activation_logging(self.model)
def rescale_gradients(self) -> Optional[float]:
return training_util.rescale_gradients(self.model, self._grad_norm)
def batch_loss(self, batch_group: List[TensorDict],
for_training: bool) -> torch.Tensor:
"""
Does a forward pass on the given batches and returns the ``loss`` value in the result.
If ``for_training`` is `True` also applies regularization penalty.
"""
if self._multiple_gpu:
output_dict = training_util.data_parallel(batch_group, self.model,
self._cuda_devices)
else:
assert len(batch_group) == 1
batch = batch_group[0]
batch = nn_util.move_to_device(batch, self._cuda_devices[0])
output_dict = self.model(**batch)
try:
loss = output_dict["loss"]
if for_training:
loss += self.model.get_regularization_penalty()
except KeyError:
if for_training:
raise RuntimeError(
"The model you are trying to optimize does not contain a"
" 'loss' key in the output of model.forward(inputs).")
loss = None
return loss
def _train_epoch(self, epoch: int) -> Dict[str, float]:
"""
Trains one epoch and returns metrics.
"""
logger.info("Epoch %d/%d", epoch, self._num_epochs - 1)
peak_cpu_usage = peak_memory_mb()
logger.info(f"Peak CPU memory usage MB: {peak_cpu_usage}")
gpu_usage = []
for gpu, memory in gpu_memory_mb().items():
gpu_usage.append((gpu, memory))
logger.info(f"GPU {gpu} memory usage MB: {memory}")
train_loss = 0.0
# Set the model to "train" mode.
self.model.train()
num_gpus = len(self._cuda_devices)
# Get tqdm for the training batches
raw_train_generator = self.iterator(self.train_data,
num_epochs=1,
shuffle=self.shuffle)
train_generator = lazy_groups_of(raw_train_generator, num_gpus)
num_training_batches = math.ceil(
self.iterator.get_num_batches(self.train_data) / num_gpus)
self._last_log = time.time()
last_save_time = time.time()
batches_this_epoch = 0
if self._batch_num_total is None:
self._batch_num_total = 0
histogram_parameters = set(
self.model.get_parameters_for_histogram_tensorboard_logging())
logger.info("Training")
train_generator_tqdm = Tqdm.tqdm(train_generator,
total=num_training_batches)
cumulative_batch_size = 0
for batch_group in train_generator_tqdm:
batches_this_epoch += 1
self._batch_num_total += 1
batch_num_total = self._batch_num_total
self.optimizer.zero_grad()
loss = self.batch_loss(batch_group, for_training=True)
if torch.isnan(loss):
raise ValueError("nan loss encountered")
loss.backward()
train_loss += loss.item()
batch_grad_norm = self.rescale_gradients()
# This does nothing if batch_num_total is None or you are using a
# scheduler which doesn't update per batch.
if self._learning_rate_scheduler:
self._learning_rate_scheduler.step_batch(batch_num_total)
if self._momentum_scheduler:
self._momentum_scheduler.step_batch(batch_num_total)
if self._tensorboard.should_log_histograms_this_batch():
# get the magnitude of parameter updates for logging
# We need a copy of current parameters to compute magnitude of updates,
# and copy them to CPU so large models won't go OOM on the GPU.
param_updates = {
name: param.detach().cpu().clone()
for name, param in self.model.named_parameters()
}
self.optimizer.step()
for name, param in self.model.named_parameters():
param_updates[name].sub_(param.detach().cpu())
update_norm = torch.norm(param_updates[name].view(-1))
param_norm = torch.norm(param.view(-1)).cpu()
self._tensorboard.add_train_scalar(
"gradient_update/" + name,
update_norm / (param_norm + 1e-7))
else:
self.optimizer.step()
# Update moving averages
if self._moving_average is not None:
self._moving_average.apply(batch_num_total)
# evaluate model performance if reaches a checkpoint.
if batches_this_epoch % self._summary_interval == 0:
# Update the description with the latest metrics
metrics = training_util.get_metrics(self.model, train_loss,
batches_this_epoch)
description = training_util.description_from_metrics(metrics)
train_generator_tqdm.set_description(description,
refresh=False)
# Log parameter values to Tensorboard
if self._tensorboard.should_log_this_batch():
self._tensorboard.log_parameter_and_gradient_statistics(
self.model, batch_grad_norm)
self._tensorboard.log_learning_rates(self.model,
self.optimizer)
self._tensorboard.add_train_scalar("loss/loss_train",
metrics["loss"])
self._tensorboard.log_metrics(
{"epoch_metrics/" + k: v
for k, v in metrics.items()})
if self._tensorboard.should_log_histograms_this_batch():
self._tensorboard.log_histograms(self.model,
histogram_parameters)
if self._log_batch_size_period:
cur_batch = sum([
training_util.get_batch_size(batch)
for batch in batch_group
])
cumulative_batch_size += cur_batch
if (batches_this_epoch - 1) % self._log_batch_size_period == 0:
average = cumulative_batch_size / batches_this_epoch
logger.info(
f"current batch size: {cur_batch} mean batch size: {average}"
)
self._tensorboard.add_train_scalar("current_batch_size",
cur_batch)
self._tensorboard.add_train_scalar("mean_batch_size",
average)
# Save model if needed.
if self._model_save_interval is not None and (
time.time() - last_save_time > self._model_save_interval):
last_save_time = time.time()
self._save_checkpoint("{0}.{1}".format(
epoch, training_util.time_to_str(int(last_save_time))))
metrics = training_util.get_metrics(self.model,
train_loss,
batches_this_epoch,
reset=True)
metrics["cpu_memory_MB"] = peak_cpu_usage
for (gpu_num, memory) in gpu_usage:
metrics["gpu_" + str(gpu_num) + "_memory_MB"] = memory
return metrics
def _save_checkpoint(self, epoch: Union[int, str]) -> None:
"""
Saves a checkpoint of the model to self._serialization_dir.
Is a no-op if self._serialization_dir is None.
Parameters
----------
epoch : Union[int, str], required.
The epoch of training. If the checkpoint is saved in the middle
of an epoch, the parameter is a string with the epoch and timestamp.
"""
# If moving averages are used for parameters, we save
# the moving average values into checkpoint, instead of the current values.
if self._moving_average is not None:
self._moving_average.assign_average_value()
# These are the training states we need to persist.
training_states = {
"metric_tracker": self._metric_tracker.state_dict(),
"optimizer": self.optimizer.state_dict(),
"batch_num_total": self._batch_num_total,
}
# If we have a learning rate or momentum scheduler, we should persist them too.
if self._learning_rate_scheduler is not None:
training_states[
"learning_rate_scheduler"] = self._learning_rate_scheduler.state_dict(
)
if self._momentum_scheduler is not None:
training_states[
"momentum_scheduler"] = self._momentum_scheduler.state_dict()
self._checkpointer.save_checkpoint(
model_state=self.model.state_dict(),
epoch=epoch,
training_states=training_states,
is_best_so_far=self._metric_tracker.is_best_so_far(),
)
# Restore the original values for parameters so that training will not be affected.
if self._moving_average is not None:
self._moving_average.restore()
def _restore_checkpoint(self) -> int:
"""
Restores the model and training state from the last saved checkpoint.
This includes an epoch count and optimizer state, which is serialized separately
from model parameters. This function should only be used to continue training -
if you wish to load a model for inference/load parts of a model into a new
computation graph, you should use the native Pytorch functions:
`` model.load_state_dict(torch.load("/path/to/model/weights.th"))``
If ``self._serialization_dir`` does not exist or does not contain any checkpointed weights,
this function will do nothing and return 0.
Returns
-------
epoch: int
The epoch at which to resume training, which should be one after the epoch
in the saved training state.
"""
model_state, training_state = self._checkpointer.restore_checkpoint()
if not training_state:
# No checkpoint to restore, start at 0
return 0
self.model.load_state_dict(model_state)
self.optimizer.load_state_dict(training_state["optimizer"])
if (self._learning_rate_scheduler is not None
and "learning_rate_scheduler" in training_state):
self._learning_rate_scheduler.load_state_dict(
training_state["learning_rate_scheduler"])
if self._momentum_scheduler is not None and "momentum_scheduler" in training_state:
self._momentum_scheduler.load_state_dict(
training_state["momentum_scheduler"])
training_util.move_optimizer_to_cuda(self.optimizer)
# Currently the ``training_state`` contains a serialized ``MetricTracker``.
if "metric_tracker" in training_state:
self._metric_tracker.load_state_dict(
training_state["metric_tracker"])
# It used to be the case that we tracked ``val_metric_per_epoch``.
elif "val_metric_per_epoch" in training_state:
self._metric_tracker.clear()
self._metric_tracker.add_metrics(
training_state["val_metric_per_epoch"])
# And before that we didn't track anything.
else:
self._metric_tracker.clear()
if isinstance(training_state["epoch"], int):
epoch_to_return = training_state["epoch"] + 1
else:
epoch_to_return = int(training_state["epoch"].split(".")[0]) + 1
# For older checkpoints with batch_num_total missing, default to old behavior where
# it is unchanged.
batch_num_total = training_state.get("batch_num_total")
if batch_num_total is not None:
self._batch_num_total = batch_num_total
return epoch_to_return
# Requires custom from_params.
@classmethod
def from_params( # type: ignore
cls,
model: Model,
serialization_dir: str,
iterator: DataIterator,
train_data: Iterable[Instance],
validation_data: Optional[Iterable[Instance]],
params: Params,
validation_iterator: DataIterator = None,
) -> "Trainer":
patience = params.pop_int("patience", None)
validation_metric = params.pop("validation_metric", "-loss")
shuffle = params.pop_bool("shuffle", True)
num_epochs = params.pop_int("num_epochs", 20)
cuda_device = parse_cuda_device(params.pop("cuda_device", -1))
grad_norm = params.pop_float("grad_norm", None)
grad_clipping = params.pop_float("grad_clipping", None)
lr_scheduler_params = params.pop("learning_rate_scheduler", None)
momentum_scheduler_params = params.pop("momentum_scheduler", None)
if isinstance(cuda_device, list):
model_device = cuda_device[0]
else:
model_device = cuda_device
if model_device >= 0:
# Moving model to GPU here so that the optimizer state gets constructed on
# the right device.
model = model.cuda(model_device)
parameters = [[n, p] for n, p in model.named_parameters()
if p.requires_grad]
optimizer = Optimizer.from_params(parameters, params.pop("optimizer"))
if "moving_average" in params:
moving_average = MovingAverage.from_params(
params.pop("moving_average"), parameters=parameters)
else:
moving_average = None
if lr_scheduler_params:
lr_scheduler = LearningRateScheduler.from_params(
optimizer, lr_scheduler_params)
else:
lr_scheduler = None
if momentum_scheduler_params:
momentum_scheduler = MomentumScheduler.from_params(
optimizer, momentum_scheduler_params)
else:
momentum_scheduler = None
if "checkpointer" in params:
if ("keep_serialized_model_every_num_seconds" in params
or "num_serialized_models_to_keep" in params):
raise ConfigurationError(
"Checkpointer may be initialized either from the 'checkpointer' key or from the "
"keys 'num_serialized_models_to_keep' and 'keep_serialized_model_every_num_seconds'"
" but the passed config uses both methods.")
checkpointer = Checkpointer.from_params(params.pop("checkpointer"))
else:
num_serialized_models_to_keep = params.pop_int(
"num_serialized_models_to_keep", 20)
keep_serialized_model_every_num_seconds = params.pop_int(
"keep_serialized_model_every_num_seconds", None)
checkpointer = Checkpointer(
serialization_dir=serialization_dir,
num_serialized_models_to_keep=num_serialized_models_to_keep,
keep_serialized_model_every_num_seconds=
keep_serialized_model_every_num_seconds,
)
model_save_interval = params.pop_float("model_save_interval", None)
summary_interval = params.pop_int("summary_interval", 100)
histogram_interval = params.pop_int("histogram_interval", None)
should_log_parameter_statistics = params.pop_bool(
"should_log_parameter_statistics", True)
should_log_learning_rate = params.pop_bool("should_log_learning_rate",
False)
log_batch_size_period = params.pop_int("log_batch_size_period", None)
params.assert_empty(cls.__name__)
return cls(
model,
optimizer,
iterator,
train_data,
validation_data,
patience=patience,
validation_metric=validation_metric,
validation_iterator=validation_iterator,
shuffle=shuffle,
num_epochs=num_epochs,
serialization_dir=serialization_dir,
cuda_device=cuda_device,
grad_norm=grad_norm,
grad_clipping=grad_clipping,
learning_rate_scheduler=lr_scheduler,
momentum_scheduler=momentum_scheduler,
checkpointer=checkpointer,
model_save_interval=model_save_interval,
summary_interval=summary_interval,
histogram_interval=histogram_interval,
should_log_parameter_statistics=should_log_parameter_statistics,
should_log_learning_rate=should_log_learning_rate,
log_batch_size_period=log_batch_size_period,
moving_average=moving_average,
)
class Trainer(TrainerBase):
def __init__(
self,
model: Model,
optimizer: torch.optim.Optimizer,
iterator: DataIterator,
train_dataset: Iterable[Instance],
validation_dataset: Optional[Iterable[Instance]] = None,
patience: Optional[int] = None,
validation_metric: str = "-loss",
validation_iterator: DataIterator = None,
shuffle: bool = True,
num_epochs: int = 20,
serialization_dir: Optional[str] = None,
num_serialized_models_to_keep: int = 20,
keep_serialized_model_every_num_seconds: int = None,
checkpointer: Checkpointer = None,
model_save_interval: float = None,
cuda_device: int = -1,
grad_norm: Optional[float] = None,
grad_clipping: Optional[float] = None,
learning_rate_scheduler: Optional[LearningRateScheduler] = None,
momentum_scheduler: Optional[MomentumScheduler] = None,
summary_interval: int = 100,
histogram_interval: int = None,
should_log_parameter_statistics: bool = True,
should_log_learning_rate: bool = False,
log_batch_size_period: Optional[int] = None,
moving_average: Optional[MovingAverage] = None,
distributed: bool = False,
local_rank: int = 0,
world_size: int = 1,
num_gradient_accumulation_steps: int = 1,
) -> None:
"""
A trainer for doing supervised learning. It just takes a labeled dataset
and a `DataIterator`, and uses the supplied `Optimizer` to learn the weights
for your model over some fixed number of epochs. You can also pass in a validation
dataset and enable early stopping. There are many other bells and whistles as well.
# Parameters
model : `Model`, required.
An AllenNLP model to be optimized. Pytorch Modules can also be optimized if
their `forward` method returns a dictionary with a "loss" key, containing a
scalar tensor representing the loss function to be optimized.
If you are training your model using GPUs, your model should already be
on the correct device. (If you use `Trainer.from_params` this will be
handled for you.)
optimizer : `torch.nn.Optimizer`, required.
An instance of a Pytorch Optimizer, instantiated with the parameters of the
model to be optimized.
iterator : `DataIterator`, required.
A method for iterating over a `Dataset`, yielding padded indexed batches.
train_dataset : `Dataset`, required.
A `Dataset` to train on. The dataset should have already been indexed.
validation_dataset : `Dataset`, optional, (default = None).
A `Dataset` to evaluate on. The dataset should have already been indexed.
patience : Optional[int] > 0, optional (default=None)
Number of epochs to be patient before early stopping: the training is stopped
after `patience` epochs with no improvement. If given, it must be `> 0`.
If None, early stopping is disabled.
validation_metric : str, optional (default="loss")
Validation metric to measure for whether to stop training using patience
and whether to serialize an `is_best` model each epoch. The metric name
must be prepended with either "+" or "-", which specifies whether the metric
is an increasing or decreasing function.
validation_iterator : `DataIterator`, optional (default=None)
An iterator to use for the validation set. If `None`, then
use the training `iterator`.
shuffle : `bool`, optional (default=True)
Whether to shuffle the instances in the iterator or not.
num_epochs : int, optional (default = 20)
Number of training epochs.
serialization_dir : str, optional (default=None)
Path to directory for saving and loading model files. Models will not be saved if
this parameter is not passed.
num_serialized_models_to_keep : `int`, optional (default=20)
Number of previous model checkpoints to retain. Default is to keep 20 checkpoints.
A value of None or -1 means all checkpoints will be kept.
keep_serialized_model_every_num_seconds : `int`, optional (default=None)
If num_serialized_models_to_keep is not None, then occasionally it's useful to
save models at a given interval in addition to the last num_serialized_models_to_keep.
To do so, specify keep_serialized_model_every_num_seconds as the number of seconds
between permanently saved checkpoints. Note that this option is only used if
num_serialized_models_to_keep is not None, otherwise all checkpoints are kept.
checkpointer : `Checkpointer`, optional (default=None)
An instance of class Checkpointer to use instead of the default. If a checkpointer is specified,
the arguments num_serialized_models_to_keep and keep_serialized_model_every_num_seconds should
not be specified. The caller is responsible for initializing the checkpointer so that it is
consistent with serialization_dir.
model_save_interval : `float`, optional (default=None)
If provided, then serialize models every `model_save_interval`
seconds within single epochs. In all cases, models are also saved
at the end of every epoch if `serialization_dir` is provided.
cuda_device : `int`, optional (default = -1)
An integer specifying the CUDA device(s) to use for this process. If -1, the CPU is used.
Data parallelism is controlled at the allennlp train level, so each trainer will have a single
GPU.
grad_norm : `float`, optional, (default = None).
If provided, gradient norms will be rescaled to have a maximum of this value.
grad_clipping : `float`, optional (default = `None`).
If provided, gradients will be clipped `during the backward pass` to have an (absolute)
maximum of this value. If you are getting `NaNs` in your gradients during training
that are not solved by using `grad_norm`, you may need this.
learning_rate_scheduler : `LearningRateScheduler`, optional (default = None)
If specified, the learning rate will be decayed with respect to
this schedule at the end of each epoch (or batch, if the scheduler implements
the `step_batch` method). If you use `torch.optim.lr_scheduler.ReduceLROnPlateau`,
this will use the `validation_metric` provided to determine if learning has plateaued.
To support updating the learning rate on every batch, this can optionally implement
`step_batch(batch_num_total)` which updates the learning rate given the batch number.
momentum_scheduler : `MomentumScheduler`, optional (default = None)
If specified, the momentum will be updated at the end of each batch or epoch
according to the schedule.
summary_interval : `int`, optional, (default = 100)
Number of batches between logging scalars to tensorboard
histogram_interval : `int`, optional, (default = `None`)
If not None, then log histograms to tensorboard every `histogram_interval` batches.
When this parameter is specified, the following additional logging is enabled:
* Histograms of model parameters
* The ratio of parameter update norm to parameter norm
* Histogram of layer activations
We log histograms of the parameters returned by
`model.get_parameters_for_histogram_tensorboard_logging`.
The layer activations are logged for any modules in the `Model` that have
the attribute `should_log_activations` set to `True`. Logging
histograms requires a number of GPU-CPU copies during training and is typically
slow, so we recommend logging histograms relatively infrequently.
Note: only Modules that return tensors, tuples of tensors or dicts
with tensors as values currently support activation logging.
should_log_parameter_statistics : `bool`, optional, (default = True)
Whether to send parameter statistics (mean and standard deviation
of parameters and gradients) to tensorboard.
should_log_learning_rate : `bool`, optional, (default = False)
Whether to send parameter specific learning rate to tensorboard.
log_batch_size_period : `int`, optional, (default = `None`)
If defined, how often to log the average batch size.
moving_average : `MovingAverage`, optional, (default = None)
If provided, we will maintain moving averages for all parameters. During training, we
employ a shadow variable for each parameter, which maintains the moving average. During
evaluation, we backup the original parameters and assign the moving averages to corresponding
parameters. Be careful that when saving the checkpoint, we will save the moving averages of
parameters. This is necessary because we want the saved model to perform as well as the validated
model if we load it later. But this may cause problems if you restart the training from checkpoint.
distributed : `bool`, optional, (default = False)
If set, PyTorch's `DistributedDataParallel` is used to train the model in multiple GPUs. This also
requires `world_size` to be greater than 1.
local_rank : `int`, optional, (default = 0)
This is the unique identifier of the `Trainer` in a distributed process group. The GPU device id is
used as the rank.
world_size : `int`, (default = 1)
The number of `Trainer` workers participating in the distributed training.
num_gradient_accumulation_steps : `int`, optional, (default = 1)
Gradients are accumulated for the given number of steps before doing an optimizer step. This can
be useful to accommodate batches that are larger than the RAM size. Refer Thomas Wolf's
[post](https://tinyurl.com/y5mv44fw) for details on Gradient Accumulation.
"""
super().__init__(serialization_dir, cuda_device, distributed,
local_rank, world_size)
# I am not calling move_to_gpu here, because if the model is
# not already on the GPU then the optimizer is going to be wrong.
self.model = model
self.iterator = iterator
self._validation_iterator = validation_iterator
self.shuffle = shuffle
self.optimizer = optimizer
self.train_data = train_dataset
self._validation_data = validation_dataset
if patience is None: # no early stopping
if validation_dataset:
logger.warning(
"You provided a validation dataset but patience was set to None, "
"meaning that early stopping is disabled")
elif (not isinstance(patience, int)) or patience <= 0:
raise ConfigurationError(
'{} is an invalid value for "patience": it must be a positive integer '
"or None (if you want to disable early stopping)".format(
patience))
# For tracking is_best_so_far and should_stop_early
self._metric_tracker = MetricTracker(patience, validation_metric)
# Get rid of + or -
self._validation_metric = validation_metric[1:]
self._num_epochs = num_epochs
if checkpointer is not None:
# We can't easily check if these parameters were passed in, so check against their default values.
# We don't check against serialization_dir since it is also used by the parent class.
if (num_serialized_models_to_keep != 20
or keep_serialized_model_every_num_seconds is not None):
raise ConfigurationError(
"When passing a custom Checkpointer, you may not also pass in separate checkpointer "
"args 'num_serialized_models_to_keep' or 'keep_serialized_model_every_num_seconds'."
)
self._checkpointer = checkpointer
else:
self._checkpointer = Checkpointer(
serialization_dir,
keep_serialized_model_every_num_seconds,
num_serialized_models_to_keep,
)
self._model_save_interval = model_save_interval
self._grad_norm = grad_norm
self._grad_clipping = grad_clipping
self._learning_rate_scheduler = learning_rate_scheduler
self._momentum_scheduler = momentum_scheduler
self._moving_average = moving_average
# We keep the total batch number as an instance variable because it
# is used inside a closure for the hook which logs activations in
# `_enable_activation_logging`.
self._batch_num_total = 0
self._tensorboard = TensorboardWriter(
get_batch_num_total=lambda: self._batch_num_total,
serialization_dir=serialization_dir,
summary_interval=summary_interval,
histogram_interval=histogram_interval,
should_log_parameter_statistics=should_log_parameter_statistics,
should_log_learning_rate=should_log_learning_rate,
)
self._log_batch_size_period = log_batch_size_period
self._last_log = 0.0 # time of last logging
self._num_gradient_accumulation_steps = num_gradient_accumulation_steps
# Enable activation logging.
if histogram_interval is not None:
self._tensorboard.enable_activation_logging(self.model)
# Using `DistributedDataParallel`(ddp) brings in a quirk wrt AllenNLP's `Model` interface and its
# usage. A `Model` object is wrapped by `ddp`, but assigning the wrapped model to `self.model`
# will break the usages such as `Model.get_regularization_penalty`, `Model.get_metrics`, etc.
#
# Hence a reference to Pytorch's object is maintained in the case of distributed training and in the
# normal case, reference to `Model` is retained. This reference is only used in
# these places: `model.__call__`, `model.train` and `model.eval`.
if self._distributed:
self._pytorch_model = DistributedDataParallel(
self.model,
device_ids=[self.cuda_device],
find_unused_parameters=True)
else:
self._pytorch_model = self.model
def rescale_gradients(self) -> Optional[float]:
return training_util.rescale_gradients(self.model, self._grad_norm)
def batch_loss(self, batch: TensorDict,
for_training: bool) -> torch.Tensor:
"""
Does a forward pass on the given batches and returns the `loss` value in the result.
If `for_training` is `True` also applies regularization penalty.
"""
batch = nn_util.move_to_device(batch, self.cuda_device)
output_dict = self._pytorch_model(**batch)
try:
loss = output_dict["loss"]
if for_training:
loss += self.model.get_regularization_penalty()
except KeyError:
if for_training:
raise RuntimeError(
"The model you are trying to optimize does not contain a"
" 'loss' key in the output of model.forward(inputs).")
loss = None
return loss
def _train_epoch(self, epoch: int) -> Dict[str, float]:
"""
Trains one epoch and returns metrics.
"""
logger.info("Epoch %d/%d", epoch, self._num_epochs - 1)
peak_cpu_usage = common_util.peak_memory_mb()
logger.info(f"Peak CPU memory usage MB: {peak_cpu_usage}")
gpu_usage = []
for gpu, memory in common_util.gpu_memory_mb().items():
gpu_usage.append((gpu, memory))
logger.info(f"GPU {gpu} memory usage MB: {memory}")
train_loss = 0.0
# Set the model to "train" mode.
self._pytorch_model.train()
# Get tqdm for the training batches
batch_generator = self.iterator(self.train_data,
num_epochs=1,
shuffle=self.shuffle)
batch_group_generator = common_util.lazy_groups_of(
batch_generator, self._num_gradient_accumulation_steps)
num_training_batches = math.ceil(
self.iterator.get_num_batches(self.train_data) /
self._num_gradient_accumulation_steps)
# Having multiple tqdm bars in case of distributed training will be a mess. Hence only the master's
# progress is shown
if self._master:
batch_group_generator_tqdm = Tqdm.tqdm(batch_group_generator,
total=num_training_batches)
else:
batch_group_generator_tqdm = batch_group_generator
self._last_log = time.time()
last_save_time = time.time()
batches_this_epoch = 0
if self._batch_num_total is None:
self._batch_num_total = 0
histogram_parameters = set(
self.model.get_parameters_for_histogram_tensorboard_logging())
logger.info("Training")
cumulative_batch_group_size = 0
done_early = False
for batch_group in batch_group_generator_tqdm:
if self._distributed:
# Check whether the other workers have stopped already (due to differing amounts of
# data in each). If so, we can't proceed because we would hang when we hit the
# barrier implicit in Model.forward. We use a IntTensor instead a BoolTensor
# here because NCCL process groups apparently don't support BoolTensor.
done = torch.tensor(0, device=self.cuda_device)
torch.distributed.all_reduce(done,
torch.distributed.ReduceOp.SUM)
if done.item() > 0:
done_early = True
logger.warning(
f"Worker {torch.distributed.get_rank()} finishing training early! "
"This implies that there is an imbalance in your training "
"data across the workers and that some amount of it will be "
"ignored. A small amount of this is fine, but a major imbalance "
"should be avoided. Note: This warning will appear unless your "
"data is perfectly balanced.")
break
batches_this_epoch += 1
self._batch_num_total += 1
batch_num_total = self._batch_num_total
self.optimizer.zero_grad()
for batch in batch_group:
loss = self.batch_loss(batch, for_training=True)
if torch.isnan(loss):
raise ValueError("nan loss encountered")
loss = loss / len(batch_group)
loss.backward()
train_loss += loss.item()
batch_grad_norm = self.rescale_gradients()
# This does nothing if batch_num_total is None or you are using a
# scheduler which doesn't update per batch.
if self._learning_rate_scheduler:
self._learning_rate_scheduler.step_batch(batch_num_total)
if self._momentum_scheduler:
self._momentum_scheduler.step_batch(batch_num_total)
if self._tensorboard.should_log_histograms_this_batch(
) and self._master:
# get the magnitude of parameter updates for logging
# We need a copy of current parameters to compute magnitude of updates,
# and copy them to CPU so large models won't go OOM on the GPU.
param_updates = {
name: param.detach().cpu().clone()
for name, param in self.model.named_parameters()
}
self.optimizer.step()
for name, param in self.model.named_parameters():
param_updates[name].sub_(param.detach().cpu())
update_norm = torch.norm(param_updates[name].view(-1))
param_norm = torch.norm(param.view(-1)).cpu()
self._tensorboard.add_train_scalar(
"gradient_update/" + name,
update_norm / (param_norm + 1e-7))
else:
self.optimizer.step()
# Update moving averages
if self._moving_average is not None:
self._moving_average.apply(batch_num_total)
# Update the description with the latest metrics
metrics = training_util.get_metrics(
self.model,
train_loss,
batches_this_epoch,
world_size=self._world_size,
cuda_device=[self.cuda_device],
)
# Updating tqdm only for the master as the trainers wouldn't have one
if self._master:
description = training_util.description_from_metrics(metrics)
batch_group_generator_tqdm.set_description(description,
refresh=False)
# Log parameter values to Tensorboard (only from the master)
if self._tensorboard.should_log_this_batch() and self._master:
self._tensorboard.log_parameter_and_gradient_statistics(
self.model, batch_grad_norm)
self._tensorboard.log_learning_rates(self.model,
self.optimizer)
self._tensorboard.add_train_scalar("loss/loss_train",
metrics["loss"])
self._tensorboard.log_metrics(
{"epoch_metrics/" + k: v
for k, v in metrics.items()})
if self._tensorboard.should_log_histograms_this_batch(
) and self._master:
self._tensorboard.log_histograms(self.model,
histogram_parameters)
if self._log_batch_size_period:
batch_group_size = sum(
training_util.get_batch_size(batch)
for batch in batch_group)
cumulative_batch_group_size += batch_group_size
if (batches_this_epoch - 1) % self._log_batch_size_period == 0:
average = cumulative_batch_group_size / batches_this_epoch
logger.info(
f"current batch size: {batch_group_size} mean batch size: {average}"
)
self._tensorboard.add_train_scalar("current_batch_size",
batch_group_size)
self._tensorboard.add_train_scalar("mean_batch_size",
average)
# Save model if needed.
if (self._model_save_interval is not None and
(time.time() - last_save_time > self._model_save_interval)
and self._master):
last_save_time = time.time()
self._save_checkpoint("{0}.{1}".format(
epoch, training_util.time_to_str(int(last_save_time))))
if self._distributed and not done_early:
logger.warning(
f"Worker {torch.distributed.get_rank()} completed its entire epoch (training)."
)
# Indicate that we're done so that any workers that have remaining data stop the epoch early.
done = torch.tensor(1, device=self.cuda_device)
torch.distributed.all_reduce(done, torch.distributed.ReduceOp.SUM)
assert done.item()
# Let all workers finish their epoch before computing
# the final statistics for the epoch.
if self._distributed:
dist.barrier()
metrics = training_util.get_metrics(
self.model,
train_loss,
batches_this_epoch,
reset=True,
world_size=self._world_size,
cuda_device=[self.cuda_device],
)
metrics["cpu_memory_MB"] = peak_cpu_usage
for (gpu_num, memory) in gpu_usage:
metrics["gpu_" + str(gpu_num) + "_memory_MB"] = memory
return metrics
def _validation_loss(self) -> Tuple[float, int]:
"""
Computes the validation loss. Returns it and the number of batches.
"""
logger.info("Validating")
self._pytorch_model.eval()
# Replace parameter values with the shadow values from the moving averages.
if self._moving_average is not None:
self._moving_average.assign_average_value()
if self._validation_iterator is not None:
val_iterator = self._validation_iterator
else:
val_iterator = self.iterator
val_generator = val_iterator(self._validation_data,
num_epochs=1,
shuffle=False)
num_validation_batches = val_iterator.get_num_batches(
self._validation_data)
val_generator_tqdm = Tqdm.tqdm(val_generator,
total=num_validation_batches)
batches_this_epoch = 0
val_loss = 0
done_early = False
for batch in val_generator_tqdm:
if self._distributed:
# Check whether the other workers have stopped already (due to differing amounts of
# data in each). If so, we can't proceed because we would hang when we hit the
# barrier implicit in Model.forward. We use a IntTensor instead a BoolTensor
# here because NCCL process groups apparently don't support BoolTensor.
done = torch.tensor(0, device=self.cuda_device)
torch.distributed.all_reduce(done,
torch.distributed.ReduceOp.SUM)
if done.item() > 0:
done_early = True
logger.warning(
f"Worker {torch.distributed.get_rank()} finishing validation early! "
"This implies that there is an imbalance in your validation "
"data across the workers and that some amount of it will be "
"ignored. A small amount of this is fine, but a major imbalance "
"should be avoided. Note: This warning will appear unless your "
"data is perfectly balanced.")
break
loss = self.batch_loss(batch, for_training=False)
if loss is not None:
# You shouldn't necessarily have to compute a loss for validation, so we allow for
# `loss` to be None. We need to be careful, though - `batches_this_epoch` is
# currently only used as the divisor for the loss function, so we can safely only
# count those batches for which we actually have a loss. If this variable ever
# gets used for something else, we might need to change things around a bit.
batches_this_epoch += 1
val_loss += loss.detach().cpu().numpy()
# Update the description with the latest metrics
val_metrics = training_util.get_metrics(
self.model,
val_loss,
batches_this_epoch,
world_size=self._world_size,
cuda_device=[self.cuda_device],
)
description = training_util.description_from_metrics(val_metrics)
val_generator_tqdm.set_description(description, refresh=False)
if self._distributed and not done_early:
logger.warning(
f"Worker {torch.distributed.get_rank()} completed its entire epoch (validation)."
)
# Indicate that we're done so that any workers that have remaining data stop validation early.
done = torch.tensor(1, device=self.cuda_device)
torch.distributed.all_reduce(done, torch.distributed.ReduceOp.SUM)
assert done.item()
# Now restore the original parameter values.
if self._moving_average is not None:
self._moving_average.restore()
return val_loss, batches_this_epoch
def train(self) -> Dict[str, Any]:
"""
Trains the supplied model with the supplied parameters.
"""
try:
epoch_counter = self._restore_checkpoint()
except RuntimeError:
traceback.print_exc()
raise ConfigurationError(
"Could not recover training from the checkpoint. Did you mean to output to "
"a different serialization directory or delete the existing serialization "
"directory?")
training_util.enable_gradient_clipping(self.model, self._grad_clipping)
logger.info("Beginning training.")
val_metrics: Dict[str, float] = {}
this_epoch_val_metric: float = None
metrics: Dict[str, Any] = {}
epochs_trained = 0
training_start_time = time.time()
metrics["best_epoch"] = self._metric_tracker.best_epoch
for key, value in self._metric_tracker.best_epoch_metrics.items():
metrics["best_validation_" + key] = value
for epoch in range(epoch_counter, self._num_epochs):
epoch_start_time = time.time()
train_metrics = self._train_epoch(epoch)
# get peak of memory usage
if "cpu_memory_MB" in train_metrics:
metrics["peak_cpu_memory_MB"] = max(
metrics.get("peak_cpu_memory_MB", 0),
train_metrics["cpu_memory_MB"])
for key, value in train_metrics.items():
if key.startswith("gpu_"):
metrics["peak_" + key] = max(metrics.get("peak_" + key, 0),
value)
if self._validation_data is not None:
with torch.no_grad():
# We have a validation set, so compute all the metrics on it.
val_loss, num_batches = self._validation_loss()
# It is safe again to wait till the validation is done. This is
# important to get the metrics right.
if self._distributed:
dist.barrier()
val_metrics = training_util.get_metrics(
self.model,
val_loss,
num_batches,
reset=True,
world_size=self._world_size,
cuda_device=[self.cuda_device],
)
# Check validation metric for early stopping
this_epoch_val_metric = val_metrics[
self._validation_metric]
self._metric_tracker.add_metric(this_epoch_val_metric)
if self._metric_tracker.should_stop_early():
logger.info("Ran out of patience. Stopping training.")
break
if self._master:
self._tensorboard.log_metrics(
train_metrics,
val_metrics=val_metrics,
log_to_console=True,
epoch=epoch + 1) # +1 because tensorboard doesn't like 0
# Create overall metrics dict
training_elapsed_time = time.time() - training_start_time
metrics["training_duration"] = str(
datetime.timedelta(seconds=training_elapsed_time))
metrics["training_start_epoch"] = epoch_counter
metrics["training_epochs"] = epochs_trained
metrics["epoch"] = epoch
for key, value in train_metrics.items():
metrics["training_" + key] = value
for key, value in val_metrics.items():
metrics["validation_" + key] = value
if self._metric_tracker.is_best_so_far():
# Update all the best_ metrics.
# (Otherwise they just stay the same as they were.)
metrics["best_epoch"] = epoch
for key, value in val_metrics.items():
metrics["best_validation_" + key] = value
self._metric_tracker.best_epoch_metrics = val_metrics
if self._serialization_dir and self._master:
common_util.dump_metrics(
os.path.join(self._serialization_dir,
f"metrics_epoch_{epoch}.json"), metrics)
# The Scheduler API is agnostic to whether your schedule requires a validation metric -
# if it doesn't, the validation metric passed here is ignored.
if self._learning_rate_scheduler:
self._learning_rate_scheduler.step(this_epoch_val_metric,
epoch)
if self._momentum_scheduler:
self._momentum_scheduler.step(this_epoch_val_metric, epoch)
if self._master:
self._save_checkpoint(epoch)
# Wait for the master to finish saving the checkpoint
if self._distributed:
dist.barrier()
epoch_elapsed_time = time.time() - epoch_start_time
logger.info("Epoch duration: %s",
datetime.timedelta(seconds=epoch_elapsed_time))
if epoch < self._num_epochs - 1:
training_elapsed_time = time.time() - training_start_time
estimated_time_remaining = training_elapsed_time * (
(self._num_epochs - epoch_counter) /
float(epoch - epoch_counter + 1) - 1)
formatted_time = str(
datetime.timedelta(seconds=int(estimated_time_remaining)))
logger.info("Estimated training time remaining: %s",
formatted_time)
epochs_trained += 1
# make sure pending events are flushed to disk and files are closed properly
self._tensorboard.close()
# Load the best model state before returning
best_model_state = self._checkpointer.best_model_state()
if best_model_state:
self.model.load_state_dict(best_model_state)
return metrics
def _save_checkpoint(self, epoch: Union[int, str]) -> None:
"""
Saves a checkpoint of the model to self._serialization_dir.
Is a no-op if self._serialization_dir is None.
# Parameters
epoch : Union[int, str], required.
The epoch of training. If the checkpoint is saved in the middle
of an epoch, the parameter is a string with the epoch and timestamp.
"""
# If moving averages are used for parameters, we save
# the moving average values into checkpoint, instead of the current values.
if self._moving_average is not None:
self._moving_average.assign_average_value()
# These are the training states we need to persist.
training_states = {
"metric_tracker": self._metric_tracker.state_dict(),
"optimizer": self.optimizer.state_dict(),
"batch_num_total": self._batch_num_total,
}
# If we have a learning rate or momentum scheduler, we should persist them too.
if self._learning_rate_scheduler is not None:
training_states[
"learning_rate_scheduler"] = self._learning_rate_scheduler.state_dict(
)
if self._momentum_scheduler is not None:
training_states[
"momentum_scheduler"] = self._momentum_scheduler.state_dict()
self._checkpointer.save_checkpoint(
model_state=self.model.state_dict(),
epoch=epoch,
training_states=training_states,
is_best_so_far=self._metric_tracker.is_best_so_far(),
)
# Restore the original values for parameters so that training will not be affected.
if self._moving_average is not None:
self._moving_average.restore()
def _restore_checkpoint(self) -> int:
"""
Restores the model and training state from the last saved checkpoint.
This includes an epoch count and optimizer state, which is serialized separately
from model parameters. This function should only be used to continue training -
if you wish to load a model for inference/load parts of a model into a new
computation graph, you should use the native Pytorch functions:
` model.load_state_dict(torch.load("/path/to/model/weights.th"))`
If `self._serialization_dir` does not exist or does not contain any checkpointed weights,
this function will do nothing and return 0.
# Returns
epoch: int
The epoch at which to resume training, which should be one after the epoch
in the saved training state.
"""
model_state, training_state = self._checkpointer.restore_checkpoint()
if not training_state:
# No checkpoint to restore, start at 0
return 0
self.model.load_state_dict(model_state)
self.optimizer.load_state_dict(training_state["optimizer"])
if (self._learning_rate_scheduler is not None
and "learning_rate_scheduler" in training_state):
self._learning_rate_scheduler.load_state_dict(
training_state["learning_rate_scheduler"])
if self._momentum_scheduler is not None and "momentum_scheduler" in training_state:
self._momentum_scheduler.load_state_dict(
training_state["momentum_scheduler"])
training_util.move_optimizer_to_cuda(self.optimizer)
# Currently the `training_state` contains a serialized `MetricTracker`.
if "metric_tracker" in training_state:
self._metric_tracker.load_state_dict(
training_state["metric_tracker"])
# It used to be the case that we tracked `val_metric_per_epoch`.
elif "val_metric_per_epoch" in training_state:
self._metric_tracker.clear()
self._metric_tracker.add_metrics(
training_state["val_metric_per_epoch"])
# And before that we didn't track anything.
else:
self._metric_tracker.clear()
if isinstance(training_state["epoch"], int):
epoch_to_return = training_state["epoch"] + 1
else:
epoch_to_return = int(training_state["epoch"].split(".")[0]) + 1
# For older checkpoints with batch_num_total missing, default to old behavior where
# it is unchanged.
batch_num_total = training_state.get("batch_num_total")
if batch_num_total is not None:
self._batch_num_total = batch_num_total
return epoch_to_return
@classmethod
def from_partial_objects(
cls,
model: Model,
serialization_dir: str,
iterator: DataIterator,
train_data: Iterable[Instance],
validation_iterator: DataIterator = None,
validation_data: Iterable[Instance] = None,
local_rank: int = 0,
patience: int = None,
validation_metric: str = "-loss",
shuffle: bool = True,
num_epochs: int = 20,
cuda_device: int = -1,
grad_norm: float = None,
grad_clipping: float = None,
model_save_interval: float = None,
summary_interval: int = 100,
histogram_interval: int = None,
should_log_parameter_statistics: bool = True,
should_log_learning_rate: bool = False,
log_batch_size_period: int = None,
distributed: bool = None,
world_size: int = 1,
num_gradient_accumulation_steps: int = 1,
no_grad: List[str] = None,
optimizer: Lazy[Optimizer] = None,
learning_rate_scheduler: Lazy[LearningRateScheduler] = None,
momentum_scheduler: Lazy[MomentumScheduler] = None,
moving_average: Lazy[MovingAverage] = None,
checkpointer: Lazy[Checkpointer] = None,
) -> "Trainer":
"""
This method exists so that we can have a documented method to construct this class using
`FromParams`. If you are not using `FromParams` or config files, you can safely ignore this
method.
The reason we can't just use `__init__` with `FromParams` here is because there are
sequential dependencies to this class's arguments. Anything that has a `Lazy[]` type
annotation needs something from one of the non-`Lazy` arguments. The `Optimizer` needs to
have the parameters from the `Model` before it's constructed, and the `Schedulers` need to
have the `Optimizer`. Because of this, the typical way we construct things `FromParams`
doesn't work, so we use `Lazy` to allow for constructing the objects sequentially.
If you're not using `FromParams`, you can just construct these arguments in the right order
yourself in your code and call the constructor directly.
"""
check_for_gpu(cuda_device)
if cuda_device >= 0:
# Moving model to GPU here so that the optimizer state gets constructed on
# the right device.
model = model.cuda(cuda_device)
if no_grad:
for name, parameter in model.named_parameters():
if any(re.search(regex, name) for regex in no_grad):
parameter.requires_grad_(False)
common_util.log_frozen_and_tunable_parameter_names(model)
parameters = [[n, p] for n, p in model.named_parameters()
if p.requires_grad]
optimizer_ = optimizer.construct(model_parameters=parameters)
if not optimizer_:
optimizer_ = Optimizer.default(parameters)
batches_per_epoch = iterator.get_num_batches(train_data)
if batches_per_epoch == 1: # get_num_batches returns 1 when it can't determine the answer
batches_per_epoch = None
moving_average_ = moving_average.construct(parameters=parameters)
learning_rate_scheduler_ = learning_rate_scheduler.construct(
optimizer=optimizer_,
num_epochs=num_epochs,
num_steps_per_epoch=batches_per_epoch)
momentum_scheduler_ = momentum_scheduler.construct(
optimizer=optimizer_)
checkpointer_ = checkpointer.construct() or Checkpointer(
serialization_dir)
return cls(
model,
optimizer_,
iterator,
train_data,
validation_data,
patience=patience,
validation_metric=validation_metric,
validation_iterator=validation_iterator,
shuffle=shuffle,
num_epochs=num_epochs,
serialization_dir=serialization_dir,
cuda_device=cuda_device,
grad_norm=grad_norm,
grad_clipping=grad_clipping,
learning_rate_scheduler=learning_rate_scheduler_,
momentum_scheduler=momentum_scheduler_,
checkpointer=checkpointer_,
model_save_interval=model_save_interval,
summary_interval=summary_interval,
histogram_interval=histogram_interval,
should_log_parameter_statistics=should_log_parameter_statistics,
should_log_learning_rate=should_log_learning_rate,
log_batch_size_period=log_batch_size_period,
moving_average=moving_average_,
distributed=distributed,
local_rank=local_rank,
world_size=world_size,
num_gradient_accumulation_steps=num_gradient_accumulation_steps,
)
class MetaTrainer(Trainer):
def __init__(
self,
model: MetaWrapper,
component_optimizers: Dict[str, ComponentOptimizer],
data_loader: DataLoader,
patience: Optional[int] = None,
validation_metric: str = "-loss",
validation_data_loader: DataLoader = None,
num_epochs: int = 20,
serialization_dir: Optional[str] = None,
checkpointer: Checkpointer = None,
cuda_device: Optional[Union[int, torch.device]] = None,
tensorboard_writer: TensorboardWriter = None,
moving_average: Optional[MovingAverage] = None,
batch_callbacks: List[BatchCallback] = None,
epoch_callbacks: List[EpochCallback] = None,
num_gradient_accumulation_steps: int = 1,
use_amp: bool = False,
):
super().__init__(serialization_dir, cuda_device)
# I am not calling move_to_gpu here, because if the model is
# not already on the GPU then the optimizer is going to be wrong.
self.model = model
self.data_loader = data_loader
self._validation_data_loader = validation_data_loader
self.component_optimizers = component_optimizers
if patience is None: # no early stopping
if validation_data_loader is not None:
logger.warning(
"You provided a validation dataset but patience was set to None, "
"meaning that early stopping is disabled")
elif (not isinstance(patience, int)) or patience <= 0:
raise ConfigurationError(
'{} is an invalid value for "patience": it must be a positive integer '
"or None (if you want to disable early stopping)".format(
patience))
# For tracking is_best_so_far and should_stop_early
self._metric_tracker = MetricTracker(patience, validation_metric)
# Get rid of + or -
self._validation_metric = validation_metric[1:]
self._num_epochs = num_epochs
if checkpointer is not None:
self._checkpointer = checkpointer
else:
self._checkpointer = Checkpointer(serialization_dir)
self._batch_callbacks = batch_callbacks or []
self._epoch_callbacks = epoch_callbacks or []
self._moving_average = moving_average
# We keep the total batch number as an instance variable because it
# is used inside a closure for the hook which logs activations in
# `_enable_activation_logging`.
self._batch_num_total = 0
self._tensorboard = tensorboard_writer or TensorboardWriter(
serialization_dir)
self._tensorboard.get_batch_num_total = lambda: self._batch_num_total
self._tensorboard.enable_activation_logging(self.model)
self._last_log = 0.0 # time of last logging
self._num_gradient_accumulation_steps = num_gradient_accumulation_steps
# Enable automatic mixed precision training.
self._scaler: Optional[amp.GradScaler] = None
self._use_amp = use_amp
if self._use_amp:
if self.cuda_device == torch.device("cpu"):
raise ValueError("Using AMP requires a cuda device")
self._scaler = amp.GradScaler()
self._is_master = True
self._pytorch_model = self.model
def train(self) -> Dict[str, Any]:
"""
Trains the supplied model with the supplied parameters.
"""
try:
epoch_counter = self._restore_checkpoint()
except RuntimeError:
traceback.print_exc()
raise ConfigurationError(
"Could not recover training from the checkpoint. Did you mean to output to "
"a different serialization directory or delete the existing serialization "
"directory?")
for optimizer in self.component_optimizers.values():
optimizer.enable_gradient_clipping()
logger.info("Beginning training.")
val_metrics: Dict[str, float] = {}
this_epoch_val_metric: float = None
metrics: Dict[str, Any] = {}
epochs_trained = 0
training_start_time = time.time()
metrics["best_epoch"] = self._metric_tracker.best_epoch
for key, value in self._metric_tracker.best_epoch_metrics.items():
metrics["best_validation_" + key] = value
for callback in self._epoch_callbacks:
callback(self, metrics={}, epoch=-1, is_master=self._master)
for epoch in range(epoch_counter, self._num_epochs):
epoch_start_time = time.time()
train_metrics = self._train_epoch(epoch)
# get peak of memory usage
for key, value in train_metrics.items():
if key.startswith("gpu_") and key.endswith("_memory_MB"):
metrics["peak_" + key] = max(metrics.get("peak_" + key, 0),
value)
elif key.startswith("worker_") and key.endswith("_memory_MB"):
metrics["peak_" + key] = max(metrics.get("peak_" + key, 0),
value)
if self._validation_data_loader is not None:
with torch.no_grad():
# We have a validation set, so compute all the metrics on it.
val_metrics = self._validation_loss(epoch)
# Check validation metric for early stopping
this_epoch_val_metric = val_metrics[
f"meta_{self._validation_metric}"]
self._metric_tracker.add_metric(this_epoch_val_metric)
if self._metric_tracker.should_stop_early():
logger.info("Ran out of patience. Stopping training.")
break
if self._master:
self._tensorboard.log_metrics(
train_metrics,
val_metrics=val_metrics,
log_to_console=True,
epoch=epoch + 1) # +1 because tensorboard doesn't like 0
# Create overall metrics dict
training_elapsed_time = time.time() - training_start_time
metrics["training_duration"] = str(
datetime.timedelta(seconds=training_elapsed_time))
metrics["training_start_epoch"] = epoch_counter
metrics["training_epochs"] = epochs_trained
metrics["epoch"] = epoch
for key, value in train_metrics.items():
metrics["training_" + key] = value
for key, value in val_metrics.items():
metrics["validation_" + key] = value
if self._metric_tracker.is_best_so_far():
# Update all the best_ metrics.
# (Otherwise they just stay the same as they were.)
metrics["best_epoch"] = epoch
for key, value in val_metrics.items():
metrics["best_validation_" + key] = value
self._metric_tracker.best_epoch_metrics = val_metrics
if self._serialization_dir and self._master:
common_util.dump_metrics(
os.path.join(self._serialization_dir,
f"metrics_epoch_{epoch}.json"), metrics)
# The Scheduler API is agnostic to whether your schedule requires a validation metric -
# if it doesn't, the validation metric passed here is ignored.
for name, sub_model in self._pytorch_model.component_models.items(
):
component_optimizer = self.component_optimizers[name]
metric = val_metrics[f"{name}_{self._validation_metric}"]
if component_optimizer._learning_rate_scheduler:
component_optimizer._learning_rate_scheduler.step(metric)
if component_optimizer._momentum_scheduler:
component_optimizer._momentum_scheduler.step(metric)
if self._master:
self._checkpointer.save_checkpoint(
epoch,
self,
is_best_so_far=self._metric_tracker.is_best_so_far())
for callback in self._epoch_callbacks:
callback(self,
metrics=metrics,
epoch=epoch,
is_master=self._master)
epoch_elapsed_time = time.time() - epoch_start_time
logger.info("Epoch duration: %s",
datetime.timedelta(seconds=epoch_elapsed_time))
if epoch < self._num_epochs - 1:
training_elapsed_time = time.time() - training_start_time
estimated_time_remaining = training_elapsed_time * (
(self._num_epochs - epoch_counter) /
float(epoch - epoch_counter + 1) - 1)
formatted_time = str(
datetime.timedelta(seconds=int(estimated_time_remaining)))
logger.info("Estimated training time remaining: %s",
formatted_time)
epochs_trained += 1
# make sure pending events are flushed to disk and files are closed properly
self._tensorboard.close()
# Load the best model state before returning
best_model_state = self._checkpointer.best_model_state()
if best_model_state:
self.model.load_state_dict(best_model_state)
return metrics
def _train_epoch(self, epoch: int) -> Dict[str, float]:
"""
Trains one epoch and returns metrics.
"""
logger.info(f"Epoch: {epoch}/{self._num_epochs - 1}")
cpu_memory_usage = []
for worker, memory in common_util.peak_memory_mb().items():
cpu_memory_usage.append((worker, memory))
logger.info(f"Worker {worker} memory usage MB: {memory}")
gpu_memory_usage = []
for gpu, memory in common_util.gpu_memory_mb().items():
gpu_memory_usage.append((gpu, memory))
logger.info(f"GPU {gpu} memory usage MB: {memory}")
for component_optimizer in self.component_optimizers.values():
component_optimizer.reset_loss('train')
self.model.train()
# Get tqdm for the training batches
batch_generator = iter(self.data_loader)
batch_group_generator = common_util.lazy_groups_of(
batch_generator, self._num_gradient_accumulation_steps)
logger.info("Training")
num_training_batches: Union[int, float]
try:
len_data_loader = len(self.data_loader)
num_training_batches = math.ceil(
len_data_loader / self._num_gradient_accumulation_steps)
except TypeError:
num_training_batches = float("inf")
batch_group_generator_tqdm = Tqdm.tqdm(batch_group_generator,
total=num_training_batches)
self._last_log = time.time()
batches_this_epoch = 0
if self._batch_num_total is None:
self._batch_num_total = 0
done_early = False
for batch_group in batch_group_generator_tqdm:
batches_this_epoch += 1
self._batch_num_total += 1
batch_num_total = self._batch_num_total
for component_optimizer in self.component_optimizers.values():
component_optimizer.zero_grad()
batch_group_metrics = []
meta_batch = deepcopy(batch_group)
# Train the Sub Models first
for name, sub_model in self._pytorch_model.component_models.items(
):
component_optimizer = self.component_optimizers[name]
batch_group_outputs, metrics = component_optimizer.process_batch_group(
batch_group, True, batch_num_total, batches_this_epoch,
True)
batch_group_metrics.append(metrics)
for i, batch_outputs in enumerate(batch_group_outputs):
component_output = batch_outputs["output"]
component_output = component_output.detach()
meta_batch[i][name] = component_output
meta_optimizer = self.component_optimizers["meta"]
meta_batch_outputs, meta_metrics = meta_optimizer.process_batch_group(
meta_batch, True, batch_num_total, batches_this_epoch, False)
# Update moving averages
if self._moving_average is not None:
self._moving_average.apply(batch_num_total)
batch_group_metrics.append(meta_metrics)
all_metrics = ChainMap(*batch_group_metrics)
description = training_util.description_from_metrics(all_metrics)
batch_group_generator_tqdm.set_description(description,
refresh=False)
for (worker, memory) in cpu_memory_usage:
metrics["worker_" + str(worker) + "_memory_MB"] = memory
for (gpu_num, memory) in gpu_memory_usage:
metrics["gpu_" + str(gpu_num) + "_memory_MB"] = memory
return all_metrics
def _validation_loss(self, epoch: int) -> Tuple[float, float, int]:
"""
Computes the validation loss. Returns it and the number of batches.
"""
logger.info("Validating")
self.model.eval()
# Replace parameter values with the shadow values from the moving averages.
if self._moving_average is not None:
self._moving_average.assign_average_value()
if self._validation_data_loader is not None:
validation_data_loader = self._validation_data_loader
else:
raise ConfigurationError(
"Validation results cannot be calculated without a validation_data_loader"
)
val_generator_tqdm = Tqdm.tqdm(validation_data_loader)
for component_optimizer in self.component_optimizers.values():
component_optimizer.reset_loss('validation')
batches_this_epoch = 0
done_early = False
for batch in val_generator_tqdm:
batches_this_epoch += 1
batch_metrics = []
batch_group = [batch]
meta_batch = deepcopy(batch_group)
# Train the Sub Models first
for name, sub_model in self._pytorch_model.component_models.items(
):
component_optimizer = self.component_optimizers[name]
batch_group_outputs, metrics = component_optimizer.process_batch_group(
batch_group,
for_training=False,
batches_this_epoch=batches_this_epoch)
batch_metrics.append(metrics)
for i, batch_outputs in enumerate(batch_group_outputs):
meta_batch[i][name] = batch_outputs["output"]
meta_optimizer = self.component_optimizers["meta"]
meta_batch_outputs, meta_metrics = meta_optimizer.process_batch_group(
meta_batch,
for_training=False,
batches_this_epoch=batches_this_epoch)
batch_metrics.append(meta_metrics)
all_metrics = ChainMap(*batch_metrics)
description = training_util.description_from_metrics(all_metrics)
val_generator_tqdm.set_description(description, refresh=False)
# Now restore the original parameter values.
if self._moving_average is not None:
self._moving_average.restore()
return all_metrics
@contextmanager
def get_checkpoint_state(
self) -> Iterator[Tuple[Dict[str, Any], Dict[str, Any]]]:
if self._moving_average is not None:
# Assigning average value to model parameters. The checkpointer will call
# `restore_state_after_checkpointing` when it is done to put this back to what it was.
self._moving_average.assign_average_value()
model_state = self.model.state_dict()
# These are the training states we need to persist.
training_states = {
"metric_tracker": self._metric_tracker.state_dict(),
"optimizers":
{k: v.get_state()
for k, v in self.component_optimizers.items()},
"batch_num_total": self._batch_num_total,
}
try:
yield model_state, training_states
finally:
if self._moving_average is not None:
self._moving_average.restore()
def _restore_checkpoint(self) -> int:
"""
Restores the model and training state from the last saved checkpoint.
This includes an epoch count and optimizer state, which is serialized separately
from model parameters. This function should only be used to continue training -
if you wish to load a model for inference/load parts of a model into a new
computation graph, you should use the native Pytorch functions:
` model.load_state_dict(torch.load("/path/to/model/weights.th"))`
If `self._serialization_dir` does not exist or does not contain any checkpointed weights,
this function will do nothing and return 0.
# Returns
epoch: `int`
The epoch at which to resume training, which should be one after the epoch
in the saved training state.
"""
model_state, training_state = self._checkpointer.restore_checkpoint()
if not training_state:
# No checkpoint to restore, start at 0
return 0
self.model.load_state_dict(model_state)
for name, cm in self.component_optimizers.items():
state = training_state["optimizers"][name]
cm._optimizer.load_state_dict(state["state"])
if (cm._learning_rate_scheduler is not None
and "learning_rate_scheduler" in state):
cm._learning_rate_scheduler.load_state_dict(
state["learning_rate_scheduler"])
if cm._momentum_scheduler is not None and "momentum_scheduler" in state:
cm._momentum_scheduler.load_state_dict(
state["momentum_scheduler"])
training_util.move_optimizer_to_cuda(cm._optimizer)
# Currently the `training_state` contains a serialized `MetricTracker`.
if "metric_tracker" in training_state:
self._metric_tracker.load_state_dict(
training_state["metric_tracker"])
# It used to be the case that we tracked `val_metric_per_epoch`.
elif "val_metric_per_epoch" in training_state:
self._metric_tracker.clear()
self._metric_tracker.add_metrics(
training_state["val_metric_per_epoch"])
# And before that we didn't track anything.
else:
self._metric_tracker.clear()
if isinstance(training_state["epoch"], int):
epoch_to_return = training_state["epoch"] + 1
else:
epoch_to_return = int(training_state["epoch"].split(".")[0]) + 1
# For older checkpoints with batch_num_total missing, default to old behavior where
# it is unchanged.
batch_num_total = training_state.get("batch_num_total")
if batch_num_total is not None:
self._batch_num_total = batch_num_total
return epoch_to_return
@classmethod
def from_partial_objects(
cls,
model: MetaWrapper,
serialization_dir: str,
data_loader: DataLoader,
validation_data_loader: DataLoader = None,
patience: int = None,
validation_metric: str = "-loss",
num_epochs: int = 20,
cuda_device: Optional[Union[int, torch.device]] = None,
num_gradient_accumulation_steps: int = 1,
use_amp: bool = False,
no_grad: List[str] = None,
component_optimizers: Dict[str, Lazy[ComponentOptimizer]] = None,
tensorboard_writer: Lazy[TensorboardWriter] = None,
moving_average: Lazy[MovingAverage] = None,
checkpointer: Lazy[Checkpointer] = None,
batch_callbacks: List[BatchCallback] = None,
epoch_callbacks: List[EpochCallback] = None,
) -> "Trainer":
if cuda_device is None:
from torch import cuda
if cuda.device_count() > 0:
cuda_device = 0
else:
cuda_device = -1
check_for_gpu(cuda_device)
if cuda_device >= 0:
# Moving model to GPU here so that the optimizer state gets constructed on
# the right device.
model = model.cuda(cuda_device)
model.meta_model = model.meta_model.cuda(cuda_device)
for name in model.component_models:
model.component_models[name] = model.component_models[
name].cuda(cuda_device)
if no_grad:
for name, parameter in model.named_parameters():
if any(re.search(regex, name) for regex in no_grad):
parameter.requires_grad_(False)
batches_per_epoch: Optional[int]
try:
batches_per_epoch = len(data_loader)
batches_per_epoch = math.ceil(batches_per_epoch /
num_gradient_accumulation_steps)
except TypeError:
batches_per_epoch = None
sub_models = model.get_all_models()
for name, sub_model in sub_models.items():
component_optimizers[name] = component_optimizers[name].construct(
name=name,
model=sub_model,
num_epochs=num_epochs,
batches_per_epoch=batches_per_epoch,
cuda_device=cuda_device)
all_parameters = [[n, p] for n, p in model.named_parameters()
if p.requires_grad]
moving_average_ = moving_average.construct(parameters=all_parameters)
checkpointer_ = checkpointer.construct() or Checkpointer(
serialization_dir)
tensorboard_writer_ = tensorboard_writer.construct(
) or TensorboardWriter(serialization_dir)
return cls(model=model,
component_optimizers=component_optimizers,
data_loader=data_loader,
patience=patience,
validation_metric=validation_metric,
validation_data_loader=validation_data_loader,
num_epochs=num_epochs,
serialization_dir=serialization_dir,
checkpointer=checkpointer_,
moving_average=moving_average_,
cuda_device=cuda_device,
tensorboard_writer=tensorboard_writer_,
batch_callbacks=batch_callbacks,
epoch_callbacks=epoch_callbacks,
use_amp=use_amp)
class Trainer(TrainerBase):
def __init__(
self,
model: Model,
optimizer: torch.optim.Optimizer,
scheduler: torch.optim.lr_scheduler,
iterator: DataIterator,
train_dataset: Iterable[Instance],
validation_dataset: Optional[Iterable[Instance]] = None,
patience: Optional[int] = None,
validation_metric: str = "-loss",
validation_iterator: DataIterator = None,
shuffle: bool = True,
num_epochs: int = 20,
accumulated_batch_count: int = 1,
serialization_dir: Optional[str] = None,
num_serialized_models_to_keep: int = 20,
keep_serialized_model_every_num_seconds: int = None,
checkpointer: Checkpointer = None,
model_save_interval: float = None,
cuda_device: Union[int, List] = -1,
grad_norm: Optional[float] = None,
grad_clipping: Optional[float] = None,
learning_rate_scheduler: Optional[LearningRateScheduler] = None,
momentum_scheduler: Optional[MomentumScheduler] = None,
summary_interval: int = 100,
histogram_interval: int = None,
should_log_parameter_statistics: bool = True,
should_log_learning_rate: bool = False,
log_batch_size_period: Optional[int] = None,
moving_average: Optional[MovingAverage] = None,
cold_step_count: int = 0,
cold_lr: float = 1e-3,
cuda_verbose_step=None,
) -> None:
"""
A trainer for doing supervised learning. It just takes a labeled dataset
and a ``DataIterator``, and uses the supplied ``Optimizer`` to learn the weights
for your model over some fixed number of epochs. You can also pass in a validation
dataset and enable early stopping. There are many other bells and whistles as well.
Parameters
----------
model : ``Model``, required.
An AllenNLP model to be optimized. Pytorch Modules can also be optimized if
their ``forward`` method returns a dictionary with a "loss" key, containing a
scalar tensor representing the loss function to be optimized.
If you are training your model using GPUs, your model should already be
on the correct device. (If you use `Trainer.from_params` this will be
handled for you.)
optimizer : ``torch.nn.Optimizer``, required.
An instance of a Pytorch Optimizer, instantiated with the parameters of the
model to be optimized.
iterator : ``DataIterator``, required.
A method for iterating over a ``Dataset``, yielding padded indexed batches.
train_dataset : ``Dataset``, required.
A ``Dataset`` to train on. The dataset should have already been indexed.
validation_dataset : ``Dataset``, optional, (default = None).
A ``Dataset`` to evaluate on. The dataset should have already been indexed.
patience : Optional[int] > 0, optional (default=None)
Number of epochs to be patient before early stopping: the training is stopped
after ``patience`` epochs with no improvement. If given, it must be ``> 0``.
If None, early stopping is disabled.
validation_metric : str, optional (default="loss")
Validation metric to measure for whether to stop training using patience
and whether to serialize an ``is_best`` model each epoch. The metric name
must be prepended with either "+" or "-", which specifies whether the metric
is an increasing or decreasing function.
validation_iterator : ``DataIterator``, optional (default=None)
An iterator to use for the validation set. If ``None``, then
use the training `iterator`.
shuffle: ``bool``, optional (default=True)
Whether to shuffle the instances in the iterator or not.
num_epochs : int, optional (default = 20)
Number of training epochs.
serialization_dir : str, optional (default=None)
Path to directory for saving and loading model files. Models will not be saved if
this parameter is not passed.
num_serialized_models_to_keep : ``int``, optional (default=20)
Number of previous model checkpoints to retain. Default is to keep 20 checkpoints.
A value of None or -1 means all checkpoints will be kept.
keep_serialized_model_every_num_seconds : ``int``, optional (default=None)
If num_serialized_models_to_keep is not None, then occasionally it's useful to
save models at a given interval in addition to the last num_serialized_models_to_keep.
To do so, specify keep_serialized_model_every_num_seconds as the number of seconds
between permanently saved checkpoints. Note that this option is only used if
num_serialized_models_to_keep is not None, otherwise all checkpoints are kept.
checkpointer : ``Checkpointer``, optional (default=None)
An instance of class Checkpointer to use instead of the default. If a checkpointer is specified,
the arguments num_serialized_models_to_keep and keep_serialized_model_every_num_seconds should
not be specified. The caller is responsible for initializing the checkpointer so that it is
consistent with serialization_dir.
model_save_interval : ``float``, optional (default=None)
If provided, then serialize models every ``model_save_interval``
seconds within single epochs. In all cases, models are also saved
at the end of every epoch if ``serialization_dir`` is provided.
cuda_device : ``Union[int, List[int]]``, optional (default = -1)
An integer or list of integers specifying the CUDA device(s) to use. If -1, the CPU is used.
grad_norm : ``float``, optional, (default = None).
If provided, gradient norms will be rescaled to have a maximum of this value.
grad_clipping : ``float``, optional (default = ``None``).
If provided, gradients will be clipped `during the backward pass` to have an (absolute)
maximum of this value. If you are getting ``NaNs`` in your gradients during training
that are not solved by using ``grad_norm``, you may need this.
learning_rate_scheduler : ``LearningRateScheduler``, optional (default = None)
If specified, the learning rate will be decayed with respect to
this schedule at the end of each epoch (or batch, if the scheduler implements
the ``step_batch`` method). If you use :class:`torch.optim.lr_scheduler.ReduceLROnPlateau`,
this will use the ``validation_metric`` provided to determine if learning has plateaued.
To support updating the learning rate on every batch, this can optionally implement
``step_batch(batch_num_total)`` which updates the learning rate given the batch number.
momentum_scheduler : ``MomentumScheduler``, optional (default = None)
If specified, the momentum will be updated at the end of each batch or epoch
according to the schedule.
summary_interval: ``int``, optional, (default = 100)
Number of batches between logging scalars to tensorboard
histogram_interval : ``int``, optional, (default = ``None``)
If not None, then log histograms to tensorboard every ``histogram_interval`` batches.
When this parameter is specified, the following additional logging is enabled:
* Histograms of model parameters
* The ratio of parameter update norm to parameter norm
* Histogram of layer activations
We log histograms of the parameters returned by
``model.get_parameters_for_histogram_tensorboard_logging``.
The layer activations are logged for any modules in the ``Model`` that have
the attribute ``should_log_activations`` set to ``True``. Logging
histograms requires a number of GPU-CPU copies during training and is typically
slow, so we recommend logging histograms relatively infrequently.
Note: only Modules that return tensors, tuples of tensors or dicts
with tensors as values currently support activation logging.
should_log_parameter_statistics : ``bool``, optional, (default = True)
Whether to send parameter statistics (mean and standard deviation
of parameters and gradients) to tensorboard.
should_log_learning_rate : ``bool``, optional, (default = False)
Whether to send parameter specific learning rate to tensorboard.
log_batch_size_period : ``int``, optional, (default = ``None``)
If defined, how often to log the average batch size.
moving_average: ``MovingAverage``, optional, (default = None)
If provided, we will maintain moving averages for all parameters. During training, we
employ a shadow variable for each parameter, which maintains the moving average. During
evaluation, we backup the original parameters and assign the moving averages to corresponding
parameters. Be careful that when saving the checkpoint, we will save the moving averages of
parameters. This is necessary because we want the saved model to perform as well as the validated
model if we load it later. But this may cause problems if you restart the training from checkpoint.
"""
super().__init__(serialization_dir, cuda_device)
# I am not calling move_to_gpu here, because if the model is
# not already on the GPU then the optimizer is going to be wrong.
self.model = model
self.iterator = iterator
self._validation_iterator = validation_iterator
self.shuffle = shuffle
self.optimizer = optimizer
self.scheduler = scheduler
self.train_data = train_dataset
self._validation_data = validation_dataset
self.accumulated_batch_count = accumulated_batch_count
self.cold_step_count = cold_step_count
self.cold_lr = cold_lr
self.cuda_verbose_step = cuda_verbose_step
if patience is None: # no early stopping
if validation_dataset:
logger.warning(
"You provided a validation dataset but patience was set to None, "
"meaning that early stopping is disabled")
elif (not isinstance(patience, int)) or patience <= 0:
raise ConfigurationError(
'{} is an invalid value for "patience": it must be a positive integer '
"or None (if you want to disable early stopping)".format(
patience))
# For tracking is_best_so_far and should_stop_early
self._metric_tracker = MetricTracker(patience, validation_metric)
# Get rid of + or -
self._validation_metric = validation_metric[1:]
self._num_epochs = num_epochs
if checkpointer is not None:
# We can't easily check if these parameters were passed in, so check against their default values.
# We don't check against serialization_dir since it is also used by the parent class.
if num_serialized_models_to_keep != 20 \
or keep_serialized_model_every_num_seconds is not None:
raise ConfigurationError(
"When passing a custom Checkpointer, you may not also pass in separate checkpointer "
"args 'num_serialized_models_to_keep' or 'keep_serialized_model_every_num_seconds'."
)
self._checkpointer = checkpointer
else:
self._checkpointer = Checkpointer(
serialization_dir,
keep_serialized_model_every_num_seconds,
num_serialized_models_to_keep,
)
self._model_save_interval = model_save_interval
self._grad_norm = grad_norm
self._grad_clipping = grad_clipping
self._learning_rate_scheduler = learning_rate_scheduler
self._momentum_scheduler = momentum_scheduler
self._moving_average = moving_average
# We keep the total batch number as an instance variable because it
# is used inside a closure for the hook which logs activations in
# ``_enable_activation_logging``.
self._batch_num_total = 0
self._tensorboard = TensorboardWriter(
get_batch_num_total=lambda: self._batch_num_total,
serialization_dir=serialization_dir,
summary_interval=summary_interval,
histogram_interval=histogram_interval,
should_log_parameter_statistics=should_log_parameter_statistics,
should_log_learning_rate=should_log_learning_rate,
)
self._log_batch_size_period = log_batch_size_period
self._last_log = 0.0 # time of last logging
# Enable activation logging.
if histogram_interval is not None:
self._tensorboard.enable_activation_logging(self.model)
def rescale_gradients(self) -> Optional[float]:
return training_util.rescale_gradients(self.model, self._grad_norm)
def batch_loss(self, batch_group: List[TensorDict],
for_training: bool) -> torch.Tensor:
"""
Does a forward pass on the given batches and returns the ``loss`` value in the result.
If ``for_training`` is `True` also applies regularization penalty.
"""
if self._multiple_gpu:
output_dict = training_util.data_parallel(batch_group, self.model,
self._cuda_devices)
else:
assert len(batch_group) == 1
batch = batch_group[0]
batch = nn_util.move_to_device(batch, self._cuda_devices[0])
output_dict = self.model(**batch) # 里面是训练过程.
try:
loss = output_dict["loss"]
if for_training:
loss += self.model.get_regularization_penalty(
) # 保持泛化性.这里面没设置,所以是惩罚系数=0
except KeyError:
if for_training:
raise RuntimeError(
"The model you are trying to optimize does not contain a"
" 'loss' key in the output of model.forward(inputs).")
loss = None
return loss
def _train_epoch(self, epoch: int) -> Dict[str, float]:
"""
Trains one epoch and returns metrics.
"""
logger.info("Epoch %d/%d", epoch, self._num_epochs - 1)
peak_cpu_usage = peak_memory_mb()
logger.info(f"Peak CPU memory usage MB: {peak_cpu_usage}")
gpu_usage = []
for gpu, memory in gpu_memory_mb().items():
gpu_usage.append((gpu, memory))
logger.info(f"GPU {gpu} memory usage MB: {memory}")
train_loss = 0.0
# Set the model to "train" mode.
self.model.train()
num_gpus = len(self._cuda_devices) # 如果没有gpu ,也返回1.
# Get tqdm for the training batches
raw_train_generator = self.iterator(self.train_data,
num_epochs=1,
shuffle=self.shuffle)
train_generator = lazy_groups_of(raw_train_generator, num_gpus)
num_training_batches = math.ceil(
self.iterator.get_num_batches(self.train_data) / num_gpus)
residue = num_training_batches % self.accumulated_batch_count
self._last_log = time.time()
last_save_time = time.time()
batches_this_epoch = 0
if self._batch_num_total is None:
self._batch_num_total = 0
histogram_parameters = set(
self.model.get_parameters_for_histogram_tensorboard_logging())
logger.info("Training")
train_generator_tqdm = Tqdm.tqdm(
train_generator, total=num_training_batches) # 打印一个进度条而已.
cumulative_batch_size = 0
self.optimizer.zero_grad()
for batch_group in train_generator_tqdm:
batches_this_epoch += 1
self._batch_num_total += 1
batch_num_total = self._batch_num_total
iter_len = self.accumulated_batch_count \
if batches_this_epoch <= (num_training_batches - residue) else residue
if self.cuda_verbose_step is not None and batch_num_total % self.cuda_verbose_step == 0:
print(
f'Before forward pass - Cuda memory allocated: {torch.cuda.memory_allocated() / 1e9}'
)
print(
f'Before forward pass - Cuda memory cached: {torch.cuda.memory_cached() / 1e9}'
)
try:
loss = self.batch_loss(
batch_group,
for_training=True) / iter_len # 输入的数据里面去除了全部都是keep的情况
except RuntimeError as e:
print(e)
for x in batch_group:
all_words = [len(y['words']) for y in x['metadata']]
print(f"Total sents: {len(all_words)}. "
f"Min {min(all_words)}. Max {max(all_words)}")
for elem in ['labels', 'd_tags']:
tt = x[elem]
print(
f"{elem} shape {list(tt.shape)} and min {tt.min().item()} and {tt.max().item()}"
)
for elem in ["bert", "mask", "bert-offsets"]:
tt = x['tokens'][elem]
print(
f"{elem} shape {list(tt.shape)} and min {tt.min().item()} and {tt.max().item()}"
)
raise e
if self.cuda_verbose_step is not None and batch_num_total % self.cuda_verbose_step == 0:
print(
f'After forward pass - Cuda memory allocated: {torch.cuda.memory_allocated() / 1e9}'
)
print(
f'After forward pass - Cuda memory cached: {torch.cuda.memory_cached() / 1e9}'
)
if torch.isnan(loss):
raise ValueError("nan loss encountered")
loss.backward()
if self.cuda_verbose_step is not None and batch_num_total % self.cuda_verbose_step == 0:
print(
f'After backprop - Cuda memory allocated: {torch.cuda.memory_allocated() / 1e9}'
)
print(
f'After backprop - Cuda memory cached: {torch.cuda.memory_cached() / 1e9}'
)
train_loss += loss.item() * iter_len
del batch_group, loss
torch.cuda.empty_cache() # 节省内存,显存
if self.cuda_verbose_step is not None and batch_num_total % self.cuda_verbose_step == 0:
print(
f'After collecting garbage - Cuda memory allocated: {torch.cuda.memory_allocated() / 1e9}'
)
print(
f'After collecting garbage - Cuda memory cached: {torch.cuda.memory_cached() / 1e9}'
)
batch_grad_norm = self.rescale_gradients()
# This does nothing if batch_num_total is None or you are using a
# scheduler which doesn't update per batch.
if self._learning_rate_scheduler:
self._learning_rate_scheduler.step_batch(batch_num_total)
if self._momentum_scheduler:
self._momentum_scheduler.step_batch(batch_num_total)
if self._tensorboard.should_log_histograms_this_batch():
# get the magnitude of parameter updates for logging
# We need a copy of current parameters to compute magnitude of updates,
# and copy them to CPU so large models won't go OOM on the GPU.
param_updates = {
name: param.detach().cpu().clone()
for name, param in self.model.named_parameters()
}
if batches_this_epoch % self.accumulated_batch_count == 0 or \
batches_this_epoch == num_training_batches:
self.optimizer.step()
self.optimizer.zero_grad()
for name, param in self.model.named_parameters():
param_updates[name].sub_(param.detach().cpu())
update_norm = torch.norm(param_updates[name].view(-1))
param_norm = torch.norm(param.view(-1)).cpu()
self._tensorboard.add_train_scalar(
"gradient_update/" + name,
update_norm / (param_norm + 1e-7))
else:
if batches_this_epoch % self.accumulated_batch_count == 0 or \
batches_this_epoch == num_training_batches:
self.optimizer.step() #多个batch才进行bp算法.
self.optimizer.zero_grad()
# Update moving averages
if self._moving_average is not None:
self._moving_average.apply(batch_num_total)
# Update the description with the latest metrics
metrics = training_util.get_metrics(self.model, train_loss,
batches_this_epoch) # 计算准确率
description = training_util.description_from_metrics(metrics)
train_generator_tqdm.set_description(description, refresh=False)
# Log parameter values to Tensorboard
if self._tensorboard.should_log_this_batch():
self._tensorboard.log_parameter_and_gradient_statistics(
self.model, batch_grad_norm)
self._tensorboard.log_learning_rates(self.model,
self.optimizer)
self._tensorboard.add_train_scalar("loss/loss_train",
metrics["loss"])
self._tensorboard.log_metrics(
{"epoch_metrics/" + k: v
for k, v in metrics.items()})
if self._tensorboard.should_log_histograms_this_batch():
self._tensorboard.log_histograms(self.model,
histogram_parameters)
if self._log_batch_size_period:
cur_batch = sum([
training_util.get_batch_size(batch)
for batch in batch_group
])
cumulative_batch_size += cur_batch
if (batches_this_epoch - 1) % self._log_batch_size_period == 0:
average = cumulative_batch_size / batches_this_epoch
logger.info(
f"current batch size: {cur_batch} mean batch size: {average}"
)
self._tensorboard.add_train_scalar("current_batch_size",
cur_batch)
self._tensorboard.add_train_scalar("mean_batch_size",
average)
# Save model if needed. 取一个间隔来存
if self._model_save_interval is not None and (
time.time() - last_save_time > self._model_save_interval):
last_save_time = time.time()
self._save_checkpoint("{0}.{1}".format(
epoch, training_util.time_to_str(int(last_save_time))))
metrics = training_util.get_metrics(self.model,
train_loss,
batches_this_epoch,
reset=True)
metrics["cpu_memory_MB"] = peak_cpu_usage
for (gpu_num, memory) in gpu_usage:
metrics["gpu_" + str(gpu_num) + "_memory_MB"] = memory
return metrics
def _validation_loss(self) -> Tuple[float, int]:
"""
Computes the validation loss. Returns it and the number of batches.
"""
logger.info("Validating")
self.model.eval()
# Replace parameter values with the shadow values from the moving averages.
if self._moving_average is not None:
self._moving_average.assign_average_value()
if self._validation_iterator is not None:
val_iterator = self._validation_iterator
else:
val_iterator = self.iterator
num_gpus = len(self._cuda_devices)
raw_val_generator = val_iterator(self._validation_data,
num_epochs=1,
shuffle=True)
val_generator = lazy_groups_of(raw_val_generator, num_gpus)
num_validation_batches = math.ceil(
val_iterator.get_num_batches(self._validation_data) / num_gpus)
val_generator_tqdm = Tqdm.tqdm(val_generator,
total=num_validation_batches)
batches_this_epoch = 0
val_loss = 0
for batch_group in val_generator_tqdm:
loss = self.batch_loss(batch_group, for_training=False)
if loss is not None:
# You shouldn't necessarily have to compute a loss for validation, so we allow for
# `loss` to be None. We need to be careful, though - `batches_this_epoch` is
# currently only used as the divisor for the loss function, so we can safely only
# count those batches for which we actually have a loss. If this variable ever
# gets used for something else, we might need to change things around a bit.
batches_this_epoch += 1
val_loss += loss.detach().cpu().numpy()
# Update the description with the latest metrics
val_metrics = training_util.get_metrics(self.model, val_loss,
batches_this_epoch)
description = training_util.description_from_metrics(val_metrics)
val_generator_tqdm.set_description(description, refresh=False)
# Now restore the original parameter values.
if self._moving_average is not None:
self._moving_average.restore()
return val_loss, batches_this_epoch
def train(self, oldmodel) -> Dict[str, Any]:
"""
Trains the supplied model with the supplied parameters.
"""#---------------通过oldmodel来加载官方模型,从而进行finetune
try:
# epoch_counter = self._restore_checkpoint()
epoch_counter = 0 # 直接进行finetune. 加速训练.
# 下面看这个地方怎么修改,保证之前的参数都赋值上去.?????????????????????????????????????之前没遇到过这个问题,需要自己克服看看.
print(oldmodel,
'旧的模型路径是这个') # 这个model里面已经涵盖了xlnet网络的参数和最后2层linear的参数.
tmp = torch.load(oldmodel, map_location=torch.device('cpu'))
out_shape = self.model.tag_labels_projection_layer._module.out_features # 把下面的东西扩充到out_shape
now_shape = tmp[
'tag_labels_projection_layer._module.weight'].shape[0]
fix_shape = out_shape - now_shape
# 通过concat
tmp['tag_labels_projection_layer._module.weight'] = torch.cat(
(tmp['tag_labels_projection_layer._module.weight'],
torch.zeros(fix_shape, 768)), 0)
tmp['tag_labels_projection_layer._module.bias'] = torch.cat(
(tmp['tag_labels_projection_layer._module.bias'],
torch.zeros(fix_shape)), 0)
# 需要补充到的数据大小:
# tmp只是一个字典而已,随便玩.
self.model.load_state_dict(
tmp) # 这次的收敛速度飞快!!!!!!!!!!!!!!!!!!!# 初步的打算是,补充shape到我们需要的,大小.
except RuntimeError:
traceback.print_exc()
raise ConfigurationError(
"Could not recover training from the checkpoint. Did you mean to output to "
"a different serialization directory or delete the existing serialization "
"directory?")
training_util.enable_gradient_clipping(self.model, self._grad_clipping)
logger.info("Beginning training.")
train_metrics: Dict[str, float] = {}
val_metrics: Dict[str, float] = {}
this_epoch_val_metric: float = None
metrics: Dict[str, Any] = {}
epochs_trained = 0
training_start_time = time.time()
if self.cold_step_count > 0: # 冷启动几个step, 这些step不用更新权重.对于embed网路都freeeze上.只更新后面自己简历的分类器网络.
base_lr = self.optimizer.param_groups[0]['lr']
for param_group in self.optimizer.param_groups:
param_group['lr'] = self.cold_lr
self.model.text_field_embedder._token_embedders[
'bert'].set_weights(freeze=True)
metrics["best_epoch"] = self._metric_tracker.best_epoch
for key, value in self._metric_tracker.best_epoch_metrics.items():
metrics["best_validation_" + key] = value
for epoch in range(epoch_counter,
self._num_epochs): # 把之前学完的epoch直接跳过.
if epoch == self.cold_step_count and epoch != 0: # 冷启动完毕,开始恢复学习率
for param_group in self.optimizer.param_groups:
param_group['lr'] = base_lr
self.model.text_field_embedder._token_embedders[
'bert'].set_weights(freeze=False) #并且把embed网络解除冻结
epoch_start_time = time.time() # 下行是训练代码
train_metrics = self._train_epoch(epoch)
# get peak of memory usage
if "cpu_memory_MB" in train_metrics:
metrics["peak_cpu_memory_MB"] = max(
metrics.get("peak_cpu_memory_MB", 0),
train_metrics["cpu_memory_MB"])
for key, value in train_metrics.items():
if key.startswith("gpu_"):
metrics["peak_" + key] = max(metrics.get("peak_" + key, 0),
value)
# clear cache before validation
torch.cuda.empty_cache()
# 在验证集上评测效果,防止过拟合.
if self._validation_data is not None:
with torch.no_grad():
# We have a validation set, so compute all the metrics on it.
val_loss, num_batches = self._validation_loss()
val_metrics = training_util.get_metrics(self.model,
val_loss,
num_batches,
reset=True)
# Check validation metric for early stopping
this_epoch_val_metric = val_metrics[
self._validation_metric]
self._metric_tracker.add_metric(this_epoch_val_metric)
if self._metric_tracker.should_stop_early():
logger.info("Ran out of patience. Stopping training.")
break
self._tensorboard.log_metrics(
train_metrics,
val_metrics=val_metrics,
log_to_console=True,
epoch=epoch + 1) # +1 because tensorboard doesn't like 0
# Create overall metrics dict
training_elapsed_time = time.time() - training_start_time
metrics["training_duration"] = str(
datetime.timedelta(seconds=training_elapsed_time))
metrics["training_start_epoch"] = epoch_counter
metrics["training_epochs"] = epochs_trained
metrics["epoch"] = epoch
for key, value in train_metrics.items():
metrics["training_" + key] = value
for key, value in val_metrics.items():
metrics["validation_" + key] = value
# if self.cold_step_count <= epoch:
self.scheduler.step(metrics['validation_loss'])
if self._metric_tracker.is_best_so_far():
# Update all the best_ metrics.
# (Otherwise they just stay the same as they were.)
metrics["best_epoch"] = epoch
for key, value in val_metrics.items():
metrics["best_validation_" + key] = value
self._metric_tracker.best_epoch_metrics = val_metrics
if self._serialization_dir:
dump_metrics(
os.path.join(self._serialization_dir,
f"metrics_epoch_{epoch}.json"), metrics)
# The Scheduler API is agnostic to whether your schedule requires a validation metric -
# if it doesn't, the validation metric passed here is ignored.
if self._learning_rate_scheduler:
self._learning_rate_scheduler.step(this_epoch_val_metric,
epoch)
if self._momentum_scheduler:
self._momentum_scheduler.step(this_epoch_val_metric, epoch)
#保存model, 只需要给定这个文件夹,那么算法自动会读取里面最新的模型.来进行finetune.很方便.
if self._num_epochs == epoch + 1: # 只存最后一个.
self._save_checkpoint(epoch) # 每一个epoch 都存,最后的磁盘占用很大.
epoch_elapsed_time = time.time() - epoch_start_time
logger.info("Epoch duration: %s",
datetime.timedelta(seconds=epoch_elapsed_time))
if epoch < self._num_epochs - 1:
training_elapsed_time = time.time() - training_start_time
estimated_time_remaining = training_elapsed_time * (
(self._num_epochs - epoch_counter) /
float(epoch - epoch_counter + 1) - 1)
formatted_time = str(
datetime.timedelta(seconds=int(estimated_time_remaining)))
logger.info("Estimated training time remaining: %s",
formatted_time)
epochs_trained += 1
# make sure pending events are flushed to disk and files are closed properly
# self._tensorboard.close()
# Load the best model state before returning # 根据路径目录找到里面存的最好的模型.
best_model_state = self._checkpointer.best_model_state()
if best_model_state:
self.model.load_state_dict(best_model_state)
return metrics
def _save_checkpoint(self, epoch: Union[int, str]) -> None:
"""
Saves a checkpoint of the model to self._serialization_dir.
Is a no-op if self._serialization_dir is None.
Parameters
----------
epoch : Union[int, str], required.
The epoch of training. If the checkpoint is saved in the middle
of an epoch, the parameter is a string with the epoch and timestamp.
"""
# If moving averages are used for parameters, we save
# the moving average values into checkpoint, instead of the current values.
if self._moving_average is not None:
self._moving_average.assign_average_value()
# These are the training states we need to persist.
training_states = {
"metric_tracker": self._metric_tracker.state_dict(),
"optimizer": self.optimizer.state_dict(),
"batch_num_total": self._batch_num_total,
}
# If we have a learning rate or momentum scheduler, we should persist them too.
if self._learning_rate_scheduler is not None:
training_states[
"learning_rate_scheduler"] = self._learning_rate_scheduler.state_dict(
)
if self._momentum_scheduler is not None:
training_states[
"momentum_scheduler"] = self._momentum_scheduler.state_dict()
if self._metric_tracker.is_best_so_far(
): # 把save改成只存一份最好的,不好的直接pass.节省磁盘空间.
print('得到最优模型,正在保存')
self._checkpointer.save_checkpoint(
model_state=self.model.state_dict(),
epoch=epoch,
training_states=training_states,
is_best_so_far=self._metric_tracker.is_best_so_far(),
)
# Restore the original values for parameters so that training will not be affected.
if self._moving_average is not None:
self._moving_average.restore()
def _restore_checkpoint(self) -> int:
"""
Restores the model and training state from the last saved checkpoint.
This includes an epoch count and optimizer state, which is serialized separately
from model parameters. This function should only be used to continue training -
if you wish to load a model for inference/load parts of a model into a new
computation graph, you should use the native Pytorch functions:
`` model.load_state_dict(torch.load("/path/to/model/weights.th"))``
If ``self._serialization_dir`` does not exist or does not contain any checkpointed weights,
this function will do nothing and return 0.
Returns
-------
epoch: int
The epoch at which to resume training, which should be one after the epoch
in the saved training state.
"""#就是去目录里面找权重文件,然后拿过来继续训练.
model_state, training_state = self._checkpointer.restore_checkpoint()
if not training_state:
# No checkpoint to restore, start at 0
return 0
self.model.load_state_dict(model_state)
self.optimizer.load_state_dict(training_state["optimizer"])
if self._learning_rate_scheduler is not None \
and "learning_rate_scheduler" in training_state:
self._learning_rate_scheduler.load_state_dict(
training_state["learning_rate_scheduler"])
if self._momentum_scheduler is not None and "momentum_scheduler" in training_state:
self._momentum_scheduler.load_state_dict(
training_state["momentum_scheduler"])
training_util.move_optimizer_to_cuda(self.optimizer)
# Currently the ``training_state`` contains a serialized ``MetricTracker``.
if "metric_tracker" in training_state:
self._metric_tracker.load_state_dict(
training_state["metric_tracker"])
# It used to be the case that we tracked ``val_metric_per_epoch``.
elif "val_metric_per_epoch" in training_state:
self._metric_tracker.clear()
self._metric_tracker.add_metrics(
training_state["val_metric_per_epoch"])
# And before that we didn't track anything.
else:
self._metric_tracker.clear()
if isinstance(training_state["epoch"], int):
epoch_to_return = training_state["epoch"] + 1
else:
epoch_to_return = int(training_state["epoch"].split(".")[0]) + 1
# For older checkpoints with batch_num_total missing, default to old behavior where
# it is unchanged.
batch_num_total = training_state.get("batch_num_total")
if batch_num_total is not None:
self._batch_num_total = batch_num_total
return epoch_to_return
# Requires custom from_params.
@classmethod
def from_params( # type: ignore
cls,
model: Model,
serialization_dir: str,
iterator: DataIterator,
train_data: Iterable[Instance],
validation_data: Optional[Iterable[Instance]],
params: Params,
validation_iterator: DataIterator = None,
) -> "Trainer":
patience = params.pop_int("patience", None)
validation_metric = params.pop("validation_metric", "-loss")
shuffle = params.pop_bool("shuffle", True)
num_epochs = params.pop_int("num_epochs", 20)
cuda_device = parse_cuda_device(params.pop("cuda_device", -1))
grad_norm = params.pop_float("grad_norm", None)
grad_clipping = params.pop_float("grad_clipping", None)
lr_scheduler_params = params.pop("learning_rate_scheduler", None)
momentum_scheduler_params = params.pop("momentum_scheduler", None)
if isinstance(cuda_device, list):
model_device = cuda_device[0]
else:
model_device = cuda_device
if model_device >= 0:
# Moving model to GPU here so that the optimizer state gets constructed on
# the right device.
model = model.cuda(model_device)
parameters = [[n, p] for n, p in model.named_parameters()
if p.requires_grad]
optimizer = Optimizer.from_params(parameters, params.pop("optimizer"))
if "moving_average" in params:
moving_average = MovingAverage.from_params(
params.pop("moving_average"), parameters=parameters)
else:
moving_average = None
if lr_scheduler_params:
lr_scheduler = LearningRateScheduler.from_params(
optimizer, lr_scheduler_params)
else:
lr_scheduler = None
if momentum_scheduler_params:
momentum_scheduler = MomentumScheduler.from_params(
optimizer, momentum_scheduler_params)
else:
momentum_scheduler = None
if "checkpointer" in params:
if "keep_serialized_model_every_num_seconds" in params \
or "num_serialized_models_to_keep" in params:
raise ConfigurationError(
"Checkpointer may be initialized either from the 'checkpointer' key or from the "
"keys 'num_serialized_models_to_keep' and 'keep_serialized_model_every_num_seconds'"
" but the passed config uses both methods.")
checkpointer = Checkpointer.from_params(params.pop("checkpointer"))
else:
num_serialized_models_to_keep = params.pop_int(
"num_serialized_models_to_keep", 20)
keep_serialized_model_every_num_seconds = params.pop_int(
"keep_serialized_model_every_num_seconds", None)
checkpointer = Checkpointer(
serialization_dir=serialization_dir,
num_serialized_models_to_keep=num_serialized_models_to_keep,
keep_serialized_model_every_num_seconds=
keep_serialized_model_every_num_seconds,
)
model_save_interval = params.pop_float("model_save_interval", None)
summary_interval = params.pop_int("summary_interval", 100)
histogram_interval = params.pop_int("histogram_interval", None)
should_log_parameter_statistics = params.pop_bool(
"should_log_parameter_statistics", True)
should_log_learning_rate = params.pop_bool("should_log_learning_rate",
False)
log_batch_size_period = params.pop_int("log_batch_size_period", None)
params.assert_empty(cls.__name__)
return cls(
model,
optimizer,
iterator,
train_data,
validation_data,
patience=patience,
validation_metric=validation_metric,
validation_iterator=validation_iterator,
shuffle=shuffle,
num_epochs=num_epochs,
serialization_dir=serialization_dir,
cuda_device=cuda_device,
grad_norm=grad_norm,
grad_clipping=grad_clipping,
learning_rate_scheduler=lr_scheduler,
momentum_scheduler=momentum_scheduler,
checkpointer=checkpointer,
model_save_interval=model_save_interval,
summary_interval=summary_interval,
histogram_interval=histogram_interval,
should_log_parameter_statistics=should_log_parameter_statistics,
should_log_learning_rate=should_log_learning_rate,
log_batch_size_period=log_batch_size_period,
moving_average=moving_average,
)
class Trainer(TrainerBase):
def __init__(
self,
model: Model,
optimizer: torch.optim.Optimizer,
iterator: DataIterator,
train_dataset: Iterable[Instance],
validation_dataset: Optional[Iterable[Instance]] = None,
train_low_dataset: Optional[Iterable[Instance]] = None,
patience: Optional[int] = None,
validation_metric: str = "-loss",
validation_iterator: DataIterator = None,
shuffle: bool = True,
num_epochs: int = 20,
serialization_dir: Optional[str] = None,
num_serialized_models_to_keep: int = 20,
keep_serialized_model_every_num_seconds: int = None,
checkpointer: Checkpointer = None,
model_save_interval: float = None,
cuda_device: Union[int, List] = -1,
grad_norm: Optional[float] = None,
grad_clipping: Optional[float] = None,
learning_rate_scheduler: Optional[LearningRateScheduler] = None,
momentum_scheduler: Optional[MomentumScheduler] = None,
summary_interval: int = 100,
histogram_interval: int = None,
should_log_parameter_statistics: bool = True,
should_log_learning_rate: bool = False,
log_batch_size_period: Optional[int] = None,
moving_average: Optional[MovingAverage] = None,
epoch_low_start: Optional[int] = None,
epoch_without_improvement_low_start: Optional[int] = None,
) -> None:
super().__init__(serialization_dir, cuda_device)
# I am not calling move_to_gpu here, because if the model is
# not already on the GPU then the optimizer is going to be wrong.
self.model = model
self.iterator = iterator
self._validation_iterator = validation_iterator
self.shuffle = shuffle
self.optimizer = optimizer
self.train_data = train_dataset
self._validation_data = validation_dataset
self._train_low_data = train_low_dataset
# set when to train with low-data only / with defaults
self._epoch_low_start = epoch_low_start or 10
self._epoch_without_improvement_low_start = epoch_without_improvement_low_start or 5
if patience is None: # no early stopping
if validation_dataset:
logger.warning(
'You provided a validation dataset but patience was set to None, '
'meaning that early stopping is disabled')
elif (not isinstance(patience, int)) or patience <= 0:
raise ConfigurationError(
'{} is an invalid value for "patience": it must be a positive integer '
'or None (if you want to disable early stopping)'.format(
patience))
# For tracking is_best_so_far and should_stop_early
self._metric_tracker = MetricTracker(patience, validation_metric)
# AX: custom parameter for reinforce trainer
self._metric_tracker.reinforce_start_with_low = None
# Get rid of + or -
self._validation_metric = validation_metric[1:]
self._num_epochs = num_epochs
if checkpointer is not None:
# We can't easily check if these parameters were passed in, so check against their default values.
# We don't check against serialization_dir since it is also used by the parent class.
if num_serialized_models_to_keep != 20 or \
keep_serialized_model_every_num_seconds is not None:
raise ConfigurationError(
"When passing a custom Checkpointer, you may not also pass in separate checkpointer "
"args 'num_serialized_models_to_keep' or 'keep_serialized_model_every_num_seconds'."
)
self._checkpointer = checkpointer
else:
self._checkpointer = Checkpointer(
serialization_dir, keep_serialized_model_every_num_seconds,
num_serialized_models_to_keep)
self._model_save_interval = model_save_interval
self._grad_norm = grad_norm
self._grad_clipping = grad_clipping
self._learning_rate_scheduler = learning_rate_scheduler
self._momentum_scheduler = momentum_scheduler
self._moving_average = moving_average
# We keep the total batch number as an instance variable because it
# is used inside a closure for the hook which logs activations in
# ``_enable_activation_logging``.
self._batch_num_total = 0
self._tensorboard = TensorboardWriter(
get_batch_num_total=lambda: self._batch_num_total,
serialization_dir=serialization_dir,
summary_interval=summary_interval,
histogram_interval=histogram_interval,
should_log_parameter_statistics=should_log_parameter_statistics,
should_log_learning_rate=should_log_learning_rate)
self._log_batch_size_period = log_batch_size_period
self._last_log = 0.0 # time of last logging
# Enable activation logging.
if histogram_interval is not None:
self._tensorboard.enable_activation_logging(self.model)
def rescale_gradients(self) -> Optional[float]:
return training_util.rescale_gradients(self.model, self._grad_norm)
def batch_loss(self, batch_group: List[TensorDict],
for_training: bool) -> torch.Tensor:
"""
Does a forward pass on the given batches and returns the ``loss`` value in the result.
If ``for_training`` is `True` also applies regularization penalty.
"""
if self._multiple_gpu:
output_dict = training_util.data_parallel(batch_group, self.model,
self._cuda_devices)
else:
assert len(batch_group) == 1
batch = batch_group[0]
batch = nn_util.move_to_device(batch, self._cuda_devices[0])
output_dict = self.model(**batch)
try:
loss = output_dict["loss"]
if for_training:
loss += self.model.get_regularization_penalty()
except KeyError:
if for_training:
raise RuntimeError(
"The model you are trying to optimize does not contain a"
" 'loss' key in the output of model.forward(inputs).")
loss = None
return loss
def _train_epoch(self, epoch: int) -> Dict[str, float]:
"""
Trains one epoch and returns metrics.
"""
logger.info("Epoch %d/%d", epoch, self._num_epochs - 1)
peak_cpu_usage = peak_memory_mb()
logger.info(f"Peak CPU memory usage MB: {peak_cpu_usage}")
gpu_usage = []
for gpu, memory in gpu_memory_mb().items():
gpu_usage.append((gpu, memory))
logger.info(f"GPU {gpu} memory usage MB: {memory}")
train_loss = 0.0
# Set the model to "train" mode.
self.model.train()
num_gpus = len(self._cuda_devices)
if not self._metric_tracker.reinforce_start_with_low and (
epoch < self._epoch_low_start
or self._metric_tracker._epochs_with_no_improvement <
self._epoch_without_improvement_low_start):
train_data = self.train_data
else:
if not self._metric_tracker.reinforce_start_with_low:
self._metric_tracker.reinforce_start_with_low = epoch
train_data = self._train_low_data
# Get tqdm for the training batches
raw_train_generator = self.iterator(train_data,
num_epochs=1,
shuffle=self.shuffle)
train_generator = lazy_groups_of(raw_train_generator, num_gpus)
num_training_batches = math.ceil(
self.iterator.get_num_batches(train_data) / num_gpus)
self._last_log = time.time()
last_save_time = time.time()
batches_this_epoch = 0
if self._batch_num_total is None:
self._batch_num_total = 0
histogram_parameters = set(
self.model.get_parameters_for_histogram_tensorboard_logging())
logger.info("Training")
train_generator_tqdm = Tqdm.tqdm(train_generator,
total=num_training_batches)
cumulative_batch_size = 0
for batch_group in train_generator_tqdm:
batches_this_epoch += 1
self._batch_num_total += 1
batch_num_total = self._batch_num_total
self.optimizer.zero_grad()
loss = self.batch_loss(batch_group, for_training=True)
if torch.isnan(loss):
raise ValueError("nan loss encountered")
loss.backward()
train_loss += loss.item()
batch_grad_norm = self.rescale_gradients()
# This does nothing if batch_num_total is None or you are using a
# scheduler which doesn't update per batch.
if self._learning_rate_scheduler:
self._learning_rate_scheduler.step_batch(batch_num_total)
if self._momentum_scheduler:
self._momentum_scheduler.step_batch(batch_num_total)
if self._tensorboard.should_log_histograms_this_batch():
# get the magnitude of parameter updates for logging
# We need a copy of current parameters to compute magnitude of updates,
# and copy them to CPU so large models won't go OOM on the GPU.
param_updates = {
name: param.detach().cpu().clone()
for name, param in self.model.named_parameters()
}
self.optimizer.step()
for name, param in self.model.named_parameters():
param_updates[name].sub_(param.detach().cpu())
update_norm = torch.norm(param_updates[name].view(-1, ))
param_norm = torch.norm(param.view(-1, )).cpu()
self._tensorboard.add_train_scalar(
"gradient_update/" + name,
update_norm / (param_norm + 1e-7))
else:
self.optimizer.step()
# Update moving averages
if self._moving_average is not None:
self._moving_average.apply(batch_num_total)
# Update the description with the latest metrics
metrics = training_util.get_metrics(self.model, train_loss,
batches_this_epoch)
description = training_util.description_from_metrics(metrics)
train_generator_tqdm.set_description(description, refresh=False)
# Log parameter values to Tensorboard
if self._tensorboard.should_log_this_batch():
self._tensorboard.log_parameter_and_gradient_statistics(
self.model, batch_grad_norm)
self._tensorboard.log_learning_rates(self.model,
self.optimizer)
self._tensorboard.add_train_scalar("loss/loss_train",
metrics["loss"])
self._tensorboard.log_metrics(
{"epoch_metrics/" + k: v
for k, v in metrics.items()})
if self._tensorboard.should_log_histograms_this_batch():
self._tensorboard.log_histograms(self.model,
histogram_parameters)
if self._log_batch_size_period:
cur_batch = sum([
training_util.get_batch_size(batch)
for batch in batch_group
])
cumulative_batch_size += cur_batch
if (batches_this_epoch - 1) % self._log_batch_size_period == 0:
average = cumulative_batch_size / batches_this_epoch
logger.info(
f"current batch size: {cur_batch} mean batch size: {average}"
)
self._tensorboard.add_train_scalar("current_batch_size",
cur_batch)
self._tensorboard.add_train_scalar("mean_batch_size",
average)
# Save model if needed.
if self._model_save_interval is not None and (
time.time() - last_save_time > self._model_save_interval):
last_save_time = time.time()
self._save_checkpoint('{0}.{1}'.format(
epoch, training_util.time_to_str(int(last_save_time))))
metrics = training_util.get_metrics(self.model,
train_loss,
batches_this_epoch,
reset=True)
metrics['cpu_memory_MB'] = peak_cpu_usage
for (gpu_num, memory) in gpu_usage:
metrics['gpu_' + str(gpu_num) + '_memory_MB'] = memory
return metrics
def _validation_loss(self) -> Tuple[float, int]:
"""
Computes the validation loss. Returns it and the number of batches.
"""
logger.info("Validating")
self.model.eval()
# Replace parameter values with the shadow values from the moving averages.
if self._moving_average is not None:
self._moving_average.assign_average_value()
if self._validation_iterator is not None:
val_iterator = self._validation_iterator
else:
val_iterator = self.iterator
num_gpus = len(self._cuda_devices)
raw_val_generator = val_iterator(self._validation_data,
num_epochs=1,
shuffle=False)
val_generator = lazy_groups_of(raw_val_generator, num_gpus)
num_validation_batches = math.ceil(
val_iterator.get_num_batches(self._validation_data) / num_gpus)
val_generator_tqdm = Tqdm.tqdm(val_generator,
total=num_validation_batches)
batches_this_epoch = 0
val_loss = 0
for batch_group in val_generator_tqdm:
loss = self.batch_loss(batch_group, for_training=False)
if loss is not None:
# You shouldn't necessarily have to compute a loss for validation, so we allow for
# `loss` to be None. We need to be careful, though - `batches_this_epoch` is
# currently only used as the divisor for the loss function, so we can safely only
# count those batches for which we actually have a loss. If this variable ever
# gets used for something else, we might need to change things around a bit.
batches_this_epoch += 1
val_loss += loss.detach().cpu().numpy()
# Update the description with the latest metrics
val_metrics = training_util.get_metrics(self.model, val_loss,
batches_this_epoch)
description = training_util.description_from_metrics(val_metrics)
val_generator_tqdm.set_description(description, refresh=False)
# Now restore the original parameter values.
if self._moving_average is not None:
self._moving_average.restore()
return val_loss, batches_this_epoch
def train(self) -> Dict[str, Any]:
"""
Trains the supplied model with the supplied parameters.
"""
try:
epoch_counter = self._restore_checkpoint()
except RuntimeError:
traceback.print_exc()
raise ConfigurationError(
"Could not recover training from the checkpoint. Did you mean to output to "
"a different serialization directory or delete the existing serialization "
"directory?")
training_util.enable_gradient_clipping(self.model, self._grad_clipping)
logger.info("Beginning training.")
train_metrics: Dict[str, float] = {}
val_metrics: Dict[str, float] = {}
this_epoch_val_metric: float = None
metrics: Dict[str, Any] = {}
epochs_trained = 0
training_start_time = time.time()
metrics['best_epoch'] = self._metric_tracker.best_epoch
for key, value in self._metric_tracker.best_epoch_metrics.items():
metrics["best_validation_" + key] = value
for epoch in range(epoch_counter, self._num_epochs):
epoch_start_time = time.time()
train_metrics = self._train_epoch(epoch)
# AX: add custom value for epoch that low-training was started
metrics[
"reinforce_start_with_low"] = self._metric_tracker.reinforce_start_with_low
# get peak of memory usage
if 'cpu_memory_MB' in train_metrics:
metrics['peak_cpu_memory_MB'] = max(
metrics.get('peak_cpu_memory_MB', 0),
train_metrics['cpu_memory_MB'])
for key, value in train_metrics.items():
if key.startswith('gpu_'):
metrics["peak_" + key] = max(metrics.get("peak_" + key, 0),
value)
if self._validation_data is not None:
with torch.no_grad():
# We have a validation set, so compute all the metrics on it.
val_loss, num_batches = self._validation_loss()
val_metrics = training_util.get_metrics(self.model,
val_loss,
num_batches,
reset=True)
# Check validation metric for early stopping
this_epoch_val_metric = val_metrics[
self._validation_metric]
self._metric_tracker.add_metric(this_epoch_val_metric)
if self._metric_tracker.should_stop_early():
logger.info("Ran out of patience. Stopping training.")
break
self._tensorboard.log_metrics(
train_metrics,
val_metrics=val_metrics,
log_to_console=True,
epoch=epoch + 1) # +1 because tensorboard doesn't like 0
# Create overall metrics dict
training_elapsed_time = time.time() - training_start_time
metrics["training_duration"] = time.strftime(
"%H:%M:%S", time.gmtime(training_elapsed_time))
metrics["training_start_epoch"] = epoch_counter
metrics["training_epochs"] = epochs_trained
metrics["epoch"] = epoch
for key, value in train_metrics.items():
metrics["training_" + key] = value
for key, value in val_metrics.items():
metrics["validation_" + key] = value
if self._metric_tracker.is_best_so_far():
# Update all the best_ metrics.
# (Otherwise they just stay the same as they were.)
metrics['best_epoch'] = epoch
for key, value in val_metrics.items():
metrics["best_validation_" + key] = value
self._metric_tracker.best_epoch_metrics = val_metrics
if self._serialization_dir:
dump_metrics(
os.path.join(self._serialization_dir,
f'metrics_epoch_{epoch}.json'), metrics)
# The Scheduler API is agnostic to whether your schedule requires a validation metric -
# if it doesn't, the validation metric passed here is ignored.
if self._learning_rate_scheduler:
self._learning_rate_scheduler.step(this_epoch_val_metric,
epoch)
if self._momentum_scheduler:
self._momentum_scheduler.step(this_epoch_val_metric, epoch)
self._save_checkpoint(epoch)
epoch_elapsed_time = time.time() - epoch_start_time
logger.info(
"Epoch duration: %s",
time.strftime("%H:%M:%S", time.gmtime(epoch_elapsed_time)))
if epoch < self._num_epochs - 1:
training_elapsed_time = time.time() - training_start_time
estimated_time_remaining = training_elapsed_time * \
((self._num_epochs - epoch_counter) / float(epoch - epoch_counter + 1) - 1)
formatted_time = str(
datetime.timedelta(seconds=int(estimated_time_remaining)))
logger.info("Estimated training time remaining: %s",
formatted_time)
epochs_trained += 1
# Load the best model state before returning
best_model_state = self._checkpointer.best_model_state()
if best_model_state:
self.model.load_state_dict(best_model_state)
return metrics
def _save_checkpoint(self, epoch: Union[int, str]) -> None:
"""
Saves a checkpoint of the model to self._serialization_dir.
Is a no-op if self._serialization_dir is None.
Parameters
----------
epoch : Union[int, str], required.
The epoch of training. If the checkpoint is saved in the middle
of an epoch, the parameter is a string with the epoch and timestamp.
"""
# If moving averages are used for parameters, we save
# the moving average values into checkpoint, instead of the current values.
if self._moving_average is not None:
self._moving_average.assign_average_value()
# These are the training states we need to persist.
training_states = {
"metric_tracker": self._metric_tracker.state_dict(),
"optimizer": self.optimizer.state_dict(),
"batch_num_total": self._batch_num_total
}
# If we have a learning rate or momentum scheduler, we should persist them too.
if self._learning_rate_scheduler is not None:
training_states[
"learning_rate_scheduler"] = self._learning_rate_scheduler.state_dict(
)
if self._momentum_scheduler is not None:
training_states[
"momentum_scheduler"] = self._momentum_scheduler.state_dict()
self._checkpointer.save_checkpoint(
model_state=self.model.state_dict(),
epoch=epoch,
training_states=training_states,
is_best_so_far=self._metric_tracker.is_best_so_far())
# Restore the original values for parameters so that training will not be affected.
if self._moving_average is not None:
self._moving_average.restore()
def _restore_checkpoint(self) -> int:
"""
Restores the model and training state from the last saved checkpoint.
This includes an epoch count and optimizer state, which is serialized separately
from model parameters. This function should only be used to continue training -
if you wish to load a model for inference/load parts of a model into a new
computation graph, you should use the native Pytorch functions:
`` model.load_state_dict(torch.load("/path/to/model/weights.th"))``
If ``self._serialization_dir`` does not exist or does not contain any checkpointed weights,
this function will do nothing and return 0.
Returns
-------
epoch: int
The epoch at which to resume training, which should be one after the epoch
in the saved training state.
"""
model_state, training_state = self._checkpointer.restore_checkpoint()
if not training_state:
# No checkpoint to restore, start at 0
return 0
self.model.load_state_dict(model_state)
self.optimizer.load_state_dict(training_state["optimizer"])
if self._learning_rate_scheduler is not None and "learning_rate_scheduler" in training_state:
self._learning_rate_scheduler.load_state_dict(
training_state["learning_rate_scheduler"])
if self._momentum_scheduler is not None and "momentum_scheduler" in training_state:
self._momentum_scheduler.load_state_dict(
training_state["momentum_scheduler"])
training_util.move_optimizer_to_cuda(self.optimizer)
# Currently the ``training_state`` contains a serialized ``MetricTracker``.
if "metric_tracker" in training_state:
self._metric_tracker.load_state_dict(
training_state["metric_tracker"])
# It used to be the case that we tracked ``val_metric_per_epoch``.
elif "val_metric_per_epoch" in training_state:
self._metric_tracker.clear()
self._metric_tracker.add_metrics(
training_state["val_metric_per_epoch"])
# And before that we didn't track anything.
else:
self._metric_tracker.clear()
if isinstance(training_state["epoch"], int):
epoch_to_return = training_state["epoch"] + 1
else:
epoch_to_return = int(training_state["epoch"].split('.')[0]) + 1
# For older checkpoints with batch_num_total missing, default to old behavior where
# it is unchanged.
batch_num_total = training_state.get('batch_num_total')
if batch_num_total is not None:
self._batch_num_total = batch_num_total
return epoch_to_return
# Requires custom from_params.
@classmethod
def from_params(
cls, # type: ignore
params: Params,
serialization_dir: str,
recover: bool = False) -> 'Trainer':
# modified for second training_data
all_datasets = datasets_from_params(params)
# copied from allennlp.training.trainer.TrainingPieces
# modified for second training_data
datasets_for_vocab_creation = set(
params.pop("datasets_for_vocab_creation", all_datasets))
if recover and os.path.exists(
os.path.join(serialization_dir, "vocabulary")):
vocab = Vocabulary.from_files(
os.path.join(serialization_dir, "vocabulary"))
params.pop("vocabulary", {})
else:
vocab = Vocabulary.from_params(params.pop(
"vocabulary", {}), (instance
for key, dataset in all_datasets.items()
for instance in dataset
if key in datasets_for_vocab_creation))
model = Model.from_params(vocab=vocab, params=params.pop('model'))
model.extend_embedder_vocab()
vocab.save_to_files(os.path.join(serialization_dir, "vocabulary"))
iterator = DataIterator.from_params(params.pop("iterator"))
iterator.index_with(model.vocab)
validation_iterator_params = params.pop("validation_iterator", None)
if validation_iterator_params:
validation_iterator = DataIterator.from_params(
validation_iterator_params)
validation_iterator.index_with(model.vocab)
else:
validation_iterator = None
train_data = all_datasets['train']
validation_data = all_datasets.get('validation')
test_data = all_datasets.get('test')
train_low_data = all_datasets.get('train_low')
trainer_params = params.pop("trainer")
no_grad_regexes = trainer_params.pop("no_grad", ())
for name, parameter in model.named_parameters():
if any(re.search(regex, name) for regex in no_grad_regexes):
parameter.requires_grad_(False)
frozen_parameter_names, tunable_parameter_names = \
get_frozen_and_tunable_parameter_names(model)
logger.info("Following parameters are Frozen (without gradient):")
for name in frozen_parameter_names:
logger.info(name)
logger.info("Following parameters are Tunable (with gradient):")
for name in tunable_parameter_names:
logger.info(name)
# END OF TrainerPieces code
params = trainer_params
# pylint: disable=arguments-differ
patience = params.pop_int("patience", None)
validation_metric = params.pop("validation_metric", "-loss")
shuffle = params.pop_bool("shuffle", True)
num_epochs = params.pop_int("num_epochs", 20)
cuda_device = parse_cuda_device(params.pop("cuda_device", -1))
grad_norm = params.pop_float("grad_norm", None)
grad_clipping = params.pop_float("grad_clipping", None)
lr_scheduler_params = params.pop("learning_rate_scheduler", None)
momentum_scheduler_params = params.pop("momentum_scheduler", None)
if isinstance(cuda_device, list):
model_device = cuda_device[0]
else:
model_device = cuda_device
if model_device >= 0:
# Moving model to GPU here so that the optimizer state gets constructed on
# the right device.
model = model.cuda(model_device)
parameters = [[n, p] for n, p in model.named_parameters()
if p.requires_grad]
optimizer = Optimizer.from_params(parameters, params.pop("optimizer"))
if "moving_average" in params:
moving_average = MovingAverage.from_params(
params.pop("moving_average"), parameters=parameters)
else:
moving_average = None
if lr_scheduler_params:
lr_scheduler = LearningRateScheduler.from_params(
optimizer, lr_scheduler_params)
else:
lr_scheduler = None
if momentum_scheduler_params:
momentum_scheduler = MomentumScheduler.from_params(
optimizer, momentum_scheduler_params)
else:
momentum_scheduler = None
if 'checkpointer' in params:
if 'keep_serialized_model_every_num_seconds' in params or \
'num_serialized_models_to_keep' in params:
raise ConfigurationError(
"Checkpointer may be initialized either from the 'checkpointer' key or from the "
"keys 'num_serialized_models_to_keep' and 'keep_serialized_model_every_num_seconds'"
" but the passed config uses both methods.")
checkpointer = Checkpointer.from_params(params.pop("checkpointer"))
else:
num_serialized_models_to_keep = params.pop_int(
"num_serialized_models_to_keep", 20)
keep_serialized_model_every_num_seconds = params.pop_int(
"keep_serialized_model_every_num_seconds", None)
checkpointer = Checkpointer(
serialization_dir=serialization_dir,
num_serialized_models_to_keep=num_serialized_models_to_keep,
keep_serialized_model_every_num_seconds=
keep_serialized_model_every_num_seconds)
model_save_interval = params.pop_float("model_save_interval", None)
summary_interval = params.pop_int("summary_interval", 100)
histogram_interval = params.pop_int("histogram_interval", None)
should_log_parameter_statistics = params.pop_bool(
"should_log_parameter_statistics", True)
should_log_learning_rate = params.pop_bool("should_log_learning_rate",
False)
log_batch_size_period = params.pop_int("log_batch_size_period", None)
epoch_low_start = params.pop_int("epoch_low_start", None)
epoch_without_improvement_low_start = params.pop_int(
"epoch_without_improvement_low_start", None)
params.assert_empty(cls.__name__)
return cls(
model,
optimizer,
iterator,
train_data,
validation_data,
train_low_dataset=train_low_data,
patience=patience,
validation_metric=validation_metric,
validation_iterator=validation_iterator,
shuffle=shuffle,
num_epochs=num_epochs,
serialization_dir=serialization_dir,
cuda_device=cuda_device,
grad_norm=grad_norm,
grad_clipping=grad_clipping,
learning_rate_scheduler=lr_scheduler,
momentum_scheduler=momentum_scheduler,
checkpointer=checkpointer,
model_save_interval=model_save_interval,
summary_interval=summary_interval,
histogram_interval=histogram_interval,
should_log_parameter_statistics=should_log_parameter_statistics,
should_log_learning_rate=should_log_learning_rate,
log_batch_size_period=log_batch_size_period,
moving_average=moving_average,
epoch_low_start=epoch_low_start,
epoch_without_improvement_low_start=
epoch_without_improvement_low_start,
)
class MetaTrainer(TrainerBase):
def __init__(
self,
model: Model,
optimizer: torch.optim.Optimizer,
iterator: DataIterator,
train_datasets: Dict[str, Iterable[Instance]],
validation_datasets: Optional[Dict[str, Iterable[Instance]]] = None,
patience: Optional[int] = None,
validation_metric: str = "-loss",
validation_iterator: DataIterator = None,
shuffle: bool = True,
num_epochs: int = 20,
serialization_dir: Optional[str] = None,
num_serialized_models_to_keep: int = 20,
save_embedder: bool = True,
keep_serialized_model_every_num_seconds: int = None,
checkpointer: Checkpointer = None,
model_save_interval: float = None,
cuda_device: int = -1,
grad_norm: Optional[float] = None,
grad_clipping: Optional[float] = None,
learning_rate_scheduler: Optional[LearningRateScheduler] = None,
momentum_scheduler: Optional[MomentumScheduler] = None,
summary_interval: int = 100,
histogram_interval: int = None,
should_log_parameter_statistics: bool = True,
should_log_learning_rate: bool = False,
log_batch_size_period: Optional[int] = None,
moving_average: Optional[MovingAverage] = None,
distributed: bool = False,
local_rank: int = 0,
world_size: int = 1,
num_gradient_accumulation_steps: int = 1,
log_grad_norm: str = "total",
wrapper: Optional[Wrapper] = None,
task_discriminator: Optional[TaskDiscriminator] = None,
discriminator_optimizer: Optional[torch.optim.Optimizer] = None,
tasks_per_step: int = 0,
writer: WandBWriter = None,
) -> None:
"""
A trainer for doing supervised learning. It just takes a labeled dataset
and a `DataIterator`, and uses the supplied `Optimizer` to learn the weights
for your model over some fixed number of epochs. You can also pass in a validation
dataset and enable early stopping. There are many other bells and whistles as well.
# Parameters
model : `Model`, required.
An AllenNLP model to be optimized. Pytorch Modules can also be optimized if
their `forward` method returns a dictionary with a "loss" key, containing a
scalar tensor representing the loss function to be optimized.
If you are training your model using GPUs, your model should already be
on the correct device. (If you use `Trainer.from_params` this will be
handled for you.)
optimizer : `torch.nn.Optimizer`, required.
An instance of a Pytorch Optimizer, instantiated with the parameters of the
model to be optimized.
iterator : `DataIterator`, required.
A method for iterating over a `Dataset`, yielding padded indexed batches.
train_dataset : `Dataset`, required.
A `Dataset` to train on. The dataset should have already been indexed.
validation_dataset : `Dataset`, optional, (default = None).
A `Dataset` to evaluate on. The dataset should have already been indexed.
patience : Optional[int] > 0, optional (default=None)
Number of epochs to be patient before early stopping: the training is stopped
after `patience` epochs with no improvement. If given, it must be `> 0`.
If None, early stopping is disabled.
validation_metric : str, optional (default="loss")
Validation metric to measure for whether to stop training using patience
and whether to serialize an `is_best` model each epoch. The metric name
must be prepended with either "+" or "-", which specifies whether the metric
is an increasing or decreasing function.
validation_iterator : `DataIterator`, optional (default=None)
An iterator to use for the validation set. If `None`, then
use the training `iterator`.
shuffle : `bool`, optional (default=True)
Whether to shuffle the instances in the iterator or not.
num_epochs : int, optional (default = 20)
Number of training epochs.
serialization_dir : str, optional (default=None)
Path to directory for saving and loading model files. Models will not be saved if
this parameter is not passed.
num_serialized_models_to_keep : `int`, optional (default=20)
Number of previous model checkpoints to retain. Default is to keep 20 checkpoints.
A value of None or -1 means all checkpoints will be kept.
keep_serialized_model_every_num_seconds : `int`, optional (default=None)
If num_serialized_models_to_keep is not None, then occasionally it's useful to
save models at a given interval in addition to the last num_serialized_models_to_keep.
To do so, specify keep_serialized_model_every_num_seconds as the number of seconds
between permanently saved checkpoints. Note that this option is only used if
num_serialized_models_to_keep is not None, otherwise all checkpoints are kept.
checkpointer : `Checkpointer`, optional (default=None)
An instance of class Checkpointer to use instead of the default. If a checkpointer is specified,
the arguments num_serialized_models_to_keep and keep_serialized_model_every_num_seconds should
not be specified. The caller is responsible for initializing the checkpointer so that it is
consistent with serialization_dir.
model_save_interval : `float`, optional (default=None)
If provided, then serialize models every `model_save_interval`
seconds within single epochs. In all cases, models are also saved
at the end of every epoch if `serialization_dir` is provided.
cuda_device : `int`, optional (default = -1)
An integer specifying the CUDA device(s) to use for this process. If -1, the CPU is used.
Data parallelism is controlled at the allennlp train level, so each trainer will have a single
GPU.
grad_norm : `float`, optional, (default = None).
If provided, gradient norms will be rescaled to have a maximum of this value.
grad_clipping : `float`, optional (default = `None`).
If provided, gradients will be clipped `during the backward pass` to have an (absolute)
maximum of this value. If you are getting `NaNs` in your gradients during training
that are not solved by using `grad_norm`, you may need this.
learning_rate_scheduler : `LearningRateScheduler`, optional (default = None)
If specified, the learning rate will be decayed with respect to
this schedule at the end of each epoch (or batch, if the scheduler implements
the `step_batch` method). If you use :class:`torch.optim.lr_scheduler.ReduceLROnPlateau`,
this will use the `validation_metric` provided to determine if learning has plateaued.
To support updating the learning rate on every batch, this can optionally implement
`step_batch(batch_num_total)` which updates the learning rate given the batch number.
momentum_scheduler : `MomentumScheduler`, optional (default = None)
If specified, the momentum will be updated at the end of each batch or epoch
according to the schedule.
summary_interval : `int`, optional, (default = 100)
Number of batches between logging scalars to tensorboard
histogram_interval : `int`, optional, (default = `None`)
If not None, then log histograms to tensorboard every `histogram_interval` batches.
When this parameter is specified, the following additional logging is enabled:
* Histograms of model parameters
* The ratio of parameter update norm to parameter norm
* Histogram of layer activations
We log histograms of the parameters returned by
`model.get_parameters_for_histogram_tensorboard_logging`.
The layer activations are logged for any modules in the `Model` that have
the attribute `should_log_activations` set to `True`. Logging
histograms requires a number of GPU-CPU copies during training and is typically
slow, so we recommend logging histograms relatively infrequently.
Note: only Modules that return tensors, tuples of tensors or dicts
with tensors as values currently support activation logging.
should_log_parameter_statistics : `bool`, optional, (default = True)
Whether to send parameter statistics (mean and standard deviation
of parameters and gradients) to tensorboard.
should_log_learning_rate : `bool`, optional, (default = False)
Whether to send parameter specific learning rate to tensorboard.
log_batch_size_period : `int`, optional, (default = `None`)
If defined, how often to log the average batch size.
moving_average : `MovingAverage`, optional, (default = None)
If provided, we will maintain moving averages for all parameters. During training, we
employ a shadow variable for each parameter, which maintains the moving average. During
evaluation, we backup the original parameters and assign the moving averages to corresponding
parameters. Be careful that when saving the checkpoint, we will save the moving averages of
parameters. This is necessary because we want the saved model to perform as well as the validated
model if we load it later. But this may cause problems if you restart the training from checkpoint.
distributed : `bool`, optional, (default = False)
If set, PyTorch's `DistributedDataParallel` is used to train the model in multiple GPUs. This also
requires `world_size` to be greater than 1.
local_rank : `int`, optional, (default = 0)
This is the unique identifier of the `Trainer` in a distributed process group. The GPU device id is
used as the rank.
world_size : `int`, (default = 1)
The number of `Trainer` workers participating in the distributed training.
num_gradient_accumulation_steps : `int`, optional, (default = 1)
Gradients are accumulated for the given number of steps before doing an optimizer step. This can
be useful to accommodate batches that are larger than the RAM size. Refer Thomas Wolf's
[post](https://tinyurl.com/y5mv44fw) for details on Gradient Accumulation.
"""
super().__init__(serialization_dir, cuda_device, distributed, local_rank, world_size)
# I am not calling move_to_gpu here, because if the model is
# not already on the GPU then the optimizer is going to be wrong.
self.model = model
self.iterator = iterator
self._validation_iterator = validation_iterator
self.shuffle = shuffle
self.optimizer = optimizer
self.train_datas = train_datasets
self._validation_datas = validation_datasets
self._save_embedder = save_embedder
if patience is None: # no early stopping
if validation_datasets:
logger.warning(
"You provided a validation dataset but patience was set to None, "
"meaning that early stopping is disabled"
)
elif (not isinstance(patience, int)) or patience <= 0:
raise ConfigurationError(
'{} is an invalid value for "patience": it must be a positive integer '
"or None (if you want to disable early stopping)".format(patience)
)
# For tracking is_best_so_far and should_stop_early
self._metric_tracker = MetricTracker(patience, validation_metric)
# Get rid of + or -
self._validation_metric = validation_metric[1:]
self._num_epochs = num_epochs
if checkpointer is not None:
# We can't easily check if these parameters were passed in, so check against their default values.
# We don't check against serialization_dir since it is also used by the parent class.
if (
num_serialized_models_to_keep != 20
or keep_serialized_model_every_num_seconds is not None
):
raise ConfigurationError(
"When passing a custom Checkpointer, you may not also pass in separate checkpointer "
"args 'num_serialized_models_to_keep' or 'keep_serialized_model_every_num_seconds'."
)
self._checkpointer = checkpointer
else:
self._checkpointer = Checkpointer(
serialization_dir,
keep_serialized_model_every_num_seconds,
num_serialized_models_to_keep,
)
self._model_save_interval = model_save_interval
self._grad_norm = grad_norm
self._grad_clipping = grad_clipping
self._learning_rate_scheduler = learning_rate_scheduler
self._momentum_scheduler = momentum_scheduler
self._moving_average = moving_average
# We keep the total batch number as an instance variable because it
# is used inside a closure for the hook which logs activations in
# `_enable_activation_logging`.
self._batch_num_total = 0
if writer is not None:
self._writer = writer
else:
self._writer = TensorboardWriter(
get_batch_num_total=lambda: self._batch_num_total,
serialization_dir=serialization_dir,
summary_interval=summary_interval,
histogram_interval=histogram_interval,
should_log_parameter_statistics=should_log_parameter_statistics,
should_log_learning_rate=should_log_learning_rate)
self._log_batch_size_period = log_batch_size_period
self._last_log = 0.0 # time of last logging
self._num_gradient_accumulation_steps = num_gradient_accumulation_steps
# Using `DistributedDataParallel`(ddp) brings in a quirk wrt AllenNLP's `Model` interface and its
# usage. A `Model` object is wrapped by `ddp`, but assigning the wrapped model to `self.model`
# will break the usages such as `Model.get_regularization_penalty`, `Model.get_metrics`, etc.
#
# Hence a reference to Pytorch's object is maintained in the case of distributed training and in the
# normal case, reference to `Model` is retained. This reference is only used in
# these places: `model.__call__`, `model.train` and `model.eval`.
if self._distributed:
self._pytorch_model = DistributedDataParallel(
self.model, device_ids=[self.cuda_device], find_unused_parameters=True
)
else:
self._pytorch_model = self.model
self._tasks_per_step = tasks_per_step if tasks_per_step > 0 else len(self.train_datas.items())
self.wrapper = wrapper
self.task_D = task_discriminator
self.optim_D = discriminator_optimizer
self.has_VIB = hasattr(self.model, 'VIB') and self.model.VIB and self.model.VIB.beta > 0
self.has_pos = hasattr(self.model, '_predict_pos') and self.model._predict_pos
self.batch_generators = {task: self.iterator(train_data, shuffle=self.shuffle)
for task, train_data in self.train_datas.items()}
self.batch_group_generators = {task: lazy_groups_of(
batch_generator, self._num_gradient_accumulation_steps
) for task, batch_generator in self.batch_generators.items()}
def update_hook(norms):
assert log_grad_norm in ["none", 'total', 'var']
if log_grad_norm in ['total', 'var']:
total_task_grad_norm = 0.0
total_summed_grad_norm = 0.0
for name, norm_list in norms.items():
if len(norm_list) == 1:
logger.info(f"{name} has no gradient; skipping")
continue
avg_task_grad_norm = (sum(norm_list[:-1]) / len(norm_list[:-1]))
total_task_grad_norm += avg_task_grad_norm
summed_grad_norm = norm_list[-1]
total_summed_grad_norm += summed_grad_norm
if log_grad_norm == 'var':
ratio = summed_grad_norm / (avg_task_grad_norm + 1e-10)
self._writer.log({f"avg_task_grad_norm_{name}": avg_task_grad_norm,
f"summed_grad_norm_{name}": summed_grad_norm,
f"task-total_norm_ratio_{name}": ratio},
step=self._batch_num_total)
avg_ratio = total_summed_grad_norm / total_task_grad_norm
self._writer.log({"avg_task-total_norm_ratio": avg_ratio,
"total_grad_norm": total_summed_grad_norm,
"total_task_grad_norm": total_task_grad_norm},
step=self._batch_num_total)
self.wrapper.update_hook = update_hook
def rescale_gradients(self) -> Optional[float]:
return training_util.rescale_gradients(self.model, self._grad_norm)
def batch_loss(self, batch: TensorDict, for_training: bool) -> torch.Tensor:
"""
Does a forward pass on the given batches and returns the `loss` value in the result.
If `for_training` is `True` also applies regularization penalty.
"""
batch = nn_util.move_to_device(batch, self.cuda_device)
output_dict = self._pytorch_model(**batch)
try:
loss = output_dict["loss"]
if for_training:
loss += self.model.get_regularization_penalty()
except KeyError:
if for_training:
raise RuntimeError(
"The model you are trying to optimize does not contain a"
" 'loss' key in the output of model.forward(inputs)."
)
loss = None
return loss
def _train_epoch(self, epoch: int) -> Dict[str, float]:
"""
Trains one epoch and returns metrics.
"""
logger.info("Epoch %d/%d", epoch, self._num_epochs)
peak_cpu_usage = peak_memory_mb()
logger.info(f"Peak CPU memory usage MB: {peak_cpu_usage}")
gpu_usage = []
for gpu, memory in gpu_memory_mb().items():
gpu_usage.append((gpu, memory))
logger.info(f"GPU {gpu} memory usage MB: {memory}")
train_loss = 0.0
# Set the model to "train" mode.
self._pytorch_model.train()
num_training_batches = [math.ceil(
self.iterator.get_num_batches(train_data) / self._num_gradient_accumulation_steps
) for task, train_data in self.train_datas.items()]
assert len(set(num_training_batches)) == 1, "num_training_batches doesn't agree"
tasks = list(self.batch_group_generators.keys())
num_tasks = len(tasks)
#if isinstance(self._learning_rate_scheduler, SlantedTriangular):
# old_num_steps_per_epoch = self._learning_rate_scheduler.num_steps_per_epoch
# self._learning_rate_scheduler.num_steps_per_epoch = num_training_batches[0]
# logger.info(f"modify num_steps_per_epoch of lr scheduler from"
# f"{old_num_steps_per_epoch} to {num_training_batches}")
self._last_log = time.time()
last_save_time = time.time()
batches_this_epoch = 0
if self._batch_num_total is None:
self._batch_num_total = 0
logger.info("Training")
cumulative_batch_group_size = 0
tqdm_bar = Tqdm.tqdm(range(num_training_batches[0]))
for _ in tqdm_bar:
randperms = torch.randperm(len(tasks)).tolist()
sampled_tasks = [tasks[idx] for idx in randperms[:self._tasks_per_step]]
sampled_task_generators = [next(self.batch_group_generators[task]) for task in sampled_tasks]
batches_this_epoch += 1
self._batch_num_total += 1
batch_num_total = self._batch_num_total
self.optimizer.zero_grad()
task_metrics = self.wrapper(tasks=sampled_task_generators, train=True, meta_train=True)
losses = [list(map(lambda x: x["loss"], metrics)) for metrics in task_metrics]
LASes = [list(map(lambda x: x["metric"]["LAS"], metrics)) for metrics in task_metrics]
names = ["loss", "LAS"]
list_values = [losses, LASes]
if self.has_VIB:
KLDivs = [list(map(lambda x: x["metric"]["kl_div"], metrics)) for metrics in task_metrics]
names.append("KLDiv")
list_values.append(KLDivs)
if self.has_pos:
pos_accs = [list(map(lambda x: x["metric"].get("pos_accuracy", 0.0), metrics)) for metrics in task_metrics]
names.append("pos_acc")
list_values.append(pos_accs)
for name, values in zip(names, list_values):
self._writer.log({f"step_{name}_{task}_{i}": value
for task, task_values in zip(sampled_tasks, values)
for i, value in enumerate(task_values)},
step=self._batch_num_total)
values_inner_steps = list(map(np.mean, zip(*values)))
self._writer.log({f"step_{name}_{i}": value for i, value in
enumerate(values_inner_steps)},
step=self._batch_num_total)
if name == "loss":
train_loss += values_inner_steps[0]
batch_grad_norm = self.rescale_gradients()
# This does nothing if batch_num_total is None or you are using a
# scheduler which doesn't update per batch.
if self._learning_rate_scheduler:
self._learning_rate_scheduler.step_batch(batch_num_total)
if self._momentum_scheduler:
self._momentum_scheduler.step_batch(batch_num_total)
# variational information bottleneck / meta-learning without memorization
if self.has_VIB:
kl_loss, kl_div, kl_div2 = ContinuousVIB.get_kl_loss(self.model, sampled_task_generators)
kl_loss.backward()
self._writer.log({"kl_loss": kl_loss.detach().item(),
"kl_div": kl_div,
"kl_div2": kl_div2},
step=self._batch_num_total)
# adversarial training
if self.task_D and self.optim_D:
# D training
self.optimizer.step()
steps_per_update = self.task_D.steps_per_update
if (batch_num_total - 1) % steps_per_update == 0:
self.optim_D.zero_grad()
hidden_states, labels, masks = self.task_D.get_hidden_states(
self.model,
sampled_task_generators
)
D_loss, _, acc = self.task_D(hidden_states, labels, masks, detach=True)
D_loss.backward()
disc_grad_norm = training_util.rescale_gradients(self.task_D, self.task_D.disc_grad_norm)
self.optim_D.step()
self._writer.log({"D_loss": D_loss.detach().item(),
"D_acc": acc},
step=self._batch_num_total)
if disc_grad_norm:
self._writer.log({"D_grad_norm": disc_grad_norm.detach().item()},
step=self._batch_num_total)
# G training
hidden_states, labels, masks = self.task_D.get_hidden_states(
self.model,
sampled_task_generators
)
_, g_loss, acc = self.task_D(hidden_states, labels, masks)
if self.task_D.weight:
alpha = self.task_D.weight
else:
alpha = self.task_D.get_alpha(self._batch_num_total,
num_training_batches[0] * self._num_epochs)
G_loss = -alpha * g_loss
G_loss.backward()
gen_grad_norm = training_util.rescale_gradients(self.model, self.task_D.gen_grad_norm)
self._writer.log({"G_loss": g_loss.detach().item(), "alpha": alpha, "G_acc": acc},
step=self._batch_num_total)
if gen_grad_norm:
self._writer.log({"G_grad_norm": gen_grad_norm.detach().item()},
step=self._batch_num_total)
self.optimizer.step()
# Update moving averages
if self._moving_average is not None:
self._moving_average.apply(batch_num_total)
# Update the description with the latest metrics
metrics = training_util.get_metrics(
self.wrapper.container,
train_loss,
batches_this_epoch,
world_size=self._world_size,
cuda_device=[self.cuda_device],
)
# Updating tqdm only for the master as the trainers wouldn't have one
if self._master:
description = training_util.description_from_metrics(metrics)
tqdm_bar.set_description(description, refresh=False)
# log learning rate.
self._writer.log({"lr": self.optimizer.param_groups[0]['lr']},
step=self._batch_num_total)
# Save model if needed.
if (
self._model_save_interval is not None
and (time.time() - last_save_time > self._model_save_interval)
and self._master
):
last_save_time = time.time()
self._save_checkpoint(
"{0}.{1}".format(epoch, training_util.time_to_str(int(last_save_time)))
)
# Let all workers finish their epoch before computing
# the final statistics for the epoch.
if self._distributed:
dist.barrier()
metrics = training_util.get_metrics(
self.wrapper.container,
train_loss,
batches_this_epoch,
reset=True,
world_size=self._world_size,
cuda_device=[self.cuda_device],
)
metrics["cpu_memory_MB"] = peak_cpu_usage
for (gpu_num, memory) in gpu_usage:
metrics["gpu_" + str(gpu_num) + "_memory_MB"] = memory
return metrics
def _validation_loss(self) -> Tuple[float, int]:
"""
Computes the validation loss. Returns it and the number of batches.
"""
logger.info("Validating")
self._pytorch_model.eval()
# Replace parameter values with the shadow values from the moving averages.
if self._moving_average is not None:
self._moving_average.assign_average_value()
if self._validation_iterator is not None:
val_iterator = self._validation_iterator
else:
val_iterator = self.iterator
batches_this_epoch = 0
val_loss = 0
val_generators = {key: val_iterator(val_data, num_epochs=1, shuffle=False)
for key, val_data in self._validation_datas.items()}
num_validation_batches = {key: val_iterator.get_num_batches(val_data)
for key, val_data in self._validation_datas.items()}
val_generators_tqdm = [Tqdm.tqdm(val_generator, total=num_validation_batches[key])
for key, val_generator in val_generators.items()]
for val_generator_tqdm in val_generators_tqdm:
for batch in val_generator_tqdm:
loss = self.batch_loss(batch, for_training=False)
if loss is not None:
# You shouldn't necessarily have to compute a loss for validation, so we allow for
# `loss` to be None. We need to be careful, though - `batches_this_epoch` is
# currently only used as the divisor for the loss function, so we can safely only
# count those batches for which we actually have a loss. If this variable ever
# gets used for something else, we might need to change things around a bit.
batches_this_epoch += 1
val_loss += loss.detach().cpu().numpy()
# Update the description with the latest metrics
val_metrics = training_util.get_metrics(
self.model,
val_loss,
batches_this_epoch,
world_size=self._world_size,
cuda_device=[self.cuda_device],
)
description = training_util.description_from_metrics(val_metrics)
for val_generator_tqdm in val_generators_tqdm:
val_generator_tqdm.set_description(description, refresh=False)
# Now restore the original parameter values.
if self._moving_average is not None:
self._moving_average.restore()
return val_loss, batches_this_epoch
def train(self) -> Dict[str, Any]:
"""
Trains the supplied model with the supplied parameters.
"""
try:
epoch_counter = self._restore_checkpoint()
except RuntimeError:
traceback.print_exc()
raise ConfigurationError(
"Could not recover training from the checkpoint. Did you mean to output to "
"a different serialization directory or delete the existing serialization "
"directory?"
)
training_util.enable_gradient_clipping(self.model, self._grad_clipping)
logger.info("Beginning training.")
train_metrics: Dict[str, float] = {}
val_metrics: Dict[str, float] = {}
this_epoch_val_metric: float = None
metrics: Dict[str, Any] = {}
epochs_trained = 0
training_start_time = time.time()
metrics["best_epoch"] = self._metric_tracker.best_epoch
for key, value in self._metric_tracker.best_epoch_metrics.items():
metrics["best_validation_" + key] = value
if self._master:
self._save_checkpoint(epoch_counter - 1)
for epoch in range(epoch_counter, self._num_epochs + 1):
epoch_start_time = time.time()
train_metrics = self._train_epoch(epoch)
# get peak of memory usage
if "cpu_memory_MB" in train_metrics:
metrics["peak_cpu_memory_MB"] = max(
metrics.get("peak_cpu_memory_MB", 0), train_metrics["cpu_memory_MB"]
)
for key, value in train_metrics.items():
if key.startswith("gpu_"):
metrics["peak_" + key] = max(metrics.get("peak_" + key, 0), value)
if self._validation_datas is not None:
with torch.no_grad():
# We have a validation set, so compute all the metrics on it.
val_loss, num_batches = self._validation_loss()
# It is safe again to wait till the validation is done. This is
# important to get the metrics right.
if self._distributed:
dist.barrier()
val_metrics = training_util.get_metrics(
self.model,
val_loss,
num_batches,
reset=True,
world_size=self._world_size,
cuda_device=[self.cuda_device],
)
# Check validation metric for early stopping
this_epoch_val_metric = val_metrics[self._validation_metric]
self._metric_tracker.add_metric(this_epoch_val_metric)
if self._metric_tracker.should_stop_early():
logger.info("Ran out of patience. Stopping training.")
break
if self._master:
self._writer.log(train_metrics, step=self._batch_num_total,
epoch=epoch, prefix="train")
self._writer.log(val_metrics, step=self._batch_num_total,
epoch=epoch, prefix="val")
# Create overall metrics dict
training_elapsed_time = time.time() - training_start_time
metrics["training_duration"] = str(datetime.timedelta(seconds=training_elapsed_time))
metrics["training_start_epoch"] = epoch_counter
metrics["training_epochs"] = epochs_trained
metrics["epoch"] = epoch
for key, value in train_metrics.items():
metrics["training_" + key] = value
for key, value in val_metrics.items():
metrics["validation_" + key] = value
if self._metric_tracker.is_best_so_far():
# Update all the best_ metrics.
# (Otherwise they just stay the same as they were.)
metrics["best_epoch"] = epoch
for key, value in val_metrics.items():
metrics["best_validation_" + key] = value
self._metric_tracker.best_epoch_metrics = val_metrics
if self._serialization_dir and self._master:
dump_metrics(
os.path.join(self._serialization_dir, f"metrics_epoch_{epoch}.json"), metrics
)
# The Scheduler API is agnostic to whether your schedule requires a validation metric -
# if it doesn't, the validation metric passed here is ignored.
if self._learning_rate_scheduler:
self._learning_rate_scheduler.step(this_epoch_val_metric, epoch)
if self._momentum_scheduler:
self._momentum_scheduler.step(this_epoch_val_metric, epoch)
if self._master:
self._save_checkpoint(epoch)
# Wait for the master to finish saving the checkpoint
if self._distributed:
dist.barrier()
epoch_elapsed_time = time.time() - epoch_start_time
logger.info("Epoch duration: %s", datetime.timedelta(seconds=epoch_elapsed_time))
if epoch < self._num_epochs:
training_elapsed_time = time.time() - training_start_time
estimated_time_remaining = training_elapsed_time * (
(self._num_epochs + 1 - epoch_counter) / float(epoch - epoch_counter + 1) - 1
)
formatted_time = str(datetime.timedelta(seconds=int(estimated_time_remaining)))
logger.info("Estimated training time remaining: %s", formatted_time)
epochs_trained += 1
# Load the best model state before returning
best_model_state = self._checkpointer.best_model_state()
if best_model_state:
self.model.load_state_dict(best_model_state)
return metrics
def _save_checkpoint(self, epoch: Union[int, str]) -> None:
"""
Saves a checkpoint of the model to self._serialization_dir.
Is a no-op if self._serialization_dir is None.
# Parameters
epoch : Union[int, str], required.
The epoch of training. If the checkpoint is saved in the middle
of an epoch, the parameter is a string with the epoch and timestamp.
"""
# If moving averages are used for parameters, we save
# the moving average values into checkpoint, instead of the current values.
if self._moving_average is not None:
self._moving_average.assign_average_value()
# These are the training states we need to persist.
training_states = {
"metric_tracker": self._metric_tracker.state_dict(),
"optimizer": self.optimizer.state_dict(),
"batch_num_total": self._batch_num_total,
}
# If we have a learning rate or momentum scheduler, we should persist them too.
if self._learning_rate_scheduler is not None:
training_states["learning_rate_scheduler"] = self._learning_rate_scheduler.state_dict()
if self._momentum_scheduler is not None:
training_states["momentum_scheduler"] = self._momentum_scheduler.state_dict()
if self.task_D is not None:
training_states["task_discriminator"] = self.task_D.state_dict()
if self.optim_D is not None:
training_states["discriminator_optimizer"] = self.optim_D.state_dict()
if self._save_embedder:
model_state = self.model.state_dict()
else:
model_state = filter_state_dict(self.model.state_dict(),
lambda k, v: 'text_field_embedder' not in k)
self._checkpointer.save_checkpoint(
model_state=model_state,
epoch=epoch,
training_states=training_states,
is_best_so_far=self._metric_tracker.is_best_so_far(),
)
# Restore the original values for parameters so that training will not be affected.
if self._moving_average is not None:
self._moving_average.restore()
def _restore_checkpoint(self) -> int:
"""
Restores the model and training state from the last saved checkpoint.
This includes an epoch count and optimizer state, which is serialized separately
from model parameters. This function should only be used to continue training -
if you wish to load a model for inference/load parts of a model into a new
computation graph, you should use the native Pytorch functions:
` model.load_state_dict(torch.load("/path/to/model/weights.th"))`
If `self._serialization_dir` does not exist or does not contain any checkpointed weights,
this function will do nothing and return 0.
# Returns
epoch: int
The epoch at which to resume training, which should be one after the epoch
in the saved training state.
"""
model_state, training_state = self._checkpointer.restore_checkpoint()
if not training_state:
# No checkpoint to restore, start at 0
return 1
missing_keys, _ = self.model.load_state_dict(model_state, strict=False)
self.optimizer.load_state_dict(training_state["optimizer"])
if (
self._learning_rate_scheduler is not None
and "learning_rate_scheduler" in training_state
):
self._learning_rate_scheduler.load_state_dict(training_state["learning_rate_scheduler"])
if self._momentum_scheduler is not None and "momentum_scheduler" in training_state:
self._momentum_scheduler.load_state_dict(training_state["momentum_scheduler"])
training_util.move_optimizer_to_cuda(self.optimizer)
if self.task_D is not None and "task_discriminator" in training_state:
self.task_D.load_state_dict(training_state["task_discriminator"])
if self.optim_D is not None and "discriminator_optimizer" in training_state:
self.optim_D.load_state_dict(training_state["discriminator_optimizer"])
# Currently the `training_state` contains a serialized `MetricTracker`.
if "metric_tracker" in training_state:
self._metric_tracker.load_state_dict(training_state["metric_tracker"])
# It used to be the case that we tracked `val_metric_per_epoch`.
elif "val_metric_per_epoch" in training_state:
self._metric_tracker.clear()
self._metric_tracker.add_metrics(training_state["val_metric_per_epoch"])
# And before that we didn't track anything.
else:
self._metric_tracker.clear()
if isinstance(training_state["epoch"], int):
epoch_to_return = training_state["epoch"] + 1
else:
epoch_to_return = int(training_state["epoch"].split(".")[0]) + 1
# For older checkpoints with batch_num_total missing, default to old behavior where
# it is unchanged.
batch_num_total = training_state.get("batch_num_total")
if batch_num_total is not None:
self._batch_num_total = batch_num_total
return epoch_to_return
# Requires custom from_params.
@classmethod
def from_params( # type: ignore
cls,
params: Params,
serialization_dir: str,
recover: bool = False,
local_rank: int = 0,
) -> "MetaTrainer":
from allennlp.training.trainer import Trainer
from src.training.trainer_pieces import MetaTrainerPieces
config = dict(as_flat_dict(params.as_dict()))
pieces = MetaTrainerPieces.from_params(params, serialization_dir, recover)
model = pieces.model
serialization_dir = serialization_dir
iterator = pieces.iterator
train_datas = pieces.train_datasets
validation_datas = pieces.validation_datasets
params = pieces.params
validation_iterator = pieces.validation_iterator
patience = params.pop_int("patience", None)
validation_metric = params.pop("validation_metric", "-loss")
shuffle = params.pop_bool("shuffle", True)
num_epochs = params.pop_int("num_epochs", 20)
cuda_device = parse_cuda_device(params.pop("cuda_device", -1))
grad_norm = params.pop_float("grad_norm", None)
grad_clipping = params.pop_float("grad_clipping", None)
lr_scheduler_params = params.pop("learning_rate_scheduler", None)
momentum_scheduler_params = params.pop("momentum_scheduler", None)
check_for_gpu(cuda_device)
if cuda_device >= 0:
# Moving model to GPU here so that the optimizer state gets constructed on
# the right device.
model = model.cuda(cuda_device)
parameters = [[n, p] for n, p in model.named_parameters() if p.requires_grad]
optimizer = Optimizer.from_params(parameters, params.pop("optimizer"))
if "moving_average" in params:
moving_average = MovingAverage.from_params(
params.pop("moving_average"), parameters=parameters
)
else:
moving_average = None
if lr_scheduler_params:
lr_scheduler = LearningRateScheduler.from_params(optimizer, lr_scheduler_params)
else:
lr_scheduler = None
if momentum_scheduler_params:
momentum_scheduler = MomentumScheduler.from_params(optimizer, momentum_scheduler_params)
else:
momentum_scheduler = None
if "checkpointer" in params:
if (
"keep_serialized_model_every_num_seconds" in params
or "num_serialized_models_to_keep" in params
):
raise ConfigurationError(
"Checkpointer may be initialized either from the 'checkpointer' key or from the "
"keys 'num_serialized_models_to_keep' and 'keep_serialized_model_every_num_seconds'"
" but the passed config uses both methods."
)
checkpointer = Checkpointer.from_params(params.pop("checkpointer"))
else:
num_serialized_models_to_keep = params.pop_int("num_serialized_models_to_keep", 20)
keep_serialized_model_every_num_seconds = params.pop_int(
"keep_serialized_model_every_num_seconds", None
)
checkpointer = Checkpointer(
serialization_dir=serialization_dir,
num_serialized_models_to_keep=num_serialized_models_to_keep,
keep_serialized_model_every_num_seconds=keep_serialized_model_every_num_seconds,
)
log_grad_norm = params.pop("log_grad_norm", "total")
save_embedder = params.pop_bool("save_embedder", True)
model_save_interval = params.pop_float("model_save_interval", None)
summary_interval = params.pop_int("summary_interval", 100)
histogram_interval = params.pop_int("histogram_interval", None)
should_log_parameter_statistics = params.pop_bool("should_log_parameter_statistics", True)
should_log_learning_rate = params.pop_bool("should_log_learning_rate", False)
log_batch_size_period = params.pop_int("log_batch_size_period", None)
distributed = params.pop_bool("distributed", False)
world_size = params.pop_int("world_size", 1)
num_gradient_accumulation_steps = params.pop("num_gradient_accumulation_steps", 1)
tasks_per_step = params.pop_int("tasks_per_step", 0)
wrapper = Wrapper.from_params(
params.pop("wrapper"),
model=model,
meta_optimizer=optimizer,
)
task_discriminator_params = params.pop("task_discriminator", None)
if task_discriminator_params:
num_tasks = model.vocab.get_vocab_size("lang_labels")
task_discriminator = TaskDiscriminator.from_params(task_discriminator_params,
num_tasks=num_tasks)
if cuda_device >= 0:
task_discriminator = task_discriminator.cuda(cuda_device)
discriminator_parameters = \
[[n, p] for n, p in task_discriminator.named_parameters() if p.requires_grad]
discriminator_optimizer = Optimizer.from_params(discriminator_parameters,
params.pop("discriminator_optimizer"))
else:
task_discriminator = None
discriminator_optimizer = None
writer = None
wandb_config = params.pop("wandb", None)
if wandb_config is not None:
writer = WandBWriter(config, wrapper.container, wandb_config)
params.assert_empty(cls.__name__)
return cls(
model,
optimizer,
iterator,
train_datas,
validation_datas,
patience=patience,
validation_metric=validation_metric,
validation_iterator=validation_iterator,
shuffle=shuffle,
num_epochs=num_epochs,
serialization_dir=serialization_dir,
save_embedder=save_embedder,
cuda_device=cuda_device,
grad_norm=grad_norm,
grad_clipping=grad_clipping,
learning_rate_scheduler=lr_scheduler,
momentum_scheduler=momentum_scheduler,
checkpointer=checkpointer,
model_save_interval=model_save_interval,
summary_interval=summary_interval,
histogram_interval=histogram_interval,
should_log_parameter_statistics=should_log_parameter_statistics,
should_log_learning_rate=should_log_learning_rate,
log_batch_size_period=log_batch_size_period,
moving_average=moving_average,
distributed=distributed,
local_rank=local_rank,
world_size=world_size,
num_gradient_accumulation_steps=num_gradient_accumulation_steps,
log_grad_norm=log_grad_norm,
wrapper=wrapper,
task_discriminator=task_discriminator,
discriminator_optimizer=discriminator_optimizer,
tasks_per_step=tasks_per_step,
writer=writer,
)
class MetaTrainer(TrainerBase):
def __init__(
self,
model: Model,
optimizer: torch.optim.Optimizer,
iterator: DataIterator,
train_datasets: List[Iterable[Instance]],
validation_datasets: List[Iterable[Instance]] = None,
patience: Optional[int] = None,
validation_metric: str = "-loss",
validation_iterator: DataIterator = None,
shuffle: bool = True,
num_epochs: int = 20,
serialization_dir: Optional[str] = None,
num_serialized_models_to_keep: int = 20,
keep_serialized_model_every_num_seconds: int = None,
checkpointer: Checkpointer = None,
model_save_interval: float = None,
cuda_device: Union[int, List] = [0, 1], #int = -1,
grad_norm: Optional[float] = None,
grad_clipping: Optional[float] = None,
learning_rate_scheduler: Optional[LearningRateScheduler] = None,
momentum_scheduler: Optional[MomentumScheduler] = None,
summary_interval: int = 100,
histogram_interval: int = None,
should_log_parameter_statistics: bool = True,
should_log_learning_rate: bool = False,
log_batch_size_period: Optional[int] = None,
moving_average: Optional[MovingAverage] = None,
# meta learner parameters
meta_batches: int = 200,
inner_steps: int = 1,
meta_batch_size: int = 3,
batch_norm=True,
) -> None:
"""
A metatrainer for doing meta-learning. It just takes a list of labeled datasets
and a ``DataIterator``, and uses the supplied meta-learner to learn the weights
for your model over some fixed number of epochs. You can also pass in a validation
datasets and enable early stopping. There are many other bells and whistles as well.
Parameters
----------
model : ``Model``, required.
"""
print('[info]============================ metatrainer.init is running')
print(
'[info] cuda_device in metatrainer.init is:{}'.format(cuda_device))
# I am not calling move_to_gpu here, because if the model is
# not already on the GPU then the optimizer is going to be wrong.
super().__init__(serialization_dir, cuda_device)
self.train_data = train_datasets
self._validation_data = validation_datasets
self.model = model
self.iterator = iterator[0]
self._validation_iterator = validation_iterator
self.shuffle = shuffle
self.optimizer = optimizer
# Meta Trainer specific params
self.meta_batches = meta_batches
self.inner_steps = inner_steps
self.innerstepsize = .001
self.meta_batch_size = meta_batch_size
self.meta_step_size = .1
self.batch_norm = batch_norm
if patience is None: # no early stopping
if validation_dataset:
logger.warning(
'You provided a validation dataset but patience was set to None, '
'meaning that early stopping is disabled')
elif (not isinstance(patience, int)) or patience <= 0:
raise ConfigurationError(
'{} is an invalid value for "patience": it must be a positive integer '
'or None (if you want to disable early stopping)'.format(
patience))
# For tracking is_best_so_far and should_stop_early
self._metric_tracker = MetricTracker(patience, validation_metric)
# Get rid of + or -
self._validation_metric = validation_metric[1:]
self._num_epochs = num_epochs
if checkpointer is not None:
# We can't easily check if these parameters were passed in, so check against their default values.
# We don't check against serialization_dir since it is also used by the parent class.
if num_serialized_models_to_keep != 20 or \
keep_serialized_model_every_num_seconds is not None:
raise ConfigurationError(
"When passing a custom Checkpointer, you may not also pass in separate checkpointer "
"args 'num_serialized_models_to_keep' or 'keep_serialized_model_every_num_seconds'."
)
self._checkpointer = checkpointer
else:
self._checkpointer = Checkpointer(
serialization_dir, keep_serialized_model_every_num_seconds,
num_serialized_models_to_keep)
self._model_save_interval = model_save_interval
self._grad_norm = grad_norm
self._grad_clipping = grad_clipping
self._learning_rate_scheduler = learning_rate_scheduler
self._momentum_scheduler = momentum_scheduler
self._moving_average = moving_average
# We keep the total batch number as an instance variable because it
# is used inside a closure for the hook which logs activations in
# ``_enable_activation_logging``.
self._batch_num_total = 0
self._tensorboard = TensorboardWriter(
get_batch_num_total=lambda: self._batch_num_total,
serialization_dir=serialization_dir,
summary_interval=summary_interval,
histogram_interval=histogram_interval,
should_log_parameter_statistics=should_log_parameter_statistics,
should_log_learning_rate=should_log_learning_rate)
self._log_batch_size_period = log_batch_size_period
self._last_log = 0.0 # time of last logging
# Enable activation logging.
if histogram_interval is not None:
self._tensorboard.enable_activation_logging(self.model)
def rescale_gradients(self) -> Optional[float]:
return training_util.rescale_gradients(self.model, self._grad_norm)
# TODO check out overriding
def batch_loss(self, batch: TensorDict,
for_training: bool) -> torch.Tensor:
"""
Does a forward pass on the given batches and returns the ``loss`` value in the result.
If ``for_training`` is `True` also applies regularization penalty.
"""
if self._multiple_gpu: #len(self.cuda_device) > 1:
# print('[info] self.cuda_device is:{}'.format(self.cuda_device))
# print('[info] batch len:{}, is:{}'.format(len(batch), batch))
output_dict = training_util.data_parallel(batch, self.model,
self._cuda_devices)
else:
batch = nn_util.move_to_device(batch, self._cuda_devices[0])
output_dict = self.model(**batch)
try:
loss = output_dict["loss"]
if for_training:
loss += self.model.get_regularization_penalty()
except KeyError:
if for_training:
raise RuntimeError(
"The model you are trying to optimize does not contain a"
" 'loss' key in the output of model.forward(inputs).")
loss = None
return loss
def reptile_inner_update(self, batch_data: TensorDict) -> float:
loss = self.batch_loss(batch_data, True)
if torch.isnan(loss):
raise ValueError("nan loss encountered")
loss.backward()
temp_loss = loss.item()
self.optimizer.step()
# This only place where vary from implementation
# for param in self.model.parameters():
# TODO add innerstepsize
# param.data -= self.innerstepsize * param.grad.data
return temp_loss
def reptile_outer_update(self, train_generators: List[Iterable],
iteration: int, num_gpus: int):
# https://github.com/farbodtm/reptile-pytorch/blob/master/reptile.py
weights_before = deepcopy(self.model.state_dict())
self.optimizer.zero_grad()
random.shuffle(train_generators)
new_weights = []
total_loss = 0.0
# for batch in train_generators[0]:
# print('[info]batch is:{}'.format(batch))
task_wrap = Tqdm.tqdm(zip(train_generators[0], train_generators[1],
train_generators[2]),
total=1)
# , train_generators[3], train_generators[4]), \
for i, batch_group in enumerate(task_wrap):
if not i:
for k in range(self.meta_batch_size): # tasks per batch
total_loss += self.reptile_inner_update(batch_group[k][0])
new_weights.append(deepcopy(self.model.state_dict()))
self.model.load_state_dict({
name: weights_before[name]
for name in weights_before
})
else:
break
weights_after = {
name: new_weights[0][name] / float(self.meta_batch_size)
for name in new_weights[0]
}
for i in range(1, self.meta_batch_size):
for name in new_weights[i]:
weights_after[name] += new_weights[i][name] / float(
self.meta_batch_size)
#They used self.step_size of 1.0 in some of their outer.
outerstepsize = self.meta_step_size * (
1 - iteration / self.meta_batches) # linear schedule
self.model.load_state_dict({
name: weights_before[name] +
(weights_after[name] - weights_before[name]) * outerstepsize
for name in weights_before
})
return total_loss / self.meta_batch_size
def _train_epoch(self, epoch: int) -> Dict[str, float]:
"""
Trains on one epoch. Differs from base trainer in that
it utilizes
"""
logger.info("Epoch %d/%d", epoch, self._num_epochs - 1)
peak_cpu_usage = peak_memory_mb()
logger.info(f"Peak CPU memory usage MB: {peak_cpu_usage}")
gpu_usage = []
for gpu, memory in gpu_memory_mb().items():
gpu_usage.append((gpu, memory))
logger.info(f"GPU {gpu} memory usage MB: {memory}")
train_loss = 0.0
# Set the model to "train" mode.
self.model.train()
num_gpus = len(self._cuda_devices)
raw_generators = []
# fix max number of batches
self._last_log = time.time()
last_save_time = time.time()
batches_this_epoch = 0
if self._batch_num_total is None:
self._batch_num_total = 0
histogram_parameters = set(
self.model.get_parameters_for_histogram_tensorboard_logging())
logger.info("Training")
cumulative_batch_size = 0
for i in range(0, self.meta_batches):
train_generators = []
for i, train_info in enumerate(self.train_data):
raw_train_generator = self.iterator(train_info,
num_epochs=1,
shuffle=self.shuffle)
train_generators.append(
lazy_groups_of(raw_train_generator, num_gpus))
loss_batch = self.reptile_outer_update(train_generators, i,
num_gpus)
# TODO figure out if is important
train_loss = loss_batch
print('[info] train_loss is:{}'.format(train_loss))
# TODO figure out BATCH NORM MAML https://openreview.net/pdf?id=HygBZnRctX
if self.batch_norm:
batch_grad_norm = self.rescale_gradients()
# This does nothing if batch_num_total is None or you are using a
# scheduler which doesn't update per batch.
# TODO investigate learning rate scheduling for meta learning
#if self._learning_rate_scheduler:
#self._learning_rate_scheduler.step_batch(batch_num_total)
#if self._momentum_scheduler:
#self._momentum_scheduler.step_batch(batch_num_total)
if self._tensorboard.should_log_histograms_this_batch():
# get the magnitude of parameter updates for logging
# We need a copy of current parameters to compute magnitude of updates,
# and copy them to CPU so large models won't go OOM on the GPU.
param_updates = {
name: param.detach().cpu().clone()
for name, param in self.model.named_parameters()
}
self.optimizer.step()
for name, param in self.model.named_parameters():
param_updates[name].sub_(param.detach().cpu())
update_norm = torch.norm(param_updates[name].view(-1, ))
param_norm = torch.norm(param.view(-1, )).cpu()
self._tensorboard.add_train_scalar(
"gradient_update/" + name,
update_norm / (param_norm + 1e-7))
else:
self.optimizer.step()
# Update moving averages
if self._moving_average is not None:
self._moving_average.apply(batch_num_total)
# Update the description with the latest metrics
metrics = training_util.get_metrics(self.model, train_loss,
batches_this_epoch)
description = training_util.description_from_metrics(metrics)
# Log parameter values to Tensorboard
if self._tensorboard.should_log_this_batch():
self._tensorboard.log_parameter_and_gradient_statistics(
self.model, batch_grad_norm)
self._tensorboard.log_learning_rates(self.model,
self.optimizer)
self._tensorboard.add_train_scalar("loss/loss_train",
metrics["loss"])
self._tensorboard.log_metrics(
{"epoch_metrics/" + k: v
for k, v in metrics.items()})
if self._tensorboard.should_log_histograms_this_batch():
self._tensorboard.log_histograms(self.model,
histogram_parameters)
if self._log_batch_size_period:
cur_batch = sum([
training_util.get_batch_size(batch)
for batch in batch_group
])
cumulative_batch_size += cur_batch
if (batches_this_epoch - 1) % self._log_batch_size_period == 0:
average = cumulative_batch_size / batches_this_epoch
logger.info(
f"current batch size: {cur_batch} mean batch size: {average}"
)
self._tensorboard.add_train_scalar("current_batch_size",
cur_batch)
self._tensorboard.add_train_scalar("mean_batch_size",
average)
# Save model if needed.
if self._model_save_interval is not None and (
time.time() - last_save_time > self._model_save_interval):
last_save_time = time.time()
self._save_checkpoint('{0}.{1}'.format(
epoch, training_util.time_to_str(int(last_save_time))))
metrics = training_util.get_metrics(self.model,
train_loss,
batches_this_epoch,
reset=True)
metrics['cpu_memory_MB'] = peak_cpu_usage
for (gpu_num, memory) in gpu_usage:
metrics['gpu_' + str(gpu_num) + '_memory_MB'] = memory
return metrics
def _validation_loss(self) -> Tuple[float, int]:
"""
Computes the validation loss. Returns it and the number of batches.
"""
logger.info("Validating")
self.model.eval()
# Replace parameter values with the shadow values from the moving averages.
if self._moving_average is not None:
self._moving_average.assign_average_value()
if self._validation_iterator is not None:
val_iterator = self._validation_iterator[0]
else:
val_iterator = self.iterator
num_gpus = len(self._cuda_devices)
valid_generators = []
for i, valid_info in enumerate(self._validation_data):
raw_val_generator = self.iterator(valid_info,
num_epochs=1,
shuffle=self.shuffle)
valid_generators.append(lazy_groups_of(raw_val_generator,
num_gpus))
num_validation_batches = min(
map(
lambda i: math.ceil(
val_iterator.get_num_batches(self._validation_data[i]) /
num_gpus), range(self.meta_batch_size)))
val_generator_tqdm = Tqdm.tqdm(zip(valid_generators[0],
valid_generators[1],
valid_generators[2]),
total=num_validation_batches)
print("val gene called")
batches_this_epoch = 0
val_loss = 0
for i, batch_group in enumerate(val_generator_tqdm):
for k in range(self.meta_batch_size): # tasks per batch
loss = self.batch_loss(batch_group[k][0], for_training=False)
if loss is not None:
# You shouldn't necessarily have to compute a loss for validation, so we allow for
# `loss` to be None. We need to be careful, though - `batches_this_epoch` is
# currently only used as the divisor for the loss function, so we can safely only
# count those batches for which we actually have a loss. If this variable ever
# gets used for something else, we might need to change things around a bit.
batches_this_epoch += 1
val_loss += loss.detach().cpu().numpy()
# Update the description with the latest metrics
val_metrics = training_util.get_metrics(self.model, val_loss,
batches_this_epoch)
description = training_util.description_from_metrics(val_metrics)
val_generator_tqdm.set_description(description, refresh=False)
# Now restore the original parameter values.
if self._moving_average is not None:
self._moving_average.restore()
return val_loss, batches_this_epoch
def train(self) -> Dict[str, Any]:
"""
Trains the supplied model with the supplied parameters.
"""
try:
epoch_counter = self._restore_checkpoint()
except RuntimeError:
traceback.print_exc()
raise ConfigurationError(
"Could not recover training from the checkpoint. Did you mean to output to "
"a different serialization directory or delete the existing serialization "
"directory?")
training_util.enable_gradient_clipping(self.model, self._grad_clipping)
logger.info("Beginning training.")
train_metrics: Dict[str, float] = {}
val_metrics: Dict[str, float] = {}
this_epoch_val_metric: float = None
metrics: Dict[str, Any] = {}
epochs_trained = 0
training_start_time = time.time()
metrics['best_epoch'] = self._metric_tracker.best_epoch
for key, value in self._metric_tracker.best_epoch_metrics.items():
metrics["best_validation_" + key] = value
for epoch in range(epoch_counter, self._num_epochs):
epoch_start_time = time.time()
train_metrics = self._train_epoch(epoch)
# get peak of memory usage
if 'cpu_memory_MB' in train_metrics:
metrics['peak_cpu_memory_MB'] = max(
metrics.get('peak_cpu_memory_MB', 0),
train_metrics['cpu_memory_MB'])
for key, value in train_metrics.items():
if key.startswith('gpu_'):
metrics["peak_" + key] = max(metrics.get("peak_" + key, 0),
value)
if self._validation_data is not None:
# We have a validation set, so compute all the metrics on it.
val_loss, num_batches = self._validation_loss()
val_metrics = training_util.get_metrics(self.model,
val_loss,
num_batches,
reset=True)
# Check validation metric for early stopping
this_epoch_val_metric = val_metrics[self._validation_metric]
self._metric_tracker.add_metric(this_epoch_val_metric)
if self._metric_tracker.should_stop_early():
logger.info("Ran out of patience. Stopping training.")
break
self._tensorboard.log_metrics(
train_metrics,
val_metrics=val_metrics,
log_to_console=True,
epoch=epoch + 1) # +1 because tensorboard doesn't like 0
# Create overall metrics dict
training_elapsed_time = time.time() - training_start_time
metrics["training_duration"] = str(
datetime.timedelta(seconds=training_elapsed_time))
metrics["training_start_epoch"] = epoch_counter
metrics["training_epochs"] = epochs_trained
metrics["epoch"] = epoch
for key, value in train_metrics.items():
metrics["training_" + key] = value
for key, value in val_metrics.items():
metrics["validation_" + key] = value
if self._metric_tracker.is_best_so_far():
# Update all the best_ metrics.
# (Otherwise they just stay the same as they were.)
metrics['best_epoch'] = epoch
for key, value in val_metrics.items():
metrics["best_validation_" + key] = value
self._metric_tracker.best_epoch_metrics = val_metrics
if self._serialization_dir:
dump_metrics(
os.path.join(self._serialization_dir,
f'metrics_epoch_{epoch}.json'), metrics)
# The Scheduler API is agnostic to whether your schedule requires a validation metric -
# if it doesn't, the validation metric passed here is ignored.
if self._learning_rate_scheduler:
self._learning_rate_scheduler.step(this_epoch_val_metric,
epoch)
if self._momentum_scheduler:
self._momentum_scheduler.step(this_epoch_val_metric, epoch)
self._save_checkpoint(epoch)
epoch_elapsed_time = time.time() - epoch_start_time
logger.info("Epoch duration: %s",
datetime.timedelta(seconds=epoch_elapsed_time))
if epoch < self._num_epochs - 1:
training_elapsed_time = time.time() - training_start_time
estimated_time_remaining = training_elapsed_time * \
((self._num_epochs - epoch_counter) / float(epoch - epoch_counter + 1) - 1)
formatted_time = str(
datetime.timedelta(seconds=int(estimated_time_remaining)))
logger.info("Estimated training time remaining: %s",
formatted_time)
epochs_trained += 1
# Load the best model state before returning
best_model_state = self._checkpointer.best_model_state()
if best_model_state:
self.model.load_state_dict(best_model_state)
return metrics
def _save_checkpoint(self, epoch: Union[int, str]) -> None:
"""
Saves a checkpoint of the model to self._serialization_dir.
Is a no-op if self._serialization_dir is None.
Parameters
----------
epoch : Union[int, str], required.
The epoch of training. If the checkpoint is saved in the middle
of an epoch, the parameter is a string with the epoch and timestamp.
"""
# If moving averages are used for parameters, we save
# the moving average values into checkpoint, instead of the current values.
if self._moving_average is not None:
self._moving_average.assign_average_value()
# These are the training states we need to persist.
training_states = {
"metric_tracker": self._metric_tracker.state_dict(),
"optimizer": self.optimizer.state_dict(),
"batch_num_total": self._batch_num_total
}
# If we have a learning rate or momentum scheduler, we should persist them too.
if self._learning_rate_scheduler is not None:
training_states[
"learning_rate_scheduler"] = self._learning_rate_scheduler.state_dict(
)
if self._momentum_scheduler is not None:
training_states[
"momentum_scheduler"] = self._momentum_scheduler.state_dict()
self._checkpointer.save_checkpoint(
model_state=self.model.state_dict(),
epoch=epoch,
training_states=training_states,
is_best_so_far=self._metric_tracker.is_best_so_far())
# Restore the original values for parameters so that training will not be affected.
if self._moving_average is not None:
self._moving_average.restore()
def _restore_checkpoint(self) -> int:
"""
Restores the model and training state from the last saved checkpoint.
This includes an epoch count and optimizer state, which is serialized separately
from model parameters. This function should only be used to continue training -
if you wish to load a model for inference/load parts of a model into a new
computation graph, you should use the native Pytorch functions:
`` model.load_state_dict(torch.load("/path/to/model/weights.th"))``
If ``self._serialization_dir`` does not exist or does not contain any checkpointed weights,
this function will do nothing and return 0.
Returns
-------
epoch: int
The epoch at which to resume training, which should be one after the epoch
in the saved training state.
"""
model_state, training_state = self._checkpointer.restore_checkpoint()
if not training_state:
# No checkpoint to restore, start at 0
return 0
self.model.load_state_dict(model_state)
self.optimizer.load_state_dict(training_state["optimizer"])
if self._learning_rate_scheduler is not None and "learning_rate_scheduler" in training_state:
self._learning_rate_scheduler.load_state_dict(
training_state["learning_rate_scheduler"])
if self._momentum_scheduler is not None and "momentum_scheduler" in training_state:
self._momentum_scheduler.load_state_dict(
training_state["momentum_scheduler"])
training_util.move_optimizer_to_cuda(self.optimizer)
# Currently the ``training_state`` contains a serialized ``MetricTracker``.
if "metric_tracker" in training_state:
self._metric_tracker.load_state_dict(
training_state["metric_tracker"])
# It used to be the case that we tracked ``val_metric_per_epoch``.
elif "val_metric_per_epoch" in training_state:
self._metric_tracker.clear()
self._metric_tracker.add_metrics(
training_state["val_metric_per_epoch"])
# And before that we didn't track anything.
else:
self._metric_tracker.clear()
if isinstance(training_state["epoch"], int):
epoch_to_return = training_state["epoch"] + 1
else:
epoch_to_return = int(training_state["epoch"].split('.')[0]) + 1
# For older checkpoints with batch_num_total missing, default to old behavior where
# it is unchanged.
batch_num_total = training_state.get('batch_num_total')
if batch_num_total is not None:
self._batch_num_total = batch_num_total
return epoch_to_return
# Requires custom from_params.
@classmethod
def from_params(
cls, # type: ignore
params: Params,
serialization_dir: str,
recover: bool = False,
cache_directory: str = None,
cache_prefix: str = None) -> 'Trainer':
# datasets = meta_dataset_from_params(params, cache_directory=cache_directory, cache_prefix=cache_prefix)
# model = Model.from_params(vocab=vocab, params=params.pop("model"))
# iterator = DataIterator.from_params(params.pop("iterator"))
# iterator.index_with(model.vocab)
pieces = MetaTrainerPieces.from_params(params, serialization_dir,
recover, cache_directory,
cache_prefix)
model = pieces.model
iterator = pieces.iterator,
# params=pieces.params,
train_data = pieces.train_dataset
validation_data = pieces.validation_dataset
validation_iterator = pieces.validation_iterator
params = pieces.params
# pylint: disable=arguments-differ
patience = params.pop_int("patience", None)
validation_metric = params.pop("validation_metric", "-loss")
shuffle = params.pop_bool("shuffle", True)
num_epochs = params.pop_int("num_epochs", 20)
cuda_device = parse_cuda_device(params.pop("cuda_device", [0, 1]))
grad_norm = params.pop_float("grad_norm", None)
grad_clipping = params.pop_float("grad_clipping", None)
lr_scheduler_params = params.pop("learning_rate_scheduler", None)
momentum_scheduler_params = params.pop("momentum_scheduler", None)
if isinstance(cuda_device, list):
model_device = cuda_device[0]
else:
model_device = cuda_device
if model_device >= 0:
# Moving model to GPU here so that the optimizer state gets constructed on
# the right device.
model = model.cuda(model_device)
parameters = [[n, p] for n, p in model.named_parameters()
if p.requires_grad]
optimizer = Optimizer.from_params(parameters, params.pop("optimizer"))
if "moving_average" in params:
moving_average = MovingAverage.from_params(
params.pop("moving_average"), parameters=parameters)
else:
moving_average = None
if lr_scheduler_params:
lr_scheduler = LearningRateScheduler.from_params(
optimizer, lr_scheduler_params)
else:
lr_scheduler = None
if momentum_scheduler_params:
momentum_scheduler = MomentumScheduler.from_params(
optimizer, momentum_scheduler_params)
else:
momentum_scheduler = None
if 'checkpointer' in params:
if 'keep_serialized_model_every_num_seconds' in params or \
'num_serialized_models_to_keep' in params:
raise ConfigurationError(
"Checkpointer may be initialized either from the 'checkpointer' key or from the "
"keys 'num_serialized_models_to_keep' and 'keep_serialized_model_every_num_seconds'"
" but the passed config uses both methods.")
checkpointer = Checkpointer.from_params(params.pop("checkpointer"))
else:
num_serialized_models_to_keep = params.pop_int(
"num_serialized_models_to_keep", 20)
keep_serialized_model_every_num_seconds = params.pop_int(
"keep_serialized_model_every_num_seconds", None)
checkpointer = Checkpointer(
serialization_dir=serialization_dir,
num_serialized_models_to_keep=num_serialized_models_to_keep,
keep_serialized_model_every_num_seconds=
keep_serialized_model_every_num_seconds)
model_save_interval = params.pop_float("model_save_interval", None)
summary_interval = params.pop_int("summary_interval", 100)
histogram_interval = params.pop_int("histogram_interval", None)
should_log_parameter_statistics = params.pop_bool(
"should_log_parameter_statistics", True)
should_log_learning_rate = params.pop_bool("should_log_learning_rate",
False)
log_batch_size_period = params.pop_int("log_batch_size_period", None)
print('[info] cuda_device in metatrainer.from_param is:{}'.format(
cuda_device))
params.assert_empty(cls.__name__)
return cls(
model,
optimizer,
iterator,
train_data,
validation_data,
patience=patience,
validation_metric=validation_metric,
validation_iterator=validation_iterator,
shuffle=shuffle,
num_epochs=num_epochs,
serialization_dir=serialization_dir,
cuda_device=cuda_device,
grad_norm=grad_norm,
grad_clipping=grad_clipping,
learning_rate_scheduler=lr_scheduler,
momentum_scheduler=momentum_scheduler,
checkpointer=checkpointer,
model_save_interval=model_save_interval,
summary_interval=summary_interval,
histogram_interval=histogram_interval,
should_log_parameter_statistics=should_log_parameter_statistics,
should_log_learning_rate=should_log_learning_rate,
log_batch_size_period=log_batch_size_period,
moving_average=moving_average,
# distributed=distributed,
# rank=local_rank,
# world_size=world_size,
# num_gradient_accumulation_steps=num_gradient_accumulation_steps,
)
class Trainer(TrainerBase):
def __init__(
self,
model: Model,
optimizer: torch.optim.Optimizer,
# scheduler 根据epoch调整学习率
scheduler: torch.optim.lr_scheduler,
iterator: DataIterator,
train_dataset: Iterable[Instance],
validation_dataset: Optional[Iterable[Instance]] = None,
patience: Optional[int] = None,
validation_metric: str = "-loss",
validation_iterator: DataIterator = None,
shuffle: bool = True,
num_epochs: int = 20,
accumulated_batch_count: int = 1,
serialization_dir: Optional[str] = None,
num_serialized_models_to_keep: int = 20,
keep_serialized_model_every_num_seconds: int = None,
checkpointer: Checkpointer = None,
model_save_interval: float = None,
cuda_device: Union[int, List] = -1,
grad_norm: Optional[float] = None,
grad_clipping: Optional[float] = None,
learning_rate_scheduler: Optional[LearningRateScheduler] = None,
momentum_scheduler: Optional[MomentumScheduler] = None,
summary_interval: int = 100,
histogram_interval: int = None,
should_log_parameter_statistics: bool = True,
should_log_learning_rate: bool = False,
log_batch_size_period: Optional[int] = None,
moving_average: Optional[MovingAverage] = None,
cold_step_count: int = 0,
cold_lr: float = 1e-3,
cuda_verbose_step=None,
) -> None:
"""
A trainer for doing supervised learning. It just takes a labeled dataset
and a ``DataIterator``, and uses the supplied ``Optimizer`` to learn the weights
for your model over some fixed number of epochs. You can also pass in a validation
dataset and enable early stopping. There are many other bells and whistles as well.
Parameters
----------
model : ``Model``, required.
An AllenNLP model to be optimized. Pytorch Modules can also be optimized if
their ``forward`` method returns a dictionary with a "loss" key, containing a
scalar tensor representing the loss function to be optimized.
If you are training your model using GPUs, your model should already be
on the correct device. (If you use `Trainer.from_params` this will be
handled for you.)
optimizer : ``torch.nn.Optimizer``, required.
An instance of a Pytorch Optimizer, instantiated with the parameters of the
model to be optimized.
iterator : ``DataIterator``, required.
A method for iterating over a ``Dataset``, yielding padded indexed batches.
train_dataset : ``Dataset``, required.
A ``Dataset`` to train on. The dataset should have already been indexed.
validation_dataset : ``Dataset``, optional, (default = None).
A ``Dataset`` to evaluate on. The dataset should have already been indexed.
patience : Optional[int] > 0, optional (default=None)
Number of epochs to be patient before early stopping: the training is stopped
after ``patience`` epochs with no improvement. If given, it must be ``> 0``.
If None, early stopping is disabled.
validation_metric : str, optional (default="loss")
Validation metric to measure for whether to stop training using patience
and whether to serialize an ``is_best`` model each epoch. The metric name
must be prepended with either "+" or "-", which specifies whether the metric
is an increasing or decreasing function.
validation_iterator : ``DataIterator``, optional (default=None)
An iterator to use for the validation set. If ``None``, then
use the training `iterator`.
shuffle: ``bool``, optional (default=True)
Whether to shuffle the instances in the iterator or not.
num_epochs : int, optional (default = 20)
Number of training epochs.
serialization_dir : str, optional (default=None)
Path to directory for saving and loading model files. Models will not be saved if
this parameter is not passed.
num_serialized_models_to_keep : ``int``, optional (default=20)
Number of previous model checkpoints to retain. Default is to keep 20 checkpoints.
A value of None or -1 means all checkpoints will be kept.
keep_serialized_model_every_num_seconds : ``int``, optional (default=None)
If num_serialized_models_to_keep is not None, then occasionally it's useful to
save models at a given interval in addition to the last num_serialized_models_to_keep.
To do so, specify keep_serialized_model_every_num_seconds as the number of seconds
between permanently saved checkpoints. Note that this option is only used if
num_serialized_models_to_keep is not None, otherwise all checkpoints are kept.
checkpointer : ``Checkpointer``, optional (default=None)
An instance of class Checkpointer to use instead of the default. If a checkpointer is specified,
the arguments num_serialized_models_to_keep and keep_serialized_model_every_num_seconds should
not be specified. The caller is responsible for initializing the checkpointer so that it is
consistent with serialization_dir.
model_save_interval : ``float``, optional (default=None)
If provided, then serialize models every ``model_save_interval``
seconds within single epochs. In all cases, models are also saved
at the end of every epoch if ``serialization_dir`` is provided.
cuda_device : ``Union[int, List[int]]``, optional (default = -1)
An integer or list of integers specifying the CUDA device(s) to use. If -1, the CPU is used.
grad_norm : ``float``, optional, (default = None).
If provided, gradient norms will be rescaled to have a maximum of this value.
grad_clipping : ``float``, optional (default = ``None``).
If provided, gradients will be clipped `during the backward pass` to have an (absolute)
maximum of this value. If you are getting ``NaNs`` in your gradients during training
that are not solved by using ``grad_norm``, you may need this.
learning_rate_scheduler : ``LearningRateScheduler``, optional (default = None)
If specified, the learning rate will be decayed with respect to
this schedule at the end of each epoch (or batch, if the scheduler implements
the ``step_batch`` method). If you use :class:`torch.optim.lr_scheduler.ReduceLROnPlateau`,
this will use the ``validation_metric`` provided to determine if learning has plateaued.
To support updating the learning rate on every batch, this can optionally implement
``step_batch(batch_num_total)`` which updates the learning rate given the batch number.
momentum_scheduler : ``MomentumScheduler``, optional (default = None)
If specified, the momentum will be updated at the end of each batch or epoch
according to the schedule.
summary_interval: ``int``, optional, (default = 100)
Number of batches between logging scalars to tensorboard
histogram_interval : ``int``, optional, (default = ``None``)
If not None, then log histograms to tensorboard every ``histogram_interval`` batches.
When this parameter is specified, the following additional logging is enabled:
* Histograms of model parameters
* The ratio of parameter update norm to parameter norm
* Histogram of layer activations
We log histograms of the parameters returned by
``model.get_parameters_for_histogram_tensorboard_logging``.
The layer activations are logged for any modules in the ``Model`` that have
the attribute ``should_log_activations`` set to ``True``. Logging
histograms requires a number of GPU-CPU copies during training and is typically
slow, so we recommend logging histograms relatively infrequently.
Note: only Modules that return tensors, tuples of tensors or dicts
with tensors as values currently support activation logging.
should_log_parameter_statistics : ``bool``, optional, (default = True)
Whether to send parameter statistics (mean and standard deviation
of parameters and gradients) to tensorboard.
should_log_learning_rate : ``bool``, optional, (default = False)
Whether to send parameter specific learning rate to tensorboard.
log_batch_size_period : ``int``, optional, (default = ``None``)
If defined, how often to log the average batch size.
moving_average: ``MovingAverage``, optional, (default = None)
If provided, we will maintain moving averages for all parameters. During training, we
employ a shadow variable for each parameter, which maintains the moving average. During
evaluation, we backup the original parameters and assign the moving averages to corresponding
parameters. Be careful that when saving the checkpoint, we will save the moving averages of
parameters. This is necessary because we want the saved model to perform as well as the validated
model if we load it later. But this may cause problems if you restart the training from checkpoint.
"""
super().__init__(serialization_dir, cuda_device)
# I am not calling move_to_gpu here, because if the model is
# not already on the GPU then the optimizer is going to be wrong.
self.model = model
self.iterator = iterator
self._validation_iterator = validation_iterator
self.shuffle = shuffle
self.optimizer = optimizer
self.scheduler = scheduler
self.train_data = train_dataset
self._validation_data = validation_dataset
self.accumulated_batch_count = accumulated_batch_count
self.cold_step_count = cold_step_count
self.cold_lr = cold_lr
self.cuda_verbose_step = cuda_verbose_step
if patience is None: # no early stopping
if validation_dataset:
logger.warning(
"You provided a validation dataset but patience was set to None, "
"meaning that early stopping is disabled")
elif (not isinstance(patience, int)) or patience <= 0:
raise ConfigurationError(
'{} is an invalid value for "patience": it must be a positive integer '
"or None (if you want to disable early stopping)".format(
patience))
# For tracking is_best_so_far and should_stop_early It mimics the PyTorch state_dict / load_state_dict
self._metric_tracker = MetricTracker(patience, validation_metric)
# Get rid of + or - 取绝对值
self._validation_metric = validation_metric[1:]
# 默认20个 epoch
self._num_epochs = num_epochs
if checkpointer is not None:
# We can't easily check if these parameters were passed in, so check against their default values.
# We don't check against serialization_dir since it is also used by the parent class.
# 默认应该有20个序列化模型 且不应该设置只保存模型一段间隔
if num_serialized_models_to_keep != 20 \
or keep_serialized_model_every_num_seconds is not None:
raise ConfigurationError(
"When passing a custom Checkpointer, you may not also pass in separate checkpointer "
"args 'num_serialized_models_to_keep' or 'keep_serialized_model_every_num_seconds'."
)
self._checkpointer = checkpointer
else:
self._checkpointer = Checkpointer(
serialization_dir,
keep_serialized_model_every_num_seconds,
num_serialized_models_to_keep,
)
self._model_save_interval = model_save_interval
self._grad_norm = grad_norm
self._grad_clipping = grad_clipping
self._learning_rate_scheduler = learning_rate_scheduler
self._momentum_scheduler = momentum_scheduler
self._moving_average = moving_average
# We keep the total batch number as an instance variable because it
# is used inside a closure for the hook which logs activations in
# ``_enable_activation_logging``.
self._batch_num_total = 0
self._tensorboard = TensorboardWriter(
get_batch_num_total=lambda: self._batch_num_total,
serialization_dir=serialization_dir,
summary_interval=summary_interval,
histogram_interval=histogram_interval,
should_log_parameter_statistics=should_log_parameter_statistics,
should_log_learning_rate=should_log_learning_rate,
)
self._log_batch_size_period = log_batch_size_period
self._last_log = 0.0 # time of last logging
# Enable activation logging.
if histogram_interval is not None:
self._tensorboard.enable_activation_logging(self.model)
# 目的是如果多个特征取值范围不一样 梯度下降收敛会慢 这里grad norm 默认是none 即no-op 无操作返回
def rescale_gradients(self) -> Optional[float]:
return training_util.rescale_gradients(self.model, self._grad_norm)
# 计算batch的loss
def batch_loss(self, batch_group: List[TensorDict],
for_training: bool) -> torch.Tensor:
"""
Does a forward pass on the given batches and returns the ``loss`` value in the result.
If ``for_training`` is `True` also applies regularization penalty.
正则化惩罚是要降低不重要特征的影响力，避免过拟合
被train epoch 和evaluate epoch复用
"""
# 处理并行， 但在gector中默认是单gpu
if self._multiple_gpu:
output_dict = training_util.data_parallel(batch_group, self.model,
self._cuda_devices)
else:
assert len(batch_group) == 1
batch = batch_group[0]
batch = nn_util.move_to_device(batch, self._cuda_devices[0])
# 前向传播
output_dict = self.model(**batch)
# 通过正则化惩罚项来计算loss
try:
loss = output_dict["loss"]
if for_training:
loss += self.model.get_regularization_penalty()
except KeyError:
if for_training:
raise RuntimeError(
"The model you are trying to optimize does not contain a"
" 'loss' key in the output of model.forward(inputs).")
loss = None
return loss
def _train_epoch(self, epoch: int) -> Dict[str, float]:
"""
Trains one epoch and returns metrics.
"""
logger.info("Epoch %d/%d", epoch, self._num_epochs - 1)
peak_cpu_usage = peak_memory_mb()
logger.info(f"Peak CPU memory usage MB: {peak_cpu_usage}")
gpu_usage = []
for gpu, memory in gpu_memory_mb().items():
gpu_usage.append((gpu, memory))
logger.info(f"GPU {gpu} memory usage MB: {memory}")
train_loss = 0.0
# Set the model to "train" mode.
self.model.train()
num_gpus = len(self._cuda_devices)
# Get tqdm for the training batches
# 使训练数据可迭代
raw_train_generator = self.iterator(self.train_data,
num_epochs=1,
shuffle=self.shuffle)
# 将可迭代的单实例批处理到list中
train_generator = lazy_groups_of(raw_train_generator, num_gpus)
# 向上取整 获取batch数 （总batch/gpu数）
num_training_batches = math.ceil(
self.iterator.get_num_batches(self.train_data) / num_gpus)
# 默认的accumulated batch count 为4，此处是求accumulate的尾巴
residue = num_training_batches % self.accumulated_batch_count
self._last_log = time.time()
last_save_time = time.time()
batches_this_epoch = 0
if self._batch_num_total is None:
self._batch_num_total = 0
histogram_parameters = set(
self.model.get_parameters_for_histogram_tensorboard_logging())
logger.info("Training")
# 训练进度条
train_generator_tqdm = Tqdm.tqdm(train_generator,
total=num_training_batches)
cumulative_batch_size = 0
# 梯度清零 常规操作
self.optimizer.zero_grad()
# 开始训练
for batch_group in train_generator_tqdm:
batches_this_epoch += 1
self._batch_num_total += 1
batch_num_total = self._batch_num_total
# 一个batch为accumulated_batch_count个iteration，梯度累积
iter_len = self.accumulated_batch_count \
if batches_this_epoch <= (num_training_batches - residue) else residue
if self.cuda_verbose_step is not None and batch_num_total % self.cuda_verbose_step == 0:
print(
f'Before forward pass - Cuda memory allocated: {torch.cuda.memory_allocated() / 1e9}'
)
print(
f'Before forward pass - Cuda memory cached: {torch.cuda.memory_cached() / 1e9}'
)
try: # 平均loss
loss = self.batch_loss(batch_group,
for_training=True) / iter_len
except RuntimeError as e:
print(e)
for x in batch_group:
all_words = [len(y['words']) for y in x['metadata']]
print(f"Total sents: {len(all_words)}. "
f"Min {min(all_words)}. Max {max(all_words)}")
for elem in ['labels', 'd_tags']:
tt = x[elem]
print(
f"{elem} shape {list(tt.shape)} and min {tt.min().item()} and {tt.max().item()}"
)
for elem in ["bert", "mask", "bert-offsets"]:
tt = x['tokens'][elem]
print(
f"{elem} shape {list(tt.shape)} and min {tt.min().item()} and {tt.max().item()}"
)
raise e
if self.cuda_verbose_step is not None and batch_num_total % self.cuda_verbose_step == 0:
print(
f'After forward pass - Cuda memory allocated: {torch.cuda.memory_allocated() / 1e9}'
)
print(
f'After forward pass - Cuda memory cached: {torch.cuda.memory_cached() / 1e9}'
)
if torch.isnan(loss):
raise ValueError("nan loss encountered")
# 反向传播
loss.backward()
if self.cuda_verbose_step is not None and batch_num_total % self.cuda_verbose_step == 0:
print(
f'After backprop - Cuda memory allocated: {torch.cuda.memory_allocated() / 1e9}'
)
print(
f'After backprop - Cuda memory cached: {torch.cuda.memory_cached() / 1e9}'
)
# 计算loss
train_loss += loss.item() * iter_len
# 删除两个变量
del batch_group, loss
# pytorch 训练时无用的临时变量可能会越来越多，导致 out of memory ，可以使用下面语句来清理这些不需要的变量。
torch.cuda.empty_cache()
if self.cuda_verbose_step is not None and batch_num_total % self.cuda_verbose_step == 0:
print(
f'After collecting garbage - Cuda memory allocated: {torch.cuda.memory_allocated() / 1e9}'
)
print(
f'After collecting garbage - Cuda memory cached: {torch.cuda.memory_cached() / 1e9}'
)
# 正则化梯度
batch_grad_norm = self.rescale_gradients()
# This does nothing if batch_num_total is None or you are using a
# scheduler which doesn't update per batch.
# lr会在epoch变大的同时予以调整，一般是逐渐变小
# momentum 动量 防止损失函数陷入局部极小值，跳出鞍点
if self._learning_rate_scheduler:
self._learning_rate_scheduler.step_batch(batch_num_total)
if self._momentum_scheduler:
self._momentum_scheduler.step_batch(batch_num_total)
if self._tensorboard.should_log_histograms_this_batch():
# copy参数 防止爆内存
# get the magnitude of parameter updates for logging
# We need a copy of current parameters to compute magnitude of updates,
# and copy them to CPU so large models won't go OOM on the GPU.
param_updates = {
name: param.detach().cpu().clone()
for name, param in self.model.named_parameters()
}
if batches_this_epoch % self.accumulated_batch_count == 0 or \
batches_this_epoch == num_training_batches:
# 自动计算梯度 optimizer.step()
self.optimizer.step()
self.optimizer.zero_grad()
for name, param in self.model.named_parameters():
param_updates[name].sub_(param.detach().cpu())
# 求l1范数
update_norm = torch.norm(param_updates[name].view(-1))
param_norm = torch.norm(param.view(-1)).cpu()
self._tensorboard.add_train_scalar(
"gradient_update/" + name,
update_norm / (param_norm + 1e-7))
else:
if batches_this_epoch % self.accumulated_batch_count == 0 or \
batches_this_epoch == num_training_batches:
self.optimizer.step()
self.optimizer.zero_grad()
# Update moving averages 在adam或SGD优化中为了平衡模型更新速度一般设置滑动平均来提高模型在测试数据上的健壮性
if self._moving_average is not None:
self._moving_average.apply(batch_num_total)
# Update the description with the latest metrics
metrics = training_util.get_metrics(self.model, train_loss,
batches_this_epoch)
description = training_util.description_from_metrics(metrics)
train_generator_tqdm.set_description(description, refresh=False)
# Log parameter values to Tensorboard
if self._tensorboard.should_log_this_batch():
self._tensorboard.log_parameter_and_gradient_statistics(
self.model, batch_grad_norm)
self._tensorboard.log_learning_rates(self.model,
self.optimizer)
self._tensorboard.add_train_scalar("loss/loss_train",
metrics["loss"])
self._tensorboard.log_metrics(
{"epoch_metrics/" + k: v
for k, v in metrics.items()})
if self._tensorboard.should_log_histograms_this_batch():
self._tensorboard.log_histograms(self.model,
histogram_parameters)
if self._log_batch_size_period:
cur_batch = sum([
training_util.get_batch_size(batch)
for batch in batch_group
])
cumulative_batch_size += cur_batch
if (batches_this_epoch - 1) % self._log_batch_size_period == 0:
average = cumulative_batch_size / batches_this_epoch
logger.info(
f"current batch size: {cur_batch} mean batch size: {average}"
)
self._tensorboard.add_train_scalar("current_batch_size",
cur_batch)
self._tensorboard.add_train_scalar("mean_batch_size",
average)
# Save model if needed.
if self._model_save_interval is not None and (
time.time() - last_save_time > self._model_save_interval):
last_save_time = time.time()
self._save_checkpoint("{0}.{1}".format(
epoch, training_util.time_to_str(int(last_save_time))))
metrics = training_util.get_metrics(self.model,
train_loss,
batches_this_epoch,
reset=True)
metrics["cpu_memory_MB"] = peak_cpu_usage
for (gpu_num, memory) in gpu_usage:
metrics["gpu_" + str(gpu_num) + "_memory_MB"] = memory
return metrics
def _validation_loss(self) -> Tuple[float, int]:
"""
Computes the validation loss. Returns it and the number of batches.
"""
logger.info("Validating")
# 与model.train()类似
self.model.eval()
# Replace parameter values with the shadow values from the moving averages.
if self._moving_average is not None:
self._moving_average.assign_average_value()
if self._validation_iterator is not None:
val_iterator = self._validation_iterator
else:
val_iterator = self.iterator
num_gpus = len(self._cuda_devices)
# 与上面train代码的流程相似
raw_val_generator = val_iterator(self._validation_data,
num_epochs=1,
shuffle=False)
val_generator = lazy_groups_of(raw_val_generator, num_gpus)
num_validation_batches = math.ceil(
val_iterator.get_num_batches(self._validation_data) / num_gpus)
val_generator_tqdm = Tqdm.tqdm(val_generator,
total=num_validation_batches)
batches_this_epoch = 0
val_loss = 0
for batch_group in val_generator_tqdm:
loss = self.batch_loss(batch_group, for_training=False)
if loss is not None:
# You shouldn't necessarily have to compute a loss for validation, so we allow for
# `loss` to be None. We need to be careful, though - `batches_this_epoch` is
# currently only used as the divisor for the loss function, so we can safely only
# count those batches for which we actually have a loss. If this variable ever
# gets used for something else, we might need to change things around a bit.
batches_this_epoch += 1
val_loss += loss.detach().cpu().numpy(
) # loss.detach().cpu().numpy()为了取出loss值
# Update the description with the latest metrics
val_metrics = training_util.get_metrics(self.model, val_loss,
batches_this_epoch)
description = training_util.description_from_metrics(val_metrics)
val_generator_tqdm.set_description(description, refresh=False)
# Now restore the original parameter values.
if self._moving_average is not None:
self._moving_average.restore()
return val_loss, batches_this_epoch
@property
def train(self) -> Dict[str, Any]:
"""
Trains the supplied model with the supplied parameters.
相关的metric字典记录的信息都在训练时产生的json文件中
"""
try:
epoch_counter = self._restore_checkpoint()
except RuntimeError:
traceback.print_exc()
raise ConfigurationError(
"Could not recover training from the checkpoint. Did you mean to output to "
"a different serialization directory or delete the existing serialization "
"directory?")
# 梯度剪裁 防止梯度爆炸跳过最优解
training_util.enable_gradient_clipping(self.model, self._grad_clipping)
logger.info("Beginning training.")
train_metrics: Dict[str, float] = {}
val_metrics: Dict[str, float] = {}
this_epoch_val_metric: float = None
metrics: Dict[str, Any] = {}
epochs_trained = 0
# ------训练开始-------
training_start_time = time.time()
# cold_step_count为只训练最后一层线性层的epoch数
# 训练阶段一，二
# 在前 cold_step_count个epoch
# 不需要训练原来的预训练模型，之后需要训练
# 阶段三直接训练预训练模型参数， 因为预训练模型的参数过多
# 同时需要注意，在cold step阶段也要使用cold lr，
# 此阶段结束后，使用base lr
if self.cold_step_count > 0:
# 1e-5
base_lr = self.optimizer.param_groups[0]['lr']
for param_group in self.optimizer.param_groups:
# 1e-3
param_group['lr'] = self.cold_lr
self.model.text_field_embedder._token_embedders[
'bert'].set_weights(freeze=True)
metrics["best_epoch"] = self._metric_tracker.best_epoch
for key, value in self._metric_tracker.best_epoch_metrics.items():
metrics["best_validation_" + key] = value
# epoch_counter = 0 if restore_checkpoint is none else continue training
for epoch in range(epoch_counter, self._num_epochs):
# 恢复正常
if epoch == self.cold_step_count and epoch != 0:
for param_group in self.optimizer.param_groups:
param_group['lr'] = base_lr
self.model.text_field_embedder._token_embedders[
'bert'].set_weights(freeze=False)
# --开始当前epoch--
epoch_start_time = time.time()
# **训练**
train_metrics = self._train_epoch(epoch)
# get peak of memory usage
if "cpu_memory_MB" in train_metrics:
metrics["peak_cpu_memory_MB"] = max(
metrics.get("peak_cpu_memory_MB", 0),
train_metrics["cpu_memory_MB"])
for key, value in train_metrics.items():
if key.startswith("gpu_"):
metrics["peak_" + key] = max(metrics.get("peak_" + key, 0),
value)
# clear cache before validation
torch.cuda.empty_cache()
# evaluate的函数说了， 不是一定需要进行验证，所以这里要做判断
if self._validation_data is not None:
# 常规操作，验证时不计算梯度，不更新参数
with torch.no_grad():
# We have a validation set, so compute all the metrics on it.
val_loss, num_batches = self._validation_loss()
val_metrics = training_util.get_metrics(self.model,
val_loss,
num_batches,
reset=True)
# Check validation metric for early stopping
# 获取性能指标--loss
this_epoch_val_metric = val_metrics[
self._validation_metric]
self._metric_tracker.add_metric(this_epoch_val_metric)
if self._metric_tracker.should_stop_early():
# 这就是为什么有的时候ckpt不足epoch个数，是因为patience耗光
# patience是配合早停机制的阈值，patience次在验证集的性能下降时，停止训练
logger.info("Ran out of patience. Stopping training.")
break
self._tensorboard.log_metrics(
train_metrics,
val_metrics=val_metrics,
log_to_console=True,
epoch=epoch + 1) # +1 because tensorboard doesn't like 0
# Create overall metrics dict
# **epoch结束**
training_elapsed_time = time.time() - training_start_time
metrics["training_duration"] = str(
datetime.timedelta(seconds=training_elapsed_time))
metrics["training_start_epoch"] = epoch_counter
metrics["training_epochs"] = epochs_trained
metrics["epoch"] = epoch
# 将train， evaluate阶段的metric记录都汇总
for key, value in train_metrics.items():
metrics["training_" + key] = value
for key, value in val_metrics.items():
metrics["validation_" + key] = value
# if self.cold_step_count <= epoch:
# step操作
self.scheduler.step(metrics['validation_loss'])
# 这些更新都在119服务器的pretraingectors目录下
if self._metric_tracker.is_best_so_far():
# Update all the best_ metrics.
# (Otherwise they just stay the same as they were.)
metrics["best_epoch"] = epoch
for key, value in val_metrics.items():
metrics["best_validation_" + key] = value
self._metric_tracker.best_epoch_metrics = val_metrics
# 以json形式存储metrics
if self._serialization_dir:
dump_metrics(
os.path.join(self._serialization_dir,
f"metrics_epoch_{epoch}.json"), metrics)
# The Scheduler API is agnostic to whether your schedule requires a validation metric -
# if it doesn't, the validation metric passed here is ignored.
if self._learning_rate_scheduler:
# step操作
self._learning_rate_scheduler.step(this_epoch_val_metric,
epoch)
if self._momentum_scheduler:
# step操作
self._momentum_scheduler.step(this_epoch_val_metric, epoch)
# 保存ckpt
self._save_checkpoint(epoch)
epoch_elapsed_time = time.time() - epoch_start_time
logger.info("Epoch duration: %s",
datetime.timedelta(seconds=epoch_elapsed_time))
if epoch < self._num_epochs - 1:
training_elapsed_time = time.time() - training_start_time
estimated_time_remaining = training_elapsed_time * (
(self._num_epochs - epoch_counter) /
float(epoch - epoch_counter + 1) - 1)
formatted_time = str(
datetime.timedelta(seconds=int(estimated_time_remaining)))
logger.info("Estimated training time remaining: %s",
formatted_time)
# 一个epoch结束
epochs_trained += 1
# make sure pending events are flushed to disk and files are closed properly
# self._tensorboard.close()
# Load the best model state before returning
best_model_state = self._checkpointer.best_model_state()
if best_model_state:
self.model.load_state_dict(best_model_state)
return metrics
def _save_checkpoint(self, epoch: Union[int, str]) -> None:
"""
主要是存储model状态和train的状态
Saves a checkpoint of the model to self._serialization_dir.
Is a no-op if self._serialization_dir is None.
Parameters
----------
epoch : Union[int, str], required.
The epoch of training. If the checkpoint is saved in the middle
of an epoch, the parameter is a string with the epoch and timestamp.
"""
# If moving averages are used for parameters, we save
# the moving average values into checkpoint, instead of the current values.
if self._moving_average is not None:
self._moving_average.assign_average_value()
# These are the training states we need to persist.
training_states = {
"metric_tracker": self._metric_tracker.state_dict(),
"optimizer": self.optimizer.state_dict(),
"batch_num_total": self._batch_num_total,
}
# If we have a learning rate or momentum scheduler, we should persist them too.
if self._learning_rate_scheduler is not None:
training_states[
"learning_rate_scheduler"] = self._learning_rate_scheduler.state_dict(
)
if self._momentum_scheduler is not None:
training_states[
"momentum_scheduler"] = self._momentum_scheduler.state_dict()
# save checkpoint
self._checkpointer.save_checkpoint(
model_state=self.model.state_dict(),
epoch=epoch,
training_states=training_states,
is_best_so_far=self._metric_tracker.is_best_so_far(),
)
# Restore the original values for parameters so that training will not be affected.
if self._moving_average is not None:
self._moving_average.restore()
def _restore_checkpoint(self) -> int:
"""
恢复上一个检查点的模型和训练状态
Restores the model and training state from the last saved checkpoint.
This includes an epoch count and optimizer state, which is serialized separately
from model parameters. This function should only be used to continue training -
if you wish to load a model for inference/load parts of a model into a new
computation graph, you should use the native Pytorch functions:
`` model.load_state_dict(torch.load("/path/to/model/weights.th"))``
If ``self._serialization_dir`` does not exist or does not contain any checkpointed weights,
this function will do nothing and return 0.
Returns
-------
epoch: int
The epoch at which to resume training, which should be one after the epoch
in the saved training state.
"""
model_state, training_state = self._checkpointer.restore_checkpoint()
if not training_state:
# No checkpoint to restore, start at 0
return 0
self.model.load_state_dict(model_state)
self.optimizer.load_state_dict(training_state["optimizer"])
if self._learning_rate_scheduler is not None \
and "learning_rate_scheduler" in training_state:
self._learning_rate_scheduler.load_state_dict(
training_state["learning_rate_scheduler"])
if self._momentum_scheduler is not None and "momentum_scheduler" in training_state:
self._momentum_scheduler.load_state_dict(
training_state["momentum_scheduler"])
training_util.move_optimizer_to_cuda(self.optimizer)
# Currently the ``training_state`` contains a serialized ``MetricTracker``.
if "metric_tracker" in training_state:
self._metric_tracker.load_state_dict(
training_state["metric_tracker"])
# It used to be the case that we tracked ``val_metric_per_epoch``.
elif "val_metric_per_epoch" in training_state:
self._metric_tracker.clear()
self._metric_tracker.add_metrics(
training_state["val_metric_per_epoch"])
# And before that we didn't track anything.
else:
self._metric_tracker.clear()
if isinstance(training_state["epoch"], int):
epoch_to_return = training_state["epoch"] + 1
else:
epoch_to_return = int(training_state["epoch"].split(".")[0]) + 1
# For older checkpoints with batch_num_total missing, default to old behavior where
# it is unchanged.
batch_num_total = training_state.get("batch_num_total")
if batch_num_total is not None:
self._batch_num_total = batch_num_total
return epoch_to_return
# Requires custom from_params.
# 通过返回和构造函数一样的参数名字，来完成实例化
@classmethod
def from_params( # type: ignore
cls,
model: Model,
serialization_dir: str,
iterator: DataIterator,
train_data: Iterable[Instance],
validation_data: Optional[Iterable[Instance]],
params: Params,
validation_iterator: DataIterator = None,
) -> "Trainer":
# 与python 字典的pop一样 返回值为对应key的value
patience = params.pop_int("patience", None)
validation_metric = params.pop("validation_metric", "-loss")
shuffle = params.pop_bool("shuffle", True)
num_epochs = params.pop_int("num_epochs", 20)
cuda_device = parse_cuda_device(params.pop("cuda_device", -1))
grad_norm = params.pop_float("grad_norm", None)
grad_clipping = params.pop_float("grad_clipping", None)
lr_scheduler_params = params.pop("learning_rate_scheduler", None)
momentum_scheduler_params = params.pop("momentum_scheduler", None)
# 单gpu
if isinstance(cuda_device, list):
model_device = cuda_device[0]
else:
model_device = cuda_device
if model_device >= 0:
# Moving model to GPU here so that the optimizer state gets constructed on
# the right device.
model = model.cuda(model_device)
parameters = [[n, p] for n, p in model.named_parameters()
if p.requires_grad]
optimizer = Optimizer.from_params(parameters, params.pop("optimizer"))
if "moving_average" in params:
moving_average = MovingAverage.from_params(
params.pop("moving_average"), parameters=parameters)
else:
moving_average = None
if lr_scheduler_params:
lr_scheduler = LearningRateScheduler.from_params(
optimizer, lr_scheduler_params)
else:
lr_scheduler = None
if momentum_scheduler_params:
momentum_scheduler = MomentumScheduler.from_params(
optimizer, momentum_scheduler_params)
else:
momentum_scheduler = None
if "checkpointer" in params:
if "keep_serialized_model_every_num_seconds" in params \
or "num_serialized_models_to_keep" in params:
raise ConfigurationError(
"Checkpointer may be initialized either from the 'checkpointer' key or from the "
"keys 'num_serialized_models_to_keep' and 'keep_serialized_model_every_num_seconds'"
" but the passed config uses both methods.")
checkpointer = Checkpointer.from_params(params.pop("checkpointer"))
else:
num_serialized_models_to_keep = params.pop_int(
"num_serialized_models_to_keep", 20)
keep_serialized_model_every_num_seconds = params.pop_int(
"keep_serialized_model_every_num_seconds", None)
checkpointer = Checkpointer(
serialization_dir=serialization_dir,
num_serialized_models_to_keep=num_serialized_models_to_keep,
keep_serialized_model_every_num_seconds=
keep_serialized_model_every_num_seconds,
)
model_save_interval = params.pop_float("model_save_interval", None)
summary_interval = params.pop_int("summary_interval", 100)
histogram_interval = params.pop_int("histogram_interval", None)
should_log_parameter_statistics = params.pop_bool(
"should_log_parameter_statistics", True)
should_log_learning_rate = params.pop_bool("should_log_learning_rate",
False)
log_batch_size_period = params.pop_int("log_batch_size_period", None)
params.assert_empty(cls.__name__)
return cls(
model,
optimizer,
iterator,
train_data,
validation_data,
patience=patience,
validation_metric=validation_metric,
validation_iterator=validation_iterator,
shuffle=shuffle,
num_epochs=num_epochs,
serialization_dir=serialization_dir,
cuda_device=cuda_device,
grad_norm=grad_norm,
grad_clipping=grad_clipping,
learning_rate_scheduler=lr_scheduler,
momentum_scheduler=momentum_scheduler,
checkpointer=checkpointer,
model_save_interval=model_save_interval,
summary_interval=summary_interval,
histogram_interval=histogram_interval,
should_log_parameter_statistics=should_log_parameter_statistics,
should_log_learning_rate=should_log_learning_rate,
log_batch_size_period=log_batch_size_period,
moving_average=moving_average,
)
class PtTrainer(TrainerBase):
def __init__(
self,
model: Model,
optimizer: torch.optim.Optimizer,
iterator: DataIterator,
train_dataset: Iterable[Instance],
validation_dataset: Optional[Iterable[Instance]] = None,
max_src_len: int = None,
patience: Optional[int] = None,
validation_metric: str = "-loss",
validation_iterator: DataIterator = None,
batch_size: int = 1,
shuffle: bool = True,
num_epochs: int = 20,
serialization_dir: Optional[str] = None,
num_serialized_models_to_keep: int = 20,
keep_serialized_model_every_num_seconds: int = None,
checkpointer: Checkpointer = None,
model_save_interval: float = None,
cuda_device: Union[int, List] = -1,
grad_norm: Optional[float] = None,
grad_clipping: Optional[float] = None,
learning_rate_scheduler: Optional[LearningRateScheduler] = None,
momentum_scheduler: Optional[MomentumScheduler] = None,
summary_interval: int = 100,
histogram_interval: int = None,
should_log_parameter_statistics: bool = True,
should_log_learning_rate: bool = False,
log_batch_size_period: Optional[int] = None,
moving_average: Optional[MovingAverage] = None) -> None:
super().__init__(serialization_dir, cuda_device)
# I am not calling move_to_gpu here, because if the model is
# not already on the GPU then the optimizer is going to be wrong.
self.model = model
self.iterator = iterator
self._validation_iterator = validation_iterator
self.shuffle = shuffle
self.optimizer = optimizer
self.train_data = train_dataset
self._validation_data = validation_dataset
self.max_src_len = max_src_len
self.batch_size = batch_size
# For tracking is_best_so_far and should_stop_early
self._metric_tracker = MetricTracker(patience, validation_metric)
# Get rid of + or -
self._validation_metric = validation_metric[1:]
self._num_epochs = num_epochs
if checkpointer is not None:
# We can't easily check if these parameters were passed in, so check against their default values.
# We don't check against serialization_dir since it is also used by the parent class.
if num_serialized_models_to_keep != 20 or \
keep_serialized_model_every_num_seconds is not None:
raise ConfigurationError(
"When passing a custom Checkpointer, you may not also pass in separate checkpointer "
"args 'num_serialized_models_to_keep' or 'keep_serialized_model_every_num_seconds'."
)
self._checkpointer = checkpointer
else:
self._checkpointer = Checkpointer(
serialization_dir, keep_serialized_model_every_num_seconds,
num_serialized_models_to_keep)
self._model_save_interval = model_save_interval
self._grad_norm = grad_norm
self._grad_clipping = grad_clipping
self._learning_rate_scheduler = learning_rate_scheduler
self._momentum_scheduler = momentum_scheduler
self._moving_average = moving_average
# We keep the total batch number as an instance variable because it
# is used inside a closure for the hook which logs activations in
# ``_enable_activation_logging``.
self._batch_num_total = 0
self._tensorboard = TensorboardWriter(
get_batch_num_total=lambda: self._batch_num_total,
serialization_dir=serialization_dir,
summary_interval=summary_interval,
histogram_interval=histogram_interval,
should_log_parameter_statistics=should_log_parameter_statistics,
should_log_learning_rate=should_log_learning_rate)
self._log_batch_size_period = log_batch_size_period
self._last_log = 0.0 # time of last logging
# Enable activation logging.
if histogram_interval is not None:
self._tensorboard.enable_activation_logging(self.model)
def rescale_gradients(self) -> Optional[float]:
return training_util.rescale_gradients(self.model, self._grad_norm)
def batch_loss(self, batch_group: List[TensorDict],
for_training: bool) -> torch.Tensor:
"""
Does a forward pass on the given batches and returns the ``loss`` value in the result.
If ``for_training`` is `True` also applies regularization penalty.
"""
if self._multiple_gpu:
output_dict = training_util.data_parallel(batch_group, self.model,
self._cuda_devices)
else:
assert len(batch_group) == 1
batch = batch_group[0]
batch = nn_util.move_to_device(batch, self._cuda_devices[0])
output_dict = self.model(**batch)
try:
loss = output_dict["loss"]
if for_training:
loss += self.model.get_regularization_penalty()
except KeyError:
if for_training:
raise RuntimeError(
"The model you are trying to optimize does not contain a"
" 'loss' key in the output of model.forward(inputs).")
loss = None
return loss
def _train_epoch(self, epoch: int) -> Dict[str, float]:
train_loss = 0.0
self.model.train()
num_gpus = len(self._cuda_devices)
if getattr(self, "train_dataset", None) is None:
self.train_dataset = DMDataSet(data=self.train_data[0],
batch_size=self.batch_size,
num_gpus=num_gpus,
shuffle=True)
self.train_dataset.set_epoch(epoch)
num_training_batches = math.ceil(
len(self.train_dataset) / self.batch_size / num_gpus)
self._last_log = time.time()
batches_this_epoch = 0
if self._batch_num_total is None:
self._batch_num_total = 0
logger.info("Training")
train_generator_tqdm = Tqdm.tqdm(self.train_dataset,
total=num_training_batches)
for batch_group in train_generator_tqdm:
# print('gpu num: ', len(batch_group))
# print('batch_size: ', len(batch_group[0]["source_tokens"]["tokens"]))
# gpu_data = batch_group[0]
# src_data = gpu_data["source_tokens"]["tokens"]
# tgt_data = gpu_data["target_tokens"]["tokens"]
# for sdata, tdata in zip(src_data, tgt_data):
# s = ''.join([self.model.vocab.get_token_from_index(x, "source_tokens") if x != 0 else '' for x in sdata.numpy()])
# t = ''.join([self.model.vocab.get_token_from_index(x, "target_tokens") if x != 0 else '' for x in tdata.numpy()])
# print(s)
# print(t)
batches_this_epoch += 1
self._batch_num_total += 1
batch_num_total = self._batch_num_total
self.optimizer.zero_grad()
loss = self.batch_loss(batch_group, for_training=True)
if torch.isnan(loss):
raise ValueError("nan loss encountered")
loss.backward()
train_loss += loss.item()
batch_grad_norm = self.rescale_gradients()
if self._learning_rate_scheduler:
self._learning_rate_scheduler.step_batch(batch_num_total)
if self._momentum_scheduler:
self._momentum_scheduler.step_batch(batch_num_total)
self.optimizer.step()
# Update moving averages
if self._moving_average is not None:
self._moving_average.apply(batch_num_total)
# Update the description with the latest metrics
metrics = training_util.get_metrics(self.model, train_loss,
batches_this_epoch)
description = training_util.description_from_metrics(metrics)
train_generator_tqdm.set_description(description, refresh=False)
# Log parameter values to Tensorboard
if self._tensorboard.should_log_this_batch():
self._tensorboard.log_learning_rates(self.model,
self.optimizer)
self._tensorboard.add_train_scalar("loss/loss_train",
metrics["loss"])
self._tensorboard.log_metrics(
{"epoch_metrics/" + k: v
for k, v in metrics.items()})
metrics = training_util.get_metrics(self.model,
train_loss,
batches_this_epoch,
reset=True)
return metrics
def _validation_loss(self) -> Tuple[float, int]:
logger.info("Validating")
self.model.eval()
# Replace parameter values with the shadow values from the moving averages.
if self._moving_average is not None:
self._moving_average.assign_average_value()
if self._validation_iterator is not None:
val_iterator = self._validation_iterator
else:
val_iterator = self.iterator
num_gpus = len(self._cuda_devices)
if getattr(self, "val_dataset", None) is None:
self.val_dataset = DMDataSet(data=self._validation_data[0],
batch_size=self.batch_size,
num_gpus=num_gpus,
shuffle=False)
num_validation_batches = math.ceil(
len(self.val_dataset) / self.batch_size / num_gpus)
val_generator_tqdm = Tqdm.tqdm(self.val_dataset,
total=num_validation_batches)
batches_this_epoch = 0
val_loss = 0
for batch_group in val_generator_tqdm:
loss = self.batch_loss(batch_group, for_training=False)
if loss is not None:
batches_this_epoch += 1
val_loss += loss.detach().cpu().numpy()
# Update the description with the latest metrics
val_metrics = training_util.get_metrics(self.model, val_loss,
batches_this_epoch)
description = training_util.description_from_metrics(val_metrics)
val_generator_tqdm.set_description(description, refresh=False)
# Now restore the original parameter values.
if self._moving_average is not None:
self._moving_average.restore()
return val_loss, batches_this_epoch
def train(self) -> Dict[str, Any]:
"""
Trains the supplied model with the supplied parameters.
"""
try:
epoch_counter = self._restore_checkpoint()
except RuntimeError:
traceback.print_exc()
raise ConfigurationError(
"Could not recover training from the checkpoint. Did you mean to output to "
"a different serialization directory or delete the existing serialization "
"directory?")
training_util.enable_gradient_clipping(self.model, self._grad_clipping)
logger.info("Beginning training.")
train_metrics: Dict[str, float] = {}
val_metrics: Dict[str, float] = {}
this_epoch_val_metric: float = None
metrics: Dict[str, Any] = {}
epochs_trained = 0
training_start_time = time.time()
metrics['best_epoch'] = self._metric_tracker.best_epoch
for key, value in self._metric_tracker.best_epoch_metrics.items():
metrics["best_validation_" + key] = value
for epoch in range(epoch_counter, self._num_epochs):
epoch_start_time = time.time()
train_metrics = self._train_epoch(epoch)
if self._validation_data is not None:
with torch.no_grad():
# We have a validation set, so compute all the metrics on it.
val_loss, num_batches = self._validation_loss()
val_metrics = training_util.get_metrics(self.model,
val_loss,
num_batches,
reset=True)
# Check validation metric for early stopping
this_epoch_val_metric = val_metrics[
self._validation_metric]
self._metric_tracker.add_metric(this_epoch_val_metric)
if self._metric_tracker.should_stop_early():
logger.info("Ran out of patience. Stopping training.")
break
self._tensorboard.log_metrics(
train_metrics,
val_metrics=val_metrics,
log_to_console=True,
epoch=epoch + 1) # +1 because tensorboard doesn't like 0
# Create overall metrics dict
training_elapsed_time = time.time() - training_start_time
metrics["training_duration"] = str(
datetime.timedelta(seconds=training_elapsed_time))
metrics["training_start_epoch"] = epoch_counter
metrics["training_epochs"] = epochs_trained
metrics["epoch"] = epoch
for key, value in train_metrics.items():
metrics["training_" + key] = value
for key, value in val_metrics.items():
metrics["validation_" + key] = value
if self._metric_tracker.is_best_so_far():
# Update all the best_ metrics.
# (Otherwise they just stay the same as they were.)
metrics['best_epoch'] = epoch
for key, value in val_metrics.items():
metrics["best_validation_" + key] = value
self._metric_tracker.best_epoch_metrics = val_metrics
if self._serialization_dir:
dump_metrics(
os.path.join(self._serialization_dir,
f'metrics_epoch_{epoch}.json'), metrics)
# The Scheduler API is agnostic to whether your schedule requires a validation metric -
# if it doesn't, the validation metric passed here is ignored.
if self._learning_rate_scheduler:
self._learning_rate_scheduler.step(this_epoch_val_metric,
epoch)
if self._momentum_scheduler:
self._momentum_scheduler.step(this_epoch_val_metric, epoch)
self._save_checkpoint(epoch)
epoch_elapsed_time = time.time() - epoch_start_time
logger.info("Epoch duration: %s",
datetime.timedelta(seconds=epoch_elapsed_time))
if epoch < self._num_epochs - 1:
training_elapsed_time = time.time() - training_start_time
estimated_time_remaining = training_elapsed_time * \
((self._num_epochs - epoch_counter) / float(epoch - epoch_counter + 1) - 1)
formatted_time = str(
datetime.timedelta(seconds=int(estimated_time_remaining)))
logger.info("Estimated training time remaining: %s",
formatted_time)
epochs_trained += 1
# make sure pending events are flushed to disk and files are closed properly
self._tensorboard.close()
# Load the best model state before returning
best_model_state = self._checkpointer.best_model_state()
if best_model_state:
self.model.load_state_dict(best_model_state)
return metrics
def _save_checkpoint(self, epoch: Union[int, str]) -> None:
"""
Saves a checkpoint of the model to self._serialization_dir.
Is a no-op if self._serialization_dir is None.
Parameters
----------
epoch : Union[int, str], required.
The epoch of training. If the checkpoint is saved in the middle
of an epoch, the parameter is a string with the epoch and timestamp.
"""
# If moving averages are used for parameters, we save
# the moving average values into checkpoint, instead of the current values.
if self._moving_average is not None:
self._moving_average.assign_average_value()
# These are the training states we need to persist.
training_states = {
"metric_tracker": self._metric_tracker.state_dict(),
"optimizer": self.optimizer.state_dict(),
"batch_num_total": self._batch_num_total
}
# If we have a learning rate or momentum scheduler, we should persist them too.
if self._learning_rate_scheduler is not None:
training_states[
"learning_rate_scheduler"] = self._learning_rate_scheduler.state_dict(
)
if self._momentum_scheduler is not None:
training_states[
"momentum_scheduler"] = self._momentum_scheduler.state_dict()
self._checkpointer.save_checkpoint(
model_state=self.model.state_dict(),
epoch=epoch,
training_states=training_states,
is_best_so_far=self._metric_tracker.is_best_so_far())
# Restore the original values for parameters so that training will not be affected.
if self._moving_average is not None:
self._moving_average.restore()
def _restore_checkpoint(self) -> int:
"""
Restores the model and training state from the last saved checkpoint.
This includes an epoch count and optimizer state, which is serialized separately
from model parameters. This function should only be used to continue training -
if you wish to load a model for inference/load parts of a model into a new
computation graph, you should use the native Pytorch functions:
`` model.load_state_dict(torch.load("/path/to/model/weights.th"))``
If ``self._serialization_dir`` does not exist or does not contain any checkpointed weights,
this function will do nothing and return 0.
Returns
-------
epoch: int
The epoch at which to resume training, which should be one after the epoch
in the saved training state.
"""
model_state, training_state = self._checkpointer.restore_checkpoint()
if not training_state:
# No checkpoint to restore, start at 0
return 0
self.model.load_state_dict(model_state)
self.optimizer.load_state_dict(training_state["optimizer"])
if self._learning_rate_scheduler is not None and "learning_rate_scheduler" in training_state:
self._learning_rate_scheduler.load_state_dict(
training_state["learning_rate_scheduler"])
if self._momentum_scheduler is not None and "momentum_scheduler" in training_state:
self._momentum_scheduler.load_state_dict(
training_state["momentum_scheduler"])
training_util.move_optimizer_to_cuda(self.optimizer)
# Currently the ``training_state`` contains a serialized ``MetricTracker``.
if "metric_tracker" in training_state:
self._metric_tracker.load_state_dict(
training_state["metric_tracker"])
# It used to be the case that we tracked ``val_metric_per_epoch``.
elif "val_metric_per_epoch" in training_state:
self._metric_tracker.clear()
self._metric_tracker.add_metrics(
training_state["val_metric_per_epoch"])
# And before that we didn't track anything.
else:
self._metric_tracker.clear()
if isinstance(training_state["epoch"], int):
epoch_to_return = training_state["epoch"] + 1
else:
epoch_to_return = int(training_state["epoch"].split('.')[0]) + 1
# For older checkpoints with batch_num_total missing, default to old behavior where
# it is unchanged.
batch_num_total = training_state.get('batch_num_total')
if batch_num_total is not None:
self._batch_num_total = batch_num_total
return epoch_to_return
# Requires custom from_params.
@classmethod
def from_params(cls,
params: Params,
serialization_dir: str,
recover: bool = False,
cache_directory: str = None,
cache_prefix: str = None) -> 'PtTrainer':
max_src_len = params.dataset_reader.get('max_src_len', None)
all_datasets = training_util.datasets_from_params(
params, cache_directory, cache_prefix)
datasets_for_vocab_creation = set(
params.pop("datasets_for_vocab_creation", all_datasets))
for dataset in datasets_for_vocab_creation:
if dataset not in all_datasets:
raise ConfigurationError(
f"invalid 'dataset_for_vocab_creation' {dataset}")
logger.info(
"From dataset instances, %s will be considered for vocabulary creation.",
", ".join(datasets_for_vocab_creation))
if recover and os.path.exists(
os.path.join(serialization_dir, "vocabulary")):
vocab = Vocabulary.from_files(
os.path.join(serialization_dir, "vocabulary"))
params.pop("vocabulary", {})
else:
vocab = Vocabulary.from_params(params.pop(
"vocabulary", {}), (instance
for key, dataset in all_datasets.items()
if key in datasets_for_vocab_creation
for instance in dataset))
model = Model.from_params(vocab=vocab, params=params.pop('model'))
# If vocab extension is ON for training, embedding extension should also be
# done. If vocab and embeddings are already in sync, it would be a no-op.
model.extend_embedder_vocab()
# Initializing the model can have side effect of expanding the vocabulary
vocab.save_to_files(os.path.join(serialization_dir, "vocabulary"))
iterator = DataIterator.from_params(params.pop("iterator"))
iterator.index_with(model.vocab)
validation_iterator_params = params.pop("validation_iterator", None)
if validation_iterator_params:
validation_iterator = DataIterator.from_params(
validation_iterator_params)
validation_iterator.index_with(model.vocab)
else:
validation_iterator = None
train_data = all_datasets['train']
validation_data = all_datasets.get('validation')
test_data = all_datasets.get('test')
trainer_params = params.pop("trainer")
no_grad_regexes = trainer_params.pop("no_grad", ())
for name, parameter in model.named_parameters():
if any(re.search(regex, name) for regex in no_grad_regexes):
parameter.requires_grad_(False)
frozen_parameter_names, tunable_parameter_names = \
get_frozen_and_tunable_parameter_names(model)
logger.info("Following parameters are Frozen (without gradient):")
for name in frozen_parameter_names:
logger.info(name)
logger.info("Following parameters are Tunable (with gradient):")
for name in tunable_parameter_names:
logger.info(name)
params = trainer_params
patience = params.pop_int("patience", None)
validation_metric = params.pop("validation_metric", "-loss")
shuffle = params.pop_bool("shuffle", True)
num_epochs = params.pop_int("num_epochs", 20)
cuda_device = parse_cuda_device(params.pop("cuda_device", -1))
grad_norm = params.pop_float("grad_norm", None)
grad_clipping = params.pop_float("grad_clipping", None)
lr_scheduler_params = params.pop("learning_rate_scheduler", None)
momentum_scheduler_params = params.pop("momentum_scheduler", None)
if isinstance(cuda_device, list):
model_device = cuda_device[0]
else:
model_device = cuda_device
if model_device >= 0:
# Moving model to GPU here so that the optimizer state gets constructed on
# the right device.
model = model.cuda(model_device)
parameters = [[n, p] for n, p in model.named_parameters()
if p.requires_grad]
optimizer = Optimizer.from_params(parameters, params.pop("optimizer"))
if "moving_average" in params:
moving_average = MovingAverage.from_params(
params.pop("moving_average"), parameters=parameters)
else:
moving_average = None
if lr_scheduler_params:
lr_scheduler = LearningRateScheduler.from_params(
optimizer, lr_scheduler_params)
else:
lr_scheduler = None
if momentum_scheduler_params:
momentum_scheduler = MomentumScheduler.from_params(
optimizer, momentum_scheduler_params)
else:
momentum_scheduler = None
if 'checkpointer' in params:
if 'keep_serialized_model_every_num_seconds' in params or \
'num_serialized_models_to_keep' in params:
raise ConfigurationError(
"Checkpointer may be initialized either from the 'checkpointer' key or from the "
"keys 'num_serialized_models_to_keep' and 'keep_serialized_model_every_num_seconds'"
" but the passed config uses both methods.")
checkpointer = Checkpointer.from_params(params.pop("checkpointer"))
else:
num_serialized_models_to_keep = params.pop_int(
"num_serialized_models_to_keep", 20)
keep_serialized_model_every_num_seconds = params.pop_int(
"keep_serialized_model_every_num_seconds", None)
checkpointer = Checkpointer(
serialization_dir=serialization_dir,
num_serialized_models_to_keep=num_serialized_models_to_keep,
keep_serialized_model_every_num_seconds=
keep_serialized_model_every_num_seconds)
model_save_interval = params.pop_float("model_save_interval", None)
summary_interval = params.pop_int("summary_interval", 100)
histogram_interval = params.pop_int("histogram_interval", None)
should_log_parameter_statistics = params.pop_bool(
"should_log_parameter_statistics", True)
should_log_learning_rate = params.pop_bool("should_log_learning_rate",
False)
log_batch_size_period = params.pop_int("log_batch_size_period", None)
return cls(
model,
optimizer,
iterator,
train_data,
validation_data,
patience=patience,
validation_metric=validation_metric,
validation_iterator=validation_iterator,
max_src_len=max_src_len,
shuffle=shuffle,
num_epochs=num_epochs,
serialization_dir=serialization_dir,
cuda_device=cuda_device,
grad_norm=grad_norm,
grad_clipping=grad_clipping,
learning_rate_scheduler=lr_scheduler,
momentum_scheduler=momentum_scheduler,
checkpointer=checkpointer,
model_save_interval=model_save_interval,
summary_interval=summary_interval,
histogram_interval=histogram_interval,
should_log_parameter_statistics=should_log_parameter_statistics,
should_log_learning_rate=should_log_learning_rate,
log_batch_size_period=log_batch_size_period,
moving_average=moving_average,
batch_size=iterator._batch_size)
class Trainer(TrainerBase):
"""
1. epoch todo
2. loss todo
"""
def __init__(self,
model: Model,
optimizer: torch.optim.Optimizer,
iterator: DataIterator,
train_dataset: Iterable[Instance],
validation_dataset: Optional[Iterable[Instance]] = None,
patience: Optional[int] = None,
validation_metric: str = "-loss",
validation_iterator: DataIterator = None,
shuffle: bool = True,
num_epochs: int = 20,
serialization_dir: Optional[str] = None,
num_serialized_models_to_keep: int = 0,
keep_serialized_model_every_num_seconds: int = None,
checkpointer: Checkpointer = None,
model_save_interval: float = None,
cuda_device: Union[int, List] = -1,
grad_norm: Optional[float] = None,
grad_clipping: Optional[float] = None,
learning_rate_scheduler: Optional[LearningRateScheduler] = None,
momentum_scheduler: Optional[MomentumScheduler] = None,
summary_interval: int = 100,
histogram_interval: int = None,
should_log_parameter_statistics: bool = True,
should_log_learning_rate: bool = False,
log_batch_size_period: Optional[int] = None,
moving_average: Optional[MovingAverage] = None,
callbacks: List[allennlp_callback.Callback]=None,
early_stopping_by_batch: bool=True,
estimator: Estimator=None,
) -> None:
"""
A trainer for doing supervised learning. It just takes a labeled dataset
and a ``DataIterator``, and uses the supplied ``Optimizer`` to learn the weights
for your model over some fixed number of epochs. You can also pass in a validation
dataset and enable early stopping. There are many other bells and whistles as well.
Parameters
----------
model : ``Model``, required.
An AllenNLP model to be optimized. Pytorch Modules can also be optimized if
their ``forward`` method returns a dictionary with a "loss" key, containing a
scalar tensor representing the loss function to be optimized.
If you are training your model using GPUs, your model should already be
on the correct device. (If you use `Trainer.from_params` this will be
handled for you.)
optimizer : ``torch.nn.Optimizer``, required.
An instance of a Pytorch Optimizer, instantiated with the parameters of the
model to be optimized.
iterator : ``DataIterator``, required.
A method for iterating over a ``Dataset``, yielding padded indexed batches.
train_dataset : ``Dataset``, required.
A ``Dataset`` to train on. The dataset should have already been indexed.
validation_dataset : ``Dataset``, optional, (default = None).
A ``Dataset`` to evaluate on. The dataset should have already been indexed.
patience : Optional[int] > 0, optional (default=None)
Number of epochs to be patient before early stopping: the training is stopped
after ``patience`` epochs with no improvement. If given, it must be ``> 0``.
If None, early stopping is disabled.
validation_metric : str, optional (default="loss")
Validation metric to measure for whether to stop training using patience
and whether to serialize an ``is_best`` model each epoch. The metric name
must be prepended with either "+" or "-", which specifies whether the metric
is an increasing or decreasing function.
validation_iterator : ``DataIterator``, optional (default=None)
An iterator to use for the validation set. If ``None``, then
use the training `iterator`.
shuffle: ``bool``, optional (default=True)
Whether to shuffle the instances in the iterator or not.
num_epochs : int, optional (default = 20)
Number of training epochs.
serialization_dir : str, optional (default=None)
Path to directory for saving and loading model files. Models will not be saved if
this parameter is not passed.
num_serialized_models_to_keep : ``int``, optional (default=20)
Number of previous model checkpoints to retain. Default is to keep 20 checkpoints.
A value of None or -1 means all checkpoints will be kept.
keep_serialized_model_every_num_seconds : ``int``, optional (default=None)
If num_serialized_models_to_keep is not None, then occasionally it's useful to
save models at a given interval in addition to the last num_serialized_models_to_keep.
To do so, specify keep_serialized_model_every_num_seconds as the number of seconds
between permanently saved checkpoints. Note that this option is only used if
num_serialized_models_to_keep is not None, otherwise all checkpoints are kept.
checkpointer : ``Checkpointer``, optional (default=None)
An instance of class Checkpointer to use instead of the default. If a checkpointer is specified,
the arguments num_serialized_models_to_keep and keep_serialized_model_every_num_seconds should
not be specified. The caller is responsible for initializing the checkpointer so that it is
consistent with serialization_dir.
model_save_interval : ``float``, optional (default=None)
If provided, then serialize models every ``model_save_interval``
seconds within single epochs. In all cases, models are also saved
at the end of every epoch if ``serialization_dir`` is provided.
cuda_device : ``Union[int, List[int]]``, optional (default = -1)
An integer or list of integers specifying the CUDA device(s) to use. If -1, the CPU is used.
grad_norm : ``float``, optional, (default = None).
If provided, gradient norms will be rescaled to have a maximum of this value.
grad_clipping : ``float``, optional (default = ``None``).
If provided, gradients will be clipped `during the backward pass` to have an (absolute)
maximum of this value. If you are getting ``NaNs`` in your gradients during training
that are not solved by using ``grad_norm``, you may need this.
learning_rate_scheduler : ``LearningRateScheduler``, optional (default = None)
If specified, the learning rate will be decayed with respect to
this schedule at the end of each epoch (or batch, if the scheduler implements
the ``step_batch`` method). If you use :class:`torch.optim.lr_scheduler.ReduceLROnPlateau`,
this will use the ``validation_metric`` provided to determine if learning has plateaued.
To support updating the learning rate on every batch, this can optionally implement
``step_batch(batch_num_total)`` which updates the learning rate given the batch number.
momentum_scheduler : ``MomentumScheduler``, optional (default = None)
If specified, the momentum will be updated at the end of each batch or epoch
according to the schedule.
summary_interval: ``int``, optional, (default = 100)
Number of batches between logging scalars to tensorboard
histogram_interval : ``int``, optional, (default = ``None``)
If not None, then log histograms to tensorboard every ``histogram_interval`` batches.
When this parameter is specified, the following additional logging is enabled:
* Histograms of model parameters
* The ratio of parameter update norm to parameter norm
* Histogram of layer activations
We log histograms of the parameters returned by
``model.get_parameters_for_histogram_tensorboard_logging``.
The layer activations are logged for any modules in the ``Model`` that have
the attribute ``should_log_activations`` set to ``True``. Logging
histograms requires a number of GPU-CPU copies during training and is typically
slow, so we recommend logging histograms relatively infrequently.
Note: only Modules that return tensors, tuples of tensors or dicts
with tensors as values currently support activation logging.
should_log_parameter_statistics : ``bool``, optional, (default = True)
Whether to send parameter statistics (mean and standard deviation
of parameters and gradients) to tensorboard.
should_log_learning_rate : ``bool``, optional, (default = False)
Whether to send parameter specific learning rate to tensorboard.
log_batch_size_period : ``int``, optional, (default = ``None``)
If defined, how often to log the average batch size.
moving_average: ``MovingAverage``, optional, (default = None)
If provided, we will maintain moving averages for all parameters. During training, we
employ a shadow variable for each parameter, which maintains the moving average. During
evaluation, we backup the original parameters and assign the moving averages to corresponding
parameters. Be careful that when saving the checkpoint, we will save the moving averages of
parameters. This is necessary because we want the saved model to perform as well as the validated
model if we load it later. But this may cause problems if you restart the training from checkpoint.
"""
super().__init__(serialization_dir, cuda_device)
# I am not calling move_to_gpu here, because if the model is
# not already on the GPU then the optimizer is going to be wrong.
self.model = model
self.iterator = iterator
self._validation_iterator = validation_iterator
self.shuffle = shuffle
self.optimizer = optimizer
self.train_data = train_dataset
self._validation_data = validation_dataset
if patience is None: # no early stopping
if validation_dataset:
logger.warning('You provided a validation dataset but patience was set to None, '
'meaning that early stopping is disabled')
elif (not isinstance(patience, int)) or patience <= 0:
raise ConfigurationError('{} is an invalid value for "patience": it must be a positive integer '
'or None (if you want to disable early stopping)'.format(patience))
# For tracking is_best_so_far and should_stop_early
self._metric_tracker = MetricTracker(patience, validation_metric)
# Get rid of + or -
self._validation_metric = validation_metric[1:]
self._num_epochs = num_epochs
if checkpointer is not None:
# We can't easily check if these parameters were passed in, so check against their default values.
# We don't check against serialization_dir since it is also used by the parent class.
if num_serialized_models_to_keep != 20 or \
keep_serialized_model_every_num_seconds is not None:
raise ConfigurationError(
"When passing a custom Checkpointer, you may not also pass in separate checkpointer "
"args 'num_serialized_models_to_keep' or 'keep_serialized_model_every_num_seconds'.")
self._checkpointer = checkpointer
else:
self._checkpointer = Checkpointer(serialization_dir,
keep_serialized_model_every_num_seconds,
num_serialized_models_to_keep)
self._model_save_interval = model_save_interval
self._grad_norm = grad_norm
self._grad_clipping = grad_clipping
self._learning_rate_scheduler = learning_rate_scheduler
self._momentum_scheduler = momentum_scheduler
self._moving_average = moving_average
# We keep the total batch number as an instance variable because it
# is used inside a closure for the hook which logs activations in
# ``_enable_activation_logging``.
self._batch_num_total = 0
self._tensorboard = TensorboardWriter(
get_batch_num_total=lambda: self._batch_num_total,
serialization_dir=serialization_dir,
summary_interval=summary_interval,
histogram_interval=histogram_interval,
should_log_parameter_statistics=should_log_parameter_statistics,
should_log_learning_rate=should_log_learning_rate)
self._log_batch_size_period = log_batch_size_period
self._last_log = 0.0 # time of last logging
# Enable activation logging.
if histogram_interval is not None:
self._tensorboard.enable_activation_logging(self.model)
self.callbacks = callbacks
self._early_stopping_by_batch = early_stopping_by_batch
self._estimator = estimator
def rescale_gradients(self) -> Optional[float]:
return training_util.rescale_gradients(self.model, self._grad_norm)
def batch_loss(self, batch_group: List[TensorDict], for_training: bool) -> torch.Tensor:
"""
Does a forward pass on the given batches and returns the ``loss`` value in the result.
If ``for_training`` is `True` also applies regularization penalty.
"""
if self._multiple_gpu:
output_dict = training_util.data_parallel(batch_group, self.model, self._cuda_devices)
else:
assert len(batch_group) == 1
batch = batch_group[0]
batch = nn_util.move_to_device(batch, self._cuda_devices[0])
output_dict = self.model(**batch)
try:
loss = output_dict["loss"]
if for_training:
loss += self.model.get_regularization_penalty()
except KeyError:
if for_training:
raise RuntimeError("The model you are trying to optimize does not contain a"
" 'loss' key in the output of model.forward(inputs).")
loss = None
return loss
def _train_epoch(self, epoch: int) -> Dict[str, float]:
"""
Trains one epoch and returns metrics.
"""
logger.info("Epoch %d/%d", epoch, self._num_epochs - 1)
peak_cpu_usage = peak_memory_mb()
logger.info(f"Peak CPU memory usage MB: {peak_cpu_usage}")
gpu_usage = []
for gpu, memory in gpu_memory_mb().items():
gpu_usage.append((gpu, memory))
logger.info(f"GPU {gpu} memory usage MB: {memory}")
train_loss = 0.0
# Set the model to "train" mode.
self.model.train()
num_gpus = len(self._cuda_devices)
# Get tqdm for the training batches
raw_train_generator = self.iterator(self.train_data,
num_epochs=1,
shuffle=self.shuffle)
train_generator = lazy_groups_of(raw_train_generator, num_gpus)
num_training_batches = math.ceil(self.iterator.get_num_batches(self.train_data)/num_gpus)
self._last_log = time.time()
last_save_time = time.time()
batches_this_epoch = 0
if self._batch_num_total is None:
self._batch_num_total = 0
histogram_parameters = set(self.model.get_parameters_for_histogram_tensorboard_logging())
logger.info("Training")
train_generator_tqdm = Tqdm.tqdm(train_generator,
total=num_training_batches)
cumulative_batch_size = 0
for batch_group in train_generator_tqdm:
self.model.train()
batches_this_epoch += 1
self._batch_num_total += 1
batch_num_total = self._batch_num_total
self.optimizer.zero_grad()
loss = self.batch_loss(batch_group, for_training=True)
if torch.isnan(loss):
raise ValueError("nan loss encountered")
loss.backward()
train_loss += loss.item()
batch_grad_norm = self.rescale_gradients()
# This does nothing if batch_num_total is None or you are using a
# scheduler which doesn't update per batch.
if self._learning_rate_scheduler:
self._learning_rate_scheduler.step_batch(batch_num_total)
if self._momentum_scheduler:
self._momentum_scheduler.step_batch(batch_num_total)
if self._tensorboard.should_log_histograms_this_batch():
# get the magnitude of parameter updates for logging
# We need a copy of current parameters to compute magnitude of updates,
# and copy them to CPU so large models won't go OOM on the GPU.
param_updates = {name: param.detach().cpu().clone()
for name, param in self.model.named_parameters()}
self.optimizer.step()
for name, param in self.model.named_parameters():
param_updates[name].sub_(param.detach().cpu())
update_norm = torch.norm(param_updates[name].view(-1, ))
param_norm = torch.norm(param.view(-1, )).cpu()
self._tensorboard.add_train_scalar("gradient_update/" + name,
update_norm / (param_norm + 1e-7))
else:
self.optimizer.step()
# Update moving averages
if self._moving_average is not None:
self._moving_average.apply(batch_num_total)
# Update the description with the latest metrics
metrics = training_util.get_metrics(self.model, train_loss, batches_this_epoch)
description = training_util.description_from_metrics(metrics)
train_generator_tqdm.set_description(description, refresh=False)
# Log parameter values to Tensorboard
if self._tensorboard.should_log_this_batch():
self._tensorboard.log_parameter_and_gradient_statistics(self.model, batch_grad_norm)
self._tensorboard.log_learning_rates(self.model, self.optimizer)
self._tensorboard.add_train_scalar("loss/loss_train", metrics["loss"])
self._tensorboard.log_metrics({"epoch_metrics/" + k: v for k, v in metrics.items()})
if self._tensorboard.should_log_histograms_this_batch():
self._tensorboard.log_histograms(self.model, histogram_parameters)
if self._log_batch_size_period:
cur_batch = sum([training_util.get_batch_size(batch) for batch in batch_group])
cumulative_batch_size += cur_batch
if (batches_this_epoch - 1) % self._log_batch_size_period == 0:
average = cumulative_batch_size/batches_this_epoch
logger.info(f"current batch size: {cur_batch} mean batch size: {average}")
self._tensorboard.add_train_scalar("current_batch_size", cur_batch)
self._tensorboard.add_train_scalar("mean_batch_size", average)
# Save model if needed.
if self._model_save_interval is not None and (
time.time() - last_save_time > self._model_save_interval
):
last_save_time = time.time()
self._save_checkpoint(
'{0}.{1}'.format(epoch, training_util.time_to_str(int(last_save_time)))
)
if self._early_stopping_by_batch and self._batch_num_total % 10 == 0:
if self._validation_data is not None:
with torch.no_grad():
# We have a validation set, so compute all the metrics on it.
val_loss, num_batches = self._validation_loss()
val_metrics = training_util.get_metrics(self.model, val_loss, num_batches, reset=True)
# Check validation metric for early stopping
this_epoch_val_metric = val_metrics[self._validation_metric]
self._metric_tracker.add_metric(this_epoch_val_metric)
if self._metric_tracker.is_best_so_far():
metrics['best_batch'] = self._batch_num_total
for key, value in val_metrics.items():
metrics["best_validation_" + key] = value
self._metric_tracker.best_epoch_metrics = val_metrics
self._save_checkpoint(self._batch_num_total)
if self.callbacks is not None:
for callback in self.callbacks:
callback.on_batch_end(self._batch_num_total)
metrics = training_util.get_metrics(self.model, train_loss, batches_this_epoch, reset=True)
metrics['cpu_memory_MB'] = peak_cpu_usage
for (gpu_num, memory) in gpu_usage:
metrics['gpu_'+str(gpu_num)+'_memory_MB'] = memory
return metrics
def _validation_loss(self) -> Tuple[float, int]:
"""
Computes the validation loss. Returns it and the number of batches.
"""
logger.info("Validating")
self.model.eval()
# Replace parameter values with the shadow values from the moving averages.
if self._moving_average is not None:
self._moving_average.assign_average_value()
if self._validation_iterator is not None:
val_iterator = self._validation_iterator
else:
val_iterator = self.iterator
num_gpus = len(self._cuda_devices)
raw_val_generator = val_iterator(self._validation_data,
num_epochs=1,
shuffle=False)
val_generator = lazy_groups_of(raw_val_generator, num_gpus)
num_validation_batches = math.ceil(val_iterator.get_num_batches(self._validation_data)/num_gpus)
val_generator_tqdm = Tqdm.tqdm(val_generator,
total=num_validation_batches)
batches_this_epoch = 0
val_loss = 0
for batch_group in val_generator_tqdm:
loss = self.batch_loss(batch_group, for_training=False)
if loss is not None:
# You shouldn't necessarily have to compute a loss for validation, so we allow for
# `loss` to be None. We need to be careful, though - `batches_this_epoch` is
# currently only used as the divisor for the loss function, so we can safely only
# count those batches for which we actually have a loss. If this variable ever
# gets used for something else, we might need to change things around a bit.
batches_this_epoch += 1
val_loss += loss.detach().cpu().numpy()
# Update the description with the latest metrics
val_metrics = training_util.get_metrics(self.model, val_loss, batches_this_epoch)
description = training_util.description_from_metrics(val_metrics)
val_generator_tqdm.set_description(description, refresh=False)
# Now restore the original parameter values.
if self._moving_average is not None:
self._moving_average.restore()
return val_loss, batches_this_epoch
def train(self) -> Dict[str, Any]:
"""
Trains the supplied model with the supplied parameters.
"""
try:
epoch_counter = self._restore_checkpoint()
except RuntimeError:
traceback.print_exc()
raise ConfigurationError("Could not recover training from the checkpoint. Did you mean to output to "
"a different serialization directory or delete the existing serialization "
"directory?")
training_util.enable_gradient_clipping(self.model, self._grad_clipping)
logger.info("Beginning training.")
train_metrics: Dict[str, float] = {}
val_metrics: Dict[str, float] = {}
this_epoch_val_metric: float = None
metrics: Dict[str, Any] = {}
epochs_trained = 0
training_start_time = time.time()
metrics['best_epoch'] = self._metric_tracker.best_epoch
for key, value in self._metric_tracker.best_epoch_metrics.items():
metrics["best_validation_" + key] = value
if self.callbacks is not None:
with torch.no_grad():
for callback in self.callbacks:
callback.on_train_begin()
for epoch in range(epoch_counter, self._num_epochs):
epoch_start_time = time.time()
if self.callbacks is not None:
with torch.no_grad():
for callback in self.callbacks:
callback.on_epoch_begin(epoch)
train_metrics = self._train_epoch(epoch)
if not self._early_stopping_by_batch:
# get peak of memory usage
if 'cpu_memory_MB' in train_metrics:
metrics['peak_cpu_memory_MB'] = max(metrics.get('peak_cpu_memory_MB', 0),
train_metrics['cpu_memory_MB'])
for key, value in train_metrics.items():
if key.startswith('gpu_'):
metrics["peak_"+key] = max(metrics.get("peak_"+key, 0), value)
if self._validation_data is not None:
with torch.no_grad():
val_metrics_temp = self._estimator.estimate(self._validation_data)
# We have a validation set, so compute all the metrics on it.
# val_loss, num_batches = self._validation_loss()
# val_metrics = training_util.get_metrics(self.model, val_loss, num_batches, reset=True)
val_metrics = {'loss': 0}
if 'sentiment_acc' in val_metrics_temp:
val_metrics['accuracy'] = val_metrics_temp['sentiment_acc']
if 'category_f1' in val_metrics_temp:
val_metrics['category_f1'] = val_metrics_temp['category_f1']['fscore']
if 'other_metrics' in val_metrics_temp and 'merge_micro_f1' in val_metrics_temp['other_metrics']:
val_metrics['merge_micro_f1'] = val_metrics_temp['other_metrics']['merge_micro_f1']
# Check validation metric for early stopping
val_metrics.update(val_metrics_temp)
this_epoch_val_metric = val_metrics[self._validation_metric]
self._metric_tracker.add_metric(this_epoch_val_metric)
if self._metric_tracker.should_stop_early():
logger.info("Ran out of patience. Stopping training.")
break
self._tensorboard.log_metrics(train_metrics,
val_metrics=val_metrics,
log_to_console=True,
epoch=epoch + 1) # +1 because tensorboard doesn't like 0
# Create overall metrics dict
training_elapsed_time = time.time() - training_start_time
metrics["training_duration"] = str(datetime.timedelta(seconds=training_elapsed_time))
metrics["training_start_epoch"] = epoch_counter
metrics["training_epochs"] = epochs_trained
metrics["epoch"] = epoch
for key, value in train_metrics.items():
metrics["training_" + key] = value
for key, value in val_metrics.items():
metrics["validation_" + key] = value
if self._metric_tracker.is_best_so_far():
# Update all the best_ metrics.
# (Otherwise they just stay the same as they were.)
metrics['best_epoch'] = epoch
for key, value in val_metrics.items():
metrics["best_validation_" + key] = value
self._metric_tracker.best_epoch_metrics = val_metrics
if self._serialization_dir:
dump_metrics(os.path.join(self._serialization_dir, f'metrics_epoch_{epoch}.json'), metrics)
# The Scheduler API is agnostic to whether your schedule requires a validation metric -
# if it doesn't, the validation metric passed here is ignored.
if self._learning_rate_scheduler:
self._learning_rate_scheduler.step(this_epoch_val_metric, epoch)
if self._momentum_scheduler:
self._momentum_scheduler.step(this_epoch_val_metric, epoch)
self._save_checkpoint(epoch)
else:
if self._metric_tracker.should_stop_early():
logger.info("Ran out of patience. Stopping training.")
break
epoch_elapsed_time = time.time() - epoch_start_time
logger.info("Epoch duration: %s", datetime.timedelta(seconds=epoch_elapsed_time))
if epoch < self._num_epochs - 1:
training_elapsed_time = time.time() - training_start_time
estimated_time_remaining = training_elapsed_time * \
((self._num_epochs - epoch_counter) / float(epoch - epoch_counter + 1) - 1)
formatted_time = str(datetime.timedelta(seconds=int(estimated_time_remaining)))
logger.info("Estimated training time remaining: %s", formatted_time)
if self.callbacks is not None:
with torch.no_grad():
for callback in self.callbacks:
callback.on_epoch_end(epoch)
epochs_trained += 1
# make sure pending events are flushed to disk and files are closed properly
# self._tensorboard.close()
# Load the best model state before returning
best_model_state = self._checkpointer.best_model_state()
if best_model_state:
self.model.load_state_dict(best_model_state)
return metrics
def _save_checkpoint(self, epoch: Union[int, str]) -> None:
"""
Saves a checkpoint of the model to self._serialization_dir.
Is a no-op if self._serialization_dir is None.
Parameters
----------
epoch : Union[int, str], required.
The epoch of training. If the checkpoint is saved in the middle
of an epoch, the parameter is a string with the epoch and timestamp.
"""
# If moving averages are used for parameters, we save
# the moving average values into checkpoint, instead of the current values.
if self._moving_average is not None:
self._moving_average.assign_average_value()
# These are the training states we need to persist.
training_states = {
"metric_tracker": self._metric_tracker.state_dict(),
"optimizer": self.optimizer.state_dict(),
"batch_num_total": self._batch_num_total
}
# If we have a learning rate or momentum scheduler, we should persist them too.
if self._learning_rate_scheduler is not None:
training_states["learning_rate_scheduler"] = self._learning_rate_scheduler.state_dict()
if self._momentum_scheduler is not None:
training_states["momentum_scheduler"] = self._momentum_scheduler.state_dict()
self._checkpointer.save_checkpoint(
model_state=self.model.state_dict(),
epoch=epoch,
training_states=training_states,
is_best_so_far=self._metric_tracker.is_best_so_far())
# Restore the original values for parameters so that training will not be affected.
if self._moving_average is not None:
self._moving_average.restore()
def _restore_checkpoint(self) -> int:
"""
Restores the model and training state from the last saved checkpoint.
This includes an epoch count and optimizer state, which is serialized separately
from model parameters. This function should only be used to continue training -
if you wish to load a model for inference/load parts of a model into a new
computation graph, you should use the native Pytorch functions:
`` model.load_state_dict(torch.load("/path/to/model/weights.th"))``
If ``self._serialization_dir`` does not exist or does not contain any checkpointed weights,
this function will do nothing and return 0.
Returns
-------
epoch: int
The epoch at which to resume training, which should be one after the epoch
in the saved training state.
"""
model_state, training_state = self._checkpointer.restore_checkpoint()
if not training_state:
# No checkpoint to restore, start at 0
return 0
self.model.load_state_dict(model_state)
self.optimizer.load_state_dict(training_state["optimizer"])
if self._learning_rate_scheduler is not None and "learning_rate_scheduler" in training_state:
self._learning_rate_scheduler.load_state_dict(training_state["learning_rate_scheduler"])
if self._momentum_scheduler is not None and "momentum_scheduler" in training_state:
self._momentum_scheduler.load_state_dict(training_state["momentum_scheduler"])
training_util.move_optimizer_to_cuda(self.optimizer)
# Currently the ``training_state`` contains a serialized ``MetricTracker``.
if "metric_tracker" in training_state:
self._metric_tracker.load_state_dict(training_state["metric_tracker"])
# It used to be the case that we tracked ``val_metric_per_epoch``.
elif "val_metric_per_epoch" in training_state:
self._metric_tracker.clear()
self._metric_tracker.add_metrics(training_state["val_metric_per_epoch"])
# And before that we didn't track anything.
else:
self._metric_tracker.clear()
if isinstance(training_state["epoch"], int):
epoch_to_return = training_state["epoch"] + 1
else:
epoch_to_return = int(training_state["epoch"].split('.')[0]) + 1
# For older checkpoints with batch_num_total missing, default to old behavior where
# it is unchanged.
batch_num_total = training_state.get('batch_num_total')
if batch_num_total is not None:
self._batch_num_total = batch_num_total
return epoch_to_return
# Requires custom from_params.
@classmethod
def from_params(cls, # type: ignore
model: Model,
serialization_dir: str,
iterator: DataIterator,
train_data: Iterable[Instance],
validation_data: Optional[Iterable[Instance]],
params: Params,
validation_iterator: DataIterator = None) -> 'Trainer':
# pylint: disable=arguments-differ
patience = params.pop_int("patience", None)
validation_metric = params.pop("validation_metric", "-loss")
shuffle = params.pop_bool("shuffle", True)
num_epochs = params.pop_int("num_epochs", 20)
cuda_device = parse_cuda_device(params.pop("cuda_device", -1))
grad_norm = params.pop_float("grad_norm", None)
grad_clipping = params.pop_float("grad_clipping", None)
lr_scheduler_params = params.pop("learning_rate_scheduler", None)
momentum_scheduler_params = params.pop("momentum_scheduler", None)
if isinstance(cuda_device, list):
model_device = cuda_device[0]
else:
model_device = cuda_device
if model_device >= 0:
# Moving model to GPU here so that the optimizer state gets constructed on
# the right device.
model = model.cuda(model_device)
parameters = [[n, p] for n, p in model.named_parameters() if p.requires_grad]
optimizer = Optimizer.from_params(parameters, params.pop("optimizer"))
if "moving_average" in params:
moving_average = MovingAverage.from_params(params.pop("moving_average"), parameters=parameters)
else:
moving_average = None
if lr_scheduler_params:
lr_scheduler = LearningRateScheduler.from_params(optimizer, lr_scheduler_params)
else:
lr_scheduler = None
if momentum_scheduler_params:
momentum_scheduler = MomentumScheduler.from_params(optimizer, momentum_scheduler_params)
else:
momentum_scheduler = None
if 'checkpointer' in params:
if 'keep_serialized_model_every_num_seconds' in params or \
'num_serialized_models_to_keep' in params:
raise ConfigurationError(
"Checkpointer may be initialized either from the 'checkpointer' key or from the "
"keys 'num_serialized_models_to_keep' and 'keep_serialized_model_every_num_seconds'"
" but the passed config uses both methods.")
checkpointer = Checkpointer.from_params(params.pop("checkpointer"))
else:
num_serialized_models_to_keep = params.pop_int("num_serialized_models_to_keep", 20)
keep_serialized_model_every_num_seconds = params.pop_int(
"keep_serialized_model_every_num_seconds", None)
checkpointer = Checkpointer(
serialization_dir=serialization_dir,
num_serialized_models_to_keep=num_serialized_models_to_keep,
keep_serialized_model_every_num_seconds=keep_serialized_model_every_num_seconds)
model_save_interval = params.pop_float("model_save_interval", None)
summary_interval = params.pop_int("summary_interval", 100)
histogram_interval = params.pop_int("histogram_interval", None)
should_log_parameter_statistics = params.pop_bool("should_log_parameter_statistics", True)
should_log_learning_rate = params.pop_bool("should_log_learning_rate", False)
log_batch_size_period = params.pop_int("log_batch_size_period", None)
params.assert_empty(cls.__name__)
return cls(model, optimizer, iterator,
train_data, validation_data,
patience=patience,
validation_metric=validation_metric,
validation_iterator=validation_iterator,
shuffle=shuffle,
num_epochs=num_epochs,
serialization_dir=serialization_dir,
cuda_device=cuda_device,
grad_norm=grad_norm,
grad_clipping=grad_clipping,
learning_rate_scheduler=lr_scheduler,
momentum_scheduler=momentum_scheduler,
checkpointer=checkpointer,
model_save_interval=model_save_interval,
summary_interval=summary_interval,
histogram_interval=histogram_interval,
should_log_parameter_statistics=should_log_parameter_statistics,
should_log_learning_rate=should_log_learning_rate,
log_batch_size_period=log_batch_size_period,
moving_average=moving_average)
class DeepspeedTrainer(Trainer):
def __init__(
self,
model: Model,
optimizer: torch.optim.Optimizer,
data_loader: DataLoader,
deepspeed_config: DeepspeedConfig,
patience: Optional[int] = None,
validation_metric: str = "-loss",
validation_data_loader: DataLoader = None,
num_epochs: int = 20,
serialization_dir: Optional[str] = None,
checkpointer: Checkpointer = None,
cuda_device: Optional[Union[int, torch.device]] = None,
grad_norm: Optional[float] = None,
grad_clipping: Optional[float] = None,
tensorboard_writer: TensorboardWriter = None,
moving_average: Optional[MovingAverage] = None,
batch_callbacks: List[BatchCallback] = None,
epoch_callbacks: List[EpochCallback] = None,
distributed: bool = False,
local_rank: int = 0,
world_size: int = 1,
num_gradient_accumulation_steps: int = 1,
use_amp: bool = False,
) -> None:
super().__init__(serialization_dir, cuda_device, distributed,
local_rank, world_size)
# I am not calling move_to_gpu here, because if the model is
# not already on the GPU then the optimizer is going to be wrong.
self.model = model
self.data_loader = data_loader
self._validation_data_loader = validation_data_loader
self.optimizer = optimizer
if patience is None: # no early stopping
if validation_data_loader is not None:
logger.warning(
"You provided a validation dataset but patience was set to None, "
"meaning that early stopping is disabled")
elif (not isinstance(patience, int)) or patience <= 0:
raise ConfigurationError(
'{} is an invalid value for "patience": it must be a positive integer '
"or None (if you want to disable early stopping)".format(
patience))
# For tracking is_best_so_far and should_stop_early
self._metric_tracker = MetricTracker(patience, validation_metric)
# Get rid of + or -
self._validation_metric = validation_metric[1:]
self._num_epochs = num_epochs
if checkpointer is not None:
self._checkpointer = checkpointer
else:
self._checkpointer = Checkpointer(serialization_dir)
self._grad_norm = grad_norm
self._grad_clipping = grad_clipping
self._moving_average = moving_average
self._batch_callbacks = batch_callbacks or []
self._epoch_callbacks = epoch_callbacks or []
# We keep the total batch number as an instance variable because it
# is used inside a closure for the hook which logs activations in
# `_enable_activation_logging`.
self._batch_num_total = 0
self._tensorboard = tensorboard_writer or TensorboardWriter(
serialization_dir)
self._tensorboard.get_batch_num_total = lambda: self._batch_num_total
self._tensorboard.enable_activation_logging(self.model)
self._last_log = 0.0 # time of last logging
self._num_gradient_accumulation_steps = num_gradient_accumulation_steps
# Enable automatic mixed precision training.
self._scaler: Optional[amp.GradScaler] = None
self._use_amp = use_amp
if self._use_amp:
if self.cuda_device == torch.device("cpu"):
raise ValueError("Using AMP requires a cuda device")
self._scaler = amp.GradScaler()
self._pytorch_model = self.model
self._ds_config = deepspeed_config
self.model_engine, self.ds_optimizer, _, _ = self._ds_config.launch(
self.model,
None, # self.optimizer,
local_rank,
serialization_dir,
self.data_loader.batch_size,
num_gradient_accumulation_steps)
def batch_outputs(self, batch: TensorDict,
for_training: bool) -> Dict[str, torch.Tensor]:
"""
Does a forward pass on the given batch and returns the output dictionary that the model
returns, after adding any specified regularization penalty to the loss (if training).
"""
# batch = nn_util.move_to_device(batch, self.cuda_device)
batch = nn_util.move_to_device(batch, self.model_engine.device)
output_dict = self.model_engine(**batch)
if for_training:
try:
assert "loss" in output_dict
regularization_penalty = self.model.get_regularization_penalty(
)
if regularization_penalty is not None:
output_dict["reg_loss"] = regularization_penalty
output_dict["loss"] += regularization_penalty
except AssertionError:
if for_training:
raise RuntimeError(
"The model you are trying to optimize does not contain a"
" 'loss' key in the output of model.forward(inputs).")
return output_dict
def _train_epoch(self, epoch: int) -> Dict[str, float]:
"""
Trains one epoch and returns metrics.
"""
logger.info("Epoch %d/%d", epoch, self._num_epochs - 1)
cpu_memory_usage = []
for worker, memory in common_util.peak_memory_mb().items():
cpu_memory_usage.append((worker, memory))
logger.info(f"Worker {worker} memory usage MB: {memory}")
gpu_memory_usage = []
for gpu, memory in common_util.gpu_memory_mb().items():
gpu_memory_usage.append((gpu, memory))
logger.info(f"GPU {gpu} memory usage MB: {memory}")
regularization_penalty = self.model.get_regularization_penalty()
train_loss = 0.0
batch_loss = 0.0
if regularization_penalty is not None:
train_reg_loss = 0.0
batch_reg_loss = 0.0
else:
train_reg_loss = None
batch_reg_loss = None
# Set the model to "train" mode.
self.model_engine.train()
# Get tqdm for the training batches
batch_generator = iter(self.data_loader)
batch_group_generator = common_util.lazy_groups_of(
batch_generator, self._num_gradient_accumulation_steps)
logger.info("Training")
num_training_batches: Union[int, float]
try:
len_data_loader = len(self.data_loader)
num_training_batches = math.ceil(
len_data_loader / self._num_gradient_accumulation_steps)
except TypeError:
num_training_batches = float("inf")
# Having multiple tqdm bars in case of distributed training will be a mess. Hence only the master's
# progress is shown
batch_group_generator_tqdm = batch_group_generator
if self._master:
batch_group_generator_tqdm = Tqdm.tqdm(batch_group_generator,
total=num_training_batches)
self._last_log = time.time()
batches_this_epoch = 0
if self._batch_num_total is None:
self._batch_num_total = 0
done_early = False
for batch_group in batch_group_generator_tqdm:
batches_this_epoch += 1
self._batch_num_total += 1
batch_num_total = self._batch_num_total
self.optimizer.zero_grad()
batch_group_outputs = []
for batch in batch_group:
with amp.autocast(self._use_amp):
batch_outputs = self.batch_outputs(batch,
for_training=True)
batch_group_outputs.append(batch_outputs)
loss = batch_outputs.get("loss")
reg_loss = batch_outputs.get("reg_loss")
if torch.isnan(loss):
raise ValueError("nan loss encountered")
loss = loss / len(batch_group)
batch_loss = loss.item()
train_loss += batch_loss
if reg_loss is not None:
reg_loss = reg_loss / len(batch_group)
batch_reg_loss = reg_loss.item()
train_reg_loss += batch_reg_loss
self.model_engine.backward(loss)
self.model_engine.step()
param_updates = None
if self._tensorboard.should_log_histograms_this_batch(
) and self._master:
# Get the magnitude of parameter updates for logging. We need to do some
# computation before and after the optimizer step, and it's expensive because of
# GPU/CPU copies (necessary for large models, and for shipping to tensorboard), so
# we don't do this every batch, only when it's requested.
param_updates = {
name: param.detach().cpu().clone()
for name, param in self.model.named_parameters()
}
if self._scaler is not None:
self._scaler.step(self.optimizer)
self._scaler.update()
else:
self.optimizer.step()
for name, param in self.model.named_parameters():
param_updates[name].sub_(param.detach().cpu())
else:
if self._scaler is not None:
self._scaler.step(self.optimizer)
self._scaler.update()
else:
self.optimizer.step()
# Update moving averages
if self._moving_average is not None:
self._moving_average.apply(batch_num_total)
# Update the description with the latest metrics
metrics = training_util.get_metrics(
self.model,
train_loss,
train_reg_loss,
batch_loss,
batch_reg_loss,
batches_this_epoch,
world_size=self._world_size,
cuda_device=self.cuda_device,
)
if self._master:
# Updating tqdm only for the master as the trainers wouldn't have one
description = training_util.description_from_metrics(metrics)
batch_group_generator_tqdm.set_description(description,
refresh=False)
self._tensorboard.log_batch(
self.model,
self.optimizer,
0., # batch_grad_norm,
metrics,
batch_group,
param_updates,
)
self._checkpointer.maybe_save_checkpoint(
self, epoch, batches_this_epoch)
for callback in self._batch_callbacks:
callback(
self,
batch_group,
batch_group_outputs,
epoch,
batches_this_epoch,
is_training=True,
is_master=self._master,
)
metrics = training_util.get_metrics(
self.model,
train_loss,
train_reg_loss,
batch_loss=None,
batch_reg_loss=None,
num_batches=batches_this_epoch,
reset=True,
world_size=self._world_size,
cuda_device=self.cuda_device,
)
for (worker, memory) in cpu_memory_usage:
metrics["worker_" + str(worker) + "_memory_MB"] = memory
for (gpu_num, memory) in gpu_memory_usage:
metrics["gpu_" + str(gpu_num) + "_memory_MB"] = memory
return metrics
def _validation_loss(self, epoch: int) -> Tuple[float, float, int]:
"""
Computes the validation loss. Returns it and the number of batches.
"""
logger.info("Validating")
self.model_engine.eval()
# Replace parameter values with the shadow values from the moving averages.
if self._moving_average is not None:
self._moving_average.assign_average_value()
if self._validation_data_loader is not None:
validation_data_loader = self._validation_data_loader
else:
raise ConfigurationError(
"Validation results cannot be calculated without a validation_data_loader"
)
regularization_penalty = self.model.get_regularization_penalty()
# Having multiple tqdm bars in case of distributed training will be a mess. Hence only the master's
# progress is shown
if self._master:
val_generator_tqdm = Tqdm.tqdm(validation_data_loader)
else:
val_generator_tqdm = validation_data_loader
batches_this_epoch = 0
val_loss = 0
val_batch_loss = 0
if regularization_penalty is not None:
val_reg_loss = 0
val_batch_reg_loss = 0
else:
val_reg_loss = None
val_batch_reg_loss = None
done_early = False
for batch in val_generator_tqdm:
if self._distributed:
# Check whether the other workers have stopped already (due to differing amounts of
# data in each). If so, we can't proceed because we would hang when we hit the
# barrier implicit in Model.forward. We use a IntTensor instead a BoolTensor
# here because NCCL process groups apparently don't support BoolTensor.
done = torch.tensor(0, device=self.cuda_device)
torch.distributed.all_reduce(done,
torch.distributed.ReduceOp.SUM)
if done.item() > 0:
done_early = True
logger.warning(
f"Worker {torch.distributed.get_rank()} finishing validation early! "
"This implies that there is an imbalance in your validation "
"data across the workers and that some amount of it will be "
"ignored. A small amount of this is fine, but a major imbalance "
"should be avoided. Note: This warning will appear unless your "
"data is perfectly balanced.")
break
with amp.autocast(self._use_amp):
batch_outputs = self.batch_outputs(batch, for_training=False)
loss = batch_outputs.get("loss")
reg_loss = batch_outputs.get("reg_loss")
if loss is not None:
# You shouldn't necessarily have to compute a loss for validation, so we allow for
# `loss` to be None. We need to be careful, though - `batches_this_epoch` is
# currently only used as the divisor for the loss function, so we can safely only
# count those batches for which we actually have a loss. If this variable ever
# gets used for something else, we might need to change things around a bit.
batches_this_epoch += 1
val_batch_loss = loss.detach().cpu().numpy()
val_loss += val_batch_loss
if reg_loss is not None:
val_batch_reg_loss = reg_loss.detach().cpu().numpy()
val_reg_loss += val_batch_reg_loss
# Update the description with the latest metrics
val_metrics = training_util.get_metrics(
self.model,
val_loss,
val_reg_loss,
val_batch_loss,
val_batch_reg_loss,
batches_this_epoch,
world_size=self._world_size,
cuda_device=self.cuda_device,
)
description = training_util.description_from_metrics(val_metrics)
if self._master:
val_generator_tqdm.set_description(description, refresh=False)
for callback in self._batch_callbacks:
callback(
self,
[batch],
[batch_outputs],
epoch,
batches_this_epoch,
is_training=False,
is_master=self._master,
)
if self._distributed and not done_early:
logger.warning(
f"Worker {torch.distributed.get_rank()} completed its entire epoch (validation)."
)
# Indicate that we're done so that any workers that have remaining data stop validation early.
done = torch.tensor(1, device=self.cuda_device)
torch.distributed.all_reduce(done, torch.distributed.ReduceOp.SUM)
assert done.item()
# Now restore the original parameter values.
if self._moving_average is not None:
self._moving_average.restore()
return val_loss, val_reg_loss, batches_this_epoch
def train(self) -> Dict[str, Any]:
"""
Trains the supplied model with the supplied parameters.
"""
try:
epoch_counter = self._restore_checkpoint()
except RuntimeError:
traceback.print_exc()
raise ConfigurationError(
"Could not recover training from the checkpoint. Did you mean to output to "
"a different serialization directory or delete the existing serialization "
"directory?")
training_util.enable_gradient_clipping(self.model, self._grad_clipping)
logger.info("Beginning training.")
val_metrics: Dict[str, float] = {}
this_epoch_val_metric: float = None
metrics: Dict[str, Any] = {}
epochs_trained = 0
training_start_time = time.time()
metrics["best_epoch"] = self._metric_tracker.best_epoch
for key, value in self._metric_tracker.best_epoch_metrics.items():
metrics["best_validation_" + key] = value
for callback in self._epoch_callbacks:
callback(self, metrics={}, epoch=-1, is_master=self._master)
for epoch in range(epoch_counter, self._num_epochs):
epoch_start_time = time.time()
train_metrics = self._train_epoch(epoch)
# get peak of memory usage
for key, value in train_metrics.items():
if key.startswith("gpu_") and key.endswith("_memory_MB"):
metrics["peak_" + key] = max(metrics.get("peak_" + key, 0),
value)
elif key.startswith("worker_") and key.endswith("_memory_MB"):
metrics["peak_" + key] = max(metrics.get("peak_" + key, 0),
value)
if self._validation_data_loader is not None:
with torch.no_grad():
# We have a validation set, so compute all the metrics on it.
val_loss, val_reg_loss, num_batches = self._validation_loss(
epoch)
# It is safe again to wait till the validation is done. This is
# important to get the metrics right.
if self._distributed:
dist.barrier()
val_metrics = training_util.get_metrics(
self.model,
val_loss,
val_reg_loss,
batch_loss=None,
batch_reg_loss=None,
num_batches=num_batches,
reset=True,
world_size=self._world_size,
cuda_device=self.cuda_device,
)
# Check validation metric for early stopping
this_epoch_val_metric = val_metrics[
self._validation_metric]
self._metric_tracker.add_metric(this_epoch_val_metric)
if self._metric_tracker.should_stop_early():
logger.info("Ran out of patience. Stopping training.")
break
if self._master:
self._tensorboard.log_metrics(
train_metrics,
val_metrics=val_metrics,
log_to_console=True,
epoch=epoch + 1) # +1 because tensorboard doesn't like 0
# Create overall metrics dict
training_elapsed_time = time.time() - training_start_time
metrics["training_duration"] = str(
datetime.timedelta(seconds=training_elapsed_time))
metrics["training_start_epoch"] = epoch_counter
metrics["training_epochs"] = epochs_trained
metrics["epoch"] = epoch
for key, value in train_metrics.items():
metrics["training_" + key] = value
for key, value in val_metrics.items():
metrics["validation_" + key] = value
if self._metric_tracker.is_best_so_far():
# Update all the best_ metrics.
# (Otherwise they just stay the same as they were.)
metrics["best_epoch"] = epoch
for key, value in val_metrics.items():
metrics["best_validation_" + key] = value
self._metric_tracker.best_epoch_metrics = val_metrics
if self._serialization_dir and self._master:
common_util.dump_metrics(
os.path.join(self._serialization_dir,
f"metrics_epoch_{epoch}.json"), metrics)
if self._master:
self._checkpointer.save_checkpoint(
epoch,
self,
is_best_so_far=self._metric_tracker.is_best_so_far())
# Wait for the master to finish saving the checkpoint
if self._distributed:
dist.barrier()
for callback in self._epoch_callbacks:
callback(self,
metrics=metrics,
epoch=epoch,
is_master=self._master)
epoch_elapsed_time = time.time() - epoch_start_time
logger.info("Epoch duration: %s",
datetime.timedelta(seconds=epoch_elapsed_time))
if epoch < self._num_epochs - 1:
training_elapsed_time = time.time() - training_start_time
estimated_time_remaining = training_elapsed_time * (
(self._num_epochs - epoch_counter) /
float(epoch - epoch_counter + 1) - 1)
formatted_time = str(
datetime.timedelta(seconds=int(estimated_time_remaining)))
logger.info("Estimated training time remaining: %s",
formatted_time)
epochs_trained += 1
# make sure pending events are flushed to disk and files are closed properly
self._tensorboard.close()
# Load the best model state before returning
best_model_state = self._checkpointer.best_model_state()
if best_model_state:
self.model.load_state_dict(best_model_state)
return metrics
@contextmanager
def get_checkpoint_state(
self) -> Iterator[Tuple[Dict[str, Any], Dict[str, Any]]]:
if self._moving_average is not None:
# Assigning average value to model parameters. The checkpointer will call
# `restore_state_after_checkpointing` when it is done to put this back to what it was.
self._moving_average.assign_average_value()
model_state = self.model.state_dict()
# These are the training states we need to persist.
training_states = {
"metric_tracker": self._metric_tracker.state_dict(),
"optimizer": self.optimizer.state_dict(),
"batch_num_total": self._batch_num_total,
}
try:
yield model_state, training_states
finally:
if self._moving_average is not None:
self._moving_average.restore()
def _restore_checkpoint(self) -> int:
"""
Restores the model and training state from the last saved checkpoint.
This includes an epoch count and optimizer state, which is serialized separately
from model parameters. This function should only be used to continue training -
if you wish to load a model for inference/load parts of a model into a new
computation graph, you should use the native Pytorch functions:
` model.load_state_dict(torch.load("/path/to/model/weights.th"))`
If `self._serialization_dir` does not exist or does not contain any checkpointed weights,
this function will do nothing and return 0.
# Returns
epoch: `int`
The epoch at which to resume training, which should be one after the epoch
in the saved training state.
"""
model_state, training_state = self._checkpointer.restore_checkpoint()
if not training_state:
# No checkpoint to restore, start at 0
return 0
self.model.load_state_dict(model_state)
self.optimizer.load_state_dict(training_state["optimizer"])
training_util.move_optimizer_to_cuda(self.optimizer)
# Currently the `training_state` contains a serialized `MetricTracker`.
if "metric_tracker" in training_state:
self._metric_tracker.load_state_dict(
training_state["metric_tracker"])
# It used to be the case that we tracked `val_metric_per_epoch`.
elif "val_metric_per_epoch" in training_state:
self._metric_tracker.clear()
self._metric_tracker.add_metrics(
training_state["val_metric_per_epoch"])
# And before that we didn't track anything.
else:
self._metric_tracker.clear()
if isinstance(training_state["epoch"], int):
epoch_to_return = training_state["epoch"] + 1
else:
epoch_to_return = int(training_state["epoch"].split(".")[0]) + 1
# For older checkpoints with batch_num_total missing, default to old behavior where
# it is unchanged.
batch_num_total = training_state.get("batch_num_total")
if batch_num_total is not None:
self._batch_num_total = batch_num_total
return epoch_to_return
@classmethod
def from_partial_objects(
cls,
model: Model,
serialization_dir: str,
data_loader: DataLoader,
validation_data_loader: DataLoader = None,
local_rank: int = 0,
patience: int = None,
validation_metric: str = "-loss",
num_epochs: int = 20,
cuda_device: Optional[Union[int, torch.device]] = None,
grad_norm: float = None,
grad_clipping: float = None,
distributed: bool = None,
world_size: int = 1,
num_gradient_accumulation_steps: int = 1,
use_amp: bool = False,
no_grad: List[str] = None,
optimizer: Lazy[Optimizer] = None,
tensorboard_writer: Lazy[TensorboardWriter] = None,
moving_average: Lazy[MovingAverage] = None,
checkpointer: Lazy[Checkpointer] = None,
batch_callbacks: List[BatchCallback] = None,
epoch_callbacks: List[EpochCallback] = None,
deepspeed_config: DeepspeedConfig = None) -> "Trainer":
"""
This method exists so that we can have a documented method to construct this class using
`FromParams`. If you are not using `FromParams` or config files, you can safely ignore this
method.
The reason we can't just use `__init__` with `FromParams` here is because there are
sequential dependencies to this class's arguments. Anything that has a `Lazy[]` type
annotation needs something from one of the non-`Lazy` arguments. The `Optimizer` needs to
have the parameters from the `Model` before it's constructed, and the `Schedulers` need to
have the `Optimizer`. Because of this, the typical way we construct things `FromParams`
doesn't work, so we use `Lazy` to allow for constructing the objects sequentially.
If you're not using `FromParams`, you can just construct these arguments in the right order
yourself in your code and call the constructor directly.
"""
if cuda_device is None:
from torch import cuda
if cuda.device_count() > 0:
cuda_device = 0
else:
cuda_device = -1
check_for_gpu(cuda_device)
if cuda_device >= 0:
# Moving model to GPU here so that the optimizer state gets constructed on
# the right device.
model = model.cuda(cuda_device)
if no_grad:
for name, parameter in model.named_parameters():
if any(re.search(regex, name) for regex in no_grad):
parameter.requires_grad_(False)
parameters = [[n, p] for n, p in model.named_parameters()
if p.requires_grad]
optimizer_ = optimizer.construct(model_parameters=parameters)
if not optimizer_:
optimizer_ = Optimizer.default(parameters)
common_util.log_frozen_and_tunable_parameter_names(model)
batches_per_epoch: Optional[int]
try:
batches_per_epoch = len(data_loader)
batches_per_epoch = math.ceil(batches_per_epoch /
num_gradient_accumulation_steps)
except TypeError:
batches_per_epoch = None
moving_average_ = moving_average.construct(parameters=parameters)
checkpointer_ = checkpointer.construct() or Checkpointer(
serialization_dir)
tensorboard_writer_ = tensorboard_writer.construct(
) or TensorboardWriter(serialization_dir)
return cls(
model,
optimizer_,
data_loader,
patience=patience,
validation_metric=validation_metric,
validation_data_loader=validation_data_loader,
num_epochs=num_epochs,
serialization_dir=serialization_dir,
cuda_device=cuda_device,
grad_norm=grad_norm,
grad_clipping=grad_clipping,
tensorboard_writer=tensorboard_writer_,
checkpointer=checkpointer_,
moving_average=moving_average_,
batch_callbacks=batch_callbacks,
epoch_callbacks=epoch_callbacks,
distributed=distributed,
local_rank=local_rank,
world_size=world_size,
num_gradient_accumulation_steps=num_gradient_accumulation_steps,
use_amp=use_amp,
deepspeed_config=deepspeed_config)
class MultiTaskTrainer(TrainerBase):
"""
A simple trainer that works in our multi-task setup.
Really the main thing that makes this task not fit into our
existing trainer is the multiple datasets.
"""
def __init__(
self,
model: Model,
serialization_dir: str,
iterator: DataIterator,
mingler: DatasetMingler,
optimizer: torch.optim.Optimizer,
datasets: Dict[str, Iterable[Instance]],
num_epochs: int = 10,
num_serialized_models_to_keep: int = 10,
) -> None:
super().__init__(serialization_dir)
self.model = model
self.iterator = iterator
self.mingler = mingler
self.optimizer = optimizer
self.datasets = datasets
self.num_epochs = num_epochs
self.checkpointer = Checkpointer(
serialization_dir,
num_serialized_models_to_keep=num_serialized_models_to_keep)
def save_checkpoint(self, epoch: int) -> None:
training_state = {
"epoch": epoch,
"optimizer": self.optimizer.state_dict()
}
self.checkpointer.save_checkpoint(epoch, self.model.state_dict(),
training_state, True)
def restore_checkpoint(self) -> int:
model_state, trainer_state = self.checkpointer.restore_checkpoint()
if not model_state and not trainer_state:
return 0
else:
self.model.load_state_dict(model_state)
self.optimizer.load_state_dict(trainer_state["optimizer"])
return trainer_state["epoch"] + 1
def train(self) -> Dict:
start_epoch = self.restore_checkpoint()
self.model.train()
for epoch in range(start_epoch, self.num_epochs):
total_loss = 0.0
batches = tqdm.tqdm(
self.iterator(self.mingler.mingle(self.datasets),
num_epochs=1))
for i, batch in enumerate(batches):
self.optimizer.zero_grad()
loss = self.model.forward(**batch)["loss"] # type: ignore
loss.backward()
total_loss += loss.item()
self.optimizer.step()
batches.set_description(
f"epoch: {epoch} loss: {total_loss / (i + 1)}")
# Save checkpoint
self.save_checkpoint(epoch)
return {}
@classmethod
def from_partial_objects(
cls,
serialization_dir: str,
train_dataset_readers: Dict[str, DatasetReader],
train_file_paths: Dict[str, str],
model: Lazy[Model],
iterator: DataIterator,
mingler: DatasetMingler,
optimizer: Lazy[Optimizer],
num_epochs: int = 10,
) -> "MultiTaskTrainer":
datasets = {
name: reader.read(train_file_paths[name])
for name, reader in train_dataset_readers.items()
}
instances = (instance for dataset in datasets.values()
for instance in dataset)
vocab = Vocabulary.from_instances(instances=instances)
model = model.construct(vocab=vocab)
iterator.index_with(vocab)
parameters = [[n, p] for n, p in model.named_parameters()
if p.requires_grad]
optimizer_ = optimizer.construct(model_parameters=parameters)
return MultiTaskTrainer(model, serialization_dir, iterator, mingler,
optimizer_, datasets, num_epochs)
class DistributeTrainer(DistributedTrainerBase):
"""
NOTE: only work in nprocess_ngpus
"""
def __init__(
self,
rank: int,
worldsize: int,
ngpus_per_node: int,
cuda_device: Union[int, List],
model: Model,
optimizer: torch.optim.Optimizer,
iterator: DataIterator,
train_dataset: Iterable[Instance],
validation_dataset: Optional[Iterable[Instance]] = None,
patience: Optional[int] = None,
validation_metric: str = "-loss",
validation_iterator: DataIterator = None,
shuffle: bool = True,
num_epochs: int = 20,
serialization_dir: Optional[str] = None,
num_serialized_models_to_keep: int = 20,
keep_serialized_model_every_num_seconds: int = None,
checkpointer: Checkpointer = None,
model_save_interval: float = None,
grad_norm: Optional[float] = None,
grad_clipping: Optional[float] = None,
learning_rate_scheduler: Optional[LearningRateScheduler] = None,
momentum_scheduler: Optional[MomentumScheduler] = None,
summary_interval: int = 100,
histogram_interval: int = None,
should_log_parameter_statistics: bool = True,
should_log_learning_rate: bool = False,
log_batch_size_period: Optional[int] = None,
moving_average: Optional[MovingAverage] = None) -> None:
super().__init__(rank, worldsize, ngpus_per_node, cuda_device,
serialization_dir)
self.model = model
self.iterator = iterator
self._validation_iterator = validation_iterator
self.shuffle = shuffle
self.optimizer = optimizer
self.train_data = train_dataset
self._validation_data = validation_dataset
self._metric_tracker = MetricTracker(patience, validation_metric)
self._validation_metric = validation_metric[1:]
self._num_epochs = num_epochs
# NOTE: although We have ckpter for everyone, only rank 0 of each node should be able to ckpt
if checkpointer is not None:
self._checkpointer = checkpointer
else:
self._checkpointer = Checkpointer(
serialization_dir, keep_serialized_model_every_num_seconds,
num_serialized_models_to_keep)
self._model_save_interval = model_save_interval
self._grad_norm = grad_norm
self._grad_clipping = grad_clipping
self._learning_rate_scheduler = learning_rate_scheduler
self._momentum_scheduler = momentum_scheduler
self._moving_average = moving_average
# We keep the total batch number as an instance variable because it
# is used inside a closure for the hook which logs activations in
# ``_enable_activation_logging``.
self._batch_num_total = 0
# NOTE: log.
serialization_dir = os.path.join(serialization_dir, str(rank))
self._tensorboard = TensorboardWriter(
get_batch_num_total=lambda: self._batch_num_total,
serialization_dir=serialization_dir,
summary_interval=summary_interval,
histogram_interval=histogram_interval,
should_log_parameter_statistics=should_log_parameter_statistics,
should_log_learning_rate=should_log_learning_rate)
self._log_batch_size_period = log_batch_size_period
self._last_log = 0.0 # time of last logging
# Enable activation logging.
if histogram_interval is not None:
self._tensorboard.enable_activation_logging(self.model)
def rescale_gradients(self) -> Optional[float]:
return training_util.rescale_gradients(self.model, self._grad_norm)
def batch_loss(self, batch_group: List[TensorDict],
for_training: bool) -> torch.Tensor:
assert len(batch_group) == 1
batch = batch_group[0]
batch = nn_util.move_to_device(batch, self._cuda_device[0])
output_dict = self.model(**batch)
try:
loss = output_dict["loss"]
if for_training:
loss += self.model.get_regularization_penalty()
except KeyError:
if for_training:
raise RuntimeError(
"The model you are trying to optimize does not contain a"
" 'loss' key in the output of model.forward(inputs).")
loss = None
return loss
def _train_epoch(self, epoch: int) -> Dict[str, float]:
"""
Trains one epoch and returns metrics.
only report system utils when we are local rank 0 at each machine.
"""
logger.info("Rank %d: Epoch %d/%d", self._rank, epoch,
self._num_epochs - 1)
peak_cpu_usage = peak_memory_mb()
if self._is_chief:
logger.info(f"Peak CPU memory usage MB: {peak_cpu_usage}")
train_loss = 0.0
# Set the model to "train" mode.
self.model.train()
# should be 1 anyway, because we are only dealing with nprocess_with_ngpus
num_gpus = len(self._cuda_device)
# TODO: Implementation of whether the generator should take into account of worldsize.
# Get tqdm for the training batches
raw_train_generator = self.iterator(self.train_data,
num_epochs=1,
shuffle=self.shuffle)
train_generator = lazy_groups_of(raw_train_generator, num_gpus)
num_training_batches = math.ceil(
self.iterator.get_num_batches(self.train_data) / num_gpus)
self._last_log = time.time()
last_save_time = time.time()
batches_this_epoch = 0
if self._batch_num_total is None:
self._batch_num_total = 0
histogram_parameters = set(
self.model.get_parameters_for_histogram_tensorboard_logging())
logger.info("Training")
train_generator_tqdm = Tqdm.tqdm(train_generator,
total=num_training_batches)
cumulative_batch_size = 0
# NOTE: only work in nprocess_ngpus
device = torch.device("cuda:%d" % self._cuda_device[0])
for batch_group in train_generator_tqdm:
batches_this_epoch += 1
self._batch_num_total += 1
batch_num_total = self._batch_num_total
self.optimizer.zero_grad()
loss = self.batch_loss(batch_group, for_training=True)
if torch.isnan(loss):
raise ValueError("nan loss encountered")
loss.backward()
train_loss += loss.item()
batch_grad_norm = self.rescale_gradients()
# This does nothing if batch_num_total is None or you are using a
# scheduler which doesn't update per batch.
if self._learning_rate_scheduler:
self._learning_rate_scheduler.step_batch(batch_num_total)
if self._momentum_scheduler:
self._momentum_scheduler.step_batch(batch_num_total)
if self._is_chief:
# only chief do tensorboard
if self._tensorboard.should_log_histograms_this_batch():
# get the magnitude of parameter updates for logging
# We need a copy of current parameters to compute magnitude of updates,
# and copy them to CPU so large models won't go OOM on the GPU.
param_updates = {
name: param.detach().cpu().clone()
for name, param in self.model.named_parameters()
}
self.optimizer.step()
for name, param in self.model.named_parameters():
param_updates[name].sub_(param.detach().cpu())
update_norm = torch.norm(param_updates[name].view(
-1, ))
param_norm = torch.norm(param.view(-1, )).cpu()
self._tensorboard.add_train_scalar(
"gradient_update/" + name,
update_norm / (param_norm + 1e-7))
else:
self.optimizer.step()
else:
self.optimizer.step()
# Update moving averages
# NOTE: not sure whether this need to be average
if self._moving_average is not None:
self._moving_average.apply(batch_num_total)
metrics = get_metrics(self.model, device, self._worldsize,
train_loss, batches_this_epoch)
description = training_util.description_from_metrics(metrics)
train_generator_tqdm.set_description(
("Rank %d: " % self._rank) + description, refresh=False)
if self._is_chief:
# Log parameter values to Tensorboard
if self._tensorboard.should_log_this_batch():
self._tensorboard.log_parameter_and_gradient_statistics(
self.model, batch_grad_norm)
self._tensorboard.log_learning_rates(
self.model, self.optimizer)
self._tensorboard.add_train_scalar("loss/loss_train",
metrics["loss"])
self._tensorboard.log_metrics(
{"epoch_metrics/" + k: v
for k, v in metrics.items()})
if self._tensorboard.should_log_histograms_this_batch():
self._tensorboard.log_histograms(self.model,
histogram_parameters)
if self._log_batch_size_period:
cur_batch = sum([
training_util.get_batch_size(batch)
for batch in batch_group
])
cumulative_batch_size += cur_batch
if (batches_this_epoch - 1) % self._log_batch_size_period == 0:
average = cumulative_batch_size / batches_this_epoch
logger.info(
f"rank {self._rank}, current batch size: {cur_batch} mean batch size: {average}"
)
if self._is_chief:
self._tensorboard.add_train_scalar(
"current_batch_size", cur_batch)
self._tensorboard.add_train_scalar(
"mean_batch_size", average)
if self._is_chief:
# Save model if needed.
if self._model_save_interval is not None and (
time.time() - last_save_time >
self._model_save_interval):
last_save_time = time.time()
self._save_checkpoint('{0}.{1}'.format(
epoch, training_util.time_to_str(int(last_save_time))))
metrics = get_metrics(self.model, device, self._worldsize,
train_loss, batches_this_epoch)
metrics['cpu_memory_MB'] = peak_cpu_usage
return metrics
def _validation_loss(self) -> Tuple[float, int]:
"""
Computes the validation loss. Returns it and the number of batches.
"""
logger.info("Rank %d Validating", self._rank)
self.model.eval()
# Replace parameter values with the shadow values from the moving averages.
if self._moving_average is not None:
self._moving_average.assign_average_value()
if self._validation_iterator is not None:
val_iterator = self._validation_iterator
else:
val_iterator = self.iterator
num_gpus = len(self._cuda_device)
raw_val_generator = val_iterator(self._validation_data,
num_epochs=1,
shuffle=False)
val_generator = lazy_groups_of(raw_val_generator, num_gpus)
num_validation_batches = math.ceil(
val_iterator.get_num_batches(self._validation_data) / num_gpus)
val_generator_tqdm = Tqdm.tqdm(val_generator,
total=num_validation_batches)
batches_this_epoch = 0
val_loss = 0
for batch_group in val_generator_tqdm:
loss = self.batch_loss(batch_group, for_training=False)
if loss is not None:
# You shouldn't necessarily have to compute a loss for validation, so we allow for
# `loss` to be None. We need to be careful, though - `batches_this_epoch` is
# currently only used as the divisor for the loss function, so we can safely only
# count those batches for which we actually have a loss. If this variable ever
# gets used for something else, we might need to change things around a bit.
batches_this_epoch += 1
val_loss += loss.detach().cpu().numpy()
# Update the description with the latest metrics
val_metrics = training_util.get_metrics(self.model, val_loss,
batches_this_epoch)
description = training_util.description_from_metrics(val_metrics)
val_generator_tqdm.set_description(description, refresh=False)
# Now restore the original parameter values.
if self._moving_average is not None:
self._moving_average.restore()
return val_loss, batches_this_epoch
def train(self) -> Dict[str, Any]:
"""
Trains the supplied model with the supplied parameters.
"""
try:
epoch_counter = self._restore_checkpoint()
except RuntimeError:
traceback.print_exc()
raise ConfigurationError(
"Could not recover training from the checkpoint. Did you mean to output to "
"a different serialization directory or delete the existing serialization "
"directory?")
training_util.enable_gradient_clipping(self.model, self._grad_clipping)
logger.info("Rank %d Beginning training.", self._rank)
train_metrics: Dict[str, float] = {}
val_metrics: Dict[str, float] = {}
this_epoch_val_metric: float = None
metrics: Dict[str, Any] = {}
epochs_trained = 0
training_start_time = time.time()
metrics['best_epoch'] = self._metric_tracker.best_epoch
for key, value in self._metric_tracker.best_epoch_metrics.items():
metrics["best_validation_" + key] = value
for epoch in range(epoch_counter, self._num_epochs):
epoch_start_time = time.time()
train_metrics = self._train_epoch(epoch)
# get peak of memory usage
if self._is_chief:
if 'cpu_memory_MB' in train_metrics:
metrics['peak_cpu_memory_MB'] = max(
metrics.get('peak_cpu_memory_MB', 0),
train_metrics['cpu_memory_MB'])
if self._validation_data is not None:
with torch.no_grad():
# We have a validation set, so compute all the metrics on it.
val_loss, num_batches = self._validation_loss()
val_metrics = training_util.get_metrics(self.model,
val_loss,
num_batches,
reset=True)
# Check validation metric for early stopping
this_epoch_val_metric = val_metrics[
self._validation_metric]
self._metric_tracker.add_metric(this_epoch_val_metric)
if self._metric_tracker.should_stop_early():
logger.info(
"Ran out of patience. Stopping training.")
break
if self._is_chief:
self._tensorboard.log_metrics(
train_metrics,
val_metrics=val_metrics,
log_to_console=True,
epoch=epoch +
1) # +1 because tensorboard doesn't like 0
# Create overall metrics dict
training_elapsed_time = time.time() - training_start_time
metrics["training_duration"] = str(
datetime.timedelta(seconds=training_elapsed_time))
metrics["training_start_epoch"] = epoch_counter
metrics["training_epochs"] = epochs_trained
metrics["epoch"] = epoch
for key, value in train_metrics.items():
metrics["training_" + key] = value
for key, value in val_metrics.items():
metrics["validation_" + key] = value
if self._metric_tracker.is_best_so_far() and self._is_chief:
# Update all the best_ metrics.
# (Otherwise they just stay the same as they were.)
metrics['best_epoch'] = epoch
for key, value in val_metrics.items():
metrics["best_validation_" + key] = value
self._metric_tracker.best_epoch_metrics = val_metrics
if self._serialization_dir and self._is_chief:
dump_metrics(
os.path.join(self._serialization_dir,
f'metrics_epoch_{epoch}.json'), metrics)
# The Scheduler API is agnostic to whether your schedule requires a validation metric -
# if it doesn't, the validation metric passed here is ignored.
if self._learning_rate_scheduler:
self._learning_rate_scheduler.step(this_epoch_val_metric,
epoch)
if self._momentum_scheduler:
self._momentum_scheduler.step(this_epoch_val_metric, epoch)
if self._is_chief:
self._save_checkpoint(epoch)
epoch_elapsed_time = time.time() - epoch_start_time
logger.info("Rank %d Epoch duration: %s", self._rank,
datetime.timedelta(seconds=epoch_elapsed_time))
if epoch < self._num_epochs - 1:
training_elapsed_time = time.time() - training_start_time
estimated_time_remaining = training_elapsed_time * \
((self._num_epochs - epoch_counter) / float(epoch - epoch_counter + 1) - 1)
formatted_time = str(
datetime.timedelta(
seconds=int(estimated_time_remaining)))
logger.info(
"Rank %d, Estimated training time remaining: %s",
self._rank, formatted_time)
epochs_trained += 1
# make sure pending events are flushed to disk and files are closed properly
self._tensorboard.close()
# Load the best model state before returning
if self._is_chief:
best_model_state = self._checkpointer.best_model_state()
if best_model_state:
self.model.load_state_dict(best_model_state)
return metrics
def _save_checkpoint(self, epoch: Union[int, str]) -> None:
"""
Saves a checkpoint of the model to self._serialization_dir.
Is a no-op if self._serialization_dir is None.
Parameters
----------
epoch : Union[int, str], required.
The epoch of training. If the checkpoint is saved in the middle
of an epoch, the parameter is a string with the epoch and timestamp.
"""
# If moving averages are used for parameters, we save
# the moving average values into checkpoint, instead of the current values.
if self._moving_average is not None:
self._moving_average.assign_average_value()
# These are the training states we need to persist.
training_states = {
"metric_tracker": self._metric_tracker.state_dict(),
"optimizer": self.optimizer.state_dict(),
"batch_num_total": self._batch_num_total
}
# If we have a learning rate or momentum scheduler, we should persist them too.
if self._learning_rate_scheduler is not None:
training_states[
"learning_rate_scheduler"] = self._learning_rate_scheduler.state_dict(
)
if self._momentum_scheduler is not None:
training_states[
"momentum_scheduler"] = self._momentum_scheduler.state_dict()
self._checkpointer.save_checkpoint(
model_state=self.model.state_dict(),
epoch=epoch,
training_states=training_states,
is_best_so_far=self._metric_tracker.is_best_so_far())
# Restore the original values for parameters so that training will not be affected.
if self._moving_average is not None:
self._moving_average.restore()
def _restore_checkpoint(self) -> int:
"""
Restores the model and training state from the last saved checkpoint.
This includes an epoch count and optimizer state, which is serialized separately
from model parameters. This function should only be used to continue training -
if you wish to load a model for inference/load parts of a model into a new
computation graph, you should use the native Pytorch functions:
`` model.load_state_dict(torch.load("/path/to/model/weights.th"))``
If ``self._serialization_dir`` does not exist or does not contain any checkpointed weights,
this function will do nothing and return 0.
Returns
-------
epoch: int
The epoch at which to resume training, which should be one after the epoch
in the saved training state.
"""
model_state, training_state = self._checkpointer.restore_checkpoint()
if not training_state:
# No checkpoint to restore, start at 0
return 0
self.model.load_state_dict(model_state)
self.optimizer.load_state_dict(training_state["optimizer"])
if self._learning_rate_scheduler is not None and "learning_rate_scheduler" in training_state:
self._learning_rate_scheduler.load_state_dict(
training_state["learning_rate_scheduler"])
if self._momentum_scheduler is not None and "momentum_scheduler" in training_state:
self._momentum_scheduler.load_state_dict(
training_state["momentum_scheduler"])
training_util.move_optimizer_to_cuda(self.optimizer)
# Currently the ``training_state`` contains a serialized ``MetricTracker``.
if "metric_tracker" in training_state:
self._metric_tracker.load_state_dict(
training_state["metric_tracker"])
# It used to be the case that we tracked ``val_metric_per_epoch``.
elif "val_metric_per_epoch" in training_state:
self._metric_tracker.clear()
self._metric_tracker.add_metrics(
training_state["val_metric_per_epoch"])
# And before that we didn't track anything.
else:
self._metric_tracker.clear()
if isinstance(training_state["epoch"], int):
epoch_to_return = training_state["epoch"] + 1
else:
epoch_to_return = int(training_state["epoch"].split('.')[0]) + 1
# For older checkpoints with batch_num_total missing, default to old behavior where
# it is unchanged.
batch_num_total = training_state.get('batch_num_total')
if batch_num_total is not None:
self._batch_num_total = batch_num_total
return epoch_to_return