Added functionality related to DP training

GANDLF/compute/inference_loop.py

@@ -89,7 +89,16 @@ def inference_loop(
  assert file_to_load != None, "The 'best_file' was not found"
 
  main_dict = torch.load(file_to_load, map_location=parameters["device"])
- model.load_state_dict(main_dict["model_state_dict"])
+ state_dict = main_dict["model_state_dict"]
+ if parameters.get("differential_privacy"):
+ # this is required for torch==1.11 and for DP inference
+ new_state_dict = {}
+ for key, val in state_dict.items():
+ new_key = key.replace("_module.", "")
+ new_state_dict[new_key] = val # remove `module.`
+ state_dict = new_state_dict
+
+ model.load_state_dict(state_dict)
  parameters["previous_parameters"] = main_dict.get("parameters", None)
  model.eval()
  elif parameters["model"]["type"].lower() == "openvino":

GANDLF/compute/training_loop.py

@@ -31,6 +31,10 @@
 from .forward_pass import validate_network
 from .generic import create_pytorch_objects
 
+from GANDLF.privacy.opacus.model_handling import empty_collate
+from GANDLF.privacy.opacus import handle_dynamic_batch_size, prep_for_opacus_training
+from opacus.utils.batch_memory_manager import wrap_data_loader
+
 # hides torchio citation request, see https://github.com/fepegar/torchio/issues/235
 os.environ["TORCHIO_HIDE_CITATION_PROMPT"] = "1"
 
@@ -91,6 +95,14 @@ def train_network(
  for batch_idx, (subject) in enumerate(
  tqdm(train_dataloader, desc="Looping over training data")
  ):
+ if params.get("differential_privacy"):
+ subject, params["batch_size"] = handle_dynamic_batch_size(
+ subject=subject, params=params
+ )
+ assert not isinstance(
+ model, torch.nn.DataParallel
+ ), "Differential privacy is not supported with DataParallel or DistributedDataParallel. Please use a single GPU or DDP with Opacus."
+
  optimizer.zero_grad()
  image = ( # 5D tensor: (B, C, H, W, D)
  torch.cat(
@@ -212,6 +224,23 @@ def train_network(
  return average_epoch_train_loss, average_epoch_train_metric
 
 
+def train_network_wrapper(model, train_dataloader, optimizer, params):
+ """
+ Wrapper Function to handle train_dataloader for benign and DP cases and pass on to train a network for a single epoch
+ """
+
+ if params.get("differential_privacy"):
+ with train_dataloader as memory_safe_data_loader:
+ epoch_train_loss, epoch_train_metric = train_network(
+ model, memory_safe_data_loader, optimizer, params
+ )
+ else:
+ epoch_train_loss, epoch_train_metric = train_network(
+ model, train_dataloader, optimizer, params
+ )
+ return epoch_train_loss, epoch_train_metric
+
+
 def training_loop(
  training_data: pd.DataFrame,
  validation_data: pd.DataFrame,
@@ -368,6 +397,7 @@ def training_loop(
  logger_csv_filename=os.path.join(output_dir, "logs_validation.csv"),
  metrics=metrics_log,
  mode="valid",
+ add_epsilon=bool(params.get("differential_privacy")),
  )
  if testingDataDefined:
  test_logger = Logger(
@@ -392,6 +422,36 @@ def training_loop(
 
  print("Using device:", device, flush=True)
 
+ if params.get("differential_privacy"):
+ print(
+ "Using Opacus to make training differentially private with respect to the training data."
+ )
+
+ model, optimizer, train_dataloader, privacy_engine = prep_for_opacus_training(
+ model=model,
+ optimizer=optimizer,
+ train_dataloader=train_dataloader,
+ params=params,
+ )
+
+ train_dataloader.collate_fn = empty_collate(train_dataloader.dataset[0])
+
+ # train_dataloader = BatchMemoryManager(
+ # data_loader=train_dataloader,
+ # max_physical_batch_size=MAX_PHYSICAL_BATCH_SIZE,
+ # optimizer=optimizer,
+ # )
+ batch_size = params["batch_size"]
+ max_physical_batch_size = params["differential_privacy"].get(
+ "physical_batch_size"
+ )
+ if max_physical_batch_size and max_physical_batch_size != batch_size:
+ train_dataloader = wrap_data_loader(
+ data_loader=train_dataloader,
+ max_batch_size=max_physical_batch_size,
+ optimizer=optimizer,
+ )
+
  # Iterate for number of epochs
  for epoch in range(start_epoch, epochs):
  if params["track_memory_usage"]:
@@ -453,6 +513,14 @@ def training_loop(
 
  patience += 1
 
+ # if training with differential privacy, print privacy epsilon
+ if params.get("differential_privacy"):
+ delta = params["differential_privacy"]["delta"]
+ this_epsilon = privacy_engine.get_epsilon(delta)
+ print(f" Epoch Final Privacy: (ε = {this_epsilon:.2f}, δ = {delta})")
+ # save for logging
+ epoch_valid_metric["epsilon"] = this_epsilon
+
  # Write the losses to a logger
  train_logger.write(epoch, epoch_train_loss, epoch_train_metric)
  valid_logger.write(epoch, epoch_valid_loss, epoch_valid_metric)

GANDLF/config_manager.py

@@ -7,6 +7,7 @@
 
 from .utils import version_check
 from GANDLF.data.post_process import postprocessing_after_reverse_one_hot_encoding
+from GANDLF.privacy.opacus import parse_opacus_params
 
 from GANDLF.metrics import surface_distance_ids
 from importlib.metadata import version
@@ -710,6 +711,10 @@ def _parseConfig(
  temp_dict["type"] = params["optimizer"]
  params["optimizer"] = temp_dict
 
+ # initialize defaults for DP
+ if params.get("differential_privacy"):
+ params = parse_opacus_params(params, initialize_key)
+
  # initialize defaults for inference mechanism
  inference_mechanism = {"grid_aggregator_overlap": "crop", "patch_overlap": 0}
  initialize_inference_mechanism = False

GANDLF/logger.py

@@ -12,14 +12,21 @@
 
 
 class Logger:
- def __init__(self, logger_csv_filename: str, metrics: List[str], mode: str) -> None:
+ def __init__(
+ self,
+ logger_csv_filename: str,
+ metrics: List[str],
+ mode: str,
+ add_epsilon: bool = False,
+ ) -> None:
  """
- Logger class to log the training and validation metrics to a csv file.
- May append to existing file if headers match; elsewise raises an error.
+ Logger class to log the training and validation metrics to a csv file. May append to existing file if headers match; elsewise raises an error.
 
  Args:
  logger_csv_filename (str): Path to a filename where the csv has to be stored.
  metrics (Dict[str, float]): The metrics to be logged.
+ mode (str): The mode of the logger, used as suffix to metric names. Normally may be `train` / `val` / `test`
+ add_epsilon (bool): Whether to log epsilon values or not (differential privacy measurement)
  """
  self.filename = logger_csv_filename
  mode = mode.lower()
@@ -28,6 +35,8 @@ def __init__(self, logger_csv_filename: str, metrics: List[str], mode: str) -> N
  new_header = ["epoch_no", f"{mode}_loss"] + [
  f"{mode}_{metric}" for metric in metrics
  ]
+ if add_epsilon:
+ new_header.append(f"{self.mode}_epsilon")
 
  # TODO: do we really need to support appending to existing files?
  if os.path.exists(self.filename):

GANDLF/models/imagenet_unet.py

@@ -252,6 +252,10 @@ def __init__(self, parameters) -> None:
  aux_params=classifier_head_parameters,
  )
 
+ # all BatchNorm should be replaced with InstanceNorm for DP experiments
+ if "differential_privacy" in parameters:
+ self.replace_batchnorm(self.model)
+
  if self.n_dimensions == 3:
  self.model = self.converter(self.model).model
 

GANDLF/privacy/opacus/__init__.py

@@ -0,0 +1,3 @@
+from .config_parsing import parse_opacus_params
+from .model_handling import opacus_model_fix, prep_for_opacus_training
+from .training_utils import handle_dynamic_batch_size

GANDLF/privacy/opacus/config_parsing.py

@@ -0,0 +1,59 @@
+from typing import Callable
+
+
+def parse_opacus_params(params: dict, initialize_key: Callable) -> dict:
+ """
+ Function to set defaults and augment the parameters related to making a trained model differentially
+ private with respect to the training data.
+
+ Args:
+ params (dict): Training parameters.
+ initialize_key (Callable): Function to fill in value for a missing key.
+
+ Returns:
+ dict: Updated training parameters.
+ """
+
+ if not isinstance(params["differential_privacy"], dict):
+ print(
+ "WARNING: Non dictionary value for the key: 'differential_privacy' was used, replacing with default valued dictionary."
+ )
+ params["differential_privacy"] = {}
+ # these are some defaults
+ params["differential_privacy"] = initialize_key(
+ params["differential_privacy"], "noise_multiplier", 10.0
+ )
+ params["differential_privacy"] = initialize_key(
+ params["differential_privacy"], "max_grad_norm", 1.0
+ )
+ params["differential_privacy"] = initialize_key(
+ params["differential_privacy"], "accountant", "rdp"
+ )
+ params["differential_privacy"] = initialize_key(
+ params["differential_privacy"], "secure_mode", False
+ )
+ params["differential_privacy"] = initialize_key(
+ params["differential_privacy"], "allow_opacus_model_fix", True
+ )
+ params["differential_privacy"] = initialize_key(
+ params["differential_privacy"], "delta", 1e-5
+ )
+ params["differential_privacy"] = initialize_key(
+ params["differential_privacy"], "physical_batch_size", params["batch_size"]
+ )
+
+ if params["differential_privacy"]["physical_batch_size"] > params["batch_size"]:
+ print(
+ f"WARNING: The physical batch size {params['differential_privacy']['physical_batch_size']} is greater"
+ f"than the batch size {params['batch_size']}, setting the physical batch size to the batch size."
+ )
+ params["differential_privacy"]["physical_batch_size"] = params["batch_size"]
+
+ # these keys need to be parsed as floats, not strings
+ for key in ["noise_multiplier", "max_grad_norm", "delta", "epsilon"]:
+ if key in params["differential_privacy"]:
+ params["differential_privacy"][key] = float(
+ params["differential_privacy"][key]
+ )
+
+ return params