Use optax's losses and schedules in mnist and imagenet flax examples. by copybara-service[bot] · Pull Request #1286 · google/flax

copybara-service · 2021-04-29T17:22:59Z

Use optax's losses and schedules in mnist and imagenet flax examples.

avital · 2021-04-30T11:58:17Z

@mtthss let's review here. LGTM, I also asked @andsteing to review as he's the person with the most holistic view on our examples (e.g. should we make sure to port all of our example to use Optax losses and schedules?)

avital · 2021-04-30T12:00:19Z

(Oh and thanks!)

examples/imagenet/train.py

examples/mnist/train.py

marcvanzee · 2021-05-03T13:11:55Z

It seems the latest release of Optax doesn't yet support linear_schedule (and perhaps other things you are using here), so could you please release a new Pypi version?

8bitmp3 · 2021-05-03T22:51:57Z

Thanks @avital 👍 (Note to self: update the Flax Linen with MNIST tutorial to use a loss from Optax)

mtthss · 2021-05-04T09:09:05Z

It seems the latest release of Optax doesn't yet support linear_schedule (and perhaps other things you are using here), so could you please release a new Pypi version?

Releasing a new version now

andsteing

Thanks for the change!

(And sorry for my delayed reply...)

andsteing · 2021-05-04T12:18:09Z

examples/imagenet/requirements.txt

@@ -1,4 +1,5 @@
 clu==0.0.1a2
 ml-collections>=0.1.0
+optax


Maybe specify version that includes optax.linear_schedule()?

andsteing · 2021-05-04T12:27:58Z

examples/imagenet/train.py

+  warmup_fn = optax.linear_schedule(
+      init_value=0., end_value=base_learning_rate,
+      transition_steps=config.warmup_epochs * steps_per_epoch)
+  cosine_epochs = max(config.num_epochs - config.warmup_epochs, 1)


Why is this max(..., 1) and not max(..., 0) ?

andsteing · 2021-05-04T12:38:52Z

examples/imagenet/train.py

-  return -jnp.sum(
-      common_utils.onehot(labels, num_classes=1000) * logits) / labels.size
+  xentropy = optax.softmax_cross_entropy(
+      logits, common_utils.onehot(labels, num_classes=NUM_CLASSES))


Could you specify these by keyword?
(It's all too easy inverting them, like thinking of H(P, Q) with P~labels and Q~softmax(logits))

andsteing · 2021-05-04T12:41:05Z

examples/imagenet/train.py

 def cross_entropy_loss(logits, labels):
-  return -jnp.sum(
-      common_utils.onehot(labels, num_classes=1000) * logits) / labels.size
+  xentropy = optax.softmax_cross_entropy(


In imagenet/models.py we return nn.log_softmax(x):

flax/examples/imagenet/models.py

Line 117 in d804b90

x = nn.log_softmax(x)

In terms of numerics, do you know if it makes any difference computing that twice?

Should we remove it from models.py in any case?

Removed nn.log_softmax(x)

andsteing · 2021-05-04T12:43:49Z

examples/mnist/requirements.txt

@@ -1,6 +1,7 @@
 clu
 flax
+optax


should we specify version?

andsteing · 2021-05-04T12:44:28Z

examples/mnist/train.py

-
 def compute_metrics(logits, labels):
-  loss = cross_entropy_loss(logits, labels)
+  loss = jnp.mean(optax.softmax_cross_entropy(logits, onehot(labels)))


(Same comments as above: should we maybe remove nn.log_softmax(x)? And would it make sense to specify arguments by keyword for extra safety?)

Removed nn.log_softmax(x)

mtthss · 2021-05-14T13:33:52Z

Addressed comments.

andsteing

There are still two minor comments open but otherwise LGTM.

andsteing · 2021-05-14T17:34:18Z

examples/mnist/train.py

-
 def compute_metrics(logits, labels):
-  loss = cross_entropy_loss(logits, labels)
+  loss = jnp.mean(optax.softmax_cross_entropy(logits, onehot(labels)))


Could you also use keyword arguments here?

codecov-commenter · 2021-06-10T13:08:42Z

Codecov Report

Merging #1286 (b0a274b) into master (8846461) will not change coverage.
The diff coverage is n/a.

@@           Coverage Diff           @@
##           master    #1286   +/-   ##
=======================================
  Coverage   82.34%   82.34%           
=======================================
  Files          65       65           
  Lines        5318     5318           
=======================================
  Hits         4379     4379           
  Misses        939      939

Continue to review full report at Codecov.

Legend - Click here to learn more
Δ = absolute <relative> (impact), ø = not affected, ? = missing data
Powered by Codecov. Last update 8846461...b0a274b. Read the comment docs.

PiperOrigin-RevId: 378640649

google-cla bot added the cla: yes label Apr 29, 2021

copybara-service bot force-pushed the test_370893380 branch from 917666c to 7e61722 Compare April 29, 2021 17:23

avital requested a review from andsteing April 30, 2021 11:57

avital self-requested a review April 30, 2021 11:58

avital approved these changes Apr 30, 2021

View reviewed changes

copybara-service bot force-pushed the test_370893380 branch from 7e61722 to 4fdbe9b Compare April 30, 2021 13:07

marcvanzee suggested changes May 3, 2021

View reviewed changes

examples/imagenet/train.py Show resolved Hide resolved

examples/mnist/train.py Show resolved Hide resolved

marcvanzee assigned andsteing May 4, 2021

marcvanzee added the Priority: P2 - no schedule Best effort response and resolution. We have no plan to work on this at the moment. label May 4, 2021

andsteing reviewed May 4, 2021

View reviewed changes

copybara-service bot force-pushed the test_370893380 branch from 4fdbe9b to 2e04d64 Compare May 4, 2021 13:18

marcvanzee approved these changes May 4, 2021

View reviewed changes

copybara-service bot force-pushed the test_370893380 branch from 2e04d64 to 162369d Compare May 4, 2021 17:23

copybara-service bot force-pushed the test_370893380 branch from 162369d to 1078af5 Compare May 14, 2021 09:01

andsteing approved these changes May 14, 2021

View reviewed changes

copybara-service bot force-pushed the test_370893380 branch 2 times, most recently from 353fab2 to 51e59fc Compare May 18, 2021 16:31

jheek added the pull ready label May 18, 2021

copybara-service bot force-pushed the test_370893380 branch 2 times, most recently from 041147f to ba04ee7 Compare May 24, 2021 16:30

copybara-service bot force-pushed the test_370893380 branch 3 times, most recently from 5246e4e to b0a274b Compare June 10, 2021 12:52

copybara-service bot force-pushed the test_370893380 branch 2 times, most recently from e4d0282 to 0fc721a Compare June 10, 2021 13:35

Use optax's losses and schedules in mnist and imagenet flax examples.

44ee6f2

PiperOrigin-RevId: 378640649

copybara-service bot force-pushed the test_370893380 branch from 0fc721a to 44ee6f2 Compare June 10, 2021 13:54

copybara-service bot merged commit 44ee6f2 into master Jun 10, 2021

copybara-service bot deleted the test_370893380 branch June 10, 2021 13:54

@@ @@ -1,6 +1,7 @@ @@
               clu
               flax
+              optax

Conversation

copybara-service bot commented Apr 29, 2021

Uh oh!

avital commented Apr 30, 2021 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

avital commented Apr 30, 2021

Uh oh!

Uh oh!

Uh oh!

marcvanzee commented May 3, 2021

Uh oh!

8bitmp3 commented May 3, 2021

Uh oh!

mtthss commented May 4, 2021

Uh oh!

andsteing left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

mtthss commented May 14, 2021

Uh oh!

andsteing left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

codecov-commenter commented Jun 10, 2021

Codecov Report

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

7 participants

avital commented Apr 30, 2021 •

edited

Loading