search

YCCS-Summer-2023-DDNMA / project

1 stars 3 forks source link

Review tutorial and implement Jax+Flax for a neural network model #77

Open rameshnatarajanus opened 1 year ago

rameshnatarajanus commented 1 year ago

Jax+Flax

You can just save all the downloads in the "ephemereal VM" in the colab notebook (that is, no need to store on the your Google drive) but note that your downloads will vanish if the runtime gets disconnected for any reason.

Also this notebook introduces tensorboard which is very useful for monitoring neural network models which you can omit that on the first pass.

Note the intended learning from this notebook (which I am reproducing here below for clarity). It will be good to record all the variations/experiments that you ran in a google doc.

Note: You can change config (which is just a dict) in the cell after the default config (which is read-only) is loaded from the file system. Or copy the default to another location, rename it, edit it to your heart's content and load config from the new location.

The Flax Notebook Workflow:

  1. Run the entire notebook end-to-end and check out the outputs.
    • This will open Python files in the right-hand editor!
    • You'll be able to interactively explore metrics in TensorBoard.
  2. Change config and train for different hyperparameters. Check out the updated TensorBoard plots.
  3. Update the code in train.py. Thanks to %autoreload, any changes you make in the file will automatically appear in the notebook. Some ideas to get you started:
    • Change the model.
    • Log some per-batch metrics during training.
    • Add new hyperparameters to configs/default.py and use them in train.py.
  4. At any time, feel free to paste code from train.py into the notebook and modify it directly there!
rameshnatarajanus commented 1 year ago

One of the learning benefits of this notebook is that the code is stored in python files and dynamically loaded into the notebook from the filesystem. This allows test code to be written for the python modules while at the same time the notebooks can be used for connecting to a runtime and executing the code, plotting data etc.

I think this will be a useful working model to adopt for code development as we go forward. Note that if you are committing your changes to the repo then you should also commit the source code in the .py files.