# Copyright (c) 2026 Guy's and St Thomas' NHS Foundation Trust & King's College London # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # http://www.apache.org/licenses/LICENSE-2.0 # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. """FLARE Client API script for the FLIP autoencoder (VAE) tutorial. Trains a KL-regularised autoencoder and its patch discriminator on 3-D brain MRI, as an ordinary single-stage FedAvg job: one ``train`` task and one ``validate`` task, dispatched on ``flare.is_train()`` / ``flare.is_evaluate()``. Data loading fetches each accession's NIfTI series and reads them through ``transforms.py``. A brain MRI study is four co-registered sequences (FLAIR, T1w, T1Gd, T2w), and **every one of them is a training sample here**: the autoencoder is deliberately modality-agnostic, so that one autoencoder compresses all four and the latent diffusion tutorial downstream can condition on which is which. ``MODALITIES`` in ``config.json`` is the switch — leave all four listed to train across sequences, or set it to ``["T1w"]`` for the single-sequence form. Nothing here reads the conditioning label; it is attached to each sample (see ``modality.py``) for the diffusion tutorial's benefit. Because the samples within one study are four views of one brain, the train/validation split is by **accession**, not by file. Splitting by file would put a subject's T1w in training and their T2w in validation, and the reported SSIM would then be measuring memorisation. A standalone job, so an autoencoder can be trained, scored and re-used on its own. The run's aggregated weights are what the latent diffusion tutorial consumes as its frozen encoder — see that tutorial's ``process_tools/extract_autoencoder.py``. """ import argparse import json import logging from pathlib import Path import einops import nibabel as nib import numpy as np import nvflare.client as flare import torch from flip import FLIP from flip.constants import ResourceType from medicalnet_perceptual import load_medicalnet_perceptual from modality import modality_of from models import get_model from monai.data import DataLoader, Dataset from monai.losses import PatchAdversarialLoss from monai.metrics import compute_ssim_and_cs from nvflare.client.tracking import SummaryWriter from torch.amp import GradScaler, autocast from transforms import get_brain_mri_transforms # plot_utils.py is a local-development helper (dev/): job.py stages it into simulator runs, and it is # never uploaded, so on a trust this import fails and debug output is simply off. try: from plot_utils import due_for_plot, save_triplanar except ModuleNotFoundError as err: if err.name != "plot_utils": raise def due_for_plot(iteration: int, config: dict) -> bool: return False def save_triplanar(*args: object, **kwargs: object) -> None: return None logger = logging.getLogger(__name__) # The perceptual loss is MedicalNet ResNet-50, a 3-D network, staged offline by # `medicalnet_perceptual` — see that module for why MONAI's own `PerceptualLoss(...)` constructor # cannot be used here (it builds its backbone through `torch.hub.load`, and FL apps never download # at run time, FLIP#1206). Being genuinely 3-D, it scores the whole volume rather than sampled # slices, so the 2-D slicing below applies only to the discriminator. def describe(name: str, tensor: torch.Tensor) -> str: """One-line summary of a tensor's finiteness and range, for a failure report.""" finite = torch.isfinite(tensor) n_nan = int(torch.isnan(tensor).sum()) n_inf = int(torch.isinf(tensor).sum()) if bool(finite.any()): lo, hi = tensor[finite].min().item(), tensor[finite].max().item() span = f"finite range [{lo:.4g}, {hi:.4g}]" else: span = "no finite values" return f"{name}: nan={n_nan} inf={n_inf} of {tensor.numel()} elements, {span}" def require_finite(tensors: dict, context: str) -> None: """Raise a report naming every non-finite tensor, rather than a bare "found NaN". The run is still fatal — a site that trains on a poisoned tensor is worse than one that stops, and silently zeroing the offending loss hides exactly the instability worth seeing. What this adds is the *evidence*: which tensor went bad, whether it was NaN or inf, and the ranges of the others at the same step. A bare "NaN in reconstruction" cannot distinguish a diverging latent from an overflowing loss, and those want opposite fixes. """ bad = [name for name, tensor in tensors.items() if not bool(torch.isfinite(tensor).all())] if not bad: return report = "\n ".join(describe(name, tensor) for name, tensor in tensors.items()) raise RuntimeError(f"Non-finite {', '.join(bad)} during {context}:\n {report}") class KLDivergenceLoss: """ Expects z_mu (B, *latent_shape) and z_sigma (B, *latent_shape). Returns per-batch mean KL (averaged over batch). Computed in float32 even when the caller is inside ``autocast``. The autoencoder hands back fp16 tensors there, and this expression squares them: fp16 tops out at 65504, so any sigma above ~256 overflows to ``inf`` — and ``z_mu**2 + inf - log(inf) - 1`` is ``inf - inf``, i.e. **NaN**. That NaN reaches the weights and the next forward produces a NaN reconstruction, which the training loop can only report as "NaN in autoencoder reconstruction" a step after the real event. Early in training sigma genuinely does reach that range (observed above 2e4 on 3-D brain MRI within ~15 steps), so this is a live failure, not a theoretical one. The cast costs one copy of the latent, which is small next to the volumes themselves. """ def __call__( self, z_mu: torch.Tensor, z_sigma: torch.Tensor, ) -> torch.Tensor: # We remove spatial notion z_mu_m = z_mu.float().flatten(start_dim=2) z_sigma_m = z_sigma.float().flatten(start_dim=2) kl_loss = 0.5 * torch.sum(z_mu_m.pow(2) + z_sigma_m.pow(2) - torch.log(z_sigma_m.pow(2)) - 1, dim=[-1]) kl_loss = torch.sum(kl_loss) / kl_loss.shape[0] return kl_loss # --------------------------------------------------------------------------- # 3-D helpers for input 3D images # --------------------------------------------------------------------------- def detect_axial_anisotropy(image_path: str, threshold: float = 1.5) -> bool: """Whether a NIfTI volume is thick-slice, i.e. much coarser through-plane than in-plane. Anisotropy decides how a 3-D autoencoder should be scored: on thick-slice data a 3-D perceptual loss compares voxels that are not comparable, so the 2-D sliced form is the better signal (see :meth:`AutoencoderTrainer.reset_perceptual_to_anisotropic`). Not called by the shipped brain-MRI loader, on purpose. ``transforms.py`` resamples every volume to isotropic 2 mm voxels, so by the time a tensor reaches a loss there is no through-plane axis to be coarse: the answer would always be False, and asking on the *raw* geometry instead (240x240x155, a ratio of 1.55) would log "found axial anisotropy" about data that is about to be made isotropic. Kept, and worth calling on the first readable volume, for a cohort that is genuinely thick-slice and resampled to a grid that preserves that. Args: image_path (str): Path to a NIfTI volume. threshold (float): In-plane to out-plane ratio above which the volume counts as anisotropic. Returns: bool: True when the volume is thick-slice. """ volume = np.asarray(nib.load(str(image_path)).dataobj) in_plane_to_out_plane = volume.shape[0] / volume.shape[-1] if in_plane_to_out_plane > threshold: logger.info(f"Found axial anisotropy (in-plane to out-plane ratio is {in_plane_to_out_plane}).") return True return False def slice_volume_for_2d(volume: torch.Tensor, slice_indices: np.ndarray) -> torch.Tensor: """Fold selected axial slices of a 3-D batch into a 2-D batch. ``(B, C, H, W, D) -> (B*len(slice_indices), C, H, W)``, so a 2-D perceptual loss or a 2-D discriminator can score a 3-D reconstruction slice by slice. Both operands of a comparison must be sliced with the SAME indices, or the loss compares unrelated anatomy. Args: volume (torch.Tensor): A 3-D batch, channel-first. slice_indices (np.ndarray): Indices along the last (through-plane) axis. Returns: torch.Tensor: The selected slices as a 2-D batch. """ return einops.rearrange(volume[..., slice_indices], "b c h w d -> (b d) c h w") # --------------------------------------------------------------------------- # Utils # --------------------------------------------------------------------------- def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser() parser.add_argument("--project_id", type=str, default="") return parser.parse_args() def load_query() -> str: """Read the cohort query from the client app config. FLARE's TaskScriptRunner does a naive whitespace split on task_script_args, so the SQL query (which can contain spaces) is plumbed via the top-level ``query`` key in ``config/config_fed_client.json`` rather than as a CLI flag. In dev/simulator mode this is ignored by ``flip.get_dataframe``. """ client_cfg = Path(__file__).parent.parent / "config" / "config_fed_client.json" if client_cfg.exists(): try: return json.loads(client_cfg.read_text()).get("query", "") except Exception: return "" return "" def load_config() -> dict: """Load the user-supplied config.json that sits next to this script.""" config_path = Path(__file__).parent.resolve() / "config.json" with open(config_path) as f: return json.load(f) def split_for_site(items: list[dict], site_name: str) -> list[dict]: """Give each simulated site half the cohort, splitting on accession rather than on items. Local/simulator runs point every client at the same DEV dataset, so without this both sites train on identical data; in production each trust's data-access API already scopes the cohort to its own. The FLARE simulator names its clients "site-1"/"site-2" — normalise so the split actually applies (as arkplus_fine_tuning's data_utils does). Production trust names never match, so real runs are returned unsplit. Splitting on the *accession* keeps a study's sequences together: halving the item list would cut through the middle of a study and give one site a brain's T1w and the other its T2w. """ normalised = site_name.replace("-", "") if normalised not in ("site1", "site2"): return items accessions = list(dict.fromkeys(item["accession_id"] for item in items)) half = len(accessions) // 2 mine = set(accessions[:half] if normalised == "site1" else accessions[half:]) return [item for item in items if item["accession_id"] in mine] def batch_accumulation_step(batch_size: int) -> int: """Accumulate gradients up to an effective batch of 8 when the configured batch is smaller.""" if batch_size < 8: return 8 // batch_size return 1 def round_lr_schedule( optimizer: torch.optim.Optimizer, lr_start: float, lr_end: float, epochs: int ) -> torch.optim.lr_scheduler.ExponentialLR: """Exponential decay from ``lr_start`` to ``lr_end`` over one round's local epochs. Built afresh each round, as the reference MSD recipe decays over its one run. A schedule built once would compound across rounds on a trust, where one trainer serves the whole job: five rounds from 1e-4 to 1e-5 would end at 1e-9. ``lr_end`` equal to ``lr_start`` (the shipped config) is a gamma of exactly 1, so the decay is opt-in by lowering ``LR_END``. """ for group in optimizer.param_groups: group["lr"] = lr_start gamma = (lr_end / lr_start) ** (1.0 / max(1, epochs)) return torch.optim.lr_scheduler.ExponentialLR(optimizer, gamma=gamma) # How often to hand cached CUDA blocks back to the driver. See the training loop for why this is # not every iteration. CACHE_CLEAR_EVERY = 50 SSIM_KERNEL_SIZE = 11 SSIM_KERNEL_SIGMA = 1.5 def foreground_ssim( reconstruction: torch.Tensor, images: torch.Tensor, kernel_size: int = SSIM_KERNEL_SIZE, kernel_sigma: float = SSIM_KERNEL_SIGMA, ) -> float: """Mean SSIM between a reconstruction and its input, over the **foreground only**. Background is excluded because it would otherwise decide the score. These volumes are skull-stripped, so everything outside the brain is one constant value — roughly 80% of the grid. Both the input and any reconstruction are near that constant there, SSIM is ~1 over all of it, and the reported number ends up measuring how much air the volume contains rather than how well the brain came back. The unmasked metric on this cohort reads ~0.72 after one round, when the reconstructions are visibly blobs. ``data_range`` is derived per volume rather than assumed. It sets SSIM's two stability constants (``c1 = (k1 * data_range)^2``, likewise ``c2``), so a wrong value silently rescales the whole metric. ``NormalizeIntensityd`` z-scores each volume, giving a span nearer 5.5 (not 1), and each volume gets its own because each is z-scored independently. The callers correspondingly pass the reconstruction **unclamped**: clamping to [0, 1] would flatten every voxel above 1, some 15% of the volume and all of the bright anatomy, to a constant. With ``data_range=1`` and a clamp, a *perfect* reconstruction scores 0.26 rather than 1.0. Two details make the mask exact rather than approximate: * ``compute_ssim_and_cs`` returns the **full SSIM map**, not a scalar, so it can be masked at all — but it is a *valid* convolution, so the map is ``kernel_size - 1`` smaller per spatial axis (96 -> 86 on the first). The mask is therefore cropped by ``(kernel_size - 1) // 2`` on each side, which lines each map voxel up with the window centred on it. * It returns ``(ssim, contrast_sensitivity)`` **in that order**. Taking the second element gives contrast sensitivity — one factor of SSIM, not SSIM — which reads higher and moves less. Args: reconstruction (torch.Tensor): Reconstructed batch, channel-first, on the inputs' own intensity scale — do not clamp or renormalise it before passing it in. images (torch.Tensor): The input batch it is compared against, channel-first. kernel_size (int): Gaussian window size per spatial axis. kernel_sigma (float): Gaussian window sigma per spatial axis. Returns: float: Mean SSIM over foreground voxels, or NaN if the batch is entirely background. """ spatial_dims = len(images.shape[2:]) per_volume = tuple(range(1, images.ndim)) # Shape (B, 1, ..., 1), which broadcasts against the SSIM map so each volume is scored on its # own range. Taken from the target, per convention: the reference defines the dynamic range. data_range = images.amax(dim=per_volume, keepdim=True) - images.amin(dim=per_volume, keepdim=True) ssim_map, _ = compute_ssim_and_cs( reconstruction, images, spatial_dims=spatial_dims, data_range=data_range, kernel_size=[kernel_size] * spatial_dims, kernel_sigma=[kernel_sigma] * spatial_dims, ) trim = (kernel_size - 1) // 2 centres = (slice(None), slice(None)) + tuple(slice(trim, size - trim) for size in images.shape[2:]) # The background is the volume's minimum, not zero: the transform chain z-scores each volume, # which maps the skull-stripped zeros to a negative constant. An affine map is monotonic, so the # originally-zero background is still exactly the minimum. (Measured against the pre-transform # mask on the shipped cohort: identical bar one voxel in 884,736, a float32 tie.) background = images.amin(dim=per_volume, keepdim=True) foreground = images[centres] > background if not bool(foreground.any()): logger.warning("SSIM: this batch is entirely background; reporting NaN rather than a score over nothing.") return float("nan") return ssim_map[foreground].mean().item() class AutoencoderTrainer: """Holds the model, losses, optimizers, and data for autoencoder training. One instance lives for the whole ``flare.is_running()`` loop, so the optimizer state persists across global rounds. """ def __init__(self, config: dict, project_id: str, query: str): self.config = config self.project_id = project_id working_dir = Path(__file__).parent.resolve() # Training parameters for the autoencoder self.params_autoencoder = { "lr_g": config["LR_G"], "lr_d": config["LR_D"], "epochs": config.get("LOCAL_ROUNDS", 5), } self.num_workers = int(config.get("DATALOADER_WORKERS", 4)) # as in Ark+ (FLIP#1434) # Model creation self.model = get_model() self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") # Losses, optimizers etc. self.losses_ae = { "reconstruction_loss": torch.nn.L1Loss(), "kld_loss": KLDivergenceLoss(), "gan_loss": PatchAdversarialLoss(criterion="least_squares"), "perceptual_loss": load_medicalnet_perceptual(working_dir), } # 3-D only tu suppport 2D adversarial / perceptual losses. self.perceptual_slices = 12 self.axial_anisotropy: bool | None = None # Reads the SAME key models.py builds the networks from (the top-level one), not the copy # inside stage_1: nothing constructs a network from that copy, so gating on it lets a 3-D # port flip one and not the other and get a 3-D network with the slicing paths switched off. self.is_volumetric = config["net_config"]["spatial_dims"] == 3 self.weights_ae = { "w_reconstruction_loss": config["w_reconstruction_loss"], "w_perceptual_loss": config["w_perceptual_loss"], "w_kl_loss": config["w_kl_loss"], "w_gan_loss": config["w_gan_loss"], } self.optimizers_ae = { "optimizer_g": torch.optim.AdamW( self.model.autoencoder.parameters(), lr=self.params_autoencoder["lr_g"], weight_decay=1e-5 ), "optimizer_d": torch.optim.AdamW( self.model.discriminator.parameters(), lr=self.params_autoencoder["lr_d"], weight_decay=1e-5, ), } # Data loading self.flip = FLIP() dataframe = self.flip.get_dataframe(project_id=project_id, query=query) self.train_items, self.val_items = self.build_datalist(dataframe) # Per-site split for local/simulator runs where every client reads the same DEV dataset; # in production each trust's data-access API already scopes the cohort to its own data. # The FLARE simulator names its clients "site-1"/"site-2"; normalise so the split actually applies # (as arkplus_fine_tuning's data_utils does). Production trust names never match, so real runs are unsplit. self.site_name = flare.get_site_name() self.train_items = split_for_site(self.train_items, self.site_name) self._train_dataset = Dataset(self.train_items, transform=get_brain_mri_transforms()) self._val_dataset = Dataset(self.val_items, transform=get_brain_mri_transforms(is_validation=True)) def build_datalist(self, dataframe) -> tuple[list, list]: """Fetch each accession's NIfTI series and split into train/validation lists. One study yields one sample per configured modality — four, unless ``MODALITIES`` narrows the run to fewer. Each sample carries the modality label and its index in ``MODALITIES`` alongside the path; ``transforms.py`` passes both through untouched, and the collate turns the index into a batch tensor. **The split is by accession.** The four series of one study are four views of the same brain, so a file-wise split would leak a subject across the boundary and make the validation score meaningless. Whole studies go to one side or the other. """ modalities = self.config["MODALITIES"] by_accession: dict[str, list[dict]] = {} for accession_id in dataframe["accession_id"]: try: accession_folder_path = self.flip.get_by_accession_number( self.project_id, accession_id, resource_type=[ ResourceType.NIFTI, ], ) except Exception as err: logger.info(f"Could not get image data folder path for {accession_id}: {err}") continue items = [] for image in sorted(accession_folder_path.rglob("input_*.nii.gz")): modality = modality_of(image, modalities) if modality is None: # A sequence this run excludes, not an error: MODALITIES is the cohort filter. continue items.append( { "image": str(image), "accession_id": str(accession_id), "modality": modality, "modality_index": modalities.index(modality), } ) if items: by_accession[str(accession_id)] = items found = sum(len(items) for items in by_accession.values()) logger.info(f"Found {found} volume(s) across {len(by_accession)} accession(s), modalities {modalities}.") if not by_accession: raise RuntimeError( f"No volumes matched MODALITIES={modalities}. The cohort's files are named " "input__.nii.gz — check that the configured labels are the ones " "the data actually carries." ) # Validation / train splits, whole accessions at a time. accessions = list(by_accession) val_accessions = accessions[: int(self.config["VAL_SPLIT"] * len(accessions))] train_accessions = accessions[len(val_accessions) :] return ( [item for accession in train_accessions for item in by_accession[accession]], [item for accession in val_accessions for item in by_accession[accession]], ) def reset_perceptual_to_anisotropic(self): """Swap the perceptual loss for its 2-D form when 3-D data turns out to be thick-slice. 3-D only, and not called by the shipped 2-D pipeline: a 3-D port's loader should call this once :func:`detect_axial_anisotropy` returns True (setting ``self.axial_anisotropy`` first). The weight is raised alongside the swap because this network reports much smaller values. """ if self.axial_anisotropy: # The perceptual loss stays 3-D regardless: MedicalNet has no 2-D counterpart we can # ship offline, and anisotropic spacing degrades its features rather than breaking # them. Only the discriminator falls back to scoring slices. logger.info("Axial anisotropy detected; the 2-D discriminator scores slices, perceptual stays 3-D.") def make_loaders(self, batch_size: int, shuffle: bool = True) -> tuple[DataLoader, DataLoader]: train_loader = DataLoader( self._train_dataset, batch_size=batch_size, shuffle=shuffle, num_workers=self.num_workers ) val_loader = DataLoader(self._val_dataset, batch_size=batch_size, shuffle=shuffle, num_workers=self.num_workers) return train_loader, val_loader def val_loader(self, batch_size: int) -> DataLoader: return DataLoader(self._val_dataset, batch_size=batch_size, shuffle=False, num_workers=self.num_workers) def load_weights(self, weights: dict[str, torch.Tensor]) -> None: self.model.load_state_dict(state_dict=weights, strict=False) def train(self, writer: SummaryWriter, global_round: int) -> int: """One ``train`` round: local AE + discriminator epochs. Returns the iteration count.""" accumulation_step = batch_accumulation_step(self.config["BATCH_SIZE"]) train_loader, val_loader = self.make_loaders(self.config["BATCH_SIZE"]) self.model.autoencoder.to(device=self.device) self.model.discriminator.to(device=self.device) self.losses_ae["perceptual_loss"].to(self.device) # Create GradScalers - ensure they don't accumulate state scaler_g = GradScaler(enabled=True) scaler_d = GradScaler(enabled=True) # Basic training self.model.autoencoder.train() self.model.discriminator.train() epochs = self.params_autoencoder["epochs"] scheduler_g = round_lr_schedule( self.optimizers_ae["optimizer_g"], self.params_autoencoder["lr_g"], self.config.get("LR_END", self.params_autoencoder["lr_g"]), epochs, ) for epoch in range(epochs): train_g_loss = 0 train_d_loss = 0 train_g_l1loss = 0 train_g_percloss = 0 train_g_ganloss = 0 train_g_klloss = 0 batch_acc_counter = 0 for ind, batch in enumerate(train_loader): images = batch["image"].to(self.device) # 3-D only: pick the axial slices a 2-D loss/discriminator will score. Both operands of # each comparison below reuse these same indices. Inert for 2-D data. slice_indices = ( np.random.randint(0, images.shape[-1], size=self.perceptual_slices) if self.is_volumetric else None ) # TRAIN GENERATOR self.model.discriminator.requires_grad_(False) with autocast(enabled=True, device_type=self.device.type): reconstruction, z_mu, z_sigma = self.model.autoencoder(images) require_finite( {"reconstruction": reconstruction, "z_mu": z_mu, "z_sigma": z_sigma}, f"the autoencoder forward pass (epoch {epoch + 1}, iteration {ind})", ) kl_loss = self.weights_ae["w_kl_loss"] * self.losses_ae["kld_loss"](z_mu, z_sigma) l1_loss = self.weights_ae["w_reconstruction_loss"] * self.losses_ae["reconstruction_loss"]( reconstruction.float(), images.float() ) # A 2-D perceptual loss on 3-D data scores slices; a 2-D image is passed straight in. if self.is_volumetric and self.losses_ae["perceptual_loss"].spatial_dims == 2: p_loss = self.losses_ae["perceptual_loss"]( slice_volume_for_2d(reconstruction, slice_indices).float(), slice_volume_for_2d(images, slice_indices).float(), ) else: p_loss = self.losses_ae["perceptual_loss"](reconstruction.float(), images.float()) p_loss = self.weights_ae["w_perceptual_loss"] * p_loss # Likewise for a 2-D discriminator judging a 3-D reconstruction. if self.is_volumetric and self.config["net_config"]["discriminator"]["spatial_dims"] == 2: logits_fake = self.model.discriminator( slice_volume_for_2d(reconstruction, slice_indices).contiguous().float() )[-1] else: logits_fake = self.model.discriminator(reconstruction.contiguous().float())[-1] gan_loss = self.weights_ae["w_gan_loss"] * self.losses_ae["gan_loss"]( logits_fake, target_is_real=True, for_discriminator=False ) require_finite( {"l1": l1_loss, "kl": kl_loss, "perceptual": p_loss, "gan": gan_loss}, f"the generator loss (epoch {epoch + 1}, iteration {ind})", ) # Loss generator train_g_loss_ = l1_loss + kl_loss + p_loss + gan_loss # print( # f"KL: {kl_loss.item()}, l1: {l1_loss.item()}, " # f"perceptual: {p_loss.item()}, gan: {gan_loss.item()}" # ) scaler_g.scale(train_g_loss_ / accumulation_step).backward() self.model.discriminator.requires_grad_(True) batch_acc_counter += 1 if batch_acc_counter % accumulation_step == 0 or ind == (len(train_loader) - 1): scaler_g.step(self.optimizers_ae["optimizer_g"]) scaler_g.update() self.optimizers_ae["optimizer_g"].zero_grad(set_to_none=True) global_step = (global_round * epochs + epoch) * len(train_loader) + ind if due_for_plot(global_step, self.config): save_triplanar( {"input": images, "reconstruction": reconstruction}, "triplanar", self.config, site_name=self.site_name, step=global_step, ) del z_mu, z_sigma, logits_fake # TRAIN DISCRIMINATOR if self.is_volumetric and self.config["net_config"]["discriminator"]["spatial_dims"] == 2: logits_fake = self.model.discriminator( slice_volume_for_2d(reconstruction.detach(), slice_indices).contiguous().float() )[-1] logits_real = self.model.discriminator( slice_volume_for_2d(images, slice_indices).contiguous().float() )[-1] else: logits_real = self.model.discriminator(images.float())[-1] logits_fake = self.model.discriminator(reconstruction.detach().contiguous().float())[-1] loss_d_fake = self.losses_ae["gan_loss"](logits_fake, target_is_real=False, for_discriminator=True) loss_d_real = self.losses_ae["gan_loss"](logits_real, target_is_real=True, for_discriminator=True) d_loss = self.weights_ae["w_gan_loss"] * (loss_d_fake + loss_d_real) * 0.5 scaler_d.scale(d_loss / accumulation_step).backward() if batch_acc_counter % accumulation_step == 0 or ind == (len(train_loader) - 1): scaler_d.step(self.optimizers_ae["optimizer_d"]) scaler_d.update() self.optimizers_ae["optimizer_d"].zero_grad(set_to_none=True) del logits_real, logits_fake, reconstruction # Aggregate train_g_loss += train_g_loss_.item() train_g_ganloss += gan_loss.item() train_g_percloss += p_loss.item() train_g_l1loss += l1_loss.item() train_g_klloss += kl_loss.item() train_d_loss += d_loss.item() del train_g_loss_, gan_loss, p_loss, l1_loss, kl_loss, d_loss, images if torch.cuda.is_available() and ind % CACHE_CLEAR_EVERY == 0: torch.cuda.empty_cache() # Aggregate at the end train_g_loss /= max(1, len(train_loader)) train_g_ganloss /= max(1, len(train_loader)) train_g_percloss /= max(1, len(train_loader)) train_g_l1loss /= max(1, len(train_loader)) train_g_klloss /= max(1, len(train_loader)) train_d_loss /= max(1, len(train_loader)) # Validation val_loss = 0 val_ssim = 0.0 debug_batch: dict[str, torch.Tensor] = {} self.model.autoencoder.eval() for batch in val_loader: images = batch["image"].to(self.device) with autocast(enabled=True, device_type=self.device.type): with torch.no_grad(): reconstruction, _, _ = self.model.autoencoder(images) if True in torch.isnan(reconstruction): break if not debug_batch: debug_batch = {"input": images, "reconstruction": reconstruction} val_loss += ( self.weights_ae["w_reconstruction_loss"] * self.losses_ae["reconstruction_loss"](reconstruction.float(), images.float()).item() ) val_ssim += foreground_ssim(reconstruction.float(), images.float()) val_loss /= max(1, len(val_loader)) val_ssim /= max(1, len(val_loader)) self.model.autoencoder.train() logger.info( f"Epoch {epoch + 1} / {epochs};\n " f"Total loss G: {train_g_loss}, GAN: {train_g_ganloss} " f"Perceptual: {train_g_percloss}, L1: {train_g_l1loss} " f"KLD: {train_g_klloss} \n" f"Total loss D: {train_d_loss},Validation loss (L1): {val_loss}, " f"Validation SSIM (foreground): {val_ssim}" ) # Send metrics to FLIP. The "@epoch" suffix names the x-axis and `step` (cumulative local # epoch) is the coordinate. scheduler_g.step() step = global_round * epochs + epoch + 1 writer.add_scalar("Train loss (G)@epoch", train_g_loss, global_step=step) writer.add_scalar("Train loss (D)@epoch", train_d_loss, global_step=step) writer.add_scalar("Perceptual loss (G)@epoch", train_g_percloss, global_step=step) writer.add_scalar("KLD loss (G)@epoch", train_g_klloss, global_step=step) writer.add_scalar("Reconstruction loss (G)@epoch", train_g_l1loss, global_step=step) writer.add_scalar("GAN loss (G)@epoch", train_g_ganloss, global_step=step) writer.add_scalar("Validation loss (L1)@epoch", val_loss, global_step=step) # Foreground-only, so it is comparable with the `val_ssim` the validate task reports and # is not dominated by the ~90% of each volume that is background. writer.add_scalar("Validation SSIM@epoch", val_ssim, global_step=step) save_triplanar(debug_batch, "reconstruction", self.config, site_name=self.site_name, step=step) return epochs * len(train_loader) def validate( model: torch.nn.Module, test_loader: DataLoader, device: torch.device, writer: SummaryWriter, config: dict, ) -> tuple[float, float]: """Score the aggregated autoencoder on the local held-out split. Args: model (torch.nn.Module): The autoencoder network with the broadcast global weights loaded. test_loader (DataLoader): Held-out data loader. device (torch.device): Device to run on. config (dict): The user app config (``config.json``); read only for the debug-sample flag. writer (SummaryWriter): FLARE Client API metrics writer (relayed to FLIP by the analytics bridge). Returns: tuple[float, float]: ``(l1_test_loss, ssim_test)`` averaged over the held-out split. The SSIM is foreground-only — see :func:`foreground_ssim`. """ reconstruction_loss = torch.nn.L1Loss() model.autoencoder.to(device=device) model.eval() l1_test_loss = 0.0 ssim_test = 0.0 debug_batch: dict[str, torch.Tensor] = {} for batch in test_loader: images = batch["image"].to(device) with torch.no_grad(): reconstruction, _, _ = model.autoencoder(images) if not debug_batch: debug_batch = {"input": images, "reconstruction": reconstruction} reconstruction = reconstruction.detach().cpu().float() images = images.detach().cpu().float() ssim_test += foreground_ssim(reconstruction, images) l1_test_loss += reconstruction_loss(reconstruction, images).item() l1_test_loss /= max(1, len(test_loader)) ssim_test /= max(1, len(test_loader)) logger.info(f"Validation loss: {l1_test_loss} SSIM: {ssim_test}") writer.add_scalar("val_l1_loss", l1_test_loss, global_step=0) writer.add_scalar("val_ssim", ssim_test, global_step=0) save_triplanar(debug_batch, "reconstruction_aggregated", config, site_name=flare.get_site_name()) return l1_test_loss, ssim_test def to_torch_weights(input_model: flare.FLModel) -> dict[str, torch.Tensor]: return {k: torch.as_tensor(v) for k, v in input_model.params.items()} def send_weight_diff(original_params: dict, model: torch.nn.Module, n_iterations: int) -> None: """Send the full-model weight diff for a completed training round. Built one tensor at a time to avoid holding a second full copy of the model in RAM (mirrors ``flip.utils.get_model_weights_diff``). """ diff = {} for k, v in model.state_dict().items(): new_arr = v.detach().cpu().numpy() diff[k] = new_arr - np.asarray(original_params[k]) del new_arr flare.send( flare.FLModel( params=diff, params_type="DIFF", meta={"NUM_STEPS_CURRENT_ROUND": n_iterations}, ) ) def main() -> None: logging.basicConfig(level=logging.INFO) args = parse_args() config = load_config() flare.init() writer = SummaryWriter() trainer = AutoencoderTrainer(config, project_id=args.project_id, query=load_query()) while flare.is_running(): input_model = flare.receive() if input_model is None: break weights = to_torch_weights(input_model) global_round = input_model.current_round or 0 if flare.is_train(): logger.info(f"[autoencoder trainer] received train task (round {global_round})") trainer.load_weights(weights) n_iterations = trainer.train(writer, global_round) send_weight_diff(input_model.params, trainer.model, n_iterations) elif flare.is_evaluate(): trainer.model.load_state_dict({k: v.to(trainer.device) for k, v in weights.items()}) test_loader = trainer.val_loader(config["BATCH_SIZE"]) val_loss, val_ssim = validate(trainer.model, test_loader, trainer.device, writer, config) logger.info(f"Validating the aggregated autoencoder on {flare.get_site_name()}'s data: SSIM {val_ssim}") # val_acc lands in the cross-validation results.json via ValidationJsonGenerator. flare.send(flare.FLModel(metrics={"val_acc": val_ssim, "val_loss": val_loss})) else: logger.warning("Received unknown task; ignoring.") if __name__ == "__main__": main()