config_train_32g.json

{
    "channel": 0,
    "spatial_dims": 2,
    "image_channels": 1,
    "latent_channels": 1,
    "sample_axis": 2,
    "autoencoder_def": {
        "_target_": "generative.networks.nets.AutoencoderKL",
        "spatial_dims": "@spatial_dims",
        "in_channels": "$@image_channels",
        "out_channels": "@image_channels",
        "latent_channels": "@latent_channels",
        "num_channels": [
            64,
            128,
            256
        ],
        "num_res_blocks": 2,
        "norm_num_groups": 32,
        "norm_eps": 1e-06,
        "attention_levels": [
            false,
            false,
            false
        ],
        "with_encoder_nonlocal_attn": true,
        "with_decoder_nonlocal_attn": true
    },
    "autoencoder_train": {
        "batch_size": 26,
        "patch_size": [240,240],
        "lr": 5e-5,
        "perceptual_weight": 1.0,
        "kl_weight": 1e-6,
        "recon_loss": "l1",
        "n_epochs": 1000,
        "val_interval": 1
    },
    "diffusion_def": {
        "_target_": "generative.networks.nets.DiffusionModelUNet",
        "spatial_dims": "@spatial_dims",
        "in_channels": "@latent_channels",
        "out_channels": "@latent_channels",
        "num_channels":[32, 64, 128, 256],
        "attention_levels":[false, true, true, true],
        "num_head_channels":[0, 32, 32, 32],
        "num_res_blocks": 2
    },
    "diffusion_train": {
        "batch_size": 80,
        "patch_size": [256,256],
        "lr": 1e-5,
        "n_epochs": 1500,
        "val_interval": 2,
        "lr_scheduler_milestones": [1000]
    },
    "NoiseScheduler": {
        "num_train_timesteps": 1000,
        "beta_start": 0.0015,
        "beta_end": 0.0195
    }
}