|
| 1 | +import argparse |
| 2 | +import os |
| 3 | +import warnings |
| 4 | +from logging import DEBUG, INFO |
| 5 | +from os.path import exists, join |
| 6 | +from pathlib import Path |
| 7 | + |
| 8 | +from fl4health.checkpointing.client_module import ClientCheckpointAndStateModule |
| 9 | +from fl4health.checkpointing.state_checkpointer import ClientStateCheckpointer |
| 10 | + |
| 11 | + |
| 12 | +with warnings.catch_warnings(): |
| 13 | + # Silence deprecation warnings from sentry sdk due to flwr and wandb |
| 14 | + # https://github.com/adap/flower/issues/4086 |
| 15 | + warnings.filterwarnings("ignore", category=DeprecationWarning) |
| 16 | + import wandb # noqa: F401 |
| 17 | + |
| 18 | +import torch |
| 19 | +from flwr.client import start_client |
| 20 | +from flwr.common.logger import log, update_console_handler |
| 21 | +from nnunetv2.dataset_conversion.convert_MSD_dataset import convert_msd_dataset |
| 22 | + |
| 23 | +from fl4health.clients.flexible.nnunet import FlexibleNnunetClient |
| 24 | +from fl4health.metrics.base_metrics import Metric |
| 25 | +from fl4health.metrics.compound_metrics import EmaMetric |
| 26 | +from fl4health.metrics.efficient_metrics import BinaryDice, MultiClassDice |
| 27 | +from fl4health.utils.load_data import load_msd_dataset |
| 28 | +from fl4health.utils.msd_dataset_sources import get_msd_dataset_enum, msd_num_labels |
| 29 | +from fl4health.utils.nnunet_utils import set_nnunet_env_and_reload_modules |
| 30 | +from fl4health.utils.random import set_all_random_seeds |
| 31 | + |
| 32 | + |
| 33 | +N_CLASSES_2D = 2 |
| 34 | + |
| 35 | + |
| 36 | +def main( |
| 37 | + dataset_path: Path, |
| 38 | + msd_dataset_name: str, |
| 39 | + server_address: str, |
| 40 | + fold: int | str, |
| 41 | + always_preprocess: bool = False, |
| 42 | + verbose: bool = True, |
| 43 | + compile: bool = True, |
| 44 | + intermediate_client_state_dir: str | None = None, |
| 45 | + client_name: str | None = None, |
| 46 | +) -> None: |
| 47 | + # Log device and server address |
| 48 | + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") |
| 49 | + log(INFO, f"Using device: {device}") |
| 50 | + log(INFO, f"Using server address: {server_address}") |
| 51 | + |
| 52 | + # Load the dataset if necessary |
| 53 | + msd_dataset_enum = get_msd_dataset_enum(msd_dataset_name) |
| 54 | + nn_unet_raw = join(dataset_path, "nnunet_raw") |
| 55 | + if not exists(join(nn_unet_raw, msd_dataset_enum.value)): |
| 56 | + log(INFO, f"Downloading and extracting {msd_dataset_enum.value} dataset") |
| 57 | + load_msd_dataset(nn_unet_raw, msd_dataset_name) |
| 58 | + |
| 59 | + # The dataset ID will be the same as the MSD Task number |
| 60 | + dataset_id = int(msd_dataset_enum.value[4:6]) |
| 61 | + nnunet_dataset_name = f"Dataset{dataset_id:03d}_{msd_dataset_enum.value.split('_')[1]}" |
| 62 | + |
| 63 | + # Convert the msd dataset if necessary |
| 64 | + if not exists(join(nn_unet_raw, nnunet_dataset_name)): |
| 65 | + log(INFO, f"Converting {msd_dataset_enum.value} into nnunet dataset") |
| 66 | + convert_msd_dataset(source_folder=join(nn_unet_raw, msd_dataset_enum.value)) |
| 67 | + |
| 68 | + # Create dice metric |
| 69 | + dice: Metric |
| 70 | + if msd_num_labels[msd_dataset_enum] > N_CLASSES_2D: |
| 71 | + dice = MultiClassDice( |
| 72 | + batch_dim=None, # Aggregate across all samples in batch/round |
| 73 | + label_dim=1, # Separate dice for each output class |
| 74 | + name="Dice", |
| 75 | + threshold=1, # Use an argmax to binarize output logits (unactivated) |
| 76 | + ignore_background=1, # Ignore background class |
| 77 | + ) |
| 78 | + else: # Background class is automatically ignored for BinaryDice |
| 79 | + dice = BinaryDice(batch_dim=None, label_dim=1, name="Dice", threshold=1) |
| 80 | + |
| 81 | + # Create EMA Dice metric |
| 82 | + ema_dice = EmaMetric(dice) |
| 83 | + |
| 84 | + # State checkpointer (being overhauled soon) |
| 85 | + if intermediate_client_state_dir is not None: |
| 86 | + checkpoint_and_state_module = ClientCheckpointAndStateModule( |
| 87 | + state_checkpointer=ClientStateCheckpointer(Path(intermediate_client_state_dir)) |
| 88 | + ) |
| 89 | + else: |
| 90 | + checkpoint_and_state_module = None |
| 91 | + |
| 92 | + # Create client |
| 93 | + client = FlexibleNnunetClient( |
| 94 | + # Args specific to nnUNetClient |
| 95 | + dataset_id=dataset_id, |
| 96 | + fold=fold, |
| 97 | + always_preprocess=always_preprocess, |
| 98 | + verbose=verbose, |
| 99 | + compile=compile, |
| 100 | + # BaseClient Args |
| 101 | + device=device, |
| 102 | + metrics=[dice, ema_dice], |
| 103 | + progress_bar=verbose, |
| 104 | + checkpoint_and_state_module=checkpoint_and_state_module, |
| 105 | + client_name=client_name, |
| 106 | + ) |
| 107 | + |
| 108 | + start_client(server_address=server_address, client=client.to_client()) |
| 109 | + |
| 110 | + # Shutdown the client |
| 111 | + client.shutdown() |
| 112 | + |
| 113 | + |
| 114 | +if __name__ == "__main__": |
| 115 | + parser = argparse.ArgumentParser( |
| 116 | + prog="nnunet_example/client.py", |
| 117 | + description="An exampled of nnUNetClient on any of the Medical \ |
| 118 | + Segmentation Decathlon (MSD) datasets. Automatically generates a \ |
| 119 | + nnunet segmentation model and trains it in a federated setting", |
| 120 | + ) |
| 121 | + |
| 122 | + # I have to use underscores instead of dashes because thats how they |
| 123 | + # defined it in run_smoke_tests |
| 124 | + parser.add_argument( |
| 125 | + "--dataset_path", |
| 126 | + type=str, |
| 127 | + required=True, |
| 128 | + help="Path to the folder in which data should be stored. This script \ |
| 129 | + will automatically create nnunet_raw, and nnunet_preprocessed \ |
| 130 | + subfolders if they don't already exist. This script will also \ |
| 131 | + attempt to download and prepare the MSD Dataset into the \ |
| 132 | + nnunet_raw folder if it does not already exist.", |
| 133 | + ) |
| 134 | + parser.add_argument( |
| 135 | + "--fold", |
| 136 | + type=str, |
| 137 | + required=False, |
| 138 | + default="0", |
| 139 | + help="[OPTIONAL] Which fold of the local client dataset to use for \ |
| 140 | + validation. nnunet defaults to 5 folds (0 to 4). Can also be set \ |
| 141 | + to 'all' to use all the data for both training and validation. \ |
| 142 | + Defaults to fold 0", |
| 143 | + ) |
| 144 | + parser.add_argument( |
| 145 | + "--msd_dataset_name", |
| 146 | + type=str, |
| 147 | + required=False, |
| 148 | + default="Task04_Hippocampus", # The smallest dataset |
| 149 | + help="[OPTIONAL] Name of the MSD dataset to use. The options are \ |
| 150 | + defined by the values of the MsdDataset enum as returned by the \ |
| 151 | + get_msd_dataset_enum function", |
| 152 | + ) |
| 153 | + parser.add_argument( |
| 154 | + "--always-preprocess", |
| 155 | + action="store_true", |
| 156 | + required=False, |
| 157 | + help="[OPTIONAL] Use this to force preprocessing the nnunet data \ |
| 158 | + even if the preprocessed data is found to already exist", |
| 159 | + ) |
| 160 | + parser.add_argument( |
| 161 | + "--server-address", |
| 162 | + type=str, |
| 163 | + required=False, |
| 164 | + default="0.0.0.0:8080", |
| 165 | + help="[OPTIONAL] The server address for the clients to communicate \ |
| 166 | + to the server through. Defaults to 0.0.0.0:8080", |
| 167 | + ) |
| 168 | + parser.add_argument( |
| 169 | + "--verbose", |
| 170 | + action="store_true", |
| 171 | + required=False, |
| 172 | + help="[OPTIONAL] Use this flag to see extra INFO logs and a progress bar", |
| 173 | + ) |
| 174 | + parser.add_argument( |
| 175 | + "--debug", |
| 176 | + help="[OPTIONAL] Include flag to print DEBUG logs", |
| 177 | + action="store_const", |
| 178 | + dest="logLevel", |
| 179 | + const=DEBUG, |
| 180 | + default=INFO, |
| 181 | + ) |
| 182 | + parser.add_argument( |
| 183 | + "--skip-compile", |
| 184 | + action="store_true", |
| 185 | + required=False, |
| 186 | + help="[OPTIONAL] Include flag to train without jit compiling the pytorch model first", |
| 187 | + ) |
| 188 | + |
| 189 | + parser.add_argument( |
| 190 | + "--intermediate-client-state-dir", |
| 191 | + type=str, |
| 192 | + required=False, |
| 193 | + default=None, |
| 194 | + help="[OPTIONAL] Directory to store client state during training. Defaults to None", |
| 195 | + ) |
| 196 | + parser.add_argument( |
| 197 | + "--client-name", |
| 198 | + type=str, |
| 199 | + required=False, |
| 200 | + default=None, |
| 201 | + help="[OPTIONAL] Name of the client used to name client state checkpoint. \ |
| 202 | + Defaults to None, in which case a random name is generated for the client", |
| 203 | + ) |
| 204 | + parser.add_argument( |
| 205 | + "--seed", |
| 206 | + action="store", |
| 207 | + type=int, |
| 208 | + help="Seed for the random number generators across python, torch, and numpy", |
| 209 | + required=False, |
| 210 | + ) |
| 211 | + |
| 212 | + args = parser.parse_args() |
| 213 | + |
| 214 | + # Set the random seed for reproducibility |
| 215 | + set_all_random_seeds(args.seed) |
| 216 | + |
| 217 | + # Set the log level |
| 218 | + update_console_handler(level=args.logLevel) |
| 219 | + |
| 220 | + # Create nnunet directory structure and set environment variables |
| 221 | + nn_unet_raw = join(args.dataset_path, "nnunet_raw") |
| 222 | + nn_unet_preprocessed = join(args.dataset_path, "nnunet_preprocessed") |
| 223 | + os.makedirs(nn_unet_raw, exist_ok=True) |
| 224 | + os.makedirs(nn_unet_preprocessed, exist_ok=True) |
| 225 | + set_nnunet_env_and_reload_modules( |
| 226 | + nnUNet_raw=nn_unet_raw, |
| 227 | + nnUNet_preprocessed=nn_unet_preprocessed, |
| 228 | + nnUNet_results=join(args.dataset_path, "nnUNet_results"), |
| 229 | + ) |
| 230 | + |
| 231 | + # Check fold argument and start main method |
| 232 | + fold: int | str = "all" if args.fold == "all" else int(args.fold) |
| 233 | + main( |
| 234 | + dataset_path=Path(args.dataset_path), |
| 235 | + msd_dataset_name=args.msd_dataset_name, |
| 236 | + server_address=args.server_address, |
| 237 | + fold=fold, |
| 238 | + always_preprocess=args.always_preprocess, |
| 239 | + verbose=args.verbose, |
| 240 | + compile=not args.skip_compile, |
| 241 | + intermediate_client_state_dir=args.intermediate_client_state_dir, |
| 242 | + client_name=args.client_name, |
| 243 | + ) |
0 commit comments