def default_callbacks( self, validation_metric: str = "-loss", patience: int = None, max_checkpoints: int = 20, checkpoint_every: int = None, model_save_interval: float = None, serialization_dir: str = "__DEFAULT__", validation_data: Iterable[Instance] = None, validation_iterator: DataIterator = None, batch_size: int = 2, ): if serialization_dir == "__DEFAULT__": serialization_dir = self.TEST_DIR checkpointer = Checkpointer(serialization_dir, checkpoint_every, max_checkpoints) tensorboard = TensorboardWriter(get_batch_num_total=lambda: None) if validation_iterator is None: validation_iterator = BasicIterator(batch_size=batch_size) validation_iterator.index_with(self.vocab) return [ LogToTensorboard(log_batch_size_period=10, tensorboard=tensorboard), Checkpoint(checkpointer, model_save_interval), Validate( validation_data=self.instances if validation_data is None else validation_data, validation_iterator=validation_iterator, ), TrackMetrics(patience, validation_metric), GradientNormAndClip(), ]
def default_callbacks(self, validation_metric: str = "-loss", patience: int = None, max_checkpoints: int = 20, checkpoint_every: int = None, serialization_dir: str = "__DEFAULT__", iterator: DataIterator = None, validation_data: Iterable[Instance] = None, validation_iterator: DataIterator = None, batch_size: int = 2): if serialization_dir == "__DEFAULT__": serialization_dir = self.TEST_DIR checkpointer = Checkpointer(serialization_dir, checkpoint_every, max_checkpoints) tensorboard = TensorboardWriter(get_batch_num_total=lambda: None) if iterator is None: iterator = BasicIterator(batch_size=batch_size) iterator.index_with(self.vocab) return [ LogToTensorboard(log_batch_size_period=10, tensorboard=tensorboard), Checkpoint(checkpointer), Validate(validation_data=self.instances if validation_data is None else validation_data, validation_iterator=iterator if validation_iterator is None else validation_iterator), TrackMetrics(patience, validation_metric), TrainSupervised(), GenerateTrainingBatches(self.instances, iterator, True) ]
def from_params( # type: ignore cls, serialization_dir: str, params: Params) -> "LogToTensorboard": log_batch_size_period = params.pop_int("log_batch_size_period", None) tensorboard = TensorboardWriter.from_params( params=params, serialization_dir=serialization_dir, get_batch_num_total=lambda: None) return LogToTensorboard(tensorboard, log_batch_size_period)
def from_params( # type: ignore cls, serialization_dir: str, params: Params, **extras) -> "LogToTensorboard": log_batch_size_period = params.pop_int("log_batch_size_period", None) tensorboard = TensorboardWriter.from_params( params=params, serialization_dir=serialization_dir, get_batch_num_total=lambda: None) # TODO(mattg): remove get_batch_num_total from TensorboardWriter, and instead just add a # method / arguments to tell the writer what batch num we're at. return LogToTensorboard(tensorboard, log_batch_size_period)
def test_trainer_can_log_learning_rates_tensorboard(self): callbacks = [cb for cb in self.default_callbacks() if not isinstance(cb, LogToTensorboard)] # The lambda: None is unfortunate, but it will get replaced by the callback. tensorboard = TensorboardWriter(lambda: None, should_log_learning_rate=True, summary_interval=2) callbacks.append(LogToTensorboard(tensorboard)) trainer = CallbackTrainer(self.model, self.optimizer, num_epochs=2, serialization_dir=self.TEST_DIR, callbacks=callbacks) trainer.train()
def test_trainer_can_log_histograms(self): # enable activation logging for module in self.model.modules(): module.should_log_activations = True callbacks = [cb for cb in self.default_callbacks() if not isinstance(cb, LogToTensorboard)] # The lambda: None is unfortunate, but it will get replaced by the callback. tensorboard = TensorboardWriter(lambda: None, histogram_interval=2) callbacks.append(LogToTensorboard(tensorboard)) trainer = CallbackTrainer(self.model, self.optimizer, num_epochs=3, serialization_dir=self.TEST_DIR, callbacks=callbacks) trainer.train()
def test_model_training(self): training_dataset = self.sample_instances if self.sample_only else self.train_instances #training_dataset = training_dataset[:500] validation_dataset = self.sample_instances if self.sample_only else self.test_instances serialization_dir = self.TEST_DATA_ROOT / "serialized_sample" if self.sample_only else "serialized" tensorboard_dir = self.TEST_DATA_ROOT / "tensorboard" batch_size = 64 train_iterator = BucketIterator(sorting_keys=[("question", "num_tokens")], padding_noise=0.0, batch_size=batch_size) val_iterator = BucketIterator(sorting_keys=[("question", "num_tokens") ], padding_noise=0.0, batch_size=batch_size) train_iterator.index_with(vocab=self.vocab) val_iterator.index_with(vocab=self.vocab) tensorboard = TensorboardWriter(get_batch_num_total=lambda: np.ceil( len(training_dataset) / batch_size), serialization_dir=tensorboard_dir, summary_interval=5, histogram_interval=5, should_log_parameter_statistics=True) trainer = CallbackTrainer( model=self.model, serialization_dir=serialization_dir, iterator=train_iterator, training_data=training_dataset, num_epochs=20, cuda_device=0, optimizer=torch.optim.Adagrad(self.model.parameters()), callbacks=[ LogToTensorboard(tensorboard), Validate(validation_data=validation_dataset, validation_iterator=val_iterator), TrackMetrics(), ResetMetricsCallback() ]) trainer.train() self.val_outputs_fp.close()
def test_model_training(self): serialization_dir = self.TEST_DATA_ROOT / "serialized_sample" tensorboard_dir = self.TEST_DATA_ROOT / "tensorboard.seq2seq" batch_size = 64 train_iterator = BucketIterator(sorting_keys=[("source_tokens", "num_tokens")], padding_noise=0.0, batch_size=batch_size) train_iterator.index_with(vocab=self.vocab) tensorboard = TensorboardWriter( get_batch_num_total=lambda: np.ceil(len(self.train_instances) / batch_size), serialization_dir=tensorboard_dir, summary_interval=5, histogram_interval=5, should_log_parameter_statistics=True) trainer = CallbackTrainer(model=self.model, serialization_dir=serialization_dir, iterator=train_iterator, training_data=self.train_instances, num_epochs=1, cuda_device=0, optimizer=torch.optim.Adam(self.model.parameters(), lr=1e-3), callbacks=[LogToTensorboard(tensorboard), Validate(validation_data=self.dev_instances, validation_iterator=train_iterator), TrackMetrics(), ResetMetricsCallback()] ) for i in range(50): print('Epoch: {}'.format(i)) trainer.train() import itertools predictor = Seq2SeqPredictor(self.model, self.reader) for instance in itertools.islice(self.dev_instances, 10): print('SOURCE:', instance.fields['source_tokens'].tokens) print('GOLD:', instance.fields['target_tokens'].tokens) print('PRED:', predictor.predict_instance(instance)['predicted_tokens'])
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 __init__(self, model: Model, task_list: List[Task], optimizer_params: Params, lr_scheduler_params: Params, patience: Optional[int] = None, num_epochs: int = 20, serialization_dir: Optional[str] = None, cuda_device: int = -1, gradient_accumulation_steps: int = 1, grad_norm: Optional[float] = None, grad_clipping: Optional[float] = None, min_lr: float = 0.00001, no_tqdm: bool = False, momentum_scheduler: Optional[MomentumScheduler] = None, summary_interval: int = 50, histogram_interval: int = 50, should_log_parameter_statistics: bool = True, should_log_learning_rate: bool = True, sampling_method: str = "proportional", moving_average: Optional[MovingAverage] = None) -> None: self._model = model self._task_list = task_list self._n_tasks = len(self._task_list) self._optimizer_params = optimizer_params self._optimizers = {} self._lr_scheduler_params = lr_scheduler_params self._schedulers = {} self._all_params = [(n, p) for n, p in self._model.named_parameters() if p.requires_grad] self._params_exclude_share_encoder = [(n, p) for n, p in self._model.named_parameters() if p.requires_grad and "_share_encoder" not in n] self._params_exclude_share_discriminator = [(n, p) for n, p in self._model.named_parameters() if p.requires_grad and "_s_domain_discriminator" not in n] for task in self._task_list: task_name = task._name self._optimizers[task_name] = {} self._optimizers[task_name]["all_params"] = Optimizer.from_params( model_parameters=self._all_params, params=deepcopy(optimizer_params) ) self._optimizers[task_name]["exclude_share_encoder"] = Optimizer.from_params( model_parameters=self._params_exclude_share_encoder, params=deepcopy(optimizer_params) ) self._optimizers[task_name]["exclude_share_discriminator"] = Optimizer.from_params( model_parameters=self._params_exclude_share_discriminator, params=deepcopy(optimizer_params) ) self._schedulers[task_name] = {} self._schedulers[task_name] = LearningRateScheduler.from_params( optimizer=self._optimizers[task_name]["all_params"], params=deepcopy(lr_scheduler_params) ) # self._schedulers[task_name]["all_params"] = LearningRateScheduler.from_params( # optimizer=self._optimizers[task_name]["all_params"], params=deepcopy(lr_scheduler_params) # ) # self._schedulers[task_name]["exclude_share_encoder"] = LearningRateScheduler.from_params( # optimizer=self._optimizers[task_name]["exclude_share_encoder"], params=deepcopy(lr_scheduler_params) # ) # self._schedulers[task_name]["exclude_share_discriminator"] = LearningRateScheduler.from_params( # optimizer=self._optimizers[task_name]["exclude_share_discriminator"], # params=deepcopy(lr_scheduler_params) # ) self._serialization_dir = serialization_dir self._cuda_device = cuda_device if self._cuda_device >= 0: check_for_gpu(self._cuda_device) self._model = self._model.cuda(self._cuda_device) self._patience = patience self._num_epochs = num_epochs self._epoch_trained = 0 self._gradient_accumulation_steps = gradient_accumulation_steps self._grad_norm = grad_norm self._grad_clipping = grad_clipping self._min_lr = min_lr self._no_tqdm = no_tqdm self._sampling_method = sampling_method self._task_infos = None self._metric_infos = None self._tr_generators = None self._global_step = 0 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._last_log = 0.0 # time of last logging # Enable activation logging. if histogram_interval is not None: self._tensorboard.enable_activation_logging(self._model)
class GanMtlTrainer: def __init__(self, model: Model, task_list: List[Task], optimizer_params: Params, lr_scheduler_params: Params, patience: Optional[int] = None, num_epochs: int = 20, serialization_dir: Optional[str] = None, cuda_device: int = -1, gradient_accumulation_steps: int = 1, grad_norm: Optional[float] = None, grad_clipping: Optional[float] = None, min_lr: float = 0.00001, no_tqdm: bool = False, momentum_scheduler: Optional[MomentumScheduler] = None, summary_interval: int = 50, histogram_interval: int = 50, should_log_parameter_statistics: bool = True, should_log_learning_rate: bool = True, sampling_method: str = "proportional", moving_average: Optional[MovingAverage] = None) -> None: self._model = model self._task_list = task_list self._n_tasks = len(self._task_list) self._optimizer_params = optimizer_params self._optimizers = {} self._lr_scheduler_params = lr_scheduler_params self._schedulers = {} self._all_params = [(n, p) for n, p in self._model.named_parameters() if p.requires_grad] self._params_exclude_share_encoder = [(n, p) for n, p in self._model.named_parameters() if p.requires_grad and "_share_encoder" not in n] self._params_exclude_share_discriminator = [(n, p) for n, p in self._model.named_parameters() if p.requires_grad and "_s_domain_discriminator" not in n] for task in self._task_list: task_name = task._name self._optimizers[task_name] = {} self._optimizers[task_name]["all_params"] = Optimizer.from_params( model_parameters=self._all_params, params=deepcopy(optimizer_params) ) self._optimizers[task_name]["exclude_share_encoder"] = Optimizer.from_params( model_parameters=self._params_exclude_share_encoder, params=deepcopy(optimizer_params) ) self._optimizers[task_name]["exclude_share_discriminator"] = Optimizer.from_params( model_parameters=self._params_exclude_share_discriminator, params=deepcopy(optimizer_params) ) self._schedulers[task_name] = {} self._schedulers[task_name] = LearningRateScheduler.from_params( optimizer=self._optimizers[task_name]["all_params"], params=deepcopy(lr_scheduler_params) ) # self._schedulers[task_name]["all_params"] = LearningRateScheduler.from_params( # optimizer=self._optimizers[task_name]["all_params"], params=deepcopy(lr_scheduler_params) # ) # self._schedulers[task_name]["exclude_share_encoder"] = LearningRateScheduler.from_params( # optimizer=self._optimizers[task_name]["exclude_share_encoder"], params=deepcopy(lr_scheduler_params) # ) # self._schedulers[task_name]["exclude_share_discriminator"] = LearningRateScheduler.from_params( # optimizer=self._optimizers[task_name]["exclude_share_discriminator"], # params=deepcopy(lr_scheduler_params) # ) self._serialization_dir = serialization_dir self._cuda_device = cuda_device if self._cuda_device >= 0: check_for_gpu(self._cuda_device) self._model = self._model.cuda(self._cuda_device) self._patience = patience self._num_epochs = num_epochs self._epoch_trained = 0 self._gradient_accumulation_steps = gradient_accumulation_steps self._grad_norm = grad_norm self._grad_clipping = grad_clipping self._min_lr = min_lr self._no_tqdm = no_tqdm self._sampling_method = sampling_method self._task_infos = None self._metric_infos = None self._tr_generators = None self._global_step = 0 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._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 _enable_gradient_clipping(self) -> None: training_util.enable_gradient_clipping(self._model, self._grad_clipping) def _train_epoch(self, total_n_tr_batches: int, sampling_prob: List, reverse=False, train_D=False) -> Dict[ str, float]: self._model.train() # Set the model to "train" mode. if reverse: logger.info("Training Generator- Begin") elif not train_D: logger.info("Training Init Generator- Begin") if train_D: logger.info("Training Discriminator- Begin") logger.info("reverse is {}, train_D is {}", reverse, train_D) ### Reset training and trained batches counter before new training epoch ### for _, task_info in self._task_infos.items(): task_info["tr_loss_cum"] = 0.0 task_info['stm_loss'] = 0.0 task_info['p_d_loss'] = 0.0 task_info['s_d_loss'] = 0.0 task_info['valid_loss'] = 0.0 task_info["n_batches_trained_this_epoch"] = 0 all_tr_metrics = {} # BUG TO COMPLETE COMMENT TO MAKE IT MORE CLEAR ### Start training epoch ### epoch_tqdm = tqdm.tqdm(range(total_n_tr_batches), total=total_n_tr_batches) histogram_parameters = set(self._model.get_parameters_for_histogram_tensorboard_logging()) for step, _ in enumerate(epoch_tqdm): task_idx = np.argmax(np.random.multinomial(1, sampling_prob)) task = self._task_list[task_idx] task_info = self._task_infos[task._name] ### One forward + backward pass ### # Call next batch to train batch = next(self._tr_generators[task._name]) self._batch_num_total += 1 task_info["n_batches_trained_this_epoch"] += 1 # Load optimizer if not train_D: optimizer = self._optimizers[task._name]["all_params"] else: optimizer = self._optimizers[task._name]["exclude_share_encoder"] # Get the loss for this batch output_dict = self._forward(tensor_batch=batch, task=task, for_training=True, reverse=reverse) # if reverse or train_D: # output_dict_fake = self._forward(tensor_batch=batch, task=task, for_training=True, reverse=True) # loss = output_dict["stm_loss"] # if train_D: # loss = (output_dict["stm_loss"] + output_dict["s_d_loss"] + output_dict_fake["stm_loss"] + # output_dict_fake["s_d_loss"]) / 2.0 # if reverse: # # loss = (output_dict["stm_loss"] + output_dict["p_d_loss"] + 0.005 * output_dict["s_d_loss"] + # # output_dict_fake["stm_loss"] + output_dict_fake["p_d_loss"] + 0.005 * output_dict_fake[ # # "s_d_loss"]) / 2.0 # loss = (output_dict['loss'] + output_dict_fake['loss']) / 2.0 loss = output_dict['loss'] if self._gradient_accumulation_steps > 1: loss /= self._gradient_accumulation_steps loss.backward() task_info["tr_loss_cum"] += loss.item() task_info['stm_loss'] += output_dict['stm_loss'].item() task_info['p_d_loss'] += output_dict['p_d_loss'].item() task_info['s_d_loss'] += output_dict['s_d_loss'].item() task_info['valid_loss'] += output_dict['valid_loss'].item() # if reverse or train_D: # task_info['stm_loss'] += output_dict_fake['stm_loss'].item() # task_info['stm_loss'] /= 2.0 # task_info['p_d_loss'] += output_dict_fake['p_d_loss'].item() # task_info['p_d_loss'] /= 2.0 # task_info['s_d_loss'] += output_dict_fake['s_d_loss'].item() # task_info['s_d_loss'] /= 2.0 # task_info['valid_loss'] += output_dict_fake['valid_loss'].item() # task_info['valid_loss'] /= 2.0 del loss if (step + 1) % self._gradient_accumulation_steps == 0: batch_grad_norm = self._rescale_gradients() if self._tensorboard.should_log_histograms_this_batch(): param_updates = {name: param.detach().cpu().clone() for name, param in self._model.named_parameters()} 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: optimizer.step() optimizer.zero_grad() ### Get metrics for all progress so far, update tqdm, display description ### task_metrics = self._get_metrics(task=task) task_metrics["loss"] = float( task_info["tr_loss_cum"] / (task_info["n_batches_trained_this_epoch"] + 0.000_001) ) task_metrics["stm_loss"] = float( task_info["stm_loss"] / (task_info["n_batches_trained_this_epoch"] + 0.000_001) ) task_metrics["p_d_loss"] = float( task_info["p_d_loss"] / (task_info["n_batches_trained_this_epoch"] + 0.000_001) ) task_metrics["s_d_loss"] = float( task_info["s_d_loss"] / (task_info["n_batches_trained_this_epoch"] + 0.000_001) ) task_metrics["valid_loss"] = float( task_info["valid_loss"] / (task_info["n_batches_trained_this_epoch"] + 0.000_001) ) description = training_util.description_from_metrics(task_metrics) epoch_tqdm.set_description(task._name + ", " + description) # 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, optimizer) self._tensorboard.log_metrics( {"epoch_metrics/" + task._name + "/" + k: v for k, v in task_metrics.items()}) if self._tensorboard.should_log_histograms_this_batch(): self._tensorboard.log_histograms(self._model, histogram_parameters) self._global_step += 1 ### Bookkeeping all the training metrics for all the tasks on the training epoch that just finished ### for task in self._task_list: task_info = self._task_infos[task._name] task_info["total_n_batches_trained"] += task_info["n_batches_trained_this_epoch"] task_info["last_log"] = time.time() task_metrics = self._get_metrics(task=task, reset=True) if task._name not in all_tr_metrics: all_tr_metrics[task._name] = {} for name, value in task_metrics.items(): all_tr_metrics[task._name][name] = value all_tr_metrics[task._name]["loss"] = float( task_info["tr_loss_cum"] / (task_info["n_batches_trained_this_epoch"] + 0.000_000_01) ) all_tr_metrics[task._name]["stm_loss"] = float( task_info["stm_loss"] / (task_info["n_batches_trained_this_epoch"] + 0.000_000_01) ) all_tr_metrics[task._name]["p_d_loss"] = float( task_info["p_d_loss"] / (task_info["n_batches_trained_this_epoch"] + 0.000_000_01) ) all_tr_metrics[task._name]["s_d_loss"] = float( task_info["s_d_loss"] / (task_info["n_batches_trained_this_epoch"] + 0.000_000_01) ) all_tr_metrics[task._name]["valid_loss"] = float( task_info["valid_loss"] / (task_info["n_batches_trained_this_epoch"] + 0.000_000_01) ) # Tensorboard - Training metrics for this epoch for metric_name, value in all_tr_metrics[task._name].items(): self._tensorboard.add_train_scalar( name="task_" + task._name + "/" + metric_name, value=value ) logger.info("Train - End") return all_tr_metrics def _validation(self, n_epoch: int) -> Tuple[float, int]: ### Begin validation of the model ### logger.info("Validation - Begin") all_val_metrics = {} self._model.eval() # Set the model into evaluation mode avg_accuracy = 0.0 for task_idx, task in enumerate(self._task_list): logger.info("Validation - Task {}/{}: {}", task_idx + 1, self._n_tasks, task._name) val_loss = 0.0 n_batches_val_this_epoch_this_task = 0 n_val_batches = self._task_infos[task._name]["n_val_batches"] scheduler = self._schedulers[task._name] # Create tqdm generator for current tasks's validation data_iterator = task._data_iterator val_generator = data_iterator(task._validation_data, num_epochs=1, shuffle=False) val_generator_tqdm = tqdm.tqdm(val_generator, total=n_val_batches) # Iterate over each validation batch for this tasks for batch in val_generator_tqdm: n_batches_val_this_epoch_this_task += 1 # Get the loss val_output_dict = self._forward(batch, task=task, for_training=False) loss = val_output_dict["stm_loss"] val_loss += loss.item() del loss # Get metrics for all progress so far, update tqdm, display description task_metrics = self._get_metrics(task=task) task_metrics["loss"] = float(val_loss / n_batches_val_this_epoch_this_task) description = training_util.description_from_metrics(task_metrics) val_generator_tqdm.set_description(description) # Get tasks validation metrics and store them in all_val_metrics task_metrics = self._get_metrics(task=task, reset=True) if task._name not in all_val_metrics: all_val_metrics[task._name] = {} for name, value in task_metrics.items(): all_val_metrics[task._name][name] = value all_val_metrics[task._name]["loss"] = float(val_loss / n_batches_val_this_epoch_this_task) avg_accuracy += task_metrics["sentiment_acc"] # Tensorboard - Validation metrics for this epoch for metric_name, value in all_val_metrics[task._name].items(): self._tensorboard.add_validation_scalar( name="task_" + task._name + "/" + metric_name, value=value ) ### Perform a patience check and update the history of validation metric for this tasks ### this_epoch_val_metric = all_val_metrics[task._name][task._val_metric] metric_history = self._metric_infos[task._name]["hist"] metric_history.append(this_epoch_val_metric) is_best_so_far, out_of_patience = self._check_history( metric_history=metric_history, cur_score=this_epoch_val_metric, should_decrease=task._val_metric_decreases, ) if is_best_so_far: logger.info("Best model found for {}.", task._name) self._metric_infos[task._name]["best"] = (n_epoch, all_val_metrics) if out_of_patience and not self._metric_infos[task._name]["is_out_of_patience"]: self._metric_infos[task._name]["is_out_of_patience"] = True logger.info("Task {} is out of patience and vote to stop the training.", task._name) # The LRScheduler API is agnostic to whether your schedule requires a validation metric - # if it doesn't, the validation metric passed here is ignored. scheduler.step(this_epoch_val_metric, n_epoch) logger.info("Validation - End") return all_val_metrics, avg_accuracy def train(self, recover: bool = False) -> Dict[str, Any]: # 1 train sentiment classifier & private classifier & domain embeddings => init G 50 epoch # 2 fix share encoder(+domain embeddings?), train share classifier(cls&real/fake) & others => train D # 3 fix share classifier, train share encoder, reverse share classifier input gradient min loss => train G training_start_time = time.time() if recover: try: n_epoch, should_stop = self._restore_checkpoint() logger.info("Loaded model from checkpoint. Starting at epoch {}", n_epoch) except RuntimeError: 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?" ) else: n_epoch, should_stop = 0, False ### Store all the necessary informations and attributes about the tasks ### task_infos = {task._name: {} for task in self._task_list} for task_idx, task in enumerate(self._task_list): task_info = task_infos[task._name] # Store statistiscs on training and validation batches data_iterator = task._data_iterator n_tr_batches = data_iterator.get_num_batches(task._train_data) n_val_batches = data_iterator.get_num_batches(task._validation_data) task_info["n_tr_batches"] = n_tr_batches task_info["n_val_batches"] = n_val_batches # Create counter for number of batches trained during the whole # training for this specific tasks task_info["total_n_batches_trained"] = 0 task_info["last_log"] = time.time() # Time of last logging self._task_infos = task_infos ### Bookkeeping the validation metrics ### metric_infos = { task._name: { "val_metric": task._val_metric, "hist": [], "is_out_of_patience": False, "min_lr_hit": False, "best": (-1, {}), } for task in self._task_list } self._metric_infos = metric_infos ### Write log ### total_n_tr_batches = 0 # The total number of training batches across all the datasets. for task_name, info in self._task_infos.items(): total_n_tr_batches += info["n_tr_batches"] logger.info("Task {}:", task_name) logger.info("\t{} training batches", info["n_tr_batches"]) logger.info("\t{} validation batches", info["n_val_batches"]) ### Create the training generators/iterators tqdm ### self._tr_generators = {} for task in self._task_list: data_iterator = task._data_iterator tr_generator = data_iterator(task._train_data, num_epochs=None) self._tr_generators[task._name] = tr_generator ### Create sampling probability distribution ### if self._sampling_method == "uniform": sampling_prob = [float(1 / self._n_tasks)] * self._n_tasks elif self._sampling_method == "proportional": sampling_prob = [float(info["n_tr_batches"] / total_n_tr_batches) for info in self._task_infos.values()] ### Enable gradient clipping ### # Only if self._grad_clipping is specified self._enable_gradient_clipping() ### Setup is ready. Training of the model can begin ### logger.info("Set up ready. Beginning training/validation.") avg_accuracies = [] best_accuracy = 0.0 ### Begin Training of the model ### while not should_stop: ### Log Infos: current epoch count and CPU/GPU usage ### logger.info("") logger.info("Epoch {}/{} - Begin", n_epoch, self._num_epochs - 1) logger.info(f"Peak CPU memory usage MB: {peak_memory_mb()}") for gpu, memory in gpu_memory_mb().items(): logger.info(f"GPU {gpu} memory usage MB: {memory}") # if n_epoch <= 10: # # init generator # all_tr_metrics = self._train_epoch(total_n_tr_batches, sampling_prob) # # train discriminator 3 epochs # # elif 10 < n_epoch < 20 or n_epoch % 2 == 0: # # all_tr_metrics = self._train_epoch(total_n_tr_batches, sampling_prob, train_D=True) # else: # train adversarial generator every 3 epoch all_tr_metrics = self._train_epoch(total_n_tr_batches, sampling_prob, reverse=True) all_val_metrics, avg_accuracy = self._validation(n_epoch) is_best = False if best_accuracy < avg_accuracy: best_accuracy = avg_accuracy logger.info("Best accuracy found --- {}", best_accuracy / self._n_tasks) is_best = True ### Print all training and validation metrics for this epoch ### logger.info("***** Epoch {}/{} Statistics *****", n_epoch, self._num_epochs - 1) for task in self._task_list: logger.info("Statistic: {}", task._name) logger.info( "\tTraining - {}: {:3d}", "Nb batches trained", self._task_infos[task._name]["n_batches_trained_this_epoch"], ) for metric_name, value in all_tr_metrics[task._name].items(): logger.info("\tTraining - {}: {:.3f}", metric_name, value) for metric_name, value in all_val_metrics[task._name].items(): logger.info("\tValidation - {}: {:.3f}", metric_name, value) logger.info("***** Average accuracy is {:.6f} *****", avg_accuracy / self._n_tasks) avg_accuracies.append(avg_accuracy / self._n_tasks) logger.info("**********") ### Check to see if should stop ### stop_tr, stop_val = True, True for task in self._task_list: # task_info = self._task_infos[tasks._name] if self._optimizers[task._name]['exclude_share_encoder'].param_groups[0]["lr"] < self._min_lr and \ self._optimizers[task._name]['exclude_share_discriminator'].param_groups[0][ "lr"] < self._min_lr: logger.info("Minimum lr hit on {}.", task._name) logger.info("Task {} vote to stop training.", task._name) metric_infos[task._name]["min_lr_hit"] = True stop_tr = stop_tr and self._metric_infos[task._name]["min_lr_hit"] stop_val = stop_val and self._metric_infos[task._name]["is_out_of_patience"] if stop_tr: should_stop = True logger.info("All tasks hit minimum lr. Stopping training.") if stop_val: should_stop = True logger.info("All metrics ran out of patience. Stopping training.") if n_epoch >= self._num_epochs - 1: should_stop = True logger.info("Maximum number of epoch hit. Stopping training.") self._save_checkpoint(n_epoch, should_stop, is_best) ### Update n_epoch ### # One epoch = doing N (forward + backward) pass where N is the total number of training batches. n_epoch += 1 self._epoch_trained = n_epoch logger.info("Max accuracy is {:.6f}", max(avg_accuracies)) ### Summarize training at the end ### logger.info("***** Training is finished *****") logger.info("Stopped training after {} epochs", n_epoch) return_metrics = {} for task_name, task_info in self._task_infos.items(): nb_epoch_trained = int(task_info["total_n_batches_trained"] / task_info["n_tr_batches"]) logger.info( "Trained {} for {} batches ~= {} epochs", task_name, task_info["total_n_batches_trained"], nb_epoch_trained, ) return_metrics[task_name] = { "best_epoch": self._metric_infos[task_name]["best"][0], "nb_epoch_trained": nb_epoch_trained, "best_epoch_val_metrics": self._metric_infos[task_name]["best"][1], } training_elapsed_time = time.time() - training_start_time return_metrics["training_duration"] = time.strftime("%d:%H:%M:%S", time.gmtime(training_elapsed_time)) return_metrics["nb_epoch_trained"] = n_epoch return return_metrics def _check_history(self, metric_history: List[float], cur_score: float, should_decrease: bool = False): patience = self._patience + 1 best_fn = min if should_decrease else max best_score = best_fn(metric_history) if best_score == cur_score: best_so_far = metric_history.index(best_score) == len(metric_history) - 1 else: best_so_far = False out_of_patience = False if len(metric_history) > patience: if should_decrease: out_of_patience = max(metric_history[-patience:]) <= cur_score else: out_of_patience = min(metric_history[-patience:]) >= cur_score if best_so_far and out_of_patience: # then something is up print("Something is up") return best_so_far, out_of_patience def _forward(self, tensor_batch: torch.Tensor, task: Task = None, for_training: bool = False, reverse=False): if task is not None: # tensor_batch = move_to_device(tensor_batch, self._cuda_device) output_dict = self._model.forward( task_name=task._name, tensor_batch=tensor_batch, reverse=reverse, for_training=for_training, epoch_trained=self._epoch_trained ) if for_training: try: # loss = output_dict["stm_loss"] output_dict["p_d_loss"] += self._model.get_regularization_penalty() except KeyError: 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 else: raise ConfigurationError("Cannot call forward through tasks `None`") def _get_metrics(self, task: Task, reset: bool = False): task_tagger = getattr(self._model, "_tagger_" + task._name) return task_tagger.get_metrics(reset) def _save_checkpoint(self, epoch: int, should_stop: bool, is_best: bool = False) -> None: ### Saving training state ### training_state = { "epoch": epoch, "should_stop": should_stop, "metric_infos": self._metric_infos, "task_infos": self._task_infos, "schedulers": {}, "optimizers": {}, } if self._optimizers is not None: for task_name, optimizers in self._optimizers.items(): training_state["optimizers"][task_name] = {} for params_name, optimizer in optimizers.items(): training_state["optimizers"][task_name][params_name] = optimizer.state_dict() if self._schedulers is not None: for task_name, scheduler in self._schedulers.items(): training_state["schedulers"][task_name] = scheduler.lr_scheduler.state_dict() training_path = os.path.join(self._serialization_dir, "training_state.th") torch.save(training_state, training_path) logger.info("Checkpoint - Saved training state to {}", training_path) ### Saving model state ### model_path = os.path.join(self._serialization_dir, "model_state.th") model_state = self._model.state_dict() torch.save(model_state, model_path) logger.info("Checkpoint - Saved model state to {}", model_path) if is_best: logger.info("Checkpoint - Best validation performance so far for all tasks") logger.info("Checkpoint - Copying weights to '{}/best_all.th'.", self._serialization_dir) shutil.copyfile(model_path, os.path.join(self._serialization_dir, "best_all.th")) ### Saving best models for each tasks ### for task_name, infos in self._metric_infos.items(): best_epoch, _ = infos["best"] if best_epoch == epoch: logger.info("Checkpoint - Best validation performance so far for {} tasks", task_name) logger.info("Checkpoint - Copying weights to '{}/best_{}.th'.", self._serialization_dir, task_name) shutil.copyfile(model_path, os.path.join(self._serialization_dir, "best_{}.th".format(task_name))) @classmethod def from_params(cls, model: Model, task_list: List[Task], serialization_dir: str, params: Params) -> "GanMtlTrainer": optimizer_params = params.pop("optimizer") lr_scheduler_params = params.pop("scheduler") patience = params.pop_int("patience", 2) num_epochs = params.pop_int("num_epochs", 20) cuda_device = params.pop_int("cuda_device", -1) gradient_accumulation_steps = params.pop_int("gradient_accumulation_steps", 1) grad_norm = params.pop_float("grad_norm", None) grad_clipping = params.pop_float("grad_clipping", None) min_lr = params.pop_float("min_lr", 0.00001) no_tqdm = params.pop_bool("no_tqdm", False) summary_interval = params.pop("summary_interval", 30) histogram_interval = params.pop("histogram_interval", 30) sampling_method = params.pop("sampling_method", "proportional") params.assert_empty(cls.__name__) return GanMtlTrainer( model=model, task_list=task_list, optimizer_params=optimizer_params, lr_scheduler_params=lr_scheduler_params, patience=patience, num_epochs=num_epochs, serialization_dir=serialization_dir, cuda_device=cuda_device, gradient_accumulation_steps=gradient_accumulation_steps, grad_norm=grad_norm, grad_clipping=grad_clipping, min_lr=min_lr, no_tqdm=no_tqdm, summary_interval=summary_interval, histogram_interval=histogram_interval, sampling_method=sampling_method )
def __init__( self, model: Model, optimizer: torch.optim.Optimizer, data_loader: torch.utils.data.DataLoader, adv_policy: adv_utils.AdvTrainingPolicy, 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, 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, batch_callbacks: List[BatchCallback] = None, epoch_callbacks: List[EpochCallback] = None, distributed: bool = False, local_rank: int = 0, world_size: int = 1, opt_level: Optional[str] = None, ) -> None: """ A trainer for doing supervised learning. 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. # 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. 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. 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. 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`. """ 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 self.adv_policy = adv_policy 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: # 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 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 = 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) # 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 test_model_training(self): training_dataset = self.sample_instances if self.sample_only else self.train_instances #training_dataset = training_dataset[:500] validation_dataset = self.sample_instances if self.sample_only else self.test_instances serialization_dir = self.TEST_DATA_ROOT / "serialized_sample" if self.sample_only else "serialized" tensorboard_dir = self.TEST_DATA_ROOT / "tensorboard.seq2seq" batch_size = 64 train_iterator = BucketIterator(sorting_keys=[("source_tokens", "num_tokens")], padding_noise=0.1, batch_size=batch_size) train_iterator.index_with(vocab=self.vocab) multiproc_iterator = MultiprocessIterator(train_iterator, num_workers=4, output_queue_size=6000) tensorboard = TensorboardWriter(get_batch_num_total=lambda: np.ceil( len(training_dataset) / batch_size), serialization_dir=tensorboard_dir, summary_interval=5, histogram_interval=5, should_log_parameter_statistics=True, should_log_learning_rate=True) optimizer = torch.optim.Adam(self.model.parameters(), lr=1e-3) scheduler = CosineWithRestarts(optimizer=optimizer, t_initial=5) trainer = CallbackTrainer( model=self.model, serialization_dir=serialization_dir, iterator=multiproc_iterator, training_data=self.train_instances, num_epochs=100, cuda_device=0, optimizer=optimizer, callbacks=[ LogToTensorboard(tensorboard), Validate(validation_data=self.test_instances, validation_iterator=multiproc_iterator), TrackMetrics(), ResetMetricsCallback(), UpdateLearningRate(scheduler), ValidationLogCallback(self.train_reader, self.test_instances) ]) # trainer = Trainer(model=self.model, # serialization_dir=serialization_dir, # iterator=train_iterator, # train_dataset=training_dataset, # num_epochs=1, # cuda_device=0, # optimizer=torch.optim.Adam(self.model.parameters(), lr=1e-3), # validation_dataset=training_dataset, # validation_iterator=train_iterator, # should_log_learning_rate=True, # learning_rate_scheduler=scheduler # ) # for i in range(50): # print('Epoch: {}'.format(i)) # trainer.train() # # import itertools # # predictor = Seq2SeqPredictor(self.model, self.train_reader) # # for instance in itertools.islice(training_dataset, 10): # print('SOURCE:', instance.fields['source_tokens'].tokens) # print('GOLD:', instance.fields['target_tokens'].tokens) # print('PRED:', predictor.predict_instance(instance)['predicted_tokens']) # # self.val_outputs_fp.close() trainer.train()
def __init__( self, model: Model, task_list: List[Task], optimizer_params: Params, lr_scheduler_params: Params, patience: Optional[int] = None, num_epochs: int = 20, serialization_dir: str = None, cuda_device: int = -1, gradient_accumulation_steps: int = 1, grad_norm: Optional[float] = None, grad_clipping: Optional[float] = None, min_lr: float = 0.00001, no_tqdm: bool = False, summary_interval: int = 10, histogram_interval: int = 10, log_parameter_statistics: bool = False, log_gradient_statistics: bool = False, ): """ 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. iterator: ``DataIterator``, required. A method for iterating over a ``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. 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. cuda_device: int, optional (default = -1) An integer specifying the CUDA device to use. If -1, the CPU is used. Multi-gpu training is not currently supported, but will be once the Pytorch DataParallel API stabilises. 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. no_tqdm : bool, optional (default=False) We use ``tqdm`` for logging, which will print a nice progress bar that updates in place after every batch. This is nice if you're running training on a local shell, but can cause problems with log files from, e.g., a docker image running on kubernetes. If ``no_tqdm`` is ``True``, we will not use tqdm, and instead log batch statistics using ``logger.info``. """ self._model = model parameters_to_train = [(n, p) for n, p in self._model.named_parameters() if p.requires_grad] self._task_list = task_list self._n_tasks = len(self._task_list) self._optimizer_params = optimizer_params self._optimizers = {} self._lr_scheduler_params = lr_scheduler_params self._schedulers = {} for task in self._task_list: task_name = task._name self._optimizers[task_name] = Optimizer.from_params( model_parameters=parameters_to_train, params=deepcopy(optimizer_params)) self._schedulers[task_name] = LearningRateScheduler.from_params( optimizer=self._optimizers[task_name], params=deepcopy(lr_scheduler_params)) self._serialization_dir = serialization_dir self._patience = patience self._num_epochs = num_epochs self._cuda_device = cuda_device if self._cuda_device >= 0: check_for_gpu(self._cuda_device) self._model = self._model.cuda(self._cuda_device) self._gradient_accumulation_steps = gradient_accumulation_steps self._grad_norm = grad_norm self._grad_clipping = grad_clipping self._min_lr = min_lr self._task_infos = None self._metric_infos = None self._tr_generators = None self._no_tqdm = no_tqdm self._summary_interval = summary_interval # num batches between logging to tensorboard # self._log_parameter_statistics = log_parameter_statistics # self._log_gradient_statistics = log_gradient_statistics self._global_step = 0 # train_log = SummaryWriter(os.path.join(self._serialization_dir, "log", "train")) # validation_log = SummaryWriter(os.path.join(self._serialization_dir, "log", "validation")) # self._tensorboard = TensorboardWriter(train_log=train_log, validation_log=validation_log) 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_learning_rate=True) # Enable activation logging. if histogram_interval is not None: self._tensorboard.enable_activation_logging(self._model)
def record_loss(outputs_dict: Dict, tensorboard: TensorboardWriter): for key, value in outputs_dict.items(): if key.endswith("loss"): tensorboard.add_train_scalar("loss/%s" % key, value)
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)
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
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, )
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, distributed: bool = False, local_rank: int = 0, world_size: int = 1, num_gradient_accumulation_steps: int = 1, opt_level: Optional[str] = None, ) -> None: """ A trainer for doing supervised learning. 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. # 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. 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`. """ 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 # 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 __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)
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 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, )
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, 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, )
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 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(Registrable): def __init__( self, model: Model, task_list: List[Task], optimizer_params: Params, lr_scheduler_params: Params, patience: Optional[int] = None, num_epochs: int = 20, serialization_dir: str = None, cuda_device: int = -1, gradient_accumulation_steps: int = 1, grad_norm: Optional[float] = None, grad_clipping: Optional[float] = None, min_lr: float = 0.00001, no_tqdm: bool = False, summary_interval: int = 10, histogram_interval: int = 10, log_parameter_statistics: bool = False, log_gradient_statistics: bool = False, ): """ 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. iterator: ``DataIterator``, required. A method for iterating over a ``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. 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. cuda_device: int, optional (default = -1) An integer specifying the CUDA device to use. If -1, the CPU is used. Multi-gpu training is not currently supported, but will be once the Pytorch DataParallel API stabilises. 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. no_tqdm : bool, optional (default=False) We use ``tqdm`` for logging, which will print a nice progress bar that updates in place after every batch. This is nice if you're running training on a local shell, but can cause problems with log files from, e.g., a docker image running on kubernetes. If ``no_tqdm`` is ``True``, we will not use tqdm, and instead log batch statistics using ``logger.info``. """ self._model = model parameters_to_train = [(n, p) for n, p in self._model.named_parameters() if p.requires_grad] self._task_list = task_list self._n_tasks = len(self._task_list) self._optimizer_params = optimizer_params self._optimizers = {} self._lr_scheduler_params = lr_scheduler_params self._schedulers = {} for task in self._task_list: task_name = task._name self._optimizers[task_name] = Optimizer.from_params( model_parameters=parameters_to_train, params=deepcopy(optimizer_params)) self._schedulers[task_name] = LearningRateScheduler.from_params( optimizer=self._optimizers[task_name], params=deepcopy(lr_scheduler_params)) self._serialization_dir = serialization_dir self._patience = patience self._num_epochs = num_epochs self._cuda_device = cuda_device if self._cuda_device >= 0: check_for_gpu(self._cuda_device) self._model = self._model.cuda(self._cuda_device) self._gradient_accumulation_steps = gradient_accumulation_steps self._grad_norm = grad_norm self._grad_clipping = grad_clipping self._min_lr = min_lr self._task_infos = None self._metric_infos = None self._tr_generators = None self._no_tqdm = no_tqdm self._summary_interval = summary_interval # num batches between logging to tensorboard # self._log_parameter_statistics = log_parameter_statistics # self._log_gradient_statistics = log_gradient_statistics self._global_step = 0 # train_log = SummaryWriter(os.path.join(self._serialization_dir, "log", "train")) # validation_log = SummaryWriter(os.path.join(self._serialization_dir, "log", "validation")) # self._tensorboard = TensorboardWriter(train_log=train_log, validation_log=validation_log) 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_learning_rate=True) # Enable activation logging. if histogram_interval is not None: self._tensorboard.enable_activation_logging(self._model) def train( self, # tasks: List[Task], # params: Params, recover: bool = False, ): raise NotImplementedError def _check_history(self, metric_history: List[float], cur_score: float, should_decrease: bool = False): """ Given a tasks, the history of the performance on that tasks, and the current score, check if current score is best so far and if out of patience. Parameters ---------- metric_history: List[float], required cur_score: float, required should_decrease: bool, default = False Wheter or not the validation metric should increase while training. For instance, the bigger the f1 score is, the better it is -> should_decrease = False Returns ------- best_so_far: bool Whether or not the current epoch is the best so far in terms of the speicified validation metric. out_of_patience: bool Whether or not the training for this specific tasks should stop (patience parameter). """ patience = self._patience + 1 best_fn = min if should_decrease else max best_score = best_fn(metric_history) if best_score == cur_score: best_so_far = metric_history.index( best_score) == len(metric_history) - 1 else: best_so_far = False out_of_patience = False if len(metric_history) > patience: if should_decrease: out_of_patience = max(metric_history[-patience:]) <= cur_score else: out_of_patience = min(metric_history[-patience:]) >= cur_score if best_so_far and out_of_patience: # then something is up print("Something is up") return best_so_far, out_of_patience def _forward(self, tensor_batch: torch.Tensor, for_training: bool = False, task: Task = None): if task is not None: # tensor_batch = tensor2HalfTensor(tensor_batch) # for k, val in tensor_batch.items(): # if isinstance(val, Dict): # for key, value in val.items(): # print(key, value.dtype) # if isinstance(val, Tensor): # print(k, val.dtype) tensor_batch = move_to_device(tensor_batch, self._cuda_device) output_dict = self._model.forward(task_name=task._name, tensor_batch=tensor_batch, for_training=for_training) if for_training: try: loss = output_dict["loss"] loss += self._model.get_regularization_penalty() except KeyError: 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 else: raise ConfigurationError( "Cannot call forward through tasks `None`") def _get_metrics(self, task: Task, reset: bool = False): task_tagger = getattr(self._model, "_tagger_" + task._name) return task_tagger.get_metrics(reset) def _description_from_metrics(self, metrics: Dict[str, float]): # pylint: disable=no-self-use return ", ".join( ["%s: %.4f" % (name, value) for name, value in metrics.items()]) + " ||" def _rescale_gradients(self) -> Optional[float]: """ Performs gradient rescaling. Is a no-op if gradient rescaling is not enabled. """ return training_util.rescale_gradients(self._model, self._grad_norm) def _enable_gradient_clipping(self) -> None: if self._grad_clipping is not None: # Pylint is unable to tell that we're in the case that _grad_clipping is not None... # pylint: disable=invalid-unary-operand-type clip_function = lambda grad: grad.clamp(-self._grad_clipping, self. _grad_clipping) for parameter in self._model.parameters(): if parameter.requires_grad: parameter.register_hook(clip_function) def _save_checkpoint(self, epoch: int, should_stop: bool) -> None: """ Save the current states (model, training, optimizers, metrics and tasks). Parameters ---------- epoch: int, required. The epoch of training. should_stop: bool, required Whether or not the training is finished. should_save_model: bool, optional (default = True) Whether or not the model state should be saved. """ ### Saving training state ### training_state = { "epoch": epoch, "should_stop": should_stop, "metric_infos": self._metric_infos, "task_infos": self._task_infos, "schedulers": {}, "optimizers": {}, } if self._optimizers is not None: for task_name, optimizer in self._optimizers.items(): training_state["optimizers"][task_name] = optimizer.state_dict( ) if self._schedulers is not None: for task_name, scheduler in self._schedulers.items(): training_state["schedulers"][ task_name] = scheduler.lr_scheduler.state_dict() training_path = os.path.join(self._serialization_dir, "training_state.th") torch.save(training_state, training_path) logger.info("Checkpoint - Saved training state to {}", training_path) ### Saving model state ### model_path = os.path.join(self._serialization_dir, "model_state.th") model_state = self._model.state_dict() torch.save(model_state, model_path) logger.info("Checkpoint - Saved model state to {}", model_path) ### Saving best models for each tasks ### for task_name, infos in self._metric_infos.items(): best_epoch, _ = infos["best"] if best_epoch == epoch: logger.info( "Checkpoint - Best validation performance so far for {} tasks", task_name) logger.info("Checkpoint - Copying weights to '{}/best_{}.th'.", self._serialization_dir, task_name) shutil.copyfile( model_path, os.path.join(self._serialization_dir, "best_{}.th".format(task_name))) def find_latest_checkpoint(self) -> Tuple[str, str]: """ Return the location of the latest model and training state files. If there isn't a valid checkpoint then return None. """ have_checkpoint = ( self._serialization_dir is not None and any("model_state" in x for x in os.listdir(self._serialization_dir)) and any("training_state" in x for x in os.listdir(self._serialization_dir))) if not have_checkpoint: return None model_path = os.path.join(self._serialization_dir, "model_state.th") training_state_path = os.path.join(self._serialization_dir, "training_state.th") return (model_path, training_state_path) def _restore_checkpoint(self): """ Restores a model from a serialization_dir to the last saved checkpoint. This includes an epoch count, optimizer state, a model state, a tasks state and a metric state. All are of which are serialized separately. 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"))`` Returns ------- epoch: int, The epoch at which to resume training. should_stop: bool Whether or not the training should already by stopped. """ latest_checkpoint = self.find_latest_checkpoint() if not self._serialization_dir: raise ConfigurationError( "`serialization_dir` not specified - cannot " "restore a model without a directory path.") if latest_checkpoint is None: raise ConfigurationError( "Cannot restore model because one of" "`model_state.th` or `training_state.th` is not in directory path." ) model_path, training_state_path = latest_checkpoint # Load the parameters onto CPU, then transfer to GPU. # This avoids potential OOM on GPU for large models that # load parameters onto GPU then make a new GPU copy into the parameter # buffer. The GPU transfer happens implicitly in load_state_dict. model_state = torch.load(model_path, map_location=device_mapping(-1)) training_state = torch.load(training_state_path, map_location=device_mapping(-1)) # Load model self._model.load_state_dict(model_state) logger.info("Checkpoint - Model loaded from {}", model_path) # Load optimizers for task_name, optimizers_state in training_state["optimizers"].items( ): self._optimizers[task_name].load_state_dict(optimizers_state) logger.info("Checkpoint - Optimizers loaded from {}", training_state_path) # Load schedulers for task_name, scheduler_state in training_state["schedulers"].items(): self._schedulers[task_name].lr_scheduler.load_state_dict( scheduler_state) logger.info("Checkpoint - Learning rate schedulers loaded from {}", training_state_path) self._metric_infos = training_state["metric_infos"] self._task_infos = training_state["task_infos"] logger.info("Checkpoint - Task infos loaded from {}", training_state_path) logger.info("Checkpoint - Metric infos loaded from {}", training_state_path) n_epoch, should_stop = training_state["epoch"], training_state[ "should_stop"] return n_epoch + 1, should_stop @classmethod def from_params(cls, model: Model, task_list: List[Task], serialization_dir: str, params: Params) -> "MultiTaskTrainer": """ Static method that constructs the multi tasks trainer described by ``params``. """ choice = params.pop_choice("type", cls.list_available()) return cls.by_name(choice).from_params( model=model, task_list=task_list, serialization_dir=serialization_dir, params=params)
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
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)