Use Jax RNG for data shuffling in MNIST example. #296
Conversation
PiperOrigin-RevId: 313495186
danielsuo
changed the title
Codecov Report
@@ Coverage Diff @@
## master #296 +/- ##
=======================================
Coverage 79.33% 79.33%
=======================================
Files 34 34
Lines 2255 2255
=======================================
Hits 1789 1789
Misses 466 466 Continue to review full report at Codecov.
|
howtos/diffs/checkpointing.diff
Outdated
| + checkpoints.save_checkpoint(FLAGS.model_dir, optimizer, epoch, keep=3) | ||
| + | ||
| + | ||
| + # Save model at every epoch |
There was a problem hiding this comment.
Is this correct? The indentation look to have changed -- doesn't this now only save the model at the end of training?
There was a problem hiding this comment.
Yes, that sounds right; fixed.
| - loss, accuracy = eval_model(optimizer.target, test_ds) | ||
| + model = jax_utils.unreplicate(optimizer.target) # Fetch from 1st device | ||
| + loss, accuracy = eval_model(model, test_ds) | ||
| + loss_accuracy = eval_model(model, test_ds) |
There was a problem hiding this comment.
Is this correct? Seem like a typo maybe? (Do HOWTO tests pass here? -- looks like yes, but maybe the tests don't yet cover this step?)
There was a problem hiding this comment.
Yes, this is a typo! The test tests train_step, but not train.
5 participants
This PR contains commit history from #293 and fixes up
howtos.