fixes PlantSeg crash on apple silicon #273

[lorenzocerrone]

I was able to reproduce this in the plantseg env (but not in a clean env) The issue is the order of import of numpy/torch, and the fact that we are reading the state dict from file.

This will fail

import numpy as np
import torch

model = torch.nn.Conv3d(64, 32, kernel_size=(3, 3, 3))

torch.save(model.state_dict(), "model.pth")
model.load_state_dict(torch.load("model.pth"))

These two will work

import torch
import numpy as np

model = torch.nn.Conv3d(64, 32, kernel_size=(3, 3, 3))

torch.save(model.state_dict(), "model.pth")
model.load_state_dict(torch.load("model.pth"))

import numpy as np
import torch

model = torch.nn.Conv3d(64, 32, kernel_size=(3, 3, 3))

state_dict = model.state_dict()
model.load_state_dict(state_dict)

Solution

For now importing torch in the __init__.py works and has no drawbacks.

[lorenzocerrone]

Looks like the pre-commit does not like this patch 🤣

[qin-yu]

Hey @lorenzocerrone

I documented the thing here: https://kreshuklab.github.io/plant-seg/chapters/getting_started/contributing/#before-submitting-a-pull-request

Hope it's not too annoying to you.

If it's a squash-and-merge commit, then just proceed. Otherwise you can do something like git rebase -i HEAD~4 and squash the commits from pre-commit.ci, or just force push a new version with the fix.

[qin-yu]

Looks like the pre-commit does not like this patch 🤣

Oh @lorenzocerrone I see what you mean. You need to do __all__ = [torch] or something to tell yourself that this import is not useless, or if you just want to tell Ruff that this is not useless, do: # noqa: F401 on this line.

[qin-yu]

Hey @lorenzocerrone,

I noticed one thing: For macOS,there is no NVIDIA GPUs and therefore should not use -c nvidia, did you install with:

mamba create -n plant-seg -c pytorch -c nvidia -c conda-forge pytorch pytorch-cuda=12.1 pyqt plant-seg --no-channel-priority

This will not work because for macOS the official PyTorch should be:

conda install pytorch::pytorch torchvision torchaudio -c pytorch

So your installation script would be:

mamba create -n plant-seg -c pytorch -c conda-forge pytorch::pytorch pytorch-cuda=12.1 pyqt plant-seg --no-channel-priority

Please try this and see if the problem is solved for you!

[lorenzocerrone]

I already used the correct mamba installation (pytorch-cuda=12.1 should also be removed). The bug persists no matter how I install plantseg, but I can not reproduce it in a clean env, so it's a bit weird.

[qin-yu]

Hey Lorenzo, since nothing fixes the environment you need (and PlantSeg is already not too fast so import torch doesn't really hurt the user experience anyways), let's squash and merge this?

[qin-yu]

I have one more question: What additional components does the unclean environment include? Or do you mean all macOS users will encounter #273 after launching it once?

[qin-yu]

I think I understood: In PlantSeg env this happens, while in a Torch + NumPy only env this doesn't.