Neural Networks in Equinox (JAX DL framework) with Optax

Equinox is a recent addition to the environment of deep learning frameworks in Python. In this video, we will use it to train a Multilayer Perceptron to mimic the sine function, also using JAX primitives jax.vmap and jax.jit. Here is the code: github.com/Ceyron/machine-lea...
👉 This educational series is supported by the world-leaders in integrating machine learning and artificial intelligence with simulation and scientific computing, Pasteur Labs and Institute for Simulation Intelligence. Check out simulation.science/ for more on their pursuit of 'Nobel-Turing' technologies (arxiv.org/abs/2112.03235 ), and for partnership or career opportunities.
Also check out the docs of Equinox: docs.kidger.site/equinox/
-------
📝 : Check out the GitHub Repository of the channel, where I upload all the handwritten notes and source-code files (contributions are very welcome): github.com/Ceyron/machine-lea...
📢 : Follow me on LinkedIn or Twitter for updates on the channel and other cool Machine Learning & Simulation stuff: / felix-koehler and / felix_m_koehler
💸 : If you want to support my work on the channel, you can become a Patreon here: / mlsim
🪙: Or you can make a one-time donation via PayPal: www.paypal.com/paypalme/Felix...
-------
⚙️ My Gear:
(Below are affiliate links to Amazon. If you decide to purchase the product or something else on Amazon through this link, I earn a small commission.)
- 🎙️ Microphone: Blue Yeti: amzn.to/3NU7OAs
- ⌨️ Logitech TKL Mechanical Keyboard: amzn.to/3JhEtwp
- 🎨 Gaomon Drawing Tablet (similar to a WACOM Tablet, but cheaper, works flawlessly under Linux): amzn.to/37katmf
- 🔌 Laptop Charger: amzn.to/3ja0imP
- 💻 My Laptop (generally I like the Dell XPS series): amzn.to/38xrABL
- 📱 My Phone: Fairphone 4 (I love the sustainability and repairability aspect of it): amzn.to/3Jr4ZmV
If I had to purchase these items again, I would probably change the following:
- 🎙️ Rode NT: amzn.to/3NUIGtw
- 💻 Framework Laptop (I do not get a commission here, but I love the vision of Framework. It will definitely be my next Ultrabook): frame.work
As an Amazon Associate I earn from qualifying purchases.
-------
Timestamps:
00:00 Intro
01:03 Imports
01:34 Hyperparameters/Constants
02:14 Generating a toy sine dataset
04:38 Setting up MLP architecture in Equinox
09:24 Initial prediction on the dataset
11:39 Defining a loss function
12:50 What is learning? Why do we need gradients?
13:27 Function transformation with autodiff
16:08 Setting up optimizer from optax
18:58 Separate function for one optimization step
20:30 Training loop
21:28 JIT compilation of the update step function
22:45 Plotting loss history
23:17 Prediction with trained parameters
24:08 Summary
26:09 Outro

Пікірлер: 4

  • @theneuralmancer
    @theneuralmancer11 ай бұрын

    Amazing work, as always! I love that you produce high quality content

  • @MachineLearningSimulation

    @MachineLearningSimulation

    11 ай бұрын

    Thanks a lot 😊 I also saw your comment under one of the previous videos, haven't yet had the time to write a detailed answer, will do so soon :).

  • @shubhampatel6908
    @shubhampatel69084 ай бұрын

    great tutorial. Thank you for this.

  • @MachineLearningSimulation

    @MachineLearningSimulation

    4 ай бұрын

    You're very welcome! 😊