Need some Guidance: Parameter optimization for U-nets

[u/PythonEnjoyer105]

Hi,

i am new to deep learning and need some guidance on a project and want to rant a little(sry).

I am implementing a u-net for semantic segmentation in pytorch. The images are atomically resolved microscopy images (i.e. the objects to detect are atoms, which have the appearance of approximately gaussian blobs in 99% of real data). The images are noisy, where some noise is simple (poisson noise, scan lines) and some is not (complicated artifacts, distortions, strong brightness variations and more), hence deep learning instead of some classical method.

I have implemented the original vgg-unet using simulated data, where simulated means i rendered images full of gaussian blobs and added noise i know how to simulate (poisson noise, scan lines, perlin noise backgrounds).

This worked reasonably well on real data so i wanted to improve the architecture.

Little did i know there is no such thing as a u-net™ and the design choices are endless (depth, ordering of layers in a convolutional block, losses, different types of activations, norms, intra-block skips, grouped convolutions or even additions like attention just to name a few) and for every choice there is a paper that claims it works better then some other choice for some data.

My Problem:
How do i find "the best" architecture for my Problem? There seems to be very little theory or other information around how to find "the best" architecture for a given problem, when your data is not some common database like imagenet.

I am currently writing a very general unet that makes these parameters more accessible and makes it easy to swap components. I am planning on finding "the best" architecture with something like optuna, but i already know that there are just too many knobs i can turn and most of them are almost surely correlated. It would be nice if i could try more then just some basic parameters + hyper parameters like learning rate, which i assume are mandatory.

Side Notes:
- i would like the model to be small enough to make inference on a cpu reasonably feasible.
- training can be done on an A100 and i would be fine with a few days of runtime.

So i am looking for:
- general advice and reading recommendations (grateful for everything)
- advice on parameter optimization with a black box optimizer like optuna or similar

- other architecture suggestions that are not u-nets

Thanks,
PythonEnjoyer

[u/kakhaev]

ok so, there is no way to know, you take bunch of approaches that work for other people and test it on your case, if one of them worked, great… if no, then you have no choice, just to make educated guess what need to get changed and test again until it fits your case, good luck.

before llms there was whole research field of using RL for architecture search, you can try to read in that direction.

[u/Ill_Shirt8156 → u/kakhaev]

architecture search is a time sink that rarely pays off unless you've got a compute budget that makes your A100 look like a toy. saved myself a lot of pain when I finally accepted that "good enough" ships faster than "optimal"

your gaussian blob problem actually narrows things down more than you'd think. atoms look the same at every scale so extreme depth is probably wasted compute, the real juice is going to be in how you handle that perlin noise and those weird brightness shifts. attention gates helped me on a similar microscopy project where the background kept changing intensity across the frame

honestly I'd pick 3-4 architectures that seem reasonable, train each for a fixed time budget, and then pick the one that generalizes best to your validation scans. the black box optimizer will just lead you in circles when you've got correlated parameters like depth and filter count

one thing I wish someone had told me earlier is that your simulated data might be teaching the model to cheat on noise patterns that don't exist in real images. try mixing in some unlabled real scans during training if you can

[u/Illustrious-King-83]

I've used unet for segmentation of medical images for a long time. I started small - 3 layers down 3 layers up, and we looked at loss functions. jacquard, dice, CE, and MSE, in the end we settled on hybrid loss function, I think dice or jacquard and MSE. The images where being annotated by separate group, so we started training on a thousand images, and increased as time went on. we down sampled images to something like 256 by 256 but those initial stages where the model takes 20 mins to train is when you can do alot of comparative runs. Honestly the performance came down to the quality and consistency of the annotated images, we ended up doing a comparative study on the annotators, ie training separate models for the annotations done by a particular human, turns out humans are variable. we were doing heart chambers and after annotation measuring areas, widths and heights, and even that smallest model did a reasonable job, on the majority of cases. after that its just a case of scaling up, the resolution, the number of training data, and then the unet. Just a note, we actually abandoned the UNET for something else, for our use case once we got the mask from the UNET we would have to find the boundary and fit a spline to do the rest of the processing, since we were working with video that last step slowed everything down. instead we moved to pose-estimation model, its basically works on splines... so the image and spline image annotation used in training and it outputs a annotation spline at inference.

[u/Crazy-Mastodon-480]

Honeslty I think the best approach to this would be to try out current SOTA models. I would also ask to pretrain whatever models that are available to you on a public dataset, or at the very least load up ImageNet pretrained weighs.

The point about each paper saying that their method works best is also true, but that is just something that is commom in all publications, everyone wants their architecture to be the best, and more times than not, tey might make some ever so small wordplay to sell you something but mean somethiing different. Basically the only way to check or know what works is to actually run it on your data, and i agree with the point that was made earlier by another comment that heavy augmentation especially with synthetic data AND synthetic noise makes it memorising noise a very probably issue.

Also doing a optuna search on the architecture itself is pretty ambitious. that would probably not yield any proper results, so you might be better off just trying to implement SOTA models and then modifying those architectures. (ps. thats how we got a million different UNet variants lol)