3D UNet implementation

[mongoose54]

I modified @FabianIsensee implementation of 2D UNet to 3D UNet using the paper's description . The implementation is here: https://gist.github.com/mongoose54/c93c113ae195188394a7b363c24e2ac0
I tried it on a 3D medical image dataset but I cannot say I am getting satisfying results. So I am not sure if my implementation is buggy or in general 3D UNet does not perform well. @FabianIsensee you mentioned you have implemented 3D UNet, what is your experience? Do you get good results?

[FabianIsensee]

Hi Alex,
like I mentioned in my other post (the one with the dice coefficient) I am currently getting on par results with BraTS 2015 state of the art. I am using a more fancy 3D segmentation network for that, but in principle 3D UNet should work as well. What data are you training on? What is your class balance? Are the voxels evenly distributed among classes? What loss do you use? What is your input patch size?
Cheers,
Fabian

[mongoose54]

Hi Fabian,

Thanks for the reply. So the data I am training is chest CT. My training set has the following shape (199, 32, 256, 256) with total_background= 380398854 total_foreground 36934394. Basically the dataset is not very biased (10-fold) given its a medical dataset. That's why I am wondering if my 3D UNet implementation is off. I am using the categorical crossentropy. The rest of the code is basically the one you have posted on the Lasagne/Recipes. The main difference is the patch. I deactivated the patch generation. I have to use a batch size of 3-5 because I am running into memory problems (annoying!). Do you think that the patch generation is the reason?

Also, when you mention you are using a "fancier 3D segmentation network" is it some network from literature? I noticed that you are using batch_norm on the Recipes/ModelZoo UNet. If I add batch_norm layers in the 3D UNet it just gives me constantly out-of-memory errors although I have tons of free GPU memory.

-Alex

[mongoose54]

@FabianIsensee I am just touching base in case you have some time to respond above. I could move the discussion to the group. Thanks a lot in advance.

-Alex

[FabianIsensee]

Hi Alex,
sorry I totally forgot about this thread. I will join the discussion in the group!
Cheers,
Fabian

[FabianIsensee]

Hey Alex,
could you provide a link to the discussion please? I did not find it in the lasagne-users google group.
Regarding your problem: Since you seem to have a 10:1 ratio between foreground and background it is possible that your model does not work because of the categorical crossentropy. I suggest sampling your data so that foreground and background classes are somewhat balanced (see forp example here: https://arxiv.org/abs/1612.08894). Alternatively use the dice loss that I posted in your other thread.
What do you mean by deactivating your patch generation?
Batch sizes of 3-5 are fine. I am sometimes even training using only one example per batch. Is you are using 3D conv's there is basically no way around it.
By 'fancier' I mean residual connections, deep supervision, different upscaling etc. Nothing that should concern you now because you should get your current UNet3D running first.
Have you manually inspected the patches and raw data before training your network? Sometimes it is the simplest things that make a network not train properly.
What do you mean by 'or in general 3D UNet does not perform well'? Do you overfit? Is the loss not decreasing? Is the network just predicting background? If you have a plot showing your learning curves (loss, acc, dice over epoch) would be helpful.
Cheers,
Fabian

[mongoose54]

Hi @FabianIsensee I managed to put an example with BRATS here: https://groups.google.com/forum/#!topic/lasagne-users/FGcOdZe6WoY As you can see in the notebook the results are not satisfactory. Can you have a quick look and let me know what I am doing wrong? I would like to share with the Lasagne community. Thanks again for the valuable help.

[mongoose54]

Hi @FabianIsensee just wanted to touch base with you if you had a chance to look into the discussion and see some results: https://groups.google.com/forum/#!topic/lasagne-users/FGcOdZe6WoY Thanks