JUMP-QC dataset

[ShataDg]

• Download 150 good images (pull_images.ipynb)
• Download 88 bad images from the annotated dataset (pull_images.ipynb)
• Test with two channels
  • DNA and AGP
  • DNA and ER
  • DNA and Mito
  • DNA, AGP, ER
  • DNA, AGP, Mito
• Train on 100 good images
• Test on 50 good images and 88 bad images
• Test --> similarity comparison / colocalization (Pearson) between raw and unmixed images

2025-12 Updated checklist

• Reduce the training set (20 images/channel)
• Complete neptune setup
• Switch to "musplit" model which does not require a noise model

[ShataDg]

1. Dowloaded the 2020_11_16_Scope1_YokogawaUS\BRO0117014-3D-2x2x1-singleplane-10x_20201026_113115\PECCU images

aws s3 cp --recursive s3://cellpainting-gallery/cpg0002-jump-scope/source_4/images/2020_11_16_Scope1_YokogawaUS/images . --no-sign-request --dryrun

2. Split the images into:

• good_images_train (100 images)
• good_images_val (50 images)
• bad_images (88 images)
  • selected_for_training_noise_model (40 images)

3. Noise model needs to be trained using the 00_noisemodels.ipynb notebook, with microscopy images that are noisy and cover the relevant range of pixel intensities that are also covered by the data you want to use

4. Noise model trained using DNA and AGP channels

5. Error message setting up the training when cpu option is used

---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
Cell In[29], line 14
      2 trainer = Trainer(
      3     max_epochs=training_config.num_epochs,
      4     # NOTE: if you are on a mac swap the accelerator to "mps"
   (...)     11     gradient_clip_algorithm=training_config.gradient_clip_algorithm,
     12 )
     13 # start the training - yay!
---> 14 trainer.fit(
     15     model=model,
     16     train_dataloaders=train_dloader,
     17     val_dataloaders=val_dloader,
     18 )

File ~/anaconda3/envs/microsplit/lib/python3.11/site-packages/pytorch_lightning/trainer/trainer.py:539, in Trainer.fit(self, model, train_dataloaders, val_dataloaders, datamodule, ckpt_path)
    537 self.state.status = TrainerStatus.RUNNING
    538 self.training = True
--> 539 call._call_and_handle_interrupt(
    540     self, self._fit_impl, model, train_dataloaders, val_dataloaders, datamodule, ckpt_path
    541 )

File ~/anaconda3/envs/microsplit/lib/python3.11/site-packages/pytorch_lightning/trainer/call.py:47, in _call_and_handle_interrupt(trainer, trainer_fn, *args, **kwargs)
     45     if trainer.strategy.launcher is not None:
     46         return trainer.strategy.launcher.launch(trainer_fn, *args, trainer=trainer, **kwargs)
---> 47     return trainer_fn(*args, **kwargs)
     49 except _TunerExitException:
     50     _call_teardown_hook(trainer)

File ~/anaconda3/envs/microsplit/lib/python3.11/site-packages/pytorch_lightning/trainer/trainer.py:575, in Trainer._fit_impl(self, model, train_dataloaders, val_dataloaders, datamodule, ckpt_path)
    568 assert self.state.fn is not None
    569 ckpt_path = self._checkpoint_connector._select_ckpt_path(
    570     self.state.fn,
    571     ckpt_path,
    572     model_provided=True,
    573     model_connected=self.lightning_module is not None,
    574 )
--> 575 self._run(model, ckpt_path=ckpt_path)
    577 assert self.state.stopped
    578 self.training = False

File ~/anaconda3/envs/microsplit/lib/python3.11/site-packages/pytorch_lightning/trainer/trainer.py:982, in Trainer._run(self, model, ckpt_path)
    977 self._signal_connector.register_signal_handlers()
    979 # ----------------------------
    980 # RUN THE TRAINER
    981 # ----------------------------
--> 982 results = self._run_stage()
    984 # ----------------------------
    985 # POST-Training CLEAN UP
    986 # ----------------------------
    987 log.debug(f"{self.__class__.__name__}: trainer tearing down")

File ~/anaconda3/envs/microsplit/lib/python3.11/site-packages/pytorch_lightning/trainer/trainer.py:1024, in Trainer._run_stage(self)
   1022 if self.training:
   1023     with isolate_rng():
-> 1024         self._run_sanity_check()
   1025     with torch.autograd.set_detect_anomaly(self._detect_anomaly):
   1026         self.fit_loop.run()

File ~/anaconda3/envs/microsplit/lib/python3.11/site-packages/pytorch_lightning/trainer/trainer.py:1053, in Trainer._run_sanity_check(self)
   1050 call._call_callback_hooks(self, "on_sanity_check_start")
   1052 # run eval step
-> 1053 val_loop.run()
   1055 call._call_callback_hooks(self, "on_sanity_check_end")
   1057 # reset logger connector

File ~/anaconda3/envs/microsplit/lib/python3.11/site-packages/pytorch_lightning/loops/utilities.py:179, in _no_grad_context.<locals>._decorator(self, *args, **kwargs)
    177     context_manager = torch.no_grad
    178 with context_manager():
--> 179     return loop_run(self, *args, **kwargs)

File ~/anaconda3/envs/microsplit/lib/python3.11/site-packages/pytorch_lightning/loops/evaluation_loop.py:144, in _EvaluationLoop.run(self)
    142     self.batch_progress.is_last_batch = data_fetcher.done
    143     # run step hooks
--> 144     self._evaluation_step(batch, batch_idx, dataloader_idx, dataloader_iter)
    145 except StopIteration:
    146     # this needs to wrap the `*_step` call too (not just `next`) for `dataloader_iter` support
    147     break

File ~/anaconda3/envs/microsplit/lib/python3.11/site-packages/pytorch_lightning/loops/evaluation_loop.py:433, in _EvaluationLoop._evaluation_step(self, batch, batch_idx, dataloader_idx, dataloader_iter)
    427 hook_name = "test_step" if trainer.testing else "validation_step"
    428 step_args = (
    429     self._build_step_args_from_hook_kwargs(hook_kwargs, hook_name)
    430     if not using_dataloader_iter
    431     else (dataloader_iter,)
    432 )
--> 433 output = call._call_strategy_hook(trainer, hook_name, *step_args)
    435 self.batch_progress.increment_processed()
    437 if using_dataloader_iter:
    438     # update the hook kwargs now that the step method might have consumed the iterator

File ~/anaconda3/envs/microsplit/lib/python3.11/site-packages/pytorch_lightning/trainer/call.py:323, in _call_strategy_hook(trainer, hook_name, *args, **kwargs)
    320     return None
    322 with trainer.profiler.profile(f"[Strategy]{trainer.strategy.__class__.__name__}.{hook_name}"):
--> 323     output = fn(*args, **kwargs)
    325 # restore current_fx when nested context
    326 pl_module._current_fx_name = prev_fx_name

File ~/anaconda3/envs/microsplit/lib/python3.11/site-packages/pytorch_lightning/strategies/strategy.py:412, in Strategy.validation_step(self, *args, **kwargs)
    410 if self.model != self.lightning_module:
    411     return self._forward_redirection(self.model, self.lightning_module, "validation_step", *args, **kwargs)
--> 412 return self.lightning_module.validation_step(*args, **kwargs)

File ~/anaconda3/envs/microsplit/lib/python3.11/site-packages/careamics/lightning/lightning_module.py:456, in VAEModule.validation_step(self, batch, batch_idx)
    454 loss = {"_".join(["val", k]): v for k, v in loss.items()}
    455 self.log_dict(loss, on_epoch=True, prog_bar=True)
--> 456 curr_psnr = self.compute_val_psnr(out, target)
    457 for i, psnr in enumerate(curr_psnr):
    458     self.log(f"val_psnr_ch{i+1}_batch", psnr, on_epoch=True)

File ~/anaconda3/envs/microsplit/lib/python3.11/site-packages/careamics/lightning/lightning_module.py:597, in VAEModule.compute_val_psnr(self, model_output, target, psnr_func)
    592     self.running_psnr[i].update(rec=recons_img[:, i], tar=target[:, i])
    594 # compute psnr for each channel in the current batch
    595 # TODO: this doesn't need do be a method of this class
    596 # and hence can be moved to a separate module
--> 597 return [
    598     psnr_func(
    599         gt=target[:, i].clone().detach().cpu().numpy(),
    600         pred=recons_img[:, i].clone().detach().cpu().numpy(),
    601     )
    602     for i in range(out_channels)
    603 ]

File ~/anaconda3/envs/microsplit/lib/python3.11/site-packages/careamics/lightning/lightning_module.py:600, in <listcomp>(.0)
    592     self.running_psnr[i].update(rec=recons_img[:, i], tar=target[:, i])
    594 # compute psnr for each channel in the current batch
    595 # TODO: this doesn't need do be a method of this class
    596 # and hence can be moved to a separate module
    597 return [
    598     psnr_func(
    599         gt=target[:, i].clone().detach().cpu().numpy(),
--> 600         pred=recons_img[:, i].clone().detach().cpu().numpy(),
    601     )
    602     for i in range(out_channels)
    603 ]

TypeError: Got unsupported ScalarType BFloat16

Error message when mps is used

---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
Cell In[12], [line 14](vscode-notebook-cell:?execution_count=12&line=14)
      2 trainer = Trainer(
      3     max_epochs=training_config.num_epochs,
      4     # NOTE: if you are on a mac swap the accelerator to "mps"
   (...)     11     gradient_clip_algorithm=training_config.gradient_clip_algorithm,
     12 )
     13 # start the training - yay!
---> [14](vscode-notebook-cell:?execution_count=12&line=14) trainer.fit(
     15     model=model,
     16     train_dataloaders=train_dloader,
     17     val_dataloaders=val_dloader,
     18 )

File ~/anaconda3/envs/microsplit/lib/python3.11/site-packages/pytorch_lightning/trainer/trainer.py:539, in Trainer.fit(self, model, train_dataloaders, val_dataloaders, datamodule, ckpt_path)
    537 self.state.status = TrainerStatus.RUNNING
    538 self.training = True
--> [539](https://file+.vscode-resource.vscode-cdn.net/Users/sdasgupt/Documents/microsplit/microsplit/MicroSplit-reproducibility/examples/2D/custom_dataset/~/anaconda3/envs/microsplit/lib/python3.11/site-packages/pytorch_lightning/trainer/trainer.py:539) call._call_and_handle_interrupt(
    540     self, self._fit_impl, model, train_dataloaders, val_dataloaders, datamodule, ckpt_path
    541 )

File ~/anaconda3/envs/microsplit/lib/python3.11/site-packages/pytorch_lightning/trainer/call.py:47, in _call_and_handle_interrupt(trainer, trainer_fn, *args, **kwargs)
     45     if trainer.strategy.launcher is not None:
     46         return trainer.strategy.launcher.launch(trainer_fn, *args, trainer=trainer, **kwargs)
---> [47](https://file+.vscode-resource.vscode-cdn.net/Users/sdasgupt/Documents/microsplit/microsplit/MicroSplit-reproducibility/examples/2D/custom_dataset/~/anaconda3/envs/microsplit/lib/python3.11/site-packages/pytorch_lightning/trainer/call.py:47)     return trainer_fn(*args, **kwargs)
     49 except _TunerExitException:
     50     _call_teardown_hook(trainer)

File ~/anaconda3/envs/microsplit/lib/python3.11/site-packages/pytorch_lightning/trainer/trainer.py:575, in Trainer._fit_impl(self, model, train_dataloaders, val_dataloaders, datamodule, ckpt_path)
    568 assert self.state.fn is not None
    569 ckpt_path = self._checkpoint_connector._select_ckpt_path(
    570     self.state.fn,
    571     ckpt_path,
    572     model_provided=True,
    573     model_connected=self.lightning_module is not None,
    574 )
--> [575](https://file+.vscode-resource.vscode-cdn.net/Users/sdasgupt/Documents/microsplit/microsplit/MicroSplit-reproducibility/examples/2D/custom_dataset/~/anaconda3/envs/microsplit/lib/python3.11/site-packages/pytorch_lightning/trainer/trainer.py:575) self._run(model, ckpt_path=ckpt_path)
    577 assert self.state.stopped
    578 self.training = False

File ~/anaconda3/envs/microsplit/lib/python3.11/site-packages/pytorch_lightning/trainer/trainer.py:982, in Trainer._run(self, model, ckpt_path)
    977 self._signal_connector.register_signal_handlers()
    979 # ----------------------------
    980 # RUN THE TRAINER
    981 # ----------------------------
--> [982](https://file+.vscode-resource.vscode-cdn.net/Users/sdasgupt/Documents/microsplit/microsplit/MicroSplit-reproducibility/examples/2D/custom_dataset/~/anaconda3/envs/microsplit/lib/python3.11/site-packages/pytorch_lightning/trainer/trainer.py:982) results = self._run_stage()
    984 # ----------------------------
    985 # POST-Training CLEAN UP
    986 # ----------------------------
    987 log.debug(f"{self.__class__.__name__}: trainer tearing down")

File ~/anaconda3/envs/microsplit/lib/python3.11/site-packages/pytorch_lightning/trainer/trainer.py:1024, in Trainer._run_stage(self)
   1022 if self.training:
   1023     with isolate_rng():
-> [1024](https://file+.vscode-resource.vscode-cdn.net/Users/sdasgupt/Documents/microsplit/microsplit/MicroSplit-reproducibility/examples/2D/custom_dataset/~/anaconda3/envs/microsplit/lib/python3.11/site-packages/pytorch_lightning/trainer/trainer.py:1024)         self._run_sanity_check()
   1025     with torch.autograd.set_detect_anomaly(self._detect_anomaly):
   1026         self.fit_loop.run()

File ~/anaconda3/envs/microsplit/lib/python3.11/site-packages/pytorch_lightning/trainer/trainer.py:1053, in Trainer._run_sanity_check(self)
   1050 call._call_callback_hooks(self, "on_sanity_check_start")
   1052 # run eval step
-> [1053](https://file+.vscode-resource.vscode-cdn.net/Users/sdasgupt/Documents/microsplit/microsplit/MicroSplit-reproducibility/examples/2D/custom_dataset/~/anaconda3/envs/microsplit/lib/python3.11/site-packages/pytorch_lightning/trainer/trainer.py:1053) val_loop.run()
   1055 call._call_callback_hooks(self, "on_sanity_check_end")
   1057 # reset logger connector

File ~/anaconda3/envs/microsplit/lib/python3.11/site-packages/pytorch_lightning/loops/utilities.py:179, in _no_grad_context.<locals>._decorator(self, *args, **kwargs)
    177     context_manager = torch.no_grad
    178 with context_manager():
--> [179](https://file+.vscode-resource.vscode-cdn.net/Users/sdasgupt/Documents/microsplit/microsplit/MicroSplit-reproducibility/examples/2D/custom_dataset/~/anaconda3/envs/microsplit/lib/python3.11/site-packages/pytorch_lightning/loops/utilities.py:179)     return loop_run(self, *args, **kwargs)

File ~/anaconda3/envs/microsplit/lib/python3.11/site-packages/pytorch_lightning/loops/evaluation_loop.py:144, in _EvaluationLoop.run(self)
    142     self.batch_progress.is_last_batch = data_fetcher.done
    143     # run step hooks
--> [144](https://file+.vscode-resource.vscode-cdn.net/Users/sdasgupt/Documents/microsplit/microsplit/MicroSplit-reproducibility/examples/2D/custom_dataset/~/anaconda3/envs/microsplit/lib/python3.11/site-packages/pytorch_lightning/loops/evaluation_loop.py:144)     self._evaluation_step(batch, batch_idx, dataloader_idx, dataloader_iter)
    145 except StopIteration:
    146     # this needs to wrap the `*_step` call too (not just `next`) for `dataloader_iter` support
    147     break

File ~/anaconda3/envs/microsplit/lib/python3.11/site-packages/pytorch_lightning/loops/evaluation_loop.py:407, in _EvaluationLoop._evaluation_step(self, batch, batch_idx, dataloader_idx, dataloader_iter)
    405     batch = trainer.precision_plugin.convert_input(batch)
    406     batch = trainer.lightning_module._on_before_batch_transfer(batch, dataloader_idx=dataloader_idx)
--> [407](https://file+.vscode-resource.vscode-cdn.net/Users/sdasgupt/Documents/microsplit/microsplit/MicroSplit-reproducibility/examples/2D/custom_dataset/~/anaconda3/envs/microsplit/lib/python3.11/site-packages/pytorch_lightning/loops/evaluation_loop.py:407)     batch = call._call_strategy_hook(trainer, "batch_to_device", batch, dataloader_idx=dataloader_idx)
    409 # the `_step` methods don't take a batch_idx when `dataloader_iter` is used, but all other hooks still do,
    410 # so we need different kwargs
    411 hook_kwargs = self._build_kwargs(
    412     batch, batch_idx, dataloader_idx if self._is_sequential and self.num_dataloaders > 1 else None
    413 )

File ~/anaconda3/envs/microsplit/lib/python3.11/site-packages/pytorch_lightning/trainer/call.py:323, in _call_strategy_hook(trainer, hook_name, *args, **kwargs)
    320     return None
    322 with trainer.profiler.profile(f"[Strategy]{trainer.strategy.__class__.__name__}.{hook_name}"):
--> [323](https://file+.vscode-resource.vscode-cdn.net/Users/sdasgupt/Documents/microsplit/microsplit/MicroSplit-reproducibility/examples/2D/custom_dataset/~/anaconda3/envs/microsplit/lib/python3.11/site-packages/pytorch_lightning/trainer/call.py:323)     output = fn(*args, **kwargs)
    325 # restore current_fx when nested context
    326 pl_module._current_fx_name = prev_fx_name

File ~/anaconda3/envs/microsplit/lib/python3.11/site-packages/pytorch_lightning/strategies/strategy.py:278, in Strategy.batch_to_device(self, batch, device, dataloader_idx)
    276 device = device or self.root_device
    277 if model is not None:
--> [278](https://file+.vscode-resource.vscode-cdn.net/Users/sdasgupt/Documents/microsplit/microsplit/MicroSplit-reproducibility/examples/2D/custom_dataset/~/anaconda3/envs/microsplit/lib/python3.11/site-packages/pytorch_lightning/strategies/strategy.py:278)     return model._apply_batch_transfer_handler(batch, device=device, dataloader_idx=dataloader_idx)
    279 return move_data_to_device(batch, device)

File ~/anaconda3/envs/microsplit/lib/python3.11/site-packages/pytorch_lightning/core/module.py:352, in LightningModule._apply_batch_transfer_handler(self, batch, device, dataloader_idx)
    348 def _apply_batch_transfer_handler(
    349     self, batch: Any, device: Optional[torch.device] = None, dataloader_idx: int = 0
    350 ) -> Any:
    351     device = device or self.device
--> [352](https://file+.vscode-resource.vscode-cdn.net/Users/sdasgupt/Documents/microsplit/microsplit/MicroSplit-reproducibility/examples/2D/custom_dataset/~/anaconda3/envs/microsplit/lib/python3.11/site-packages/pytorch_lightning/core/module.py:352)     batch = self._call_batch_hook("transfer_batch_to_device", batch, device, dataloader_idx)
    353     batch = self._call_batch_hook("on_after_batch_transfer", batch, dataloader_idx)
    354     return batch

File ~/anaconda3/envs/microsplit/lib/python3.11/site-packages/pytorch_lightning/core/module.py:341, in LightningModule._call_batch_hook(self, hook_name, *args)
    338     else:
    339         trainer_method = call._call_lightning_datamodule_hook
--> [341](https://file+.vscode-resource.vscode-cdn.net/Users/sdasgupt/Documents/microsplit/microsplit/MicroSplit-reproducibility/examples/2D/custom_dataset/~/anaconda3/envs/microsplit/lib/python3.11/site-packages/pytorch_lightning/core/module.py:341)     return trainer_method(trainer, hook_name, *args)
    342 hook = getattr(self, hook_name)
    343 return hook(*args)

File ~/anaconda3/envs/microsplit/lib/python3.11/site-packages/pytorch_lightning/trainer/call.py:171, in _call_lightning_module_hook(trainer, hook_name, pl_module, *args, **kwargs)
    168 pl_module._current_fx_name = hook_name
    170 with trainer.profiler.profile(f"[LightningModule]{pl_module.__class__.__name__}.{hook_name}"):
--> [171](https://file+.vscode-resource.vscode-cdn.net/Users/sdasgupt/Documents/microsplit/microsplit/MicroSplit-reproducibility/examples/2D/custom_dataset/~/anaconda3/envs/microsplit/lib/python3.11/site-packages/pytorch_lightning/trainer/call.py:171)     output = fn(*args, **kwargs)
    173 # restore current_fx when nested context
    174 pl_module._current_fx_name = prev_fx_name

File ~/anaconda3/envs/microsplit/lib/python3.11/site-packages/pytorch_lightning/core/hooks.py:611, in DataHooks.transfer_batch_to_device(self, batch, device, dataloader_idx)
    564 def transfer_batch_to_device(self, batch: Any, device: torch.device, dataloader_idx: int) -> Any:
    565     """Override this hook if your :class:`~torch.utils.data.DataLoader` returns tensors wrapped in a custom data
    566     structure.
    567 
   (...)    609 
    610     """
--> [611](https://file+.vscode-resource.vscode-cdn.net/Users/sdasgupt/Documents/microsplit/microsplit/MicroSplit-reproducibility/examples/2D/custom_dataset/~/anaconda3/envs/microsplit/lib/python3.11/site-packages/pytorch_lightning/core/hooks.py:611)     return move_data_to_device(batch, device)

File ~/anaconda3/envs/microsplit/lib/python3.11/site-packages/lightning_fabric/utilities/apply_func.py:110, in move_data_to_device(batch, device)
    107     # user wrongly implemented the `_TransferableDataType` and forgot to return `self`.
    108     return data
--> [110](https://file+.vscode-resource.vscode-cdn.net/Users/sdasgupt/Documents/microsplit/microsplit/MicroSplit-reproducibility/examples/2D/custom_dataset/~/anaconda3/envs/microsplit/lib/python3.11/site-packages/lightning_fabric/utilities/apply_func.py:110) return apply_to_collection(batch, dtype=_TransferableDataType, function=batch_to)

File ~/anaconda3/envs/microsplit/lib/python3.11/site-packages/lightning_utilities/core/apply_func.py:68, in apply_to_collection(data, dtype, function, wrong_dtype, include_none, allow_frozen, *args, **kwargs)
     66     return function(data, *args, **kwargs)
     67 if data.__class__ is list and all(isinstance(x, dtype) for x in data):  # 1d homogeneous list
---> [68](https://file+.vscode-resource.vscode-cdn.net/Users/sdasgupt/Documents/microsplit/microsplit/MicroSplit-reproducibility/examples/2D/custom_dataset/~/anaconda3/envs/microsplit/lib/python3.11/site-packages/lightning_utilities/core/apply_func.py:68)     return [function(x, *args, **kwargs) for x in data]
     69 if data.__class__ is tuple and all(isinstance(x, dtype) for x in data):  # 1d homogeneous tuple
     70     return tuple(function(x, *args, **kwargs) for x in data)

File ~/anaconda3/envs/microsplit/lib/python3.11/site-packages/lightning_utilities/core/apply_func.py:68, in <listcomp>(.0)
     66     return function(data, *args, **kwargs)
     67 if data.__class__ is list and all(isinstance(x, dtype) for x in data):  # 1d homogeneous list
---> [68](https://file+.vscode-resource.vscode-cdn.net/Users/sdasgupt/Documents/microsplit/microsplit/MicroSplit-reproducibility/examples/2D/custom_dataset/~/anaconda3/envs/microsplit/lib/python3.11/site-packages/lightning_utilities/core/apply_func.py:68)     return [function(x, *args, **kwargs) for x in data]
     69 if data.__class__ is tuple and all(isinstance(x, dtype) for x in data):  # 1d homogeneous tuple
     70     return tuple(function(x, *args, **kwargs) for x in data)

File ~/anaconda3/envs/microsplit/lib/python3.11/site-packages/lightning_fabric/utilities/apply_func.py:104, in move_data_to_device.<locals>.batch_to(data)
    102 if isinstance(data, Tensor) and isinstance(device, torch.device) and device.type not in _BLOCKING_DEVICE_TYPES:
    103     kwargs["non_blocking"] = True
--> [104](https://file+.vscode-resource.vscode-cdn.net/Users/sdasgupt/Documents/microsplit/microsplit/MicroSplit-reproducibility/examples/2D/custom_dataset/~/anaconda3/envs/microsplit/lib/python3.11/site-packages/lightning_fabric/utilities/apply_func.py:104) data_output = data.to(device, **kwargs)
    105 if data_output is not None:
    106     return data_output

TypeError: Cannot convert a MPS Tensor to float64 dtype as the MPS framework doesn't support float64. Please use float32 instead.

[ShataDg]

@bethac07

Troubleshooting the "MPS framework doesn't support float64 error"

I deleted the noise models and tried running the notebook, but that did not work because the noise models were required by the get_microsplit_parameters function. I copied the noise models from the HT_LIF24 example and now its running with mps accelerator! So the retrained noise models were likely the problem.
I am currently running with 2 epochs, but if this succeeds, I will try with a higher number of epochs. I have in the meantime pushed my environment config as microsplit_environment.yml.

[bethac07]

Awesome, congrats!

[ShataDg]

@bethac07
The following results are from training with 2 epochs - it ran without errors but the results are off. The next step would be to increase the epochs?
Another thing I am trying to figure out - the HT_LIF24 example produces a metrics.csv in the lightning_logs folder. This is not produced by the custom_dataset notebook. The metrics.csv allows plotting of the reconstruction loss, kl loss etc., which is not too relevant when I am training with 2 epochs, but might be relevant for more serious training. Below is the output from the HT_LIF24 example.

Training results from minimal number of epochs

(I used the noise models from he HT_LIF24 example)

Input patches

Predictions

Crops of predictions

[bethac07]

Yes, I would not have assumed good results with only 2 epochs, so this is not surprising! Out of curiosity, how long does 2 epochs take?

[ShataDg]

Around 2-3 hrs

[bethac07]

🤞

[ShataDg]

from pandas import read_csv
from microsplit_reproducibility.notebook_utils.custom_dataset_2D import find_recent_metrics, plot_metrics
df = read_csv(find_recent_metrics())
plot_metrics(df)

produces

---------------------------------------------------------------------------
AssertionError                            Traceback (most recent call last)
Cell In[14], [line 3](vscode-notebook-cell:?execution_count=14&line=3)
      1 from pandas import read_csv
      2 from microsplit_reproducibility.notebook_utils.custom_dataset_2D import find_recent_metrics, plot_metrics
----> [3](vscode-notebook-cell:?execution_count=14&line=3) df = read_csv(find_recent_metrics())
      4 plot_metrics(df)

File ~/Documents/microsplit/microsplit/MicroSplit-reproducibility/src/microsplit_reproducibility/notebook_utils/custom_dataset_2D.py:145, in find_recent_metrics()
    143 def find_recent_metrics():
    144     last = ((Path("lightning_logs") / Path(sorted(os.listdir("lightning_logs")[:10])[-1])).absolute() / "metrics.csv")
--> [145](https://file+.vscode-resource.vscode-cdn.net/Users/sdasgupt/Documents/microsplit/microsplit/MicroSplit-reproducibility/examples/2D/custom_dataset/~/Documents/microsplit/microsplit/MicroSplit-reproducibility/src/microsplit_reproducibility/notebook_utils/custom_dataset_2D.py:145)     assert last.exists(), f"File {last} does not exist. Please remove the lightning_logs folder and re-run the training."
    146     return last

AssertionError: File /Users/sdasgupt/Documents/microsplit/microsplit/MicroSplit-reproducibility/examples/2D/custom_dataset/lightning_logs/version_9/metrics.csv does not exist. Please remove the lightning_logs folder and re-run the training.

[ShataDg]

I moved beyond the LightningModule/ VAEModule related error and was able to get some graphing in Neptune, but I don't think it is graphing the right thing and it is just a flat line.
The experiment params keys are:

dict_keys(['algorithm', 'loss_type', 'img_size', 'encoder_conv_strides', 'decoder_conv_strides', 'target_channels', 'multiscale_count', 'predict_logvar', 'nm_paths', 'kl_type', 'batch_size', 'lr', 'lr_scheduler_patience', 'earlystop_patience', 'num_epochs', 'num_workers', 'mmse_count', 'grid_size', 'data_stats'])

Do you have any suggestions on how to fix this?

[bethac07]

Indeed it is not graphing the right thing, you passed it the lr, not any actual parameter being learned. Looking into some stuff now.

[bethac07]

Here is the diff between before you started adding neptune and after. I tried to look for a commit where you were getting the error you showed, but I didn't see one (though I might have missed it)

IF the issue was neptune, but I am not conviced it is:

The major thing I see here as suspicious is you're just importing neptune:

Whereas the documentation all says to use the pytorch version

That import is there, but commented out, and clearly not every run is committed (and not all of them have to be, but it's nice to commit things you're bringing up in discussions/issues so we can link a particular error to a particular full code state), so I can't really tell if you used it at any point or not.

So it may indeed just be as simple as trading your import statement, but let's plan to go through this together in checkin tomorrow, because...

I'm not sure this has to do with Neptune at all, necessarily

Nothing in the stack trace you're showing actually looks like it has anything to do with Neptune or logging at all - it looks like it's trying and failing to import an existing model. I'm not saying it ISN'T neptune, but I'm not convinced it IS until you can take a successful run, change only the logger, and come up with the same error you were trying to debug before.

---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
Cell In[14], [line 15](vscode-notebook-cell:?execution_count=14&line=15)
      2 trainer = Trainer(
      3     max_epochs=training_config.num_epochs,
      4     # NOTE: if you are on a mac swap the accelerator to "mps"
   (...)     12 
     13 )
     14 # start the training - yay!
---> [15](vscode-notebook-cell:?execution_count=14&line=15) trainer.fit(
     16     model=model,
     17     train_dataloaders=train_dloader,
     18     val_dataloaders=val_dloader,
     19 )

File ~/anaconda3/envs/microsplit/lib/python3.11/site-packages/lightning/pytorch/trainer/trainer.py:577, in Trainer.fit(self, model, train_dataloaders, val_dataloaders, datamodule, ckpt_path, weights_only)
    522 def fit(
    523     self,
    524     model: "pl.LightningModule",
   (...)    529     weights_only: Optional[bool] = None,
    530 ) -> None:
    531     r"""Runs the full optimization routine.
    532 
    533     Args:
   (...)    575 
    576     """
--> [577](https://file+.vscode-resource.vscode-cdn.net/Users/sdasgupt/Documents/microsplit/microsplit/MicroSplit-reproducibility/examples/2D/custom_dataset/~/anaconda3/envs/microsplit/lib/python3.11/site-packages/lightning/pytorch/trainer/trainer.py:577)     model = _maybe_unwrap_optimized(model)
    578     self.strategy._lightning_module = model
    579     _verify_strategy_supports_compile(model, self.strategy)

File ~/anaconda3/envs/microsplit/lib/python3.11/site-packages/lightning/pytorch/utilities/compile.py:111, in _maybe_unwrap_optimized(model)
    109     return model
    110 _check_mixed_imports(model)
--> [111](https://file+.vscode-resource.vscode-cdn.net/Users/sdasgupt/Documents/microsplit/microsplit/MicroSplit-reproducibility/examples/2D/custom_dataset/~/anaconda3/envs/microsplit/lib/python3.11/site-packages/lightning/pytorch/utilities/compile.py:111) raise TypeError(
    112     f"`model` must be a `LightningModule` or `torch._dynamo.OptimizedModule`, got `{type(model).__qualname__}`"
    113 )

TypeError: `model` must be a `LightningModule` or `torch._dynamo.OptimizedModule`, got `VAEModule`

[ShataDg]

Nothing in the stack trace you're showing actually looks like it has anything to do with Neptune or logging at all - it looks like it's trying and failing to import an existing model. I'm not saying it ISN'T neptune, but I'm not convinced it IS until you can take a successful run, change only the logger, and come up with the same error you were trying to debug before.

Indeed..this (per stack overflow) was a Pytorch lightning related issue, for which I needed to change the import statement.

from pytorch_lightning import Trainer

But I was thinking that neptune_logger was not compatible with the VAEModule and was doing the init manually. But I think now (latest version has been committed), it is running and producing graphs (which might still be incorrect) I haven't seen before!

[bethac07]

Sorry, not totally following, that statement was already in the code before

There's now a second one that was commented out, are you saying that what happened was the second import was messing things up?

In any case, congrats on it now working!

[ShataDg]

Connecting jupyter notebook to neptune

Neptune's docs say

and a search for neptune-notebooks does not return anything in VS Code

Should I follow the steps in Neptune docs regarding going back to a different version of neptune-client or let this be?
I currently have neptune-client==1.14.0.

[bethac07]

I think fine to instead use the other method of incoporating code; there was something in the Neptune app i recall, but you also might ping Kamal and/or look at the SIXCE repo to see what he was doing there, because I remember it had a nice thing where it also highlighted differences since last git-committed code when run. Not sure that's super helpful in a notebook context, but might be nice.

[ShataDg]

@bethac07
I am trying to determine if neptune logging is being set up ok. Based on the following, it seems it is logging the diff. Would you mind checking if that is indeed the case?
(I will work on adding images next and then moving on to testing the other non-noise model requiring algorithms)

Current committed version

Version I ran

Diff I can see on Neptune