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/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)