File size: 5,481 Bytes
d4700c7 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 | # 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.
"""Autoencoder network for the FLIP autoencoder (VAE) tutorial.
The network pairs a KL-regularised autoencoder with a patch discriminator that adversarially
sharpens its reconstructions. Both are trained and aggregated together — unlike the latent
diffusion tutorial, nothing here is frozen.
The submodule holding the autoencoder is deliberately named ``autoencoder``, so its parameters
are keyed ``autoencoder.*`` in the state dict. That name is a **contract with the latent
diffusion tutorial**: a run of this tutorial produces the checkpoint that one loads as its frozen
encoder (via ``SERVER_CHECKPOINT``), and the keys line up only because both networks name the
submodule identically. See ``process_tools/extract_autoencoder.py`` there. Renaming this
attribute silently breaks that handoff — the load is ``strict=False``, so a mismatch is tolerated
rather than raised, and the diffusion model then trains against a randomly-initialised encoder.
"""
import json
from collections.abc import Mapping
from pathlib import Path
from typing import Any
from monai.networks.nets import AutoencoderKL, PatchDiscriminator
from torch import nn
def load_net_config():
"""Read the ``net_config`` block from the ``config.json`` next to this file."""
config_path = Path(__file__).parent / "config.json"
with open(config_path, "r") as f:
config = json.load(f)
net_config = config.get("net_config", {})
print(f"Loaded network config: {net_config}")
return net_config
def autoencoder_kwargs(net_config: Mapping[str, Any]) -> dict[str, Any]:
"""The exact keyword arguments this tutorial builds its ``AutoencoderKL`` with.
This is the sole definition of the autoencoder's architecture, and the reason it is a named
function rather than inline in ``__init__``: the latent diffusion tutorial's
``process_tools/extract_autoencoder.py`` calls it to record the architecture into the YAML that
travels with an extracted checkpoint, so the two tutorials cannot disagree about what a set of
weights means.
Every argument is passed explicitly, including the ones that happen to equal MONAI's defaults.
``norm_num_groups`` is why that matters: it changes how GroupNorm computes without changing a
single parameter *shape*, so a network built with one value loads a checkpoint trained with
another under ``strict=True`` and reports nothing. The only symptom is silently wrong latents.
Leaving it to the default meant ``config.json`` could name a value that steered nothing.
"""
stage_1 = net_config["stage_1"]
return {
"spatial_dims": net_config["spatial_dims"],
"in_channels": stage_1["in_channels"],
"out_channels": stage_1["out_channels"],
"num_res_blocks": stage_1["num_res_blocks"],
"channels": stage_1["channels"],
"attention_levels": stage_1["attention_levels"],
"latent_channels": stage_1["latent_channels"],
"norm_num_groups": stage_1["norm_num_groups"],
"norm_eps": stage_1["norm_eps"],
"with_encoder_nonlocal_attn": stage_1["with_encoder_nonlocal_attn"],
"with_decoder_nonlocal_attn": stage_1["with_decoder_nonlocal_attn"],
}
class AutoencoderNetwork(nn.Module):
"""Creates the autoencoder training network, containing a:
- Variational Autoencoder (compresses inputs into a latent space and reconstructs them)
- Discriminator (adversarially trains the autoencoder)
"""
def __init__(self):
super().__init__()
net_config = load_net_config()
self.autoencoder = AutoencoderKL(**autoencoder_kwargs(net_config))
self.discriminator = PatchDiscriminator(
spatial_dims=net_config["discriminator"]["spatial_dims"],
in_channels=net_config["discriminator"]["in_channels"],
channels=net_config["discriminator"]["channels"],
out_channels=net_config["discriminator"]["out_channels"],
num_layers_d=net_config["discriminator"]["num_layers_d"],
)
def forward_ae(self, x):
return self.autoencoder(x)
def discriminate(self, x):
return self.discriminator(x)
def load_autoencoder_dict(self, state_dict: Mapping[str, Any], strict: bool = True):
self.autoencoder.load_state_dict(state_dict, strict=strict)
def load_discriminator_dict(self, state_dict: Mapping[str, Any], strict: bool = True):
self.discriminator.load_state_dict(state_dict, strict=strict)
_net = AutoencoderNetwork()
def get_model() -> nn.Module:
"""
Returns the model defined in this file.
NOTE: This function needs to exist and cannot take any input arguments. If you would like to parameterize the
configuration of your model, for example loaded from a config file, do it when instantiating the model above.
"""
return _net
|