reproducing pretrained model

[dojoh]

Hey,
I'm currently trying to reproduce the results from the pretrained model "bact_phase_omni".
According to https://omnipose.readthedocs.io/training.html I should be able to train my own model with

omnipose --train --use_gpu --dir .../data/bact_phase_train --mask_filter _masks \ --n_epochs 4000 --pretrained_model None --learning_rate 0.1 --diameter 0 \ --batch_size 16 --RAdam --nclasses 3

I'm currently predicting the instances from inside python with

parser = get_arg_parser()

args = parser.parse_args()
args.omni = True
args.dir = ".../data/bact_phase_test_pre"
args.use_gpu = True
args.pretrained_model = ".../data/bact_phase_train/models/cellpose_residual_on_style_on_concatenation_off_omni_abstract_nclasses_3_nchan_1_dim_2_bact_phase_train_2024_01_09_09_18_53.248908"
#or args.pretrained_model = "bact_phase_omni"
args.save_flows = True
args.save_outlines = True
args.save_tif = True
args.in_folders = True
args.nchan = 1
args.nclasses = 3
cellpose_omni_main(args)

Unfortunately, the model I trained on my own gives inferior results. What am I doing wrong?

Cheers
Ole

[kevinjohncutler]

@dojoh What version are you using? I can say that my recent phase model trained with the bact_phase dataset (supplemented with some other images) works just as well as the original model, and it does look like your training command is correct. Could you post a picture of the inferior results?

[dojoh]

Hey,
I'm using version 1.0.6 provided through pypi. I can try the current repro version.

I used the code you provided to compute the "JI vs IoU matching threshold" plot. Dashed lines correspond to your suggested weighted scheme:

Do you think this difference could be explained by the additional supplemented images? Is there any chance to get access to these images?

[kevinjohncutler]

@dojoh I have seen pretty similar dips in the performance plot in recent versions of Omnipose, and I should first say that these differences are pretty minor, and are at the level of typical fluctuations that you see when comparing models from different epochs in a training run. I can get near or better performance as long as I choose among several of the models near the end of the training run (--save_each). Is that what your dashed vs solid lines are, early vs late models?

Still, just to debug, please find some example images where the new model is clearly doing worse.

I have done a bit of tweaking to the loss function of Omnipose in order to get better flow field predictions at cell-cell interfaces. It is possible that this makes the training a bit more unstable (torchvf loss in particular), but I have not done a comparison to earlier releases. You'll want to use a PyPi version such as 0.4.4 (corresponding to the paper publication) to fully reproduce the bact_phase_omni model.