google / flax

Flax is a neural network library for JAX that is designed for flexibility.

Home Page:https://flax.readthedocs.io

Geek Repo:Geek Repo

Github PK Tool:Github PK Tool

Create Colab Snippets

andsteing opened this issue · comments

I find myself often copying setup code (e.g. for a simple MNIST-like example) when trying out new features.

It would be very handy to have specially crafted Snippets for Flax development for this purpose. We could provide these snippets as a normal Colab notebook that lives on Github and they would automatically get updated when we change the source. Users imply add them to their config as described in e.g. so/48760503

Great idea @andsteing.

A lot of Flax code is quite modular.

Thanks @chiamp for working on this 👍

Reassigning this to @chiamp since he is currently working on this.

@chiamp another idea for a snippet: for OS users it might be useful to be able to import our examples for Colab. @andsteing does this in the imagenet example and I think it is really convenient since you don't have to copy over all code: https://github.com/google/flax/blob/main/examples/imagenet/imagenet.ipynb

https://github.com/google/flax/blob/main/examples/imagenet/imagenet.ipynb

From the UX point of view, this ResNet notebook may reflect the user experience of someone running the code in a Python IDE, rather than a Jupyter notebook. Lots of boilerplate. It may be tough to follow especially if you're new to Flax.

Maybe we can make the ResNet (and other models) snippets available on Flax RTD, since they are very readable + 👍 for researchers/future research papers? Users wouldn't have to go to another site outside of Flax RTD to learn how to create a ResNet, etc.

This code is copy-pastable (for snippets):

  1. Flax - https://github.com/google/flax/blob/main/examples/imagenet/models.py - you can learn how to build a ResNet yourself by looking at the source code and copy-paste it, modify it. Can we port/copy this code as code snippets for the new Colab?

Also:

  1. Keras - tf.keras.applications.resnet source code https://github.com/keras-team/keras/blob/master/keras/applications/resnet.py (you can learn how to build one yourself by looking at the source code and copy-paste it).

  2. Haiku hk.nets ResNets source code https://github.com/deepmind/dm-haiku/blob/main/haiku/_src/nets/resnet.py (you can learn how to build one yourself), followed by training in https://github.com/deepmind/dm-haiku/blob/main/examples/imagenet/train.py (using hk.nets.ResNetX).

Examples 1 and 2 above may not be the most efficient implementations, but they are readable, good for learning.