Privacy

def after_training(self, params: Dict, renaming: bool = True) -> Dict:
    """DP actions after the training.

    Args:
        params: Contains model parameters after training with DP
        renaming: whether to modify parameters names
    Returns:
        `params` fixed model parameters after applying differential privacy
    """
    if self._is_active:
        logger.debug(
            "Applying DP after training: dp_type=%s parameter_count=%d",
            self._dp_args.get("type"),
            len(params),
        )
        params = self._postprocess_dp(params, renaming)
    return params

def before_training(
    self, optimizer: NativeTorchOptimizer, loader: DataLoader
) -> Tuple[NativeTorchOptimizer, DPDataLoader]:
    """DP action before starting training.

    Args:
        optimizer: NativeTorchOptimizer for training
        loader: Data loader for training

    Returns:
        Differential privacy applied Optimizer and data loader
    """
    # Before training is called once per node per round.
    # The effects remain throughout the batch steps inside that round
    #
    # dp_type=local: Per-batch dp, inside the local training loop of a single round
    # This is done per batch thanks to the Opacus `PrivacyEngine` `make_private` function that is attached to the model, optimizer and data loader.
    #
    # dp_type=central: Noise is added once after training, in the `after_training` function.
    # before_training() still runs once at the start of the round, but with sigma=0.0 after normalization, so the training-time noise path is disabled, and it only does clipping.

    if self._is_active:
        logger.debug(
            "Applying DP before training: optimizer_type=%s loader_type=%s dp_type=%s",
            type(optimizer.optimizer).__name__,
            type(loader).__name__,
            self._dp_args.get("type"),
        )
        if not isinstance(optimizer.optimizer, torch.optim.Optimizer):
            raise FedbiomedDPControllerError(
                f"{ErrorNumbers.FB616.value}: "
                f"Optimizer must be an instance of torch.optim.Optimizer, but got {optimizer}"
                "\nDeclearn optimizers are not yet compatible with Differential Privacy"
            )
        if not isinstance(loader, DataLoader):
            raise FedbiomedDPControllerError(
                f"{ErrorNumbers.FB616.value}: "
                "Data loader must be an instance of torch.utils.data.DataLoader"
            )
        try:
            optimizer._model.model, optimizer.optimizer, loader = (
                self._privacy_engine.make_private(
                    module=optimizer._model.model,
                    optimizer=optimizer.optimizer,
                    data_loader=loader,
                    noise_multiplier=float(self._dp_args["sigma"]),
                    max_grad_norm=float(self._dp_args["clip"]),
                )
            )
            logger.debug(
                "DP privacy engine attached successfully: dp_type=%s loader_type=%s",
                self._dp_args.get("type"),
                type(loader).__name__,
            )
        except Exception as e:
            raise FedbiomedDPControllerError(
                f"{ErrorNumbers.FB616.value}: "
                f"Error while running privacy engine: {e}"
            ) from e
    return optimizer, loader

def rename_params(self, params: Dict) -> Tuple[Dict, Dict[str, str]]:
    """Rename the model parameters modified by Opacus

    Args:
        params: Contains model parameters after training with DP

    Returns:
        renamed_params: parameters with Opacus prefix stripped.
        rename_mapping: maps new names to original Opacus names,
            only for parameters whose names were actually modified.
    """
    renamed_params = {}
    rename_mapping = {}
    for key, param in params.items():
        new_key = key.replace("_module.", "")
        renamed_params[new_key] = param
        if new_key != key:
            rename_mapping[new_key] = key
    return renamed_params, rename_mapping

def revert_rename_params(
    self, params: Dict, rename_mapping: Dict[str, str]
) -> Dict:
    """Restore the original Opacus parameter names using a rename mapping.

    Args:
        params: Contains model parameters with stripped names.
        rename_mapping: Maps stripped names back to original Opacus names,
            as returned by rename_params.
    Returns:
        Parameters with original Opacus names restored.
    """
    return {rename_mapping.get(key, key): param for key, param in params.items()}

def validate_and_fix_model(self, model: Module) -> Module:
    """Validate and Fix model to be DP-compliant.

    Args:
        model: An instance of [`Module`][torch.nn.Module]

    Returns:
        Fixed or validated model
    """
    if self._is_active and not ModuleValidator.is_valid(model):
        logger.debug("DP model validation failed, attempting automatic fix")
        try:
            model = ModuleValidator.fix(model)
        except Exception as e:
            raise FedbiomedDPControllerError(
                f"{ErrorNumbers.FB616.value}: "
                f"Error while making model DP-compliant: {e}"
            ) from e
        logger.debug("DP model automatic fix applied successfully")
    elif self._is_active:
        logger.debug("DP model validation passed without modification")
    return model