Model Not Converging (Issue with ControlNet Fine-Tuning Script)

[MaybeRichard]

Description:

Hello, I am encountering an issue with the fine-tuning script for ControlNet provided in this repository. Specifically, when fine-tuning the model on both the KiTS dataset and my custom dataset, the model seems unable to converge.

To Reproduce
Steps to reproduce the behavior:

1. Go to 'generation/maisi'
2. Run commands 'python -m scripts.train_controlnet -c ./configs/config_maisi.json -t ./configs/config_maisi_controlnet_train.json -e ./configs/environment_maisi_controlnet_train.json -g 1'

Problem Background:

• Datasets: I am using the KiTS dataset as well as a custom medical image dataset.
• Fine-tuning Script: I followed the steps outlined in the documentation for using the fine-tuning script.
• Training Configuration: I used the default configuration and also tried different learning rates, batch sizes and weighted_loss.
• Results: During training, the loss does not significantly decrease over multiple epochs(5k-8kepochs), and the model outputs show little to no variation.

Screenshots

Output
[2024-12-14 17:52:38.415] INFO - load trained controlnet model from ./models/controlnet-20datasets-e20wl100fold0bc_noi_dia_fsize_current.pt
[2024-12-14 17:52:38.429] INFO - total number of training steps: 600.0.
[2024-12-14 17:52:38.430] INFO - apply weighted loss = 100 on labels: [129]
[Epoch 1/100] [Batch 1/6] [LR: 0.00000997] [loss: 0.0594] ETA: 0:00:57.509962
[Epoch 1/100] [Batch 2/6] [LR: 0.00000993] [loss: 0.0397] ETA: 0:00:05.898367
[Epoch 1/100] [Batch 3/6] [LR: 0.00000990] [loss: 0.0117] ETA: 0:00:04.425231
[Epoch 1/100] [Batch 4/6] [LR: 0.00000987] [loss: 0.0154] ETA: 0:00:02.951793
[Epoch 1/100] [Batch 5/6] [LR: 0.00000983] [loss: 0.0096] ETA: 0:00:01.475626
[Epoch 1/100] [Batch 6/6] [LR: 0.00000980] [loss: 0.0155] ETA: 0:00:00
[2024-12-14 17:52:58.030] INFO - best loss -> 0.02520664595067501.
[Epoch 2/100] [Batch 1/6] [LR: 0.00000977] [loss: 0.1253] ETA: 0:00:31.438770
[Epoch 2/100] [Batch 2/6] [LR: 0.00000974] [loss: 0.5812] ETA: 0:00:05.911908
[Epoch 2/100] [Batch 3/6] [LR: 0.00000970] [loss: 0.7666] ETA: 0:00:04.445841
[Epoch 2/100] [Batch 4/6] [LR: 0.00000967] [loss: 0.2461] ETA: 0:00:02.972570
[Epoch 2/100] [Batch 5/6] [LR: 0.00000964] [loss: 0.0693] ETA: 0:00:01.480636
[Epoch 2/100] [Batch 6/6] [LR: 0.00000960] [loss: 0.0123] ETA: 0:00:00
[Epoch 3/100] [Batch 1/6] [LR: 0.00000957] [loss: 0.3098] ETA: 0:00:32.303094
[Epoch 3/100] [Batch 2/6] [LR: 0.00000954] [loss: 0.4730] ETA: 0:00:06.003017
[Epoch 3/100] [Batch 3/6] [LR: 0.00000951] [loss: 0.2315] ETA: 0:00:04.474122
[Epoch 3/100] [Batch 4/6] [LR: 0.00000947] [loss: 0.2388] ETA: 0:00:02.972777
[Epoch 3/100] [Batch 5/6] [LR: 0.00000944] [loss: 0.4927] ETA: 0:00:01.482379
[Epoch 3/100] [Batch 6/6] [LR: 0.00000941] [loss: 0.6689] ETA: 0:00:00
[Epoch 4/100] [Batch 1/6] [LR: 0.00000938] [loss: 0.0097] ETA: 0:00:32.200603
[Epoch 4/100] [Batch 2/6] [LR: 0.00000934] [loss: 0.2147] ETA: 0:00:05.911577
[Epoch 4/100] [Batch 3/6] [LR: 0.00000931] [loss: 0.1359] ETA: 0:00:04.428170
[Epoch 4/100] [Batch 4/6] [LR: 0.00000928] [loss: 0.1607] ETA: 0:00:02.957243
[Epoch 4/100] [Batch 5/6] [LR: 0.00000925] [loss: 0.2544] ETA: 0:00:01.484895
[Epoch 4/100] [Batch 6/6] [LR: 0.00000922] [loss: 0.1566] ETA: 0:00:00
[Epoch 5/100] [Batch 1/6] [LR: 0.00000918] [loss: 0.3608] ETA: 0:00:39.206282
[Epoch 5/100] [Batch 2/6] [LR: 0.00000915] [loss: 0.3654] ETA: 0:00:05.935111
[Epoch 5/100] [Batch 3/6] [LR: 0.00000912] [loss: 0.1226] ETA: 0:00:04.445168
[Epoch 5/100] [Batch 4/6] [LR: 0.00000909] [loss: 0.4570] ETA: 0:00:02.968678
[Epoch 5/100] [Batch 5/6] [LR: 0.00000906] [loss: 0.1441] ETA: 0:00:01.491980
[Epoch 5/100] [Batch 6/6] [LR: 0.00000903] [loss: 0.2887] ETA: 0:00:00
[Epoch 6/100] [Batch 1/6] [LR: 0.00000899] [loss: 0.3012] ETA: 0:00:33.762065
[Epoch 6/100] [Batch 2/6] [LR: 0.00000896] [loss: 0.7926] ETA: 0:00:06.078450
[Epoch 6/100] [Batch 3/6] [LR: 0.00000893] [loss: 0.0092] ETA: 0:00:04.465131
[Epoch 6/100] [Batch 4/6] [LR: 0.00000890] [loss: 0.2825] ETA: 0:00:02.993035
[Epoch 6/100] [Batch 5/6] [LR: 0.00000887] [loss: 0.1342] ETA: 0:00:01.493944
[Epoch 6/100] [Batch 6/6] [LR: 0.00000884] [loss: 0.0131] ETA: 0:00:00
[Epoch 7/100] [Batch 1/6] [LR: 0.00000880] [loss: 0.6823] ETA: 0:00:36.323832
[Epoch 7/100] [Batch 2/6] [LR: 0.00000877] [loss: 0.0610] ETA: 0:00:06.089784
[Epoch 7/100] [Batch 3/6] [LR: 0.00000874] [loss: 0.5895] ETA: 0:00:04.482495
[Epoch 7/100] [Batch 4/6] [LR: 0.00000871] [loss: 0.0163] ETA: 0:00:02.972429
[Epoch 7/100] [Batch 5/6] [LR: 0.00000868] [loss: 0.7321] ETA: 0:00:01.492317
[Epoch 7/100] [Batch 6/6] [LR: 0.00000865] [loss: 0.6678] ETA: 0:00:00
[Epoch 8/100] [Batch 1/6] [LR: 0.00000862] [loss: 0.1816] ETA: 0:00:30.358437
[Epoch 8/100] [Batch 2/6] [LR: 0.00000859] [loss: 0.3500] ETA: 0:00:05.981302
[Epoch 8/100] [Batch 3/6] [LR: 0.00000856] [loss: 0.1950] ETA: 0:00:04.491425
[Epoch 8/100] [Batch 4/6] [LR: 0.00000853] [loss: 0.0150] ETA: 0:00:02.980718
[Epoch 8/100] [Batch 5/6] [LR: 0.00000849] [loss: 0.6160] ETA: 0:00:01.494416
[Epoch 8/100] [Batch 6/6] [LR: 0.00000846] [loss: 0.7288] ETA: 0:00:00
[Epoch 9/100] [Batch 1/6] [LR: 0.00000843] [loss: 0.2774] ETA: 0:00:30.263555
[Epoch 9/100] [Batch 2/6] [LR: 0.00000840] [loss: 0.2248] ETA: 0:00:05.925602
[Epoch 9/100] [Batch 3/6] [LR: 0.00000837] [loss: 0.1062] ETA: 0:00:04.446892
[Epoch 9/100] [Batch 4/6] [LR: 0.00000834] [loss: 0.7962] ETA: 0:00:02.997850
[Epoch 9/100] [Batch 5/6] [LR: 0.00000831] [loss: 0.3089] ETA: 0:00:01.493284
[Epoch 9/100] [Batch 6/6] [LR: 0.00000828] [loss: 0.5279] ETA: 0:00:00
[Epoch 10/100] [Batch 1/6] [LR: 0.00000825] [loss: 0.1299] ETA: 0:00:35.395999
[Epoch 10/100] [Batch 2/6] [LR: 0.00000822] [loss: 0.5649] ETA: 0:00:05.970950
[Epoch 10/100] [Batch 3/6] [LR: 0.00000819] [loss: 0.3622] ETA: 0:00:04.472051
[Epoch 10/100] [Batch 4/6] [LR: 0.00000816] [loss: 0.5737] ETA: 0:00:02.999292
[Epoch 10/100] [Batch 5/6] [LR: 0.00000813] [loss: 0.1142] ETA: 0:00:01.490247
[Epoch 10/100] [Batch 6/6] [LR: 0.00000810] [loss: 0.0369] ETA: 0:00:00
[Epoch 11/100] [Batch 1/6] [LR: 0.00000807] [loss: 0.7214] ETA: 0:00:29.923871
[Epoch 11/100] [Batch 2/6] [LR: 0.00000804] [loss: 0.0399] ETA: 0:00:06.025214
[Epoch 11/100] [Batch 3/6] [LR: 0.00000801] [loss: 0.0103] ETA: 0:00:04.477165
[Epoch 11/100] [Batch 4/6] [LR: 0.00000798] [loss: 0.3072] ETA: 0:00:02.993378
[Epoch 11/100] [Batch 5/6] [LR: 0.00000795] [loss: 0.5458] ETA: 0:00:01.499980
[Epoch 11/100] [Batch 6/6] [LR: 0.00000792] [loss: 0.2169] ETA: 0:00:00
[Epoch 12/100] [Batch 1/6] [LR: 0.00000789] [loss: 0.0542] ETA: 0:00:35.199655
[Epoch 12/100] [Batch 2/6] [LR: 0.00000786] [loss: 0.1887] ETA: 0:00:06.034812
[Epoch 12/100] [Batch 3/6] [LR: 0.00000783] [loss: 0.0109] ETA: 0:00:04.507118
[Epoch 12/100] [Batch 4/6] [LR: 0.00000780] [loss: 0.2449] ETA: 0:00:03.011293
[Epoch 12/100] [Batch 5/6] [LR: 0.00000777] [loss: 0.7797] ETA: 0:00:01.508624
[Epoch 12/100] [Batch 6/6] [LR: 0.00000774] [loss: 0.1621] ETA: 0:00:00
[Epoch 73/100] [Batch 4/6] [LR: 0.00000075] [loss: 0.0447] ETA: 0:00:03.022383
[Epoch 73/100] [Batch 5/6] [LR: 0.00000074] [loss: 0.4118] ETA: 0:00:01.515217
[Epoch 73/100] [Batch 6/6] [LR: 0.00000073] [loss: 0.4299] ETA: 0:00:00
[Epoch 74/100] [Batch 1/6] [LR: 0.00000072] [loss: 0.0393] ETA: 0:00:31.598071
[Epoch 74/100] [Batch 2/6] [LR: 0.00000071] [loss: 0.0467] ETA: 0:00:06.056095
[Epoch 74/100] [Batch 3/6] [LR: 0.00000070] [loss: 0.7438] ETA: 0:00:04.549044
[Epoch 74/100] [Batch 4/6] [LR: 0.00000069] [loss: 0.7230] ETA: 0:00:03.040020
[Epoch 74/100] [Batch 5/6] [LR: 0.00000068] [loss: 0.0777] ETA: 0:00:01.519922
[Epoch 74/100] [Batch 6/6] [LR: 0.00000068] [loss: 0.0102] ETA: 0:00:00
[Epoch 75/100] [Batch 1/6] [LR: 0.00000067] [loss: 0.5960] ETA: 0:00:32.156808
[Epoch 75/100] [Batch 2/6] [LR: 0.00000066] [loss: 0.6143] ETA: 0:00:06.073716
[Epoch 75/100] [Batch 3/6] [LR: 0.00000065] [loss: 0.1679] ETA: 0:00:04.548900
[Epoch 75/100] [Batch 4/6] [LR: 0.00000064] [loss: 0.3282] ETA: 0:00:03.040982

Environment (please complete the following information):

• OS: Ubuntu 22.04
• Python version: 3.11
• MONAI version: 1.4.0
• CUDA/cuDNN version: 12.3
• GPU models and configuration: 4090 and A100(80G)

[KumoLiu]

Hi @MaybeRichard, thanks for your interest here!

I recommend starting by checking your data to ensure the dataset is loaded correctly and contains meaningful samples. It’s a good idea to visualize a few examples (inputs and their corresponding labels/masks) to verify proper alignment.

Next, confirm that the label masks are binary (0 and 1) and appropriately normalized.

Finally, begin with a low learning rate to stabilize the training process.

cc @guopengf for additional suggestions.

[MaybeRichard]

Thanks for your suggestions! I download the KiTS dataset and corresponding json file from NGC Catalog, and I thought the data was already processed. I will try your suggestion, thanks again!

[MaybeRichard]

Sorry to bother you again. I tried the solution you provided, but the network still does not converge. Specifically:

1. I visualized the dataset, including both the labels and the image_emb after VAE encoding.
2. I checked the code in scripts/diff_model_create_training_data.py, where the intensity range of [-1000, 1000] was selected and normalized to [0, 1] before performing VAE encoding on the original images.
3. I adjusted the learning rates, including 2e-6 and 1e-7, but the results still do not converge.
4. The label data used is the mask_combined_label, which is a combination of the pseudo labels generated by VISTA 3D (from the tutorial in this link) and the original labels, so it does not only have values of 0 and 1.

Note: I am using the KiTS dataset and corresponding JSON files provided from the NGC Catalog.

[wxywhu]

Hi, I have the same problem the network does not converge and I didn't use the same dataset as you. Did you solve the problem?

[guopengf]

Hi, the loss value is not informative for DDPM model training. We usually cannot tell model convergence from the loss value, especially when you fine-tune the pre-trained weights. If you train a model from scratch, you probably will observe a quick decrease from a big value (e.g., about 1.0) to a certain value (e.g., about 0.3) at the beginning of training, then the loss value will oscillate around this value.

The loss curve that you posted is quite common for ControlNet fine-tuning. Visualization will be the most straightforward solution if you want to check whether the model converges.

[wxywhu]

Thanks for your response. The loss value did oscillate around a certain value. But I am confused with the visualization. For example, when I finetune controlnet of the MAISI, I visualize the noise, the denoised latent at the current timestep. I still don't know when to stop the training process since the latent output is hard to determine whether it is good now and the ultimate output is always generated by resampling 1000 steps during inference. Do you have any suggestions about which one should be visualized in LDM fine-tuning?

[Can-Zhao]

What we do is to generate one image every N epochs and see if the quality visually looks correct