LVAE losses tests #180

CatEek · 2024-07-11T14:32:21Z

Description

What: Moved losses/NMs/LL/etc out of lightning module
Why: To be able to train LVAE based algorithms in CAREamics
How: Coded with my elbows, unfortunately

Changes Made

Added:

Some amount of very questionable quality code which needs to be refactored right away
Very basic tests which just ensure loss is not nan

Known problems

Noise models doesn't work because even the logic is actually not there yet
Training logic for NMs isn't clear
Most of the parameters should be documented, some should be refactored or removed
The whole pipeline isn't complete

Please ensure your PR meets the following requirements:

Code builds and passes tests locally, including doctests
New tests have been added (for bug fixes/features)
Pre-commit passes
PR to the documentation exists (for bug fixes / features)

jdeschamps

I think we should stick with the original plan of only using the new algorithms (HDN, denoiSplit) with the LightningAPI and not make them available through the CAREamist AP.

We are very close to a 0.1.0 release, which means that we might get an influx of users and incorporating code that does not work or might break can create a lot of issues.

If the modules are in the code base, they can be instantiated independently with the Ligthning API, that will be more than enough for Ashesh's paper. Once we are confident the quality of the code (integration, testing, documentation) is high enough, then we can make it available through the CAREamist.

So in the meantime, I would suggest keeping the configuration clean (removing LVAE, NoiseModels, Likelihood, and the new entries in the supported enums) and adding placeholder NotImplementedErrors in CAREamist.

For now, NoiseModel and Likelihood could be passed directly as parameters along the algorithm configuration. We can also discuss alternatives (e.g. allowing additional parameters in the Pydantic model via ConfigDict, but that would require a lot of asserts in the LVAE related code).

jdeschamps · 2024-07-12T07:17:38Z

src/careamics/careamist.py

+                self.cfg.algorithm_config.model.architecture
+                == SupportedArchitecture.LVAE
+            ):  # TODO fuckin mypy thinks model "already defined on line 158"
+                self.model: VAEModule = VAEModule(


I believe what you can do about this is to define a type at the top of the module:

Module = Union[FCNModule, VAEModule]

And use that type on line 158.

But mypy isn't complaining about type ?

One needs to try, I haven't. 😄

src/careamics/careamist.py

jdeschamps · 2024-07-12T07:20:42Z

src/careamics/careamist.py

+                self.model: VAEModule = VAEModule(
+                    algorithm_config=self.cfg.algorithm_config,
+                )
+            else:


The configuration allows custom models that are neither Unet nor LVAE (Custom I believe). This breaks that possibility.

We could potentially consider that they should be FCNModule by default. Not sure what it entails.

Let's discuss

Now that the configuration has only custom and UNet models, these statements should go.

Once the LVAE is in the configuration, we will need to reassess how to deal with custom models.

src/careamics/config/algorithm_model.py

src/careamics/losses/lvae/losses.py

src/careamics/models/lvae/likelihoods.py

jdeschamps · 2024-07-12T08:07:03Z

src/careamics/models/lvae/lvae.py


-import ml_collections
+import ml_collections  # TODO: refactor this out


I don't see it used except in the declaration of config: ml_collections and then this config is not used (as far as I can tell). So kick it out as we discussed to allow the tests from passing (hopefully)!

src/careamics/models/lvae/noise_models.py

…t/lvae_losses

melisande-c · 2024-07-15T12:27:59Z

So in the meantime, I would suggest keeping the configuration clean (removing LVAE, NoiseModels, Likelihood, and the new entries in the supported enums) and adding placeholder NotImplementedErrors in CAREamist.

I will agree and say let's make the minimum number of modifications possible to the rest of the code until the LVAE pipeline is more complete, maybe not even adding not-implemented-errors to CAREamist. It might become clearer how we add the new components into the rest of CAREamist and the configuration once we understand how they work. I think it will also make the changes to the code easier to follow.

jdeschamps · 2024-07-15T21:54:29Z

It seems there is still a hdn hanging around in one of the tests!

…t/lvae_losses

melisande-c

So I still think there are some changes here that are out of scope for the PR (also the current PR title is not descriptive of the changes). I see though that it's pretty hard to make the pydantic model when the LVAE class does need some refactoring to extract components that don't necessarily need to be contained within the class. Maybe the splitting of the configuration you mentioned will help.

melisande-c · 2024-07-18T15:24:54Z

src/careamics/careamist.py

-                self.model = CAREamicsModule(
-                    algorithm_config=self.cfg.algorithm_config,
-                )
+                if (
+                    self.cfg.algorithm_config.model.architecture
+                    == SupportedArchitecture.UNET
+                ):
+                    self.model = FCNModule(
+                        algorithm_config=self.cfg.algorithm_config,
+                    )
+                else:
+                    raise NotImplementedError("Architecture not supported.")


Repetition of inner if-else statement can be removed through separation of config initialisation and model initialisation. i.e.

if isinstance(source, Path) and source.is_file(): if source.suffix == ".yaml" or source.suffix == ".yml": self.cfg = load_configuration(source) self.model = None else: self.model, self.cfg = load_pretrained(source) elif isinstance(source, Configuration): self.cfg = source self.model = None else: raise ValueError("appropriate message") # instantiate model if self.model is None: if ( self.cfg.algorithm_config.model.architecture == SupportedArchitecture.UNET ): self.model = FCNModule( algorithm_config=self.cfg.algorithm_config, ) elif ( self.cfg.algorithm_config.model.architecture == SupportedArchitecture.LVAE ): raise NotImplementedError( "LVAE is currently only supported through lightning api" ) else: raise NotImplementedError("Architecture not supported.")

^ just an example can probably be improved further. Extracting to smaller testable sub functions might be good too.

Those things are gone in the config refac pr

melisande-c · 2024-07-18T15:31:19Z

src/careamics/config/algorithm_model.py

    # Mandatory fields
-    algorithm: Literal["n2v", "care", "n2n", "custom"]  # defined in SupportedAlgorithm
-    """Name of the algorithm, as defined in SupportedAlgorithm."""
-
+    # defined in SupportedAlgorithm
+    algorithm: Literal["n2v", "care", "n2n"]
    loss: Literal["n2v", "mae", "mse"]
-    """Loss function to use, as defined in SupportedLoss."""
-
-    model: Union[UNetModel, VAEModel, CustomModel] = Field(discriminator="architecture")
-    """Model architecture to use, defined in SupportedArchitecture."""
+    model: Union[UNetModel, CustomModel] = Field(discriminator="architecture")

    # Optional fields
    optimizer: OptimizerModel = OptimizerModel()
    """Optimizer to use, defined in SupportedOptimizer."""

    lr_scheduler: LrSchedulerModel = LrSchedulerModel()
-    """Learning rate scheduler to use, defined in SupportedScheduler."""


Some docstrings have been deleted, probably by mistake.

src/careamics/lightning/lightning_module.py

src/careamics/config/likelihood_model.py

melisande-c · 2024-07-18T15:57:11Z

src/careamics/config/support/supported_architectures.py

    CUSTOM = (
-        "Custom"  # TODO all the others tags are small letters, except the architect
+        "custom"  # TODO all the others tags are small letters, except the architect


can remove TODO

melisande-c · 2024-07-18T16:26:42Z

src/careamics/models/lvae/lvae.py

+            self.likelihood_gm = GaussianLikelihood(
+                self.decoder_n_filters,
+                self.target_ch,
+                predict_logvar=self.predict_logvar,
+                logvar_lowerbound=self.logvar_lowerbound,
+                conv2d_bias=self.topdown_conv2d_bias,
+            )


When is self.likelihood_gm used? as far as I can see it doesn't seem to be used by any of theLVAE class' methods. If it is only used by external functions maybe it can be removed from the LVAE class and passed to those functions explicitly.

tests/config/architectures/test_lvae_model.py

melisande-c · 2024-07-18T16:29:59Z

tests/config/architectures/test_lvae_model.py

+        assert act in activations
+
+
+def test_activation_wrong_values():


maybe rename test_activation_unknown_value

I don't know what that function even doing

Just testing an error is raised if "nonlinearity" is not a supported activation, no?

tests/config/test_algorithm_model.py

tests/conftest.py

jdeschamps · 2024-07-18T17:24:33Z

I pushed a fix for the doctests as promised.

jdeschamps · 2024-07-19T12:13:01Z

Comment https://github.com/CAREamics/careamics/pull/180/files#r1675427899 still needs to be fixed regarding the custom algorithms and models.

jdeschamps · 2024-07-19T12:13:47Z

src/careamics/careamist.py

+                self.model: VAEModule = VAEModule(
+                    algorithm_config=self.cfg.algorithm_config,
+                )
+            else:


Now that the configuration has only custom and UNet models, these statements should go.

Once the LVAE is in the configuration, we will need to reassess how to deal with custom models.

src/careamics/lightning/lightning_module.py

jdeschamps · 2024-07-19T12:19:13Z

src/careamics/lightning/lightning_module.py

+        target = aux[0]
+        self.loss_parameters.current_epoch = self.current_epoch
+        self.loss_parameters.inputs = x
+        self.loss_parameters.mask = ~((target == 0).reshape(len(target), -1).all(dim=1))


Please add comments about what is happening in these steps.

melisande-c · 2024-07-19T12:26:08Z

src/careamics/config/algorithm_model.py


    loss: Literal["n2v", "mae", "mse"]
-    """Loss function to use, as defined in SupportedLoss."""
+    """Name of the algorithm, as defined in SupportedAlgorithm."""


Algorithm docs instead of loss docs.

melisande-c · 2024-07-19T12:27:35Z

src/careamics/config/algorithm_model.py

-            raise ValueError("VAE are currently not implemented.")
+        if (self.algorithm == "custom") != (self.model.architecture == "custom"):
+            raise ValueError(
+                "Algorithm and model architecture must be both `custom` or not."


grammar could be clearer. not what?

I guess it was meant to be "None"!

CatEek and others added 25 commits June 18, 2024 17:18

draft + test

ab97671

types, docstrings

c9639c8

target stats, test

8bd270e

mypy fix

a127edf

stats calculation as a class refac

4ccdc25

Merge branch 'main' into iz/feat/welford_norm

e0fbf5f

Merge branch 'main' into iz/feat/welford_norm

f239e40

fix: fix calls to parameters in tests

7edcba2

rename test, minor refac

d812971

pydantic model draft + tests

383e4c3

Merge branch 'main' into iz/feat/lvae_pydantic_model

94cd81a

loss refac init

195f1d2

tests drafts

f244b91

loss params dataclass

3a29fe0

tests draft

76adbee

model refac to make init work

9711c28

necessary model refac

7fab256

loss refac up to ll module

eab4355

added ll config + model

dcf5d4c

musplit test works

f827e77

Merge remote-tracking branch 'origin/main' into iz/feat/lvae_losses

14f77b8

vae loss basic tests pass

f5b3211

tests pass

6415121

format + remove extra

313fb95

fixed some mypy

446d175

CatEek requested review from ashesh-0, jdeschamps, veegalinova, melisande-c and federico-carrara July 11, 2024 14:32

style(pre-commit.ci): auto fixes [...]

fb86481

jdeschamps requested changes Jul 12, 2024

View reviewed changes

CatEek and others added 4 commits July 12, 2024 12:34

rm ml_collections, minor fixes

bdae3d2

Merge remote-tracking branch 'origin/main' into iz/feat/lvae_losses

99213a2

Merge remote-tracking branch 'origin/iz/feat/lvae_losses' into iz/fea…

3d20c5f

…t/lvae_losses

style(pre-commit.ci): auto fixes [...]

fe2d0e2

jdeschamps self-requested a review July 15, 2024 21:54

CatEek and others added 3 commits July 17, 2024 16:00

removed registration, vae refs in config, only loss tests

0e2e9a8

Merge remote-tracking branch 'origin/iz/feat/lvae_losses' into iz/fea…

dd9d580

…t/lvae_losses

style(pre-commit.ci): auto fixes [...]

544b194

CatEek changed the title ~~Towards LVAE training~~ LVAE losses tests Jul 17, 2024

melisande-c reviewed Jul 18, 2024

View reviewed changes

fix: fix doctests and some inconsistencies

64f7381

jdeschamps added 6 commits July 18, 2024 19:33

chore: invert change to pre-commit, to do another PR

e7c759c

fix: add back custom algo and add tests

9aa978b

fix: solve pydantic warning

32b100e

fix: solve pydantic warning

56621d5

fix: skip test to allow others to pass

3827777

refac: refactor custom model

b33799c

jdeschamps requested changes Jul 19, 2024

View reviewed changes

melisande-c reviewed Jul 19, 2024

View reviewed changes

CatEek closed this Aug 22, 2024

jdeschamps deleted the iz/feat/lvae_losses branch December 6, 2024 14:33

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

LVAE losses tests #180

LVAE losses tests #180

CatEek commented Jul 11, 2024

jdeschamps left a comment

jdeschamps Jul 12, 2024

CatEek Jul 12, 2024

jdeschamps Jul 15, 2024

jdeschamps Jul 12, 2024

CatEek Jul 12, 2024

jdeschamps Jul 19, 2024

jdeschamps Jul 12, 2024

melisande-c commented Jul 15, 2024

jdeschamps commented Jul 15, 2024

melisande-c left a comment

melisande-c Jul 18, 2024

CatEek Jul 19, 2024

melisande-c Jul 18, 2024

melisande-c Jul 18, 2024

melisande-c Jul 18, 2024

melisande-c Jul 18, 2024

CatEek Jul 19, 2024

melisande-c Jul 19, 2024

jdeschamps commented Jul 18, 2024 •

edited

Loading

jdeschamps commented Jul 19, 2024 •

edited

Loading

jdeschamps Jul 19, 2024

jdeschamps Jul 19, 2024

melisande-c Jul 19, 2024

melisande-c Jul 19, 2024

jdeschamps Jul 22, 2024


		import ml_collections
		import ml_collections # TODO: refactor this out

		assert act in activations


		def test_activation_wrong_values():

LVAE losses tests #180

LVAE losses tests #180

Conversation

CatEek commented Jul 11, 2024

Description

Changes Made

Known problems

jdeschamps left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

melisande-c commented Jul 15, 2024

jdeschamps commented Jul 15, 2024

melisande-c left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

jdeschamps commented Jul 18, 2024 • edited Loading

jdeschamps commented Jul 19, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

jdeschamps commented Jul 18, 2024 •

edited

Loading

jdeschamps commented Jul 19, 2024 •

edited

Loading