SSL crappify to SEMANTIC_SEG transfer: model outputs full-image binarization despite SKIP_UNMATCHED_LAYERS=true

[Ranim-Benkhlifa]

Hi BiaPy developers,

I am working on a 3D semantic segmentation task in BiaPy v3.6.8 and would like to ask whether SSL (crappify) to SEMANTIC_SEG transfer is fully supported in this version.

I pretrained a resunet with crappify (SSL converges well, reconstructions look meaningful) and tried to use the checkpoint as initialization for the downstream task, with the same architecture and same normalization. My labels are binary masks of the target class only.

The downstream model trains but the final prediction is a binarization of the entire image (all cellular structures as foreground) instead of segmenting only the target class. The same config trained from scratch produces a correct segmentation so the issue is specifically the checkpoint loading.

Following issue #132 and inspection of config.py and misc.py for v3.6.8, I use:

MODEL:
LOAD_CHECKPOINT: true
LOAD_CHECKPOINT_ONLY_WEIGHTS: true
LOAD_MODEL_FROM_CHECKPOINT: true # required for SKIP_UNMATCHED_LAYERS to activate
SKIP_UNMATCHED_LAYERS: true

Is this the intended combination, or am I missing something? I also noticed master replaces these flags with MODEL.ITEMS_TO_LOAD_FROM_CHECKPOINT , is that newer mechanism required for this use case, suggesting v3.6.8 may not fully support it?

Environment: BiaPy v3.6.8, Python 3.11.14, PyTorch 2.9.1, Apptainer/Singularity on HPC.

Thanks a lot!

[danifranco]

Hi,

MODEL.ITEMS_TO_LOAD_FROM_CHECKPOINT was added to have more flexibility on what to load from the checkpoint. In the upcoming release LOAD_MODEL_FROM_CHECKPOINT: true will be translated into MODEL.ITEMS_TO_LOAD_FROM_CHECKPOINT: ["weights", "norm", "model_arch"] automatically. MODEL.SKIP_UNMATCHED_LAYERS will still be there.

Regarding the issue, have you checked how the predictions in "per_image" folder look like? can you send me the results folder and the log of the pretraining and finetunning so I can check them? Other option should be to do the experiments with the latest commit, as lots of things were changed for this release.

Best,

[Ranim-Benkhlifa]

Hi Dani,

Thank you for your quick answer.

Regarding per_image/ for the downstream run, I checked it and the probability maps essentially reproduce the entire input image (all cellular structures appear with high values not specifically the lipid droplets "target class"). the model seems to have inherited the reconstruction behavior from the SSL crappify pretraining instead of learning the segmentation task.

I am sharing the results folders and logs from both the pretext task (SSL pretraining) and the downstream finetuning. The archives are about 1.5 GB in total, so I uploaded them via WeTransfer :

https://we.tl/t-VzK56RQBV9b6wnz9

Note: the WeTransfer link expires in 3 days, please let me know if you need me to re-upload.
Thanks again for taking the time to look into this!

[danifranco]

Hi,

I don't see the output log of BiaPy (actually, we should be storing it always so I just make some changes to do it always starting from the nest release: 62a46d1), but with all the files in the folders I have a rough idea of the task you are doing.

I still don't understand why the segmentation is not learning anything valuable after the pretraining. You were saying that when you are applying just a semantic segmentation workflow from scratch it is working, right? It's strange. How many epochs you were setting when training from scratch? Because something that I've seen is that lots of training patches are almost empty (in terms of class information) so maybe the models didn't have enough time to learn? So, something you can try is filtering those samples not containing labels, e.g. pure background patches, so to ensure every patch fed into the model have some information. This normally helps the training when the labels are sparse. You can do it like that (more info here):

DATA:
    TRAIN:
        FILTER_SAMPLES:
          ENABLE: True
          PROPS: [['foreground']]
          VALUES: [[0.0000001]]
          SIGNS: [["lt"]]

Hope it helps!!