Fine-tuning 3d model

[ElpadoCan]

Hello, first of all, congratulations on the software tool. It is very easy to use, and the code is super clean.

I am having some issues with fine-tuning the 3d model. I have 170 training z-stack volumes and 10 validation volumes. I fine-tuned with the API as explained in the documentation (see below) starting from the synth_3d model.

The synth_3d model had a very high recall but low precision. However, the fine-tuned model has the opposite problem, it achieves 100% precision, but crazy low recall. What do you think I should try? Training from scratch? Thank you very much and have a good day!

pretrained_model_name = "synth_3d"
model = Spotiflow.from_pretrained(
    pretrained_model_name, 
    inference_mode=False
)
model.fit(
    train_imgs,
    train_spots,
    val_imgs,
    val_spots,
    device='cuda',
    save_dir=model_path,
)

[AlbertDominguez]

Hey Francesco,

Thanks for the kind words! Hmm... after the finetuning procedure, which is the probability threshold that was obtained through the automatic optimization that happens after training finishes? You can find it either printed at the end of the run or in the thresholds.yaml file in the model_path folder. It might be that there is a bug in the optimization and that the threshold is too high/too low. In this direction too, which is the crop size you are using for fine-tuning and how large are the validation images (pre-cropping)?

[ElpadoCan]

Hi Albert, thanks for the quick reply.

which is the probability threshold that was obtained through the automatic optimization that happens after training finishes?

In the thresholds.yaml I have these values

prob_thresh_best:
- 0.696
prob_thresh_last:
- 0.5879999999999999

Are they too high? Should I try 0.5?

which is the crop size you are using for fine-tuning and how large are the validation images (pre-cropping)?

I didn't set any crop size, but in the train_config.yaml I see crop_size : 512. My images are all with (43, 377, 377) shape, where 43 is the number of z-slices. I don't know if it can help, but the pixel size is 0.073 micrometers.

Thanks for the help!

[AlbertDominguez]

So the default one used is the best, which indeed seems to be too high. It would be nice to try something smaller and see if it improves. If that's the case, it might be a sampling bug (well, more of an artifact) during validation.

You can maybe try the following snippet to recompute the probability threshold on the whole validation set automatically and see if it differs a lot from the current 0.696 (and also if the precision/recall/f1 seems better then):

from spotiflow.utils import points_matching_dataset
from spotiflow.model import Spotiflow
from tqdm.auto import tqdm

import numpy as np

def get_ds_f1(model, imgs, gt_spots, prob_thresh=0.5):
    preds = tuple(model.predict(x, prob_thresh=prob_thresh, verbose=False)[0] for x in imgs)
    f1 = points_matching_dataset(gt_spots, preds).f1
    return f1

val_imgs, val_spots = ...
model = Spotiflow.from_folder(..., inference_mode=True) # load the trained model
best_score, best_p = -1, -1
for p in tqdm(np.arange(0.3, 0.705, 0.05), desc="Sweeping thresholds"):
    f1 = get_ds_f1(model, val_imgs, val_spots, prob_thresh=p)
    if f1 > best_score:
        best_score, best_p = f1, p
print(f"Best F1={best_score:.3f} (@ threshold={best_p:.2f})")

This doesnt change the model threshold that's stored (although that's possible, lmk if you want to do that!), but you can then pass the threshold as an argument to the .predict() method for trying it e.g. on the test set. If the score/thr are substantially different it's highly likely that there's some bug in the optimization of the threshold after training, I'll start taking a look just in case!

[ElpadoCan]

I just ran your script and I got this

Best F1=0.809 (@ threshold=0.35)

I guess I should try running my fine-tuned model with this right?

prob_thresh_best:
- 0.35

[AlbertDominguez]

Yes, I'd suggest doing that! So it really does seem to be very off from the automatic one. I'll hunt the bug later today and let you know any updates in here, so I'll keep this open. Let me know if with this change the predictions look better!
Thx a lot for reporting the issue :)

[ElpadoCan]

The lower threshold value did the trick, thanks!

As a suggestion, it would be nice to be able to pass the threshold value as a parameter to predict instead of changing a YAML file, what do you think?

[AlbertDominguez]

Happy to hear it works! You can change the threshold with the prob_thresh argument of the predict method, e.g. model.predict(img, prob_thresh=0.5)

[ElpadoCan]

Oh yeah, I'm stupid, I knew I could. Sorry about that :D