Diffusion Model Tutorials - Parameters for AutoEncoderKL and DiffusionModelUNet

[agokli]

I am new and trying to understand parameters for AutoEncoderKL and DiffusionModelUNet (Tutorial - 3D Latent Diffusion Model - 3d_ldm_tutorial.ipynb). This tutorial uses Task01_BrainTumour dataset.

Here are questions:

(1) AutoEncoderKL

The AutoEnCoderKL is defined as:

autoencoder = AutoencoderKL(
spatial_dims=3,
in_channels=1,
out_channels=1,
channels=(32, 64, 64),
latent_channels=3,
num_res_blocks=1,
norm_num_groups=16,
attention_levels=(False, False, True),
)
autoencoder.to(device)

in above, the in_channels = 1 and out_channels = 1. are defined and I am assuming that only one modality (out of four is selected and therefore in_channels = 1 is defined accordingly. I am also assuming that AutoEnCoderKL requires in_channels = out_channels and there fore out_channels=1 is defined.

Is my thinking correct? If yes, by default, which modality will be selected? Can you select specific modality?

(2) Diffusion Model

The Diffusion Model is defined as:

unet = DiffusionModelUNet(
spatial_dims=3,
in_channels=3,
out_channels=3,
num_res_blocks=1,
channels=(32, 64, 64),
attention_levels=(False, True, True),
num_head_channels=(0, 64, 64),
)
unet.to(device)

**(a) I am assuming that out_channels represet three output classes - Edema, Enhancing Tumor and Nectrotic Core. Am I thinking correctly?

(b) But, how is in_channels = 3 derived?**

Will appreciate for feedback.

Thanks

[Can-Zhao]

1. Yes AutoEnCoderKL requires in_channels = out_channels. In your case, if you have four modality, you can set it as 4.
2. in_channels of Diffusion Model is the latent_channels in AutoEnCoderKL