Update requirements, restructure files and fix formatting for VAE example #3046

canyon289 · 2023-04-21T15:27:17Z

What does this PR do?

Addressing a portion of the issues in #573
Putting up for review early to ensure I'm trending in the right direction and what else would be needed for a merge

Fixes

File structure - copied from WMT Example
Google style fixes - 99% done, some small errors may remain
Updates Requirements

I can fix up colab, readme, clu, model configs, and adding tests in separate PRs if you don't mind

canyon289 · 2023-04-21T18:40:56Z

I see the linting issues. Will get them fixed

8bitmp3 · 2023-04-21T22:55:42Z

Thanks @canyon289

@andsteing @cgarciae

andsteing

This looks great, thanks!

Some minor comments inline.

examples/vae/input_pipeline.py

andsteing · 2023-04-24T13:16:19Z

examples/vae/main.py

+  # Make sure tf does not allocate gpu memory.
+  tf.config.experimental.set_visible_devices([], 'GPU')
+  train.train_and_evaluate(FLAGS)
+


nit: we usually have two empty lines before this block.

What I meant: we usually have single empty lines within functions (lines 55-56, 59-60 should both be a single line), but we have double empty lines between "top-level codeblocks" (line 65 should be two empty lines)

examples/vae/train.py

examples/vae/input_pipeline.py

examples/vae/train.py

examples/vae/input_pipeline.py

examples/vae/models.py

examples/vae/train.py

andsteing · 2023-04-24T13:29:35Z

It's fine to clean-up README in a second step, but could you already update it to run python main.py instead of python train.py ?

andsteing · 2023-04-24T13:30:03Z

Tested that everything installs und runs (with expected loss) in a fresh Colab CPU runtime:
https://colab.research.google.com/drive/1MDrchufz1eUoB03znEnRV-Z8yQtg0Qau#revisionId=0B6Hoz8j9CmpCNlBmRENvTW5lSnpjZ1RHV1JxZGY2MVQyS1dRPQ

canyon289 · 2023-04-24T15:08:33Z

Thanks @andsteing for all the comments. I'll address them all

codecov-commenter · 2023-04-25T02:14:20Z

Codecov Report

Merging #3046 (38c972b) into main (d8708ed) will not change coverage.
The diff coverage is n/a.

@@           Coverage Diff           @@
##             main    #3046   +/-   ##
=======================================
  Coverage   81.97%   81.97%           
=======================================
  Files          55       55           
  Lines        6031     6031           
=======================================
  Hits         4944     4944           
  Misses       1087     1087

Help us with your feedback. Take ten seconds to tell us how you rate us. Have a feature suggestion? Share it here.

canyon289 · 2023-04-25T03:42:27Z

examples/vae/train.py

-  return VAE(latents=FLAGS.latents)
-
-
-@jax.jit


A question for understanding, should the train step be jitted? I looked at other examples and that didn't seem to be the case.

https://github.com/google/flax/blob/main/examples/wmt/train.py#L166

it's jitted here:

flax/examples/wmt/train.py

Line 528 in df66c81

p_train_step = jax.pmap(

(the pmap() transform also compiles the code like jit(), but at the same time parallelizes it onto multiple devices)

should the train step be jitted?

to answer you question: yes, you should always compile the largest possble code block; usually that's the train_step().

Ah I missed it that, was looking for the decorator. Thank you @andsteing

canyon289 · 2023-05-04T02:34:09Z

Do let me know if there's anything I missed. Happy to keep working on this

8bitmp3 · 2023-05-22T20:26:09Z

cc @andsteing @chiamp

canyon289 · 2023-05-27T01:39:21Z

Im still committed to finishing this!

andsteing · 2023-05-30T07:25:49Z

@canyon289 if you would like reviewers to have another look at the PR, you can click on the Re-quest review button in the Github UI

(otherwise busy reviewers might not read all the individual updates and miss the PR until that button is pressed)

andsteing · 2023-05-30T08:11:34Z

Tested that everything installs und runs (with expected loss) in a fresh Colab CPU runtime:
https://colab.research.google.com/drive/1MDrchufz1eUoB03znEnRV-Z8yQtg0Qau#revisionId=0B6Hoz8j9CmpCMlZIdlN0N1MwTitLS21FL05UZUtCNXFIU3dNPQ&scrollTo=1nCF9d0DX0Kk
(runs 30 epochs in 31 minutes, final loss 100.73, 103.4 after 10 eps)

andsteing

Looks good from my side.
(one open comment about empty lines, otherwise good to submit)

andsteing · 2023-05-30T07:27:12Z

examples/vae/main.py

+  # Make sure tf does not allocate gpu memory.
+  tf.config.experimental.set_visible_devices([], 'GPU')
+  train.train_and_evaluate(FLAGS)
+


What I meant: we usually have single empty lines within functions (lines 55-56, 59-60 should both be a single line), but we have double empty lines between "top-level codeblocks" (line 65 should be two empty lines)

andsteing · 2023-05-30T08:14:26Z

@cgarciae @8bitmp3 should we wait for your review?
(since you're both listed as reviewers)

canyon289 · 2023-05-30T15:25:18Z

What I meant: we usually have single empty lines within functions (lines 55-56, 59-60 should both be a single line), but we have double empty lines between "top-level codeblocks" (line 65 should be two empty lines)

You're right, updated! Thank you

canyon289 · 2023-05-30T15:30:16Z

@canyon289 if you would like reviewers to have another look at the PR, you can click on the Re-quest review button in the Github UI

(otherwise busy reviewers might not read all the individual updates and miss the PR until that button is pressed)

I dont have the rights in this repo to rerequest review from pending reviewers for some reason, I wouldnt mind another review if someone wants to, but dont want to obligate anybody. Happy to have this merged and keep moving on from there

canyon289 · 2023-05-30T15:40:38Z

Squashed to single commit and force pushed

8bitmp3 · 2023-06-26T17:22:57Z

@marcvanzee and @levskaya PTAL when you have time please

andsteing · 2023-06-27T05:45:05Z

@canyon289 if you would like reviewers to have another look at the PR, you can click on the Re-quest review button in the Github UI
(otherwise busy reviewers might not read all the individual updates and miss the PR until that button is pressed)

I dont have the rights in this repo to rerequest review from pending reviewers for some reason, I wouldnt mind another review if someone wants to, but dont want to obligate anybody. Happy to have this merged and keep moving on from there

oh, that's good to know! I must have missed your message previously, since nobody else commented on the PR, let's move forward.

canyon289 changed the title ~~Update requirements, restructure files and fix formatting~~ Update requirements, restructure files and fix formatting for VAE example Apr 21, 2023

8bitmp3 requested review from cgarciae, andsteing and 8bitmp3 April 21, 2023 22:55

andsteing requested changes Apr 24, 2023

View reviewed changes

canyon289 force-pushed the update_main branch from 93ae40f to fc7b968 Compare April 25, 2023 01:32

canyon289 commented Apr 25, 2023

View reviewed changes

andsteing approved these changes May 30, 2023

View reviewed changes

Update VAE example

db97574

canyon289 force-pushed the update_main branch from b06961f to db97574 Compare May 30, 2023 15:40

andsteing removed request for cgarciae and 8bitmp3 June 27, 2023 05:43

andsteing added the pull ready label Jun 27, 2023

copybara-service bot merged commit fe93a7e into google:main Jun 27, 2023
19 checks passed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Update requirements, restructure files and fix formatting for VAE example #3046

Update requirements, restructure files and fix formatting for VAE example #3046

canyon289 commented Apr 21, 2023 •

edited

canyon289 commented Apr 21, 2023

8bitmp3 commented Apr 21, 2023

andsteing left a comment

andsteing Apr 24, 2023

canyon289 Apr 25, 2023

andsteing May 30, 2023

andsteing commented Apr 24, 2023

andsteing commented Apr 24, 2023 •

edited

canyon289 commented Apr 24, 2023

codecov-commenter commented Apr 25, 2023

canyon289 Apr 25, 2023

andsteing May 30, 2023

andsteing May 30, 2023

canyon289 May 30, 2023

canyon289 commented May 4, 2023

8bitmp3 commented May 22, 2023

canyon289 commented May 27, 2023

andsteing commented May 30, 2023

andsteing commented May 30, 2023

andsteing left a comment

andsteing May 30, 2023

andsteing commented May 30, 2023

canyon289 commented May 30, 2023

canyon289 commented May 30, 2023

canyon289 commented May 30, 2023

8bitmp3 commented Jun 26, 2023

andsteing commented Jun 27, 2023 •

edited by 8bitmp3

Update requirements, restructure files and fix formatting for VAE example #3046

Update requirements, restructure files and fix formatting for VAE example #3046

Conversation

canyon289 commented Apr 21, 2023 • edited

What does this PR do?

canyon289 commented Apr 21, 2023

8bitmp3 commented Apr 21, 2023

andsteing 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

andsteing commented Apr 24, 2023

andsteing commented Apr 24, 2023 • edited

canyon289 commented Apr 24, 2023

codecov-commenter commented Apr 25, 2023

Codecov Report

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

canyon289 commented May 4, 2023

8bitmp3 commented May 22, 2023

canyon289 commented May 27, 2023

andsteing commented May 30, 2023

andsteing commented May 30, 2023

andsteing left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

andsteing commented May 30, 2023

canyon289 commented May 30, 2023

canyon289 commented May 30, 2023

canyon289 commented May 30, 2023

8bitmp3 commented Jun 26, 2023

andsteing commented Jun 27, 2023 • edited by 8bitmp3

canyon289 commented Apr 21, 2023 •

edited

andsteing commented Apr 24, 2023 •

edited

andsteing commented Jun 27, 2023 •

edited by 8bitmp3