Rework Flax and Paxml training tutorials

[awolant]

Category:

Refactoring

Description:

Rework training examples for Flax and Paxml to use data_iterator.

Additional information:

Flax and Paxml training examples were adjusted to the new API with data_iterator.

Affected modules and functionalities:

Flax and Paxml training examples.

Key points relevant for the review:

Tests:

• Existing tests apply
• New tests added
  • Python tests
  • GTests
  • Benchmark
  • Other
• N/A

Checklist

Documentation

• Existing documentation applies
• Documentation updated
  • Docstring
  • Doxygen
  • RST
  • Jupyter
  • Other
• N/A

DALI team only

Requirements

• Implements new requirements
• Affects existing requirements
• N/A

REQ IDs: N/A

JIRA TASK: DALI-3711

[review-notebook-app]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB

[mzient]

(...) create an iterator (...)

Also, I'm not sure that listing of all the operators adds much here.

Reply via ReviewNB

[awolant]

Done. I left the list for now.

[mzient]

Line #27.            labels = fn.one_hot(labels, num_classes=num_classes)

How does one_hot encoding depend on random shuffling? Is that a bug or exploiting some coincidence?

If it's about training vs validation, then perhaps is_training parameter would be less confusing.

Reply via ReviewNB

[awolant]

Done

[mzient]

(...) creating ~a~an iterator definition function. It is a slightly modified (...)

Simpler / less verbose version:

Note tThe new arguments passed to fn.readers.caffe2:, (comma instead of a colon) num_shards and shard_id~. They are used to~, (comma) control sharding:

We add devices argument to the decorator to indicate to the iterator specify which devices we want to use. (...)

Reply via ReviewNB

[awolant]

Done with some changes.

[mzient]

It produces a log compatible with tensorboard under in /tmp/dali_pax_logs/. To read this log in console we create a helper function that prints training accuracy from the logs. We use a helper function that reads training accuracy from the logs and prints it in the terminal.

Reply via ReviewNB

[awolant]

Done

[awolant]

!build

[dali-automaton]

CI MESSAGE: [11169573]: BUILD STARTED

[dali-automaton]

CI MESSAGE: [11169573]: BUILD FAILED

[dali-automaton]

CI MESSAGE: [11169573]: BUILD PASSED