TrainingPlans

def increment_sample_counters(self, n_samples: int):
    """Increments internal counter for numbers of observed samples"""
    self.num_samples_observed_in_epoch += n_samples
    self.num_samples_observed_in_total += n_samples

def iterate_batches(self):
    """Returns an instance of a batches iterator."""
    return MiniBatchTrainingIterationsAccountant.BatchIter(self)

def iterate_epochs(self):
    """Returns an instance of an epochs iterator."""
    return MiniBatchTrainingIterationsAccountant.EpochsIter(self)

def num_batches_in_this_epoch(self) -> int:
    """Returns the number of iterations to be performed in the current epoch"""
    if self.cur_epoch == self.epochs:
        return self.num_batches_in_last_epoch
    else:
        return self.num_batches_per_epoch

def reporting_on_epoch(self) -> Optional[int]:
    """Returns the optional index of the current epoch, for reporting."""
    if self._training_plan.training_args()['num_updates'] is not None:
        return None
    else:
        return self.cur_epoch

def reporting_on_num_iter(self) -> Tuple[int, int]:
    """Outputs useful reporting information about the number of iterations

    If the researcher specified num_updates, then the iteration number will be the cumulated total, and
    similarly the maximum number of iterations will be equal to the requested number of updates.
    If the researcher specified epochs, then the iteration number will be the batch index in the current epoch,
    while the maximum number of iterations will be computed specifically for the current epoch.

    Returns:
        the iteration number
        the maximum number of iterations to be reported
    """
    if self._training_plan.training_args()['num_updates'] is not None:
        num_iter = (self.cur_epoch - 1) * self.num_batches_per_epoch + self.cur_batch
        total_batches_to_be_observed = (self.epochs - 1) * self.num_batches_per_epoch + \
            self.num_batches_in_last_epoch
        num_iter_max = total_batches_to_be_observed
    else:
        num_iter = self.cur_batch
        num_iter_max = self.num_batches_per_epoch
    return num_iter, num_iter_max

def reporting_on_num_samples(self) -> Tuple[int, int]:
    """Outputs useful reporting information about the number of observed samples

    If the researcher specified num_updates, then the number of observed samples will be the grand total, and
    similarly the maximum number of samples will be the grand total over all iterations.
    If the researcher specified epochs, then both values will be specific to the current epoch.

    Returns:
        the number of samples observed until the current iteration
        the maximum number of samples to be observed
    """
    # get batch size
    if 'batch_size' in self._training_plan.loader_args():
        batch_size = self._training_plan.loader_args()['batch_size']
    else:
        raise FedbiomedUserInputError('Missing required key `batch_size` in `loader_args`.')
    # compute number of observed samples
    if self._training_plan.training_args()['num_updates'] is not None:
        num_samples = self.num_samples_observed_in_total
        total_batches_to_be_observed = (self.epochs - 1) * self.num_batches_per_epoch + \
            self.num_batches_in_last_epoch
        total_n_samples_to_be_observed = batch_size * total_batches_to_be_observed
        num_samples_max = total_n_samples_to_be_observed
    else:
        num_samples = self.num_samples_observed_in_epoch
        num_samples_max = batch_size*self.num_batches_in_this_epoch() if \
            self.cur_batch < self.num_batches_in_this_epoch() else num_samples
    return num_samples, num_samples_max

def should_log_this_batch(self) -> bool:
    """Whether the current batch should be logged or not.

    A batch shall be logged if at least one of the following conditions is True:

        - the cumulative batch index is a multiple of the logging interval
        - the dry_run condition was specified by the researcher
        - it is the last batch of the epoch
        - it is the first batch of the epoch
    """
    current_iter = (self.cur_epoch - 1) * self.num_batches_per_epoch + self.cur_batch
    return (current_iter % self._training_plan.training_args()['log_interval'] == 0 or
            self._training_plan.training_args()['dry_run'] or
            self.cur_batch >= self.num_batches_in_this_epoch() or  # last batch
            self.cur_batch == 1)  # first batch