Issue with training 3D U-Net using ZeroCostDL4Mic: validation set

[jbogomolovas2]

First of all, I want to express my gratitude towards ZeroCostDL4Mic. It is an amazing tool that I use for almost every image analysis task in my work. However, I am currently facing an issue with training a 3D U-Net to recognize bright cell clones in light-sheet microscopy images.

To train the model, I have provided 22 labeled 128x128x128 pixel images. I have varied the validation set (20-50%), batch size (1-5), and patch size (64-128x64-128x8-32), but the validation loss function does not converge. As seen in the attached image, the validation set behaves way worse than the training set.

I also came across a similar issue reported on the image.sc forum:
https://forum.image.sc/t/u-net-3d-zerocostdl4mic-training/79457

To investigate the issue further, I ran a sanity check on the dataset by repeating the same 128x128x128 labeled image 10 times and using a 50% validation fraction. However, even in this case, the validation set behaved differently than the training set. :

I understand that loss functions might differ because of regularization and at what point loss is measured, but shouldn't dice coefficient be identical in this test case? This suggests that there might be a bug in the script.

I would appreciate any help or suggestions to resolve this issue. Thank you in advance!

[jing00011]

I have similar graph before and adding elastic deformation + choosing default parameters helped with the issue.

[jbogomolovas2]

@jing00011 Thank you for suggestion. So activate elastic deformation and keep the rest default?