TrainingPlans

def add_preprocess(
        self,
        method: Callable,
        process_type: ProcessTypes
    ) -> None:
    """Register a pre-processing method to be executed on training data.

    Args:
        method: Pre-processing method to be run before training.
        process_type: Type of pre-processing that will be run.
            The expected signature of `method` and the arguments
            passed to it depend on this parameter.
    """
    if not callable(method):
        msg = (
            f"{ErrorNumbers.FB605.value}: error while adding "
            "preprocess, `method` should be callable."
        )
        logger.critical(msg)
        raise FedbiomedTrainingPlanError(msg)
    if not isinstance(process_type, ProcessTypes):
        msg = (
            f"{ErrorNumbers.FB605.value}: error while adding "
            "preprocess, `process_type` should be an instance "
            "of `fedbiomed.common.constants.ProcessTypes`."
        )
        logger.critical(msg)
        raise FedbiomedTrainingPlanError(msg)
    # NOTE: this may be revised into a list rather than OrderedDict
    self.pre_processes[method.__name__] = {
        'method': method,
        'process_type': process_type
    }

def after_training_params(
    self,
    flatten: bool = False,
) -> Union[Dict[str, Any], List[float]]:
    """Return the wrapped model's parameters for aggregation.

    This method returns a dict containing parameters that need to be
    reported back and aggregated in a federated learning setting.

    It may also implement post-processing steps to make these parameters
    suitable for sharing with the researcher after training - hence its
    being used over `get_model_params` at the end of training rounds.

    Returns:
        The trained parameters to aggregate.
    """
    exclude_buffers = not self._training_args['share_persistent_buffers']
    if flatten:
        return self._model.flatten(exclude_buffers=exclude_buffers)
    return self.get_model_params(exclude_buffers=exclude_buffers)

def export_model(self, filename: str) -> None:
    """Export the wrapped model to a dump file.

    Args:
        filename: path to the file where the model will be saved.

    Raises:
        FedBiomedTrainingPlanError: raised if model has not be initialized through the
        `post_init` method. If you need to export the model, you must do it through
        [`Experiment`][`fedbiomed.researcher.federated_workflows.Experiment`]'s interface.

    !!! info "Notes":
        This method is designed to save the model to a local dump
        file for easy re-use by the same user, possibly outside of
        Fed-BioMed. It is not designed to produce trustworthy data
        dumps and is not used to exchange models and their weights
        as part of the federated learning process.

        To save the model parameters for sharing as part of the FL process,
        use the `after_training_params` method (or `get_model_params` one
        outside of a training context) and export results using
        [`Serializer`][fedbiomed.common.serializer.Serializer].
    """
    if self._model is None:
        raise FedbiomedTrainingPlanError(self._error_msg_import_model % "export")
    self._model.export(filename)

def get_model_params(self,
                     only_trainable: bool = False,
                     exclude_buffers: bool = True) -> Dict[str, Any]:
    """Return a copy of the model's trainable weights.

    The type of data structure used to store weights depends on the actual
    framework of the wrapped model.

    Args:
        only_trainable: Whether to ignore non-trainable model parameters
            from outputs (e.g. frozen neural network layers' parameters),
            or include all model parameters (the default).
        exclude_buffers: Whether to ignore buffers (the default), or
            include them.

    Returns:
        Model weights, as a dictionary mapping parameters' names to their value.
    """
    return self._model.get_weights(only_trainable=only_trainable, exclude_buffers=exclude_buffers)

def get_model_wrapper_class(self) -> Optional[Model]:
    """Gets training plan's model wrapper class.

    Returns:
        the wrapper class for the model, or None
        if model is not instantiated.
    """
    return self._model

def import_model(self, filename: str) -> None:
    """Import and replace the wrapped model from a dump file.

    Args:
        filename: path to the file where the model has been exported.

    Raises:
        FedBiomedTrainingPlanError: raised if model has not be initialized through the
        `post_init` method. If you need to export the model from the Training Plan, you
        must do it through [`Experiment`][`fedbiomed.researcher.federated_workflows.Experiment`]'s
        interface.

    !!! info "Notes":
        This method is designed to load the model from a local dump
        file, that might not be in a trustworthy format. It should
        therefore only be used to re-load data exported locally and
        not received from someone else, including other FL peers.

        To load model parameters shared as part of the FL process, use the
        [`Serializer`][fedbiomed.common.serializer.Serializer] to read the
        network-exchanged file, and the `set_model_params` method to assign
        the loaded values into the wrapped model.
    """
    if self._model is None:
        raise FedbiomedTrainingPlanError(self._error_msg_import_model % "import")
    try:
        self._model.reload(filename)
    except FedbiomedModelError as exc:
        msg = (
            f"{ErrorNumbers.FB304.value}: failed to import a model from "
            f"a dump file: {exc}"
        )
        logger.critical(msg)
        raise FedbiomedTrainingPlanError(msg) from exc

def init_dependencies(self) -> List[str]:
    """Default method where dependencies are returned

    Returns:
        Empty list as default
    """
    return []

@abstractmethod
def init_optimizer(self) -> Any:
    """Method for declaring optimizer by default

    Returns:
        either framework specific optimizer (or None) or
        FedBiomed [`Optimizers`][`fedbiomed.common.optimizers.Optimizer`]
    """

def loader_args(self) -> Dict[str, Any]:
    """Retrieve loader arguments

    Returns:
        Loader arguments
    """
    return self._loader_args

@abstractmethod
def model(self):
    """Gets model instance of the training plan"""

def model_args(self) -> Dict[str, Any]:
    """Retrieve model arguments.

    Returns:
        Model arguments
    """
    return self._model_args

def optimizer(self) -> Optional[BaseOptimizer]:
    """Get the BaseOptimizer wrapped by this training plan.

    Returns:
        BaseOptimizer wrapped by this training plan, or None if
        it has not been initialized yet.
    """
    # FUTURE: return `self._optimizer.optimizer` instead?
    # Currently, the legacy Scaffold implem. needs the BaseOptimizer,
    # but IMHO it really should remain a private backend component.
    return self._optimizer

def optimizer_args(self) -> Dict[str, Any]:
    """Retrieves optimizer arguments

    Returns:
        Optimizer arguments
    """
    return self._optimizer_args

@abstractmethod
def post_init(
        self,
        model_args: Dict[str, Any],
        training_args: Dict[str, Any],
        aggregator_args: Optional[Dict[str, Any]] = None,
        initialize_optimizer: bool = True
    ) -> None:
    """Process model, training and optimizer arguments.

    Args:
        model_args: Arguments defined to instantiate the wrapped model.
        training_args: Arguments that are used in training routines
            such as epoch, dry_run etc.
            Please see [`TrainingArgs`][fedbiomed.common.training_args.TrainingArgs]
        aggregator_args: Arguments managed by and shared with the
            researcher-side aggregator.
        initialize_optimizer: whether to initialize the optimizer or not. Defaults
            to True.
    """

    # Store various arguments provided by the researcher
    self._model_args = model_args
    self._aggregator_args = aggregator_args or {}
    self._optimizer_args = training_args.optimizer_arguments() or {}
    self._loader_args = training_args.loader_arguments() or {}
    self._training_args = training_args.pure_training_arguments()

    # Set random seed: the seed can be either None or an int provided by the researcher.
    # when it is None, both random.seed and np.random.seed rely on the OS to generate a random seed.
    rseed = training_args['random_seed']
    random.seed(rseed)
    np.random.seed(rseed)

def save_code(self, filepath: str, from_code: Union[str, None] = None) -> None:
    """Saves the class source/codes of the training plan class that is created byuser.

    Args:
        filepath: path to the destination file

    Raises:
        FedbiomedTrainingPlanError: raised when source of the model class cannot be assessed
        FedbiomedTrainingPlanError: raised when model file cannot be created/opened/edited
    """
    if from_code is None:
        content = self.source()
    else:
        if not isinstance(from_code, str):
            raise FedbiomedTrainingPlanError(f"{ErrorNumbers.FB605}: Expected type str for `from_code`, "
                                             "got: {type(from_code)}")
        content = from_code

    try:
        # should we write it in binary (for the sake of space optimization)?
        with open(filepath, "w", encoding="utf-8") as file:
            file.write(content)
        logger.debug("Model file has been saved: " + filepath)
    except PermissionError as exc:
        _msg = ErrorNumbers.FB605.value + f" : Unable to read {filepath} due to unsatisfactory privileges" + \
               ", can't write the model content into it"
        logger.critical(_msg)
        raise FedbiomedTrainingPlanError(_msg) from exc
    except MemoryError as exc:
        _msg = ErrorNumbers.FB605.value + f" : Can't write model file on {filepath}: out of memory!"
        logger.critical(_msg)
        raise FedbiomedTrainingPlanError(_msg) from exc
    except OSError as exc:
        _msg = ErrorNumbers.FB605.value + f" : Can't open file {filepath} to write model content"
        logger.critical(_msg)
        raise FedbiomedTrainingPlanError(_msg) from exc

def set_aggregator_args(self, aggregator_args: Dict[str, Any]):
    raise FedbiomedTrainingPlanError("method not implemented and needed")

def set_data_loaders(
        self,
        train_data_loader: Union[DataLoader, NPDataLoader, None],
        test_data_loader: Union[DataLoader, NPDataLoader, None]
    ) -> None:
    """Sets data loaders

    Args:
        train_data_loader: Data loader for training routine/loop
        test_data_loader: Data loader for validation routine
    """
    self.training_data_loader = train_data_loader
    self.testing_data_loader = test_data_loader

def set_dataset_path(self, dataset_path: str) -> None:
    """Dataset path setter for TrainingPlan

    Args:
        dataset_path: The path where data is saved on the node.
            This method is called by the node that executes the training.
    """
    self.dataset_path = dataset_path
    logger.debug(f"Dataset path has been set as {self.dataset_path}")

def set_model_params(self, params: Dict[str, Any]) -> None:
    """Assign new values to the model's trainable parameters.

    The type of data structure used to store weights depends on the actual
    framework of the wrapped model.

    Args:
        params: model weights, as a dictionary mapping parameters' names
            to their value.
    """
    self._model.set_weights(params)

def set_optimizer(self, optimizer: BaseOptimizer):
    self._optimizer = optimizer

def source(self) -> str:

    try:
        class_source = get_class_source(self.__class__)
    except FedbiomedError as exc:
        msg = f"{ErrorNumbers.FB605.value}: error while getting source of the model class: {exc}"
        logger.critical(msg)
        raise FedbiomedTrainingPlanError(msg) from exc

    # Preparing content of the module
    content = "\n".join(self._dependencies)
    content += "\n"
    content += class_source

    return content

def testing_routine(
        self,
        metric: Optional[MetricTypes],
        metric_args: Dict[str, Any],
        history_monitor: Optional['HistoryMonitor'],
        before_train: bool,
    ) -> None:
    """Evaluation routine, to be called once per round.

    !!! info "Note"
        If the training plan implements a `testing_step` method
        (the signature of which is func(data, target) -> metrics)
        then it will be used rather than the input metric.

    Args:
        metric: The metric used for validation.
            If None, use MetricTypes.ACCURACY.
        metric_args: dicitonary containing additinal arguments for setting up metric,
            that maps <argument_name; argument_value> ad that will be passed to the
            metric function as positinal arguments.
        history_monitor: HistoryMonitor instance,
            used to record computed metrics and communicate them to
            the researcher (server).
        before_train: Whether the evaluation is being performed
            before local training occurs, of afterwards. This is merely
            reported back through `history_monitor`.
    """
    # TODO: Add preprocess option for testing_data_loader.
    if self.testing_data_loader is None:
        msg = f"{ErrorNumbers.FB605.value}: no validation dataset was set."
        logger.critical(msg)
        raise FedbiomedTrainingPlanError(msg)

    n_samples = len(self.testing_data_loader.dataset)
    n_batches = max(len(self.testing_data_loader) , 1)

    # Set up a batch-wise metrics-computation function.
    # Either use an optionally-implemented custom training routine.
    if hasattr(self, "testing_step"):
        evaluate = getattr(self, "testing_step")
        metric_name = "Custom"
    # Or use the provided `metric` (or its default value).
    else:
        if metric is None:
            metric = MetricTypes.ACCURACY
        metric_controller = Metrics()
        def evaluate(data, target):
            nonlocal metric, metric_args, metric_controller
            output = self._model.predict(data)
            if isinstance(target, torch.Tensor):
                target = target.numpy()

            return metric_controller.evaluate(
                target, output, metric=metric, **metric_args
            )
        metric_name = metric.name
    # Iterate over the validation dataset and run the defined routine.
    num_samples_observed_till_now: int = 0


    for idx, (data, target) in enumerate(self.testing_data_loader, 1):

        num_samples_observed_till_now += self._infer_batch_size(data)
        # Run the evaluation step; catch and raise exceptions.
        try:
            m_value = evaluate(data, target)
        except Exception as exc:
            msg = (
                f"{ErrorNumbers.FB605.value}: An error occurred "
                f"while computing the {metric_name} metric: {exc}"
            )
            logger.critical(msg)
            raise FedbiomedTrainingPlanError(msg) from exc
        # Log the computed value.
        # Reporting

        if idx % self.training_args()['log_interval'] == 0 or idx == 1 or idx == n_batches:
            logger.debug(
                f"Validation: Batch {idx}/{n_batches} "
                f"| Samples {num_samples_observed_till_now}/{n_samples} "
                f"| Metric[{metric_name}]: {m_value}"
            )
            # Further parse, and report it (provided a monitor is set).
            if history_monitor is not None:
                m_dict = self._create_metric_result_dict(m_value, metric_name)
                history_monitor.add_scalar(
                    metric=m_dict,
                    iteration=idx,
                    epoch=None,
                    test=True,
                    test_on_local_updates=(not before_train),
                    test_on_global_updates=before_train,
                    total_samples=n_samples,
                    batch_samples=num_samples_observed_till_now,
                    num_batches=n_batches
                )

def training_args(self) -> Dict[str, Any]:
    """Retrieve training arguments

    Returns:
        Training arguments
    """
    return self._training_args

def training_data(self):
    """All subclasses must provide a training_data routine the purpose of this actual code is to detect
    that it has been provided

    Raises:
        FedbiomedTrainingPlanError: if called and not inherited
    """
    msg = f"{ErrorNumbers.FB303.value}: training_data must be implemented"
    logger.critical(msg)
    raise FedbiomedTrainingPlanError(msg)

@abstractmethod
def training_routine(
        self,
        history_monitor: Optional['HistoryMonitor'] = None,
        node_args: Optional[Dict[str, Any]] = None
    ) -> None:
    """Training routine, to be called once per round.

    Args:
        history_monitor: optional HistoryMonitor
            instance, recording training metadata.
        node_args: Command line arguments for node.
            These arguments can specify GPU use; however, this is not
            supported for scikit-learn models and thus will be ignored.
    """
    return None

def post_init(
        self,
        model_args: Dict[str, Any],
        training_args: Dict[str, Any],
        aggregator_args: Optional[Dict[str, Any]] = None,
        **kwargs
    ) -> None:
    # get default values of Perceptron model (different from SGDClassifier model default values)
    perceptron_default_values = Perceptron().get_params()
    sgd_classifier_default_values = SGDClassifier().get_params()
    # make sure loss used is perceptron loss - can not be changed by user
    model_args["loss"] = "perceptron"
    super().post_init(model_args, training_args)
    self._model.set_params(loss="perceptron")

    # collect default values of Perceptron and set it to the model FedPerceptron
    model_hyperparameters = self._model.get_params()
    for hyperparameter_name, val in perceptron_default_values.items():
        if model_hyperparameters[hyperparameter_name] == sgd_classifier_default_values[hyperparameter_name]:
            # this means default parameter of SGDClassifier has not been changed by user
            self._model.set_params(**{hyperparameter_name: val})

def init_optimizer(self) -> Optional[FedOptimizer]:
    """Creates and configures optimizer. By default, returns None (meaning native inner scikit
    learn optimization SGD based will be used).

    In the case a Declearn Optimizer is used, this method should be overridden in the Training Plan and return
    a Fedbiomed [`Optimizer`][fedbiomed.common.optimizers.optimizer.Optimizer]"""
    pass

def model(self) -> Optional[BaseEstimator]:
    """Retrieve the wrapped scikit-learn model instance.

    Returns:
        Scikit-learn model instance
    """
    if self._model is not None:
        return self._model.model
    else:
        return self._model

def post_init(
        self,
        model_args: Dict[str, Any],
        training_args: TrainingArgs,
        aggregator_args: Optional[Dict[str, Any]] = None,
        initialize_optimizer: bool = True
) -> None:
    """Process model, training and optimizer arguments.

    Args:
        model_args: Arguments defined to instantiate the wrapped model.
        training_args: Arguments that are used in training routines
            such as epoch, dry_run etc.
            Please see [`TrainingArgs`][fedbiomed.common.training_args.TrainingArgs]
        aggregator_args: Arguments managed by and shared with the
            researcher-side aggregator.
        initialize_optimizer: Unused.
    """
    model_args.setdefault("verbose", 1)
    super().post_init(model_args, training_args, aggregator_args)
    self._model = SkLearnModel(self._model_cls)
    self._batch_maxnum = self._training_args.get('batch_maxnum', self._batch_maxnum)
    self._warn_about_training_args()

    # configure optimizer (if provided in the TrainingPlan)
    self._configure_optimizer()

    # FIXME: should we do that in `_configure_optimizer`
    # from now on, `self._optimizer`` is not None
    # Override default model parameters based on `self._model_args`.
    params = {
        key: model_args.get(key, val)
        for key, val in self._model.get_params().items()
    }
    self._model.set_params(**params)
    # Set up additional parameters (normally created by `self._model.fit`).
    self._model.set_init_params(model_args)

def set_data_loaders(
        self,
        train_data_loader: Union[DataLoader, NPDataLoader, None],
        test_data_loader: Union[DataLoader, NPDataLoader, None]
) -> None:
    """Sets data loaders

    Args:
        train_data_loader: Data loader for training routine/loop
        test_data_loader: Data loader for validation routine
    """
    args = (train_data_loader, test_data_loader)
    if not all(isinstance(data, NPDataLoader) for data in args):
        msg = (
            f"{ErrorNumbers.FB310.value}: SKLearnTrainingPlan expects "
            "NPDataLoader instances as training and testing data "
            f"loaders, but received {type(train_data_loader)} "
            f"and {type(test_data_loader)} respectively."
        )
        logger.error(msg)
        raise FedbiomedTrainingPlanError(msg)
    self.training_data_loader = train_data_loader
    self.testing_data_loader = test_data_loader

def testing_routine(
        self,
        metric: Optional[MetricTypes],
        metric_args: Dict[str, Any],
        history_monitor: Optional['HistoryMonitor'],
        before_train: bool
) -> None:
    """Evaluation routine, to be called once per round.

    !!! info "Note"
        If the training plan implements a `testing_step` method
        (the signature of which is func(data, target) -> metrics)
        then it will be used rather than the input metric.

    Args:
        metric: The metric used for validation.
            If None, use MetricTypes.ACCURACY.
        history_monitor: HistoryMonitor instance,
            used to record computed metrics and communicate them to
            the researcher (server).
        before_train: Whether the evaluation is being performed
            before local training occurs, of afterwards. This is merely
            reported back through `history_monitor`.
    """
    # Check that the testing data loader is of proper type.
    if not isinstance(self.testing_data_loader, NPDataLoader):
        msg = (
            f"{ErrorNumbers.FB310.value}: SKLearnTrainingPlan cannot be "
            "evaluated without a NPDataLoader as `testing_data_loader`."
        )
        logger.error(msg)
        raise FedbiomedTrainingPlanError(msg)
    # If required, make up for the lack of specifications regarding target
    # classification labels.
    if self._model.is_classification and not hasattr(self.model(), 'classes_'):
        classes = self._classes_from_concatenated_train_test()
        setattr(self.model(), 'classes_', classes)
    # If required, select the default metric (accuracy or mse).
    if metric is None:
        if self._model.is_classification:
            metric = MetricTypes.ACCURACY
        else:
            metric = MetricTypes.MEAN_SQUARE_ERROR
    # Delegate the actual evalation routine to the parent class.
    super().testing_routine(
        metric, metric_args, history_monitor, before_train
    )

def training_routine(
        self,
        history_monitor: Optional['HistoryMonitor'] = None,
        node_args: Optional[Dict[str, Any]] = None
) -> None:
    """Training routine, to be called once per round.

    Args:
        history_monitor: optional HistoryMonitor
            instance, recording training metadata. Defaults to None.
        node_args: command line arguments for node.
            These arguments can specify GPU use; however, this is not
            supported for scikit-learn models and thus will be ignored.
    """
    if self._optimizer is None:
        raise FedbiomedTrainingPlanError('Optimizer is None, please run `post_init` beforehand')

    # Run preprocesses
    self._preprocess()

    if not isinstance(self.training_data_loader, NPDataLoader):
        msg = (
            f"{ErrorNumbers.FB310.value}: SKLearnTrainingPlan cannot "
            "be trained without a NPDataLoader as `training_data_loader`."
        )
        logger.critical(msg)
        raise FedbiomedTrainingPlanError(msg)

    # Warn if GPU-use was expected (as it is not supported).
    if node_args is not None and node_args.get('gpu_only', False):
        logger.warning(
            'Node would like to force GPU usage, but sklearn training '
            'plan does not support it. Training on CPU.'
        )
    # Run the model-specific training routine.
    try:
        return self._training_routine(history_monitor)
    except Exception as exc:
        msg = (
            f"{ErrorNumbers.FB605.value}: error while fitting "
            f"the model: {exc}"
        )
        logger.critical(msg)
        raise FedbiomedTrainingPlanError(msg)

def type(self) -> TrainingPlans:
    """Getter for training plan type """
    return self.__type

def after_training_params(self, flatten: bool = False) -> Dict[str, torch.Tensor]:
    """Return the wrapped model's parameters for aggregation.

    This method returns a dict containing parameters that need to be
    reported back and aggregated in a federated learning setting.

    If the `postprocess` method exists (i.e. has been defined by end-users)
    it is called in the context of this method. DP-required adjustments are
    also set to happen as part of this method.

    If the researcher specified `share_persistent_buffers: False` in the
    training arguments, then we return only the output of
    [Model.get_weights][fedbiomed.common.models.TorchModel.get_weights],
    which considers only the trainable parameters.
    Otherwise, the default behaviour is to return the complete `state_dict`.

    Returns:
        The trained parameters to aggregate.
    """
    # Either include non-parameter buffers to the outputs or not.
    # Note: this is mostly about sharing statistics from BatchNorm layers.
    params = super().after_training_params()
    # Check whether postprocess method exists, and use it.
    if hasattr(self, 'postprocess'):
        logger.debug("running model.postprocess() method")
        try:
            params = self.postprocess(self._model.model.state_dict())  # Post process
        except Exception as e:
            raise FedbiomedTrainingPlanError(f"{ErrorNumbers.FB605.value}: Error while running post-process "
                                             f"{e}") from e

    # Run (optional) DP controller adjustments as well.
    params = self._dp_controller.after_training(params)
    if flatten:
        params = self._model.flatten(exclude_buffers=not self._share_persistent_buffers)
    return params

@abstractmethod
def init_model(self):
    """Abstract method where model should be defined."""

def init_optimizer(self) -> Union[FedOptimizer, torch.optim.Optimizer]:
    """Abstract method for declaring optimizer by default """
    try:
        self._optimizer = torch.optim.Adam(self._model.model.parameters(), **self._optimizer_args)
    except AttributeError as e:
        raise FedbiomedTrainingPlanError(f"{ErrorNumbers.FB605.value}: Invalid argument for default "
                                         f"optimizer Adam. Error: {e}") from e

    return self._optimizer

def initial_parameters(self) -> Dict:
    """Returns initial parameters without DP or training applied

    Returns:
        State dictionary of torch Module
    """
    return self._model.init_params

def model(self) -> Optional[torch.nn.Module]:
    if self._model is None:
        return None
    return self._model.model

def optimizer_args(self) -> Dict[str, Any]:
    """Retrieves optimizer arguments

    Returns:
        Optimizer arguments
    """
    self.update_optimizer_args()  # update `optimizer_args` (eg after training)
    return super().optimizer_args()

def post_init(
        self,
        model_args: Dict[str, Any],
        training_args: TrainingArgs,
        aggregator_args: Optional[Dict[str, Any]] = None,
        initialize_optimizer: bool = True
) -> None:
    """Process model, training and optimizer arguments.

    Args:
        model_args: Arguments defined to instantiate the wrapped model.
        training_args: Arguments that are used in training routines
            such as epoch, dry_run etc.
            Please see [`TrainingArgs`][fedbiomed.common.training_args.TrainingArgs]
        aggregator_args: Arguments managed by and shared with the
            researcher-side aggregator.
        initialize_optimizer: If True, configures optimizer. It has to be True for node
            side configuration to prepare optimizer for the training.
    Raises:
        FedbiomedTrainingPlanError: If the provided arguments do not
            match expectations, or if the optimizer, model and dependencies
            configuration goes wrong.
    """
    super().post_init(model_args, training_args, aggregator_args)
    # Assign scalar attributes.
    self._use_gpu = self._training_args.get('use_gpu')
    self._batch_maxnum = self._training_args.get('batch_maxnum')
    self._log_interval = self._training_args.get('log_interval')
    self._epochs = self._training_args.get('epochs')
    self._num_updates = self._training_args.get('num_updates', 1)
    self._dry_run = self._training_args.get('dry_run')
    self._share_persistent_buffers = training_args.get('share_persistent_buffers', True)
    # Set random seed (Pytorch-specific)
    rseed = training_args['random_seed']
    rseed = rseed if rseed is not None else torch.seed()
    torch.manual_seed(rseed)
    # Optionally set up differential privacy.
    self._dp_controller = DPController(training_args.dp_arguments() or None)
    # Configure aggregator-related arguments
    # TODO: put fedprox mu inside strategy_args
    self._fedprox_mu = self._training_args.get('fedprox_mu')
    self.set_aggregator_args(aggregator_args or {})
    # Configure the model and optimizer.

    self._configure_model_and_optimizer(initialize_optimizer)

def send_to_device(self,
                   to_send: Union[torch.Tensor, list, tuple, dict],
                   device: torch.device
                   ):
    """Send inputs to correct device for training.

    Recursively traverses lists, tuples and dicts until it meets a torch Tensor, then sends the Tensor
    to the specified device.

    Args:
        to_send: the data to be sent to the device.
        device: the device to send the data to.

    Raises:
       FedbiomedTrainingPlanError: when to_send is not the correct type
    """
    if isinstance(to_send, torch.Tensor):
        return to_send.to(device)
    elif isinstance(to_send, dict):
        return {key: self.send_to_device(val, device) for key, val in to_send.items()}
    elif isinstance(to_send, tuple):
        return tuple(self.send_to_device(d, device) for d in to_send)
    elif isinstance(to_send, list):
        return [self.send_to_device(d, device) for d in to_send]
    else:
        raise FedbiomedTrainingPlanError(f'{ErrorNumbers.FB310.value} cannot send data to device. '
                                         f'Data must be a torch Tensor or a list, tuple or dict '
                                         f'ultimately containing Tensors.')

def set_aggregator_args(self, aggregator_args: Dict[str, Any]):
    """Handles and loads aggregators arguments to received from the researcher

    Args:
        aggregator_args (Dict[str, Any]): dictionary mapping aggregator argument
            name with its value (eg 'aggregator_correction' with correction states)
    """

    self.aggregator_name = aggregator_args.get('aggregator_name', self.aggregator_name)
    # FIXME: this is too specific to Scaffold. Should be redesigned, or handled
    # by an aggregator handler that contains all keys for all strategies
    # implemented in fedbiomed
    # here we ae loading all args that have been sent from file exchange system
    for arg_name, aggregator_arg in aggregator_args.items():
        if arg_name == 'aggregator_correction':
            if not isinstance(aggregator_arg, dict):
                raise FedbiomedTrainingPlanError(
                    f"{ErrorNumbers.FB309.value}: TorchTrainingPlan received "
                    "invalid 'aggregator_correction' aggregator args."
                )
            self.correction_state = aggregator_arg

def testing_routine(
        self,
        metric: Optional[MetricTypes],
        metric_args: Dict[str, Any],
        history_monitor: Optional['HistoryMonitor'],
        before_train: bool
) -> None:
    """Evaluation routine, to be called once per round.

    !!! info "Note"
        If the training plan implements a `testing_step` method
        (the signature of which is func(data, target) -> metrics)
        then it will be used rather than the input metric.

    Args:
        metric: The metric used for validation.
            If None, use MetricTypes.ACCURACY.
        history_monitor: HistoryMonitor instance,
            used to record computed metrics and communicate them to
            the researcher (server).
        before_train: Whether the evaluation is being performed
            before local training occurs, of afterwards. This is merely
            reported back through `history_monitor`.
    """
    if not isinstance(self.model(), torch.nn.Module):
        msg = (
            f"{ErrorNumbers.FB320.value}: model should be a torch "
            f"nn.Module, but is of type {type(self.model())}"
        )
        logger.critical(msg)
        raise FedbiomedTrainingPlanError(msg)
    try:

        with torch.no_grad():
            super().testing_routine(
                metric, metric_args, history_monitor, before_train
            )
    finally:
        self.model().train()  # restore training behaviors

@abstractmethod
def training_data(self):
    """Abstract method to return training data"""

def training_routine(self,
                     history_monitor: Any = None,
                     node_args: Union[dict, None] = None,
                     ) -> int:
    # FIXME: add betas parameters for ADAM solver + momentum for SGD
    # FIXME 2: remove parameters specific for validation specified in the
    # training routine
    """Training routine procedure.

    End-user should define;

    - a `training_data()` function defining how sampling / handling data in node's dataset is done. It should
        return a generator able to output tuple (batch_idx, (data, targets)) that is iterable for each batch.
    - a `training_step()` function defining how cost is computed. It should output loss values for backpropagation.

    Args:
        history_monitor: Monitor handler for real-time feed. Defined by the Node and can't be overwritten
        node_args: command line arguments for node. Can include:
            - `gpu (bool)`: propose use a GPU device if any is available. Default False.
            - `gpu_num (Union[int, None])`: if not None, use the specified GPU device instead of default
                GPU device if this GPU device is available. Default None.
            - `gpu_only (bool)`: force use of a GPU device if any available, even if researcher
                doesn't request for using a GPU. Default False.
    Returns:
        Total number of samples observed during the training.
    """

    #self.model().train()  # pytorch switch for training
    self._optimizer.init_training()
    # set correct type for node args
    node_args = {} if not isinstance(node_args, dict) else node_args

    # send all model to device, ensures having all the requested tensors
    self._set_device(self._use_gpu, node_args)
    self._model.send_to_device(self._device)

    # Run preprocess when everything is ready before the training
    self._preprocess()

    # # initial aggregated model parameters
    # self._init_params = deepcopy(list(self.model().parameters()))

    # DP actions
    self._optimizer, self.training_data_loader = \
        self._dp_controller.before_training(optimizer= self._optimizer, loader=self.training_data_loader)

    # set number of training loop iterations
    iterations_accountant = MiniBatchTrainingIterationsAccountant(self)

    # Training loop iterations
    for epoch in iterations_accountant.iterate_epochs():
        training_data_iter: Iterator = iter(self.training_data_loader)

        for batch_idx in iterations_accountant.iterate_batches():
            # retrieve data and target
            data, target = next(training_data_iter)

            # update accounting for number of observed samples
            batch_size = self._infer_batch_size(data)
            iterations_accountant.increment_sample_counters(batch_size)

            # handle training on accelerator devices
            data, target = self.send_to_device(data, self._device), self.send_to_device(target, self._device)

            # train this batch
            corrected_loss, loss = self._train_over_batch(data, target)

            # Reporting
            if iterations_accountant.should_log_this_batch():
                # Retrieve reporting information: semantics differ whether num_updates or epochs were specified
                num_samples, num_samples_max = iterations_accountant.reporting_on_num_samples()
                num_iter, num_iter_max = iterations_accountant.reporting_on_num_iter()
                epoch_to_report = iterations_accountant.reporting_on_epoch()

                logger.debug('Train {}| '
                             'Iteration {}/{} | '
                             'Samples {}/{} ({:.0f}%)\tLoss: {:.6f}'.format(
                                 f'Epoch: {epoch_to_report} ' if epoch_to_report is not None else '',
                                 num_iter,
                                 num_iter_max,
                                 num_samples,
                                 num_samples_max,
                                 100. * num_iter / num_iter_max,
                                 loss.item())
                             )

                # Send scalar values via general/feedback topic
                if history_monitor is not None:
                    # the researcher only sees the average value of samples observed until now
                    history_monitor.add_scalar(metric={'Loss': loss.item()},
                                               iteration=num_iter,
                                               epoch=epoch_to_report,
                                               train=True,
                                               num_samples_trained=num_samples,
                                               num_batches=num_iter_max,
                                               total_samples=num_samples_max,
                                               batch_samples=batch_size)

            # Handle dry run mode
            if self._dry_run:
                self._model.send_to_device(self._device_init)
                torch.cuda.empty_cache()
                return iterations_accountant.num_samples_observed_in_total

    # release gpu usage as much as possible though:
    # - it should be done by deleting the object
    # - and some gpu memory remains used until process (cuda kernel ?) finishes

    self._model.send_to_device(self._device_init)
    torch.cuda.empty_cache()

    # # test (to be removed)
    # assert id(self._optimizer.model.model) == id(self._model.model)

    # assert id(self._optimizer.model) == id(self._model)

    # for (layer, val), (layer2, val2) in zip(self._model.model.state_dict().items(), self._optimizer.model.model.state_dict().items()):
    #     assert layer == layer2
    #     print(val, layer2)
    #     assert torch.isclose(val, val2).all()

    # for attributes, values in self._model.__dict__.items():
    #     print("ATTRIBUTES", values)
    #     assert values == getattr(self._optimizer.model, attributes)

    # for attributes, values in self._model.model.__dict__.items():
    #     print("ATTRIBUTES", values)
    #     assert values == getattr(self._optimizer.model.model, attributes)
    return iterations_accountant.num_samples_observed_in_total

@abstractmethod
def training_step(self):
    """Abstract method, all subclasses must provide a training_step."""

def type(self) -> TrainingPlans.TorchTrainingPlan:
    """ Gets training plan type"""
    return self.__type

def update_optimizer_args(self) -> Dict:
    """
    Updates `_optimizer_args` variable. Can prove useful
    to retrieve optimizer parameters after having trained a
    model, parameters which may have changed during training (eg learning rate).

    Updated arguments:
     - learning_rate

    Returns:
        Dict: updated `_optimizer_args`
    """
    if self._optimizer_args is None:
        self._optimizer_args = {}
    if self.aggregator_name is not None and self.aggregator_name.lower() == "scaffold":
        self._optimizer_args['lr'] = self._optimizer.get_learning_rate()
    return self._optimizer_args