Microsplit - train noise models

[ShataDg]

@bethac07
I tried training the model a few times after reducing the number of epochs and number of workers to one, but that did not fix the issue. I kept getting Your system has ran out of application memory even though I had no other applications open, and abot ~200GB available on the disc.

I have opened an issue on the Microsplit repo.

[ShataDg]

I tried running on CPU and downgrading PyTorch but that did not fix the issue.

[ShataDg]

@bethac07

The run completed and the denoised images were still blank, and there were no error messages.
I nevertheless ran the next block of code

for channel_idx in range(train_data.shape[-1]):
    
    # train Noise Model for current channel
    print(f"Training noise model for channel {channel_idx}")
    channel_data = train_data[..., channel_idx]
    channel_prediction = np.concatenate(prediction)[:, channel_idx]
    noise_model_config = GaussianMixtureNMConfig(model_type="GaussianMixtureNoiseModel",
                                                 min_signal=channel_data.min(),
                                                 max_signal=channel_data.max(),
                                                 n_coeff=4,
                                                 n_gaussian=6)
    noise_model = GaussianMixtureNoiseModel(noise_model_config)
    noise_model.fit(signal=channel_data, observation=channel_prediction, n_epochs=1000)

    # save result on disk for later re-use
    noise_model.save(path=f"noise_models/{EXPOSURE_DURATION}/", name=f"noise_model_Ch{channel_idx}")

    # show the result
    histogram = create_histogram(bins=100, min_val=channel_data.min(), max_val=channel_data.max(), signal=channel_data, observation=channel_prediction)
    plot_probability_distribution(noise_model, signalBinIndex=50, histogram=histogram[0], channel=channel_idx)

and it gave me

Training noise model for channel 0
[GaussianMixtureNoiseModel] min_sigma: 200.0
0 nan
NaN or Inf detected in the weights. Aborting training at epoch:  1
===================

The trained parameters (noise_model_Ch0) is saved at location: noise_models/5ms/
[/Users/sdasgupt/anaconda3/envs/microsplit/lib/python3.9/site-packages/microsplit_reproducibility/utils/utils.py:82](https://file+.vscode-resource.vscode-cdn.net/Users/sdasgupt/anaconda3/envs/microsplit/lib/python3.9/site-packages/microsplit_reproducibility/utils/utils.py:82): UserWarning: To copy construct from a tensor, it is recommended to use sourceTensor.clone().detach() or sourceTensor.clone().detach().requires_grad_(True), rather than torch.tensor(sourceTensor).
  noise_model.min_signal = torch.tensor(
[/Users/sdasgupt/anaconda3/envs/microsplit/lib/python3.9/site-packages/microsplit_reproducibility/utils/utils.py:85](https://file+.vscode-resource.vscode-cdn.net/Users/sdasgupt/anaconda3/envs/microsplit/lib/python3.9/site-packages/microsplit_reproducibility/utils/utils.py:85): UserWarning: To copy construct from a tensor, it is recommended to use sourceTensor.clone().detach() or sourceTensor.clone().detach().requires_grad_(True), rather than torch.tensor(sourceTensor).
  noise_model.max_signal = torch.tensor(
[/Users/sdasgupt/anaconda3/envs/microsplit/lib/python3.9/site-packages/microsplit_reproducibility/utils/utils.py:90](https://file+.vscode-resource.vscode-cdn.net/Users/sdasgupt/anaconda3/envs/microsplit/lib/python3.9/site-packages/microsplit_reproducibility/utils/utils.py:90): UserWarning: To copy construct from a tensor, it is recommended to use sourceTensor.clone().detach() or sourceTensor.clone().detach().requires_grad_(True), rather than torch.tensor(sourceTensor).
  observations=queryObservations, signals=torch.tensor(querySignal)
Training noise model for channel 1
[GaussianMixtureNoiseModel] min_sigma: 200.0
0 nan
NaN or Inf detected in the weights. Aborting training at epoch:  1
===================

The trained parameters (noise_model_Ch1) is saved at location: noise_models/5ms/
Training noise model for channel 2
[GaussianMixtureNoiseModel] min_sigma: 200.0
0 nan
NaN or Inf detected in the weights. Aborting training at epoch:  1
===================

The trained parameters (noise_model_Ch2) is saved at location: noise_models/5ms/

I shared all the outputs and the contents of my environment with the microsplit team.

[ShataDg]

Finally the some non-blank denoised images!

What might have been the issue:
I had followed the instructions for Mac and created an env in Python 3.9 using

CONDA_SUBDIR=osx-arm64 conda create -n microsplit python=3.9

My original environment had the following (because I think later versions of torch were not compatible with Python 3.9)

careamics==0.0.12
careamics-portfolio==0.0.15
pytorch-lightning==2.2.0
torch==2.2.0
torchmetrics==1.8.2
torchvision==0.17.0

Folks from the microsplit team suggested I create the env with Python 3.11 and also mentioned the specific versions of torch and torchvision that are working for them

I created a new environment with

CONDA_SUBDIR=osx-arm64 conda create -n microsplit python=3.11

Then

pip install torch torchvision

And after re-cloning the repo

pip install -e .

The environment thereby created had

pytorch-lightning==2.5.5
torch==2.8.0
torchmetrics==1.8.2
torchvision==0.23.0

careamics==0.0.16
careamics-portfolio==0.0.15

But when I tried running the notebook (HT_LI24), while running

from microsplit_reproducibility.configs.data.HT_LIF24 import get_data_configs
from microsplit_reproducibility.datasets.HT_LIF24 import get_train_val_data
from microsplit_reproducibility.utils.utils import plot_probability_distribution

I got

---------------------------------------------------------------------------
ImportError                               Traceback (most recent call last)
Cell In[1], [line 15](vscode-notebook-cell:?execution_count=1&line=15)
     12 from careamics.config import create_n2v_configuration
     14 # Dataset specific imports...
---> [15](vscode-notebook-cell:?execution_count=1&line=15) from microsplit_reproducibility.configs.data.HT_LIF24 import get_data_configs
     16 from microsplit_reproducibility.datasets.HT_LIF24 import get_train_val_data
     17 from microsplit_reproducibility.utils.utils import plot_probability_distribution

File ~/Documents/microsplit/microsplit/MicroSplit-reproducibility/src/microsplit_reproducibility/configs/data/HT_LIF24.py:3
      1 from typing import Literal
----> [3](https://file+.vscode-resource.vscode-cdn.net/Users/sdasgupt/Documents/microsplit/microsplit/MicroSplit-reproducibility/examples/2D/HT_LIF24/~/Documents/microsplit/microsplit/MicroSplit-reproducibility/src/microsplit_reproducibility/configs/data/HT_LIF24.py:3) from careamics.lvae_training.dataset import DatasetConfig, DataSplitType, DataType
      5 CH_IDX_LIST = [1, 2, 3, 17]
      8 class HTLIF24DataConfig(DatasetConfig):

ImportError: cannot import name 'DatasetConfig' from 'careamics.lvae_training.dataset' (/Users/sdasgupt/anaconda3/envs/microsplit/lib/python3.11/site-packages/careamics/lvae_training/dataset/__init__.py)

I uninstalled torch and torchvision and installed the exact versions suggested, and then my environment had

pytorch-lightning==2.5.0.post0
torch==2.6.0
torchmetrics==1.8.2
torchvision==0.21.0
careamics==0.0.12
careamics-portfolio==0.0.15

I used the following configuration to train the model and got the expected results!

config = create_n2v_configuration(
    experiment_name="ht_lif24_noise_models_n2v",
    data_type="array",
    axes="SYXC",
    n_channels=len(TARGET_CHANNEL_IDX_LIST),
    patch_size=(64, 64),
    batch_size=64,
    num_epochs=5,
    train_dataloader_params={"num_workers": get_num_workers()},
    val_dataloader_params={"num_workers": get_num_workers()},
)

[ShataDg]

I don't think I have access to these repos:

Master parquet file: https://github.com/jump-cellpainting/jump-qc/blob/main/data/qc_features_annotations.parquet
Batch we know has at least a couple blurry and debris: https://github.com/jump-cellpainting/jump-qc/blob/main/data/broad_annotations/Classifier_Batch11_Erin.csv

[bethac07]

I just sent an invite, lmk when you accept, and then I'll put the repo back in archive mode

[ShataDg]

Done!