Add monge gap #361

theouscidda6 · 2023-05-04T16:25:21Z

This pull request proposes to add the MongeGap class to ott-jax, by creating a new losses.py file in the ott/src/ott/solvers/nn folder. It includes the addition of a tutorial showing how to recode, from scratch, a simple solver to fit a neural OT map using the Monge gap.

review-notebook-app · 2023-05-04T16:25:27Z

Check out this pull request on

See visual diffs & provide feedback on Jupyter Notebooks.

Powered by ReviewNB

codecov-commenter · 2023-05-04T19:25:29Z

Codecov Report

Merging #361 (ff9832a) into main (b2b7ebb) will increase coverage by 0.04%.
The diff coverage is 88.88%.

❗ Your organization is not using the GitHub App Integration. As a result you may experience degraded service beginning May 15th. Please install the Github App Integration for your organization. Read more.

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #361      +/-   ##
==========================================
+ Coverage   87.95%   88.00%   +0.04%     
==========================================
  Files          52       54       +2     
  Lines        5679     5769      +90     
  Branches      841      857      +16     
==========================================
+ Hits         4995     5077      +82     
- Misses        561      566       +5     
- Partials      123      126       +3

Impacted Files	Coverage Δ
src/ott/geometry/geometry.py	`92.69% <ø> (ø)`
src/ott/tools/map_estimator.py	`87.01% <87.01%> (ø)`
src/ott/solvers/nn/losses.py	`100.00% <100.00%> (ø)`

... and 1 file with indirect coverage changes

marcocuturi · 2023-05-14T18:01:13Z

hi Théo, can you fix the linter issues?

you can check first the CONTRIBUTING.md file and install pre-commit

Thanks!

marcocuturi

This is fantastic Théo! thanks so much, I am sure this will be very useful.

I would suggest moving the estimation pipeline in nn/solvers, I am sure other people will benefit from it.

src/ott/solvers/nn/losses.py

marcocuturi · 2023-05-18T15:41:44Z

src/ott/solvers/nn/losses.py

+from typing import Any, Mapping, Callable
+from types import MappingProxyType
+
+class MongeGap:


Sounds like this is missing a simple test.

This could be running the Monge gap on the gradient of a convex function on a small batch of 10 points.

The test should be put in the same location as losses.py but in the test folder, and be losses_test.py.

Agree with Marco, this function needs some tests.

marcocuturi · 2023-05-18T15:49:26Z

docs/tutorials/notebooks/Monge_Gap.ipynb

@@ -0,0 +1,734 @@
+{


Feels like the first "big" equation (in the begin{equation} env) is not rendering in this NB.

squared-Euclidean (capital letter)

Explain in what sense this provides an alternative to the solver = Sinkhorn(norm_error=1) solver, and, ideally, we might want to run it in this solver too?

When defining the monge gap, mention to cost_fn is not provided.

Reply via ReviewNB

docs/tutorials/notebooks/Monge_Gap.ipynb

src/ott/solvers/nn/losses.py

michalk8 · 2023-05-25T14:01:47Z

src/ott/solvers/nn/losses.py

+        self.sinkhorn_kwargs = sinkhorn_kwargs
+
+    def __call__(
+        self, samples: jnp.ndarray, T: Callable[[jnp.ndarray], jnp.ndarray]


I don't think this function should receive the function computing the transport, as it limits what can be passed (e.g., classes implementing __call__). I'd rather you pass the transported samples.

I think it might be worth revisiting this choice.

In fact, I think the notion of Monge gap only has meaning when it is passed with callable (the Monge gap is defined as a function defined on maps in the paper), so I think the initial implementation was better conceptually (i.e. I would pass callable T and source points as arguments, just to instantiate target as being T run on samples.

We can have a function with source and target points passed directly, but this should be called something different (e.g. id_coupling_gap).

michalk8 · 2023-05-25T14:12:29Z

src/ott/solvers/nn/losses.py

+from typing import Any, Mapping, Callable
+from types import MappingProxyType
+
+class MongeGap:


Agree with Marco, this function needs some tests.

src/ott/solvers/nn/losses.py

docs/tutorials/notebooks/Monge_Gap.ipynb

michalk8 · 2023-05-25T15:09:28Z

docs/tutorials/notebooks/Monge_Gap.ipynb

@@ -0,0 +1,734 @@
+{


Line #163. if self.plot_fitted_map:
I feel like plotting shouldn't be part of the training loop - it creates a lot of arguments and pollutes the class (+ it's simple enough to be added after the map has been fitted).

Reply via ReviewNB

michalk8 · 2023-05-25T15:09:28Z

docs/tutorials/notebooks/Monge_Gap.ipynb

@@ -0,0 +1,734 @@
+{


... instantiated for the squared Euclidean cost.
... that will be bot -> ... that will be both

Please use numbered lists instead of unnumbered lists + (i)/(ii).
The rendering of "... class. For its instantiation, we need ..." seems to have strike-through enabled.
Please add link to the MLP class as {class}....

TODO(michalk8): add a section header for MapEstimator and use it when referring to it.

Reply via ReviewNB

docs/tutorials/notebooks/Monge_Gap.ipynb

michalk8 · 2023-05-25T15:09:29Z

docs/tutorials/notebooks/Monge_Gap.ipynb

@@ -0,0 +1,734 @@
+{


Line #35. dict_monge_gaps = {
I'd slightly prefer to defined 2 variables for the 2 different Monge gaps and later use:
for key, monge_gap in zip(["no_monge_gap", ...], [None, ]):
...
But if you prefer this, let's keep it!

Reply via ReviewNB

michalk8 · 2023-05-25T15:09:29Z

docs/tutorials/notebooks/Monge_Gap.ipynb

@@ -0,0 +1,734 @@
+{


Line #1. # neural network and optimizer
Please add above a cell that briefly explains what you're doing in this cell.

Reply via ReviewNB

docs/tutorials/notebooks/Monge_Gap.ipynb

michalk8 · 2023-05-25T15:09:39Z

docs/tutorials/notebooks/Monge_Gap.ipynb

@@ -0,0 +1,734 @@
+{


... instantiated for the squared Euclidean cost.
... that will be bot -> ... that will be both

Please use numbered lists instead of unnumbered lists + (i)/(ii).
The rendering of "... class. For its instantiation, we need ..." seems to have strike-through enabled.
Please add link to the MLP class as {class}....

TODO(michalk8): add a section header for MapEstimator and use it when referring to it.

Reply via ReviewNB

docs/tutorials/notebooks/Monge_Gap.ipynb

michalk8 · 2023-05-25T15:09:39Z

docs/tutorials/notebooks/Monge_Gap.ipynb

@@ -0,0 +1,734 @@
+{


Line #35. dict_monge_gaps = {
I'd slightly prefer to defined 2 variables for the 2 different Monge gaps and later use:
for key, monge_gap in zip(["no_monge_gap", ...], [None, ]):
...
But if you prefer this, let's keep it!

Reply via ReviewNB

docs/tutorials/notebooks/Monge_Gap.ipynb

michalk8 · 2023-06-15T08:11:33Z

Hi @theouscidda6 , any progress on this?

theouscidda6 · 2023-06-15T15:23:18Z

Thank you very much for your help and comments! Here is a revised pull request.

docs/references.bib

src/ott/solvers/nn/losses.py

src/ott/tools/map_estimator.py

make branch up to date

michalk8

Thanks @theouscidda6 , I think the tests are shaping up nicely, just some minor comments! Further, I think we should test for the following:

case when the source/target is the same
differentiating through the MongeGap
this comment by @marcocuturi : Add monge gap #361 (comment)

tests/solvers/nn/icnn_test.py

tests/solvers/nn/losses_test.py

marcocuturi · 2023-07-04T14:24:01Z

This LGTM Théo, thanks a lot for all this work!!

* add monge gap and corresponding tutorial * add monge gap and corresponding tutorial * update bib * revising the pull-request * revising the pull request * revising the pull request * revising the pull request * fix linter issues * fix linter issues * fix linter issues * fix linter issues * fix linter issues * fix linter issues * fix linter issues * fix linter issues * fix linter issues * fix linter issues * fix linter issues * fix linter issues * fix lint cod issue * update losses * update output * update doc * update doc * update doc * Update docs v3 * Fix docs linter * Refer to shapes * Add a skeleton test * update rng type for icnn test for consistentency * add monge gap test * add map_estimator test * add_new_tests * add_new_tests * Clean tests * Reduce number of tests * Further reduce number of tests * Pin numpy * CHange k-means threshold * Fix missing fixture --------- Co-authored-by: Michal Klein <[email protected]> Co-authored-by: Marco Cuturi <[email protected]>

theouscidda6 added 2 commits May 4, 2023 18:11

add monge gap and corresponding tutorial

b196ce1

add monge gap and corresponding tutorial

d4693d1

update bib

405fa4f

michalk8 assigned theouscidda6 May 5, 2023

michalk8 added the enhancement New feature or request label May 5, 2023

michalk8 self-requested a review May 11, 2023 13:32

marcocuturi reviewed May 18, 2023

View reviewed changes

michalk8 requested changes May 25, 2023

View reviewed changes

michalk8 reviewed May 25, 2023

View reviewed changes

theouscidda6 added 4 commits June 15, 2023 10:59

revising the pull-request

b039367

revised pull request

d09e9a2

revising the pull request

830bdde

revising the pull request

aacfa31

theouscidda6 added 8 commits June 15, 2023 17:32

revising the pull request

d1d56dd

fix linter issues

4dde18a

fix linter issues

2c8278b

fix linter issues

0efe106

fix linter issues

3d8ca61

fix linter issues

ea0b39c

fix linter issues

8e1b0d1

fix linter issues

de9e036

michalk8 requested changes Jun 26, 2023

View reviewed changes

theouscidda6 and others added 12 commits June 27, 2023 17:33

Merge pull request #1 from ott-jax/main

6fda8b2

make branch up to date

update losses

734065b

update output

e36d96a

update doc

3789d8e

update doc

0ced88d

update doc

b46dfc7

Update docs v3

ef80498

Fix docs linter

4c0364d

Refer to shapes

47915fa

Add a skeleton test

968cc6c

update rng type for icnn test for consistentency

eca8ad0

add monge gap test

7774484

michalk8 requested changes Jun 28, 2023

View reviewed changes

theouscidda6 and others added 10 commits June 28, 2023 15:45

add map_estimator test

ccc928d

add_new_tests

92984b0

add_new_tests

92eb8bb

Clean tests

2adfaee

Reduce number of tests

98ec92e

Further reduce number of tests

9a2a27b

Pin numpy

b5acfcb

CHange k-means threshold

0dd2132

Fix missing fixture

e84bb25

Merge remote-tracking branch 'upstream/main' into add_monge_gap

ff9832a

Merge branch 'main' into add_monge_gap

a1b8eb2

marcocuturi merged commit b47aaaf into ott-jax:main Jul 4, 2023

marcocuturi mentioned this pull request Jul 5, 2023

changes in definition of Monge gap #389

Merged

michalk8 mentioned this pull request Jul 21, 2023

Add Monge Gap #268

Closed

3 tasks

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add monge gap #361

Add monge gap #361

theouscidda6 commented May 4, 2023

review-notebook-app bot commented May 4, 2023

codecov-commenter commented May 4, 2023 •

edited

Loading

marcocuturi commented May 14, 2023

marcocuturi left a comment

marcocuturi May 18, 2023

michalk8 May 25, 2023

marcocuturi May 18, 2023 •

edited

Loading

michalk8 May 25, 2023

marcocuturi Jul 4, 2023

michalk8 May 25, 2023

michalk8 May 25, 2023

michalk8 May 25, 2023

michalk8 May 25, 2023

michalk8 May 25, 2023

michalk8 May 25, 2023

michalk8 May 25, 2023

michalk8 commented Jun 15, 2023

theouscidda6 commented Jun 15, 2023

michalk8 left a comment

marcocuturi commented Jul 4, 2023

Add monge gap #361

Add monge gap #361

Conversation

theouscidda6 commented May 4, 2023

review-notebook-app bot commented May 4, 2023

codecov-commenter commented May 4, 2023 • edited Loading

Codecov Report

marcocuturi commented May 14, 2023

marcocuturi left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

marcocuturi May 18, 2023 • 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

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

michalk8 commented Jun 15, 2023

theouscidda6 commented Jun 15, 2023

michalk8 left a comment

Choose a reason for hiding this comment

marcocuturi commented Jul 4, 2023

codecov-commenter commented May 4, 2023 •

edited

Loading

marcocuturi May 18, 2023 •

edited

Loading