Adds "full eval" HOWTO.

[andsteing]

Adds a HOWTO explaining the problem and possible solutions for "full dataset processing", which is especially relevant for computing evaluation metrics.

The added function flax.jax_utils.pad_shard_unpad() is copied verbatim from big_vision and was created by @lucasb-eyer - thanks!

For reviewing this PR, see:

https://flax--2111.org.readthedocs.build/en/2111/howtos/full_eval.html

https://colab.research.google.com/github/andsteing/flax/blob/doc/docs/notebooks/full_eval.ipynb

Note that for running the Colab you'll need to replace the line

!pip install -q git+https://github.com/google/flax

with

!pip install -q git+https://github.com/andsteing/flax@doc

[review-notebook-app]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB

[codecov-commenter]

Codecov Report

Merging #2111 (d0822db) into main (7c56ba9) will increase coverage by 0.19%.
The diff coverage is 96.87%.

@@            Coverage Diff             @@
##             main    #2111      +/-   ##
==========================================
+ Coverage   74.91%   75.10%   +0.19%     
==========================================
  Files          59       59              
  Lines        5042     5094      +52     
==========================================
+ Hits         3777     3826      +49     
- Misses       1265     1268       +3     

Impacted Files	Coverage Δ
flax/jax_utils.py	66.66% <96.87%> (+10.28%)	⬆️
flax/linen/__init__.py	100.00% <0.00%> (ø)
flax/linen/activation.py	100.00% <0.00%> (ø)
flax/linen/partitioning.py	81.49% <0.00%> (+0.01%)	⬆️
flax/traverse_util.py	98.97% <0.00%> (+0.03%)	⬆️

---

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 7c56ba9...d0822db. Read the comment docs.

[marcvanzee]

Looks great!

[marcvanzee]

It seems the problem is two-fold: double compilation (both in training and eval), and incorrect metric results (in eval). Is this correct? Maybe it is worth emphasizing this. Currently you state that is especially important during eval but you don't explain why.

[andsteing]

I see it more like this

1. In eval, we care about processing the full dataset because otherwise the metrics are off. (in training we usually do multiple epochs and using some examples 1x less for training does not matter)
2. When we want to avoid loosing data at the end, we run into other problems (like e.g. multiple compliations)

I think 1. is mentioned in the first paragraph "Especially when evaluating a model, it is important that we process all examples", and 2. is mentioned further down as disadvantage of some solutions.

[cgarciae]

One case that might be worth discussing is what to do if you update the metric state inside a compiled train_step/test_step. It would seem that you would have to use padding plus some form of masking?

[andsteing]

One case that might be worth discussing is what to do if you update the metric state inside a compiled train_step/test_step. It would seem that you would have to use padding plus some form of masking?

@cgarciae, please check out the added section about using a eval_step() with the new argument pad_shard_unpad(static_return)

[jheek]

another solution is to let each host process indepedently and do the pmean(metrics) add the very end in a seperate pmapped program

[andsteing]

yes, this is actually very similar to what I do with count_correct_p() below

note that the above two sections Adding "infinite padding" and Computing metrics in eval_step() can easily be combined

so I think the different usecases are covered with the subsections, but if you feel there is a specific combination that should be added, feel free to add some more specific comments and I'll write it up.

[jheek]

Ah indeed a parallel stoppinig criteria is not so different when you use padding. I was thinking more towards a note on what to do when you don't want to pad. You might not want to include that route for simplicity though

[andsteing]

but without padding different hosts would have different batch sizes? and this would also lead to re-compilations?

(in that case I think we should keep it simple with "solution=padding" since it seems superior, at the cost of a little more complicated code, but that added complication is quite small, especially in the case where one compute the metrics in the main eval loop)

[cgarciae]

Left a new comment about static_return which is not urgent.
A previous comment that was marked as resolved says that maybe you were going to change the name vs_p to variables but I still see vs_p. Maybe I misunderstood, just wanted to clarify that in case its not intended.

[cgarciae]

LGTM @andsteing! Enjoyed the read, also made me realize all clu and jax_metrics metrics should support masking.

[andsteing]

@cgarciae yes exactly, ideally all metrics support a single mask feature that is expected exactly for this reason; how the mask is incorporated for intermediate value updates then depends on the individual metrics ...