Conversion of bact_fluor_omnitorch_0 model to .onnx throws error

[JensWendt]

Hello,

we are trying to import your models as custom models into the Arivis Vision4D Cellpose segmenter(first paragraph) module for some bacteria segmentation. For that we have to convert your models into the .onnx format, for which we use their custom python script which basically calls torch.onnx.export() Converting your model bact_fluor_cptorch_0 works smoothly, but for bact_fluor_omnitorch_0 i get the following error:

Traceback (most recent call last):
  File "cellpose_to_onnx.py", line 92, in <module>
    main()
  File "cellpose_to_onnx.py", line 82, in main
    convert_to_ONNX(model_path=args.model_path, output_directory=args.output_directory, diam_mean=args.mean_diameter)
  File "cellpose_to_onnx.py", line 19, in convert_to_ONNX
    model.load_model(model_path)
  File "C:\ProgramData\anaconda3\envs\cellpose\lib\site-packages\cellpose\resnet_torch.py", line 236, in load_model
    self.load_state_dict(dict([(name, param) for name, param in state_dict.items()]), strict=False)
  File "C:\ProgramData\anaconda3\envs\cellpose\lib\site-packages\torch\nn\modules\module.py", line 2152, in load_state_dict
    raise RuntimeError('Error(s) in loading state_dict for {}:\n\t{}'.format(
RuntimeError: Error(s) in loading state_dict for CPnet:
        size mismatch for output.2.weight: copying a param with shape torch.Size([4, 32, 1, 1]) from checkpoint, the shape in current model is torch.Size([3, 32, 1, 1]).
        size mismatch for output.2.bias: copying a param with shape torch.Size([4]) from checkpoint, the shape in current model is torch.Size([3]).

I think it is generally beneficial to have the models as "open" and FAIR as possible. Thus a conversion to .onnx would be nice. Maybe you can look into what went wrong there?

[kevinjohncutler]

@JensWendt I don't currently have the time to look into this, but I suspect the reason it was failing is that the cp models have 3 prediction classes and the published omni models have 4. Please let me know if you get a new conversion script sorted out so that future models can be converted as well.

[kevinjohncutler]

@JensWendt Archiving for now, please reopen if you have any updates.