Migrate haiku_image_classif.py to notebook format: resnet_haiku.ipynb

Follows a similar pattern to the resnet_flax.ipynb notebook, but using haiku instead of flax. Also gets SOTA results on CIFAR-10 and takes around 13 minutes to run on a single GPU (same as the flax version). Interestingly, this notebook is shorter than the flax one since we don't need to define Resnets, we can just import them from Haiku.
google · Sep 5, 2023 · 2ca3a20 · 2ca3a20
1 parent 5bf8d3d
commit 2ca3a20
Show file tree

Hide file tree

Showing 6 changed files with 1,066 additions and 203 deletions.
diff --git a/docs/changelog.rst b/docs/changelog.rst
@@ -1,6 +1,18 @@
 Changelog
 =========
 
+Version 0.9
+-----------
+
+
+Bug fixes and enhancements
+~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+- Improved example :doc:`/notebooks/deep_learning/resnet_flax` (better accuracy in less time) by Fabian Pedregosa.
+- Improved example :doc:`/notebooks/deep_learning/resnet_haiku`, by Fabian Pedregosa.
+- Documentation is built as part of the continuous integration workflow, by Fabian Pedregosa
+
+
 Version 0.8
 -----------
 
@@ -392,7 +404,7 @@ Examples
 
 - :ref:`sphx_glr_auto_examples_constrained_binary_kernel_svm_with_intercept.py`
 - :ref:`sphx_glr_auto_examples_deep_learning_flax_image_classif.py`
-- :ref:`sphx_glr_auto_examples_deep_learning_haiku_image_classif.py`
+- Image Classification with Haiku and JAXopt
 - :ref:`sphx_glr_auto_examples_deep_learning_haiku_vae.py`
 - :ref:`sphx_glr_auto_examples_implicit_diff_lasso_implicit_diff.py`
 - :ref:`sphx_glr_auto_examples_constrained_multiclass_linear_svm.py`

diff --git a/docs/notebooks/deep_learning/resnet_haiku.ipynb b/docs/notebooks/deep_learning/resnet_haiku.ipynb