Machine Learning with Flax - From Zero to Hero
❤️ Become The AI Epiphany Patreon ❤️
https://www.patreon.com/theaiepiphany

👨‍👩‍👧‍👦 Join our Discord community 👨‍👩‍👧‍👦
https://discord.gg/peBrCpheKE

❤️ Deepnote referral link: https://deepnote.com/referral?token=823d18856ad5 - get 20h of free "Pro machine" + support the channel at 0 cost to you! ❤️

In this video I cover Flax - a JAX-based machine learning library. It's a part of my machine learning with JAX series of videos! You'll learn everything you need to get you started building ML models in Flax.

▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬ 
✅ Flax notebook (do follow along!): https://github.com/gordicaleksa/get-started-with-JAX/blob/main/Tutorial_4_Flax_Zero2Hero_Colab.ipynb
✅ Flax GitHub: https://github.com/google/flax
✅ Flax docs: https://flax.readthedocs.io/en/latest/
✅ HuggingFace Flax community week: https://github.com/huggingface/transformers/tree/master/examples/research_projects/jax-projects
▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬ 

⌚️ Timetable: 
00:00:00 Intro - Flax is performant and reproducible
00:02:00 Deepnote walk-through (sponsored)
00:04:57 Flax basics
00:14:05 Flax vs Haiku
00:17:25 Benchmarking Flax
00:18:28 Linear regression toy example
00:27:05 Introducing Optax (Adam state example)
00:33:15 Creating custom models
00:39:30 self.param example
00:45:55 self.variable example
00:57:15 Handling dropout, BatchNorm, etc.
01:02:25 CNN on MNIST example
01:11:00 TrainState source code
01:13:00 CNN dropout modification
01:15:05 Outro and summary

▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬ 
💰 BECOME A PATREON OF THE AI EPIPHANY ❤️ 

If these videos, GitHub projects, and blogs help you,
consider helping me out by supporting me on Patreon! 

The AI Epiphany - https://www.patreon.com/theaiepiphany
One-time donation - https://www.paypal.com/paypalme/theaiepiphany 

Huge thank you to these AI Epiphany patreons:
Eli Mahler
Petar Veličković

▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬ 

💼 LinkedIn - https://www.linkedin.com/in/aleksagordic/ 
🐦 Twitter - https://twitter.com/gordic_aleksa 
👨‍👩‍👧‍👦 Discord - https://discord.gg/peBrCpheKE

📺 YouTube - https://www.youtube.com/c/TheAIEpiphany/
📚 Medium - https://gordicaleksa.medium.com/ 
💻 GitHub - https://github.com/gordicaleksa 
📢 AI Newsletter - https://aiepiphany.substack.com/

▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬

#flax #jax #deeplearning
 What's cracking you guys? In this video I'm resuming my Jax series of videos and in this one I'm gonna cover Flex. Basically a neural network library built on top of Jax and you probably are already aware of frameworks such as TensorFlow, Keras, PyTorch. So then like why is there a need for yet another deep learning framework? And the answer would be like twofold probably. One of the most important things I think is the performance component. We're gonna later see some numbers on certain benchmarks and Flex kinda was the best, the fastest. I.e. in general like Jax is faster than the other alternatives. And B. you get better reproducibility because of the way that Jax and thus Flex and some other frameworks such as Haiku which is also built on top of Jax. How they handle state. And basically because they're using the functional programming paradigm, they handle state externally. They deal with pure functions and so basically even the random number generators handle state externally and all of that leads to better reproducibility. So those are some whys behind learning Flex and yeah. So before you start watching this video you should have some knowledge of Jax already and if you don't know what Jax is, if you've never done any Jax, I strongly encourage you to check out my Machine Learning with Jax series of videos. You can see here the first one from zero to hero which is going to teach you the basics. Then we have from hero to hero pro plus which is going to go beyond and teach you like a nitty-gritty details of Jax in general. And finally I have a video where I'm covering where I'm coding a neural network from scratch in Pure Jax. So yeah, if you don't know anything about Jax, maybe go and check out at least the tutorial number one before you watch this one. Okay, now you maybe notice something different in this video. I'm not using colab anymore and what this is is deep note and I first want to give up like a huge thank you to the deep note team for deciding to sponsor this video. So I want to briefly mention that like I have so far decided to decline like multiple sponsorship opportunities and the reason is I really want to promote something that I truly believe in myself. So that means I should either use the product like on a daily or on a weekly basis myself or I think that the product has a lot of value and deep note definitely checks the marks. And so yeah, here I am. So what's the thing with deep note? If you take a look at it, it looks like a regular Jupyter notebook like your colab notebook. So what's the difference? So you can see here like you've got your markdown cells. You can click shift enter and render the markdown cell. You have the code cells. Well, the one of the most important things about deep note is the collaborative feature. So that basically means you can share this notebook with multiple collaborators with a team and you can real time collaborate on the notebook and you can basically see other people like typing real time commenting, etc. It's a very good collaborative tool. And unfortunately, I cannot show you right now those features because I'm so here. Secondly, it integrates which bunch of different data sources you can check out those yourself. But like drive, Amazon s3, Redshift, BigQuery, a bunch of other applications with GitHub as well. They also provide templates which can get you started like you can maybe pick this A B testing. Let me show you this one. So it's got a very cool widgets, cool dashboards, and you can even create deep note apps. So let me show you briefly how this thing looks like. You can see there is like a lot of these widgets. You can kind of create these interactive plots. And if you go and hit this publishing editor, you can even go and create a dashboard. And so that means no code, which means you can kind of can basically use this as a data science notebook and present the results to your like business partners or whoever is not tech savvy, but wants to see the results and get some insights. So yeah, that's very cool. So let me see what else I think I've covered. The main things you can check it out in your own pace. You've even got terminals. You can kind of edit the environment if for whatever reason you need to do that. You can you can track the common history of the whole team. In general, the whole history of what happened is fairly nicely tracked. So and finally, let me get back to my notebook. You can see I've installed some of the packages here. You don't have to do that. You can create a fully custom Docker container. They have a support for that. I even created some of the Docker images myself here, but I won't be using that. I'm going to be using this 3.7, which is the same one. I recommend you use if you want to start with this specific notebook. And finally, I've got a referral link for you guys down in the video description. If you want to support the channel and also if you want to get 20 additional hours of these pro machines on deep node, do click on the referral link and do check out deep node. So so yeah, let's get started. The point of this video is to learn to get some basic understanding of flex so that you can build train your own neural networks and then you can take it from there and start reading the documentation, which is probably the best source of information right now because the community is not as strong as like say pie torts. Even though it's growing and it's being adopted by hugging face, they even had a community week where they were coded model in flex. I'm going to link some of those in the in the video description so you can check those out. But yeah, let's get started here. Basically, I'm going to import here some some some like Jax and certain libraries, which is non-pie from Jax. So Jax, Numbay, and then we're going to import the flex, which is the library we're going to learn today. And it's a it's a library developed by Google Research, in particular, Google Brain, and it's a the name itself is a contraction between like you have flexibility, you have Jax, and then you kind of get flex very creative. I'm going to import some functions here. They have this leaning API, which is I think the second iteration of their API. And you can see that the convention is to import the API itself as an N, which you may be familiar with if you come from the pie torch background. It's a fairly common notation. Okay, train state is just a wrapper function that's going to help us make things a bit cleaner. We're going to see that in a moment. Yeah, I think I mentioned Haiku. It's it's also like a library built on top of Jax that comes from DeepMind. And the only reason I have it here is for comparison purposes. So we won't be learning Haiku in this video, although I'll cover it in the next one. So stay tuned for that. I'm importing Optex here, which is like another library built on top of Jax for Optimizers. And the rationale here is the following. You want to have in this Jax ecosystem, every one of these libraries takes care of a specific vertical. So you have Optex for Optimizers, you've got Jax, you've got Flex and Haiku for like neural network building and construction. You've got checks for testing the networks. You've got libraries for reinforcement learning, libraries for graph neural networks, etc. So that's like kind of the paradigm, the mindset of the Jax ecosystem, so to speak. Finally, I'm going to use a pie torch data loading functionalities here, because if you recall from my Jax tutorials, they did not want to reinvent the wheel here. They think that data loading functionality is already good. And so they're just kind of utilizing either pie torch or tensorflow data sets, etc. I'm going to use a pie torch because that's what I'm the most familiar with. Finally, I'm importing some functions here and some common modules. I already ran these cells because, yeah, just to save some time. But now starting from here, we're going to run every single cell and see what is going on. So let's start by building the most simple model, which is like just a simple, this is like a feed forward layer. You can also think about it as a linear regression because we don't have any nonlinear activation functions. And so this is just a linear regression model. So let me run this. So you just kind of click shift enter the same as in Colab. And the notation is dense instead of linear. And you can see that the dense class inherits from this module class, which is the same analogous to pie torch. And the idea is even though this is a functional programming like environment, like in its core, they tried and built like an objective oriented paradigm wrapper around it. And that's what Flex is. And that's what Haiku is. So that's you as a programmer are kind of programming in your usual using your usual mental model like the object oriented programming. OK, let's continue here. So let's see how we can do inference using this model. And the answer would be two steps. We've got this in that function and we've got the apply function. The point of init is to initialize the initial parameters because as you recall, Jax and thus Flex by transitivity property is handling like state externally. So that's why we need init and apply. And let's see what we do here. Basically, hopefully you can see the text reasonably well. I'm going to zoom in a little bit more for just in case. But the UI is a bit uglier. Hopefully you'll see the text, which is the point here. OK, we have some random seed here. Your regular Jax thing here. We are externally handling the random state of the random number generator and we split it into two keys. Then we just instantiate a random like data point here using the key one is just a 10 dimensional vector. And then what we do is we take the model, we call this init function and we pass the key. And the key is used obviously because initialization itself is random process. So we want to have key for reproducibility reasons. We pass the input X and basically we get back the parameters and the output in the case of using this function here. In it with output, not that it's that important. I'm going to later just using it. I'm going to show in a second. But yeah, then you can just we can visualize the shape of the parameters and see that it matches our expectation for a simple linear model. So let me run the cell and you can see here that the output is exactly what we expect. Bias term has is the dimensionality is five because the as we call here, we have five output neurons and we have 10 to five for the weight for the kernel. And you can always a couple of things here. So first is the automatic shape inference because I haven't specified 10 anywhere here. I've just specified the output dimensionality. And so the shape is automatically inferred through this input variable X. So that's what happened there. The second thing you can notice here is this immutable structure because as you recall again, Jax loves dealing with immutable objects. So what flex has done here is just made it easier for Jax to not introduce any box here by returning immutable structures. So that's why we have the frozen dict here. Finally, I'm using in it with output here, but you usually use just in it. You don't care about the output initially. So I can just run in it here and then I'm going to delete the Y here. I'm going to comment this out and we'll get the same result. So yeah, usually you'll see the in it function again. To recap here, what this thing does here is it's iterating through the supremes of pie tree and a pie tree object, which is a term from from Jax again. And basically you're iterating. You're kind of traversing the tree and for every single leaf you're applying this function here, which means you'll be taking like, let's say you have bias and you have some like a Jax lump array in there. You're going to take it. You're going to take its shape and you're going to change. You're basically going to map the tree of parameters into tree of shapes. So that's what you what you get here. So yeah, just a small, small recap there. Okay, so that should be fairly familiar. So the new thing is the in it function and the second thing is this apply function. So let's see what happens here. So we have model we call apply and we pass the parameters and the input X. So this is the only thing that you have to kind of that you have to adjust your mental model to this new paradigm and it may be very weird initially, but like believe me, I started playing with Jax in November last year 2021 and I'm already feeling very, very comfortable with it with this much. I feel much better than when I first encountered this this kind of signature. So you call apply you get the output and we can run this and get the results as expected five outputs here. So again, you cannot call the model like this. So this is what you're used to in other frameworks, but here this is not going to work. It's going to throw an exception. We're going to catch it and print something can call compact methods on unbound modules. So the details themselves are not that important here. Okay, so that's the very, very basics of flex. Now I want to do a small coding exercise here and contrast flex with Haiku and in my personal opinion Haiku is a bit cleaner than flex. I like it better me personally, although it may happen the flex is more powerful because of the additional like notation differences you'll see soon see. So yeah, I'm not very strongly opinionated about this, but so far for me it seems Haiku is even nicer than flex even though flex is way more popular in the community. I mentioned hugging face community week. I mentioned they have even have like 5,000 plus models on hugging face hub, which is fairly, fairly nice amount of models for a relatively new framework compared to TensorFlow or like PyTorch. So let me show you how you would do the same thing in Haiku now. It's going to be very brief. So let's create a model. So Haiku uses this HK as a, as a, as a like a notation then linear and then five. Now if I were to run this, it's going to fail. So let me try and run this. And as you can see here, all of these modules must be initialized inside an HK Haiku transform function. So what the trick is Haiku is very explicit about when you want to, you take a function like linear and you, in order to have the init and apply. So it's using the same notation as flex. You have to pass it into this HQ transform function. So let me do that. Transform. And what we do here, we create a Lambda anonymous function here and we just simply call this, this, this, this linear model. And that's it. Now this will work. So if I were to run this, it will not fail. But now let me show you what else happens. So let me just copy paste these parts here because they remain the same, the seeds and the inputs will remain the same. Now let me initialize the parameters. So we get model init again, we pass the key, we pass the X and we get parameters. So so far so good. It looks very, very similar to two flex. And finally, let's get the output by calling model apply. We pass the parameters here. We pass none because we don't have any stochasticity in the forward pass. And for now it may be, it may look as high because even more complex and a kind of uglier in the notation compared to two, two, two flex, but it's not the case. You'll see that in the next video. So you'll do it like this. And by the way, I made a small mistake here. You don't want to reuse keys. So this should be key two. So I should be able to run this right now and everything should work. So this is, this is high Q very similar. You just have this transform function and, and yeah, so that's cool. I had the solutions here. So let me kind of show this code and hopefully this looks the same. So yeah, we have the model, we form the data here, we call in it, we call apply and we print. And again, we have this, I forgot to do this. Basically it just shows that high Q follows the same paradigm as flex. It's inheriting from this module class as well. So let me just hide this code again. So hide the code and that's it. Cool. That was section number one. Hopefully you're not too overwhelmed because as I said, you will have to kind of get accustomed with the, with this functional programming paradigm. But yeah, it gets easier after, after, after a quick. So as I said, flex and high Q offer a performance boost. So here's some numbers from the hugging face team. You can see here using a single GPU and using a flex you get compared to PyTorch, you get certain improvements. Like in all of these, I think these are, all of these are NLP tasks. So you have Cola, you have the natural language inference task. You have, yeah, semantic, I guess, yeah, whatever. So in any case, you get a fairly nice improvements here speed wise. And in total, you can see that here you've kind of saved more than an hour using a single GPU, which is not a small thing, especially when you scale these two big language, large language models. Okay. Let's get going and let's introduce a single, a simple toy example where we're going to just form like a two dimensional data set and try and train a simple linear regression model using the NN dense model, which is an implementation and instantiation of the linear regression model. So I'm going to have 150 data points. The X dimension is two. The output is going to be, we are regressing like a scalar just to make things a plottable, if that's even a word, because yeah, if we had multiple dimensions here, I wouldn't be able to plot. I'd have to do some dimensionality reduction. So I'm going to skip that. Again, we're going to form some keys here being explicit about the random state as usual. We are going to form the ground truth weights and biases here. So basically we'll be trying to learn this weight and this bias later during the training. And now what we do here is we form a structure. That's, that's something that flex is, is used to. So you have the freeze is supposed to create a frozen dict, a frozen dictionary out of the normal Python native dictionary here. As you can recall from, from this section, we had frozen dict here. So let me get back here. And the structure is such that we have this params collection as flex calls it. And then we have bias B and kernel W and that's it. Now we, we've used this ground truth model and generate the data. So again, we split the key, which we created here, we split it and we get the key that we're going to use to generate the axis, the domain of our data. So we have 150 samples of, so 150 to two doubles. So that means just 150 two dimensional data points. And then we're going to use it, use the ground truth model and regress the, the YS here. And then we're just going to do a superposition of noise here using this amplitude. That's pretty much it. Let me run this now. And you can see the shapes, 152 and 151. And now let's plot the, the, the data set we just created. And you can see it here. Now I was lazy and so I didn't try and create an interactive plot here. And so let's be smart about it. And let's try and visualize where this data lies in. And you probably already know the answer. If not, let's do a simple exercise and try to approach it. So first one is we know how we created, how we generated the data. So basically what we've done here is we had Y is equal to W X plus some bias. And because X is two dimensional, this basically turns into W one X one plus W two X two plus B. And you can see that this is basically your, your analytical formulation of a plane, of a 2D plane embedded into 3D space. So that would be one approach. The second approach is data driven as a data driven. I mean, what I'm going to do here is increase from 150 to 1500. Okay, let's run this cell now and see what we get. So we get 1500 data points and let's run this again. And hopefully now you can see that it lies, it lies on the, on the 2D plane here. And yeah, the reason we added noise, hopefully that's, that's clear enough is to kind of hide the actual, the ground to the, the actual model that, that lies in the background that generated this data. And now we're going to try next up to, to kind of like a regress and infer the model that generated the data. So first let's create this, this loss function. This make MSC is just going to return this MSC loss function, which is a jitted version of the MSC. So just kind of trying to optimize things a little bit here. So squared error itself takes a single data point X and Y. We do inference, we just do a prediction. We pass in the single data point, we get a prediction and return the mean squared error here. You may be confused by the inner function here and inner is used just because so as to be able to generalize two cases when like Y was a multi-dimensional vector and not a scalar like it is right now. So this is just a generalization to you of our mean squared error loss. Finally we take the square error and we call the V map. So that's the transform from Jack's that makes this thing run on a batch of data. So we pass now like a batch of our data, like the, let me, whoops, let me actually reduce from thousand five hundred. No, I'm going to leave it at thousand five hundred. It's still like a small data set. So we're going to run this on the batch of data. So this will give us like a thousand five hundred by one vector. And then we're going to do mean over the X is equal zero. So that means we're going to kind of average out all of those, of those loss losses for all of these singular data points to get the final loss. And that's pretty much it. That's how we get our MSC loss here. And finally, I kind of wrap it into this value and grad transform function from Jack's to get the, yeah. That's how we calculate gradients and you should be familiar with this already. So now this will be our model. So we're kind of cheating here a little bit because this is the exact same model that generated the data. So we, we usually wouldn't have this advantage of knowing the model, but here for the sake of just demoing this, I'm going to suppose we, I'm going to assume we know the exact model. So we're going to do the usual stuff here, initialize the model. We get initial parameters and then we're going to do a simple training here, a training loop. We're going to have 20 epochs, just an arbitrary number. You can tweak it, but like it's enough for this toy problem. We called, we called the, we just passed the params and we get the gradients and we get the loss. So this may be like a dark magic a little bit, but the trick is this value and grad function already contains the MSC loss, which already contains our data because we had this wrapper function make MSC loss and the data excess and YS is already stored internally. So that's why we don't have to even pass anything here. So what this does is for the current set of parameters and for the car, for the dataset, which is static, it figures out the current gradients and the current loss for those parameters. And then we just do a simple stochastic grid in the send here. I'm being very explicit here. You usually will not be doing it like this. You're going to be using Optex, but here I'm being explicit and doing it in pure Jax way. By the way, I'm having some internet connection problems. Hopefully that will not interfere. I think not, but yeah. So you can see here, we do the tree multimap from Jax. We basically have two pie trees here, params and grads. We are going and traversing the leaves and we do the following simple manipulation from the current parameter. We subtract the learning rate times the gradient, which is a definition of your simplest stochastic gradient descent. And that's it. So this is actually not stochastic gradient descent because we're using the whole data. So this will be a batch gradient descent. So it's exact. We're not being stochastic here. And finally, we are logging every now and then. So every five epochs, I just print the epoch and then print the loss. And yeah, if I were to run this shift enter, we can get, we can see the initial parameters which were initialized using this key, which is different compared to the way we've initialized the model here. We use the specialized keys here and then we kind of wrap those into this freeze dictionary, et cetera. Okay. So those were the initial parameters. Then we had loss going down towards zero. And finally we get like, as you can see here, the final parameters, the learn parameters are here and you can see that they are pretty much the same. We get two five four here, two five, very similar numbers, which is also indicated by the loss. So yeah, that's how the training loop in Flex looks like. Almost in Flex. As I said, I'm using more, this is more like a puregex type of syntax here. So let's now use OptX, which is a dedicated library for this purpose in particular, and let's initialize it. So from OptX we call the SGD, we pass in the learning rates and then we have to again initialize the OptX because state has to be externally manipulated. And here we get a state and we kind of print the state. So let's run this cell and let's see what happens. And as you can see, we've got a tuple of two empty states and that may not make sense if you're not familiar with how optimizers work, but if you are, this makes perfect sense because SGD doesn't have any state. Compare that to Atom, which has to kind of keep the various statistics of gradients, like the variance of gradients in memory in state and that's why they're like, we have to like kind of explicitly manipulate it. So if I were to change this to Atom, and yeah, I won't be changing the variable name, I'm just gonna change and use Atom. Let me try and run this. You can see that now we have like a meaningful state here and that's something you have to kind of pass externally here in Flex, in Jax in general. So yeah, you can see what happens here is we have two set of parameters here. First is the mu parameters and then we have the nu parameters. And these are just your Greek alphabet letters that are denoting the same quantities as in the original Atom paper. So if I were to Google here Atom paper, blah, blah, blah, and if I were to open this paper on archive, we're going to see the formulas are using the exact same variable name. So I think it's here, is this the Atom? Yeah. And this is the nu, I don't know whether I'm pronouncing it correctly, probably not, but basically those are the statistics, the gradient statistics that Atom is kind of keeping in the background in order to do the update. So here is the Atom update. It's way more complicated compared to, let's say, SGD because yeah, you can see here we're using the mu and the nu here to form the novel parameter, to form the update. Okay, so that's it. Let me return back to SGD so that we can have comparable results compared to the pure Jax approach we used here. And so the tree multimap. So let me go back here and we're going to have the same exact thing replicated. The only thing that's different here is we are using instead of explicitly handling the optimization part, we're going to use OptX. And you can see in this example, it actually takes two lines to implement the code that we implemented here. I think in, yeah, we had just a single line here. So you may ask yourself why would we use OptX then? And then, I mean, it's pretty obvious. I mean, when you have some more complex optimizers, you'd have in pure Jax, you'd have like 10, 15 lines or more. Here you always have two lines and that's very cool. So this is how it looks like. So you have the optimizer, the SGD optimizer here. We pass the gradients and we pass its state and outcomes, the novel state and outcome, the updates, which are the modified gradients. So in the case of Atom, depending on the internal state, we'd be modifying these updates and then we apply, we just call this apply updates on the parameters and we get the novel parameters and that's it. So this should not work. Let me try and run this and see what we get. And you can see here that we get like a loss. I think it's the same loss. It's 512 here and if we take a look here, it's also 512. The last two digits are 85. The last two digits are 85. We get the same loss because I mean, this implements the SGD perfectly. Okay. So that's cool. And later, I have a note here. Later we'll see how we can kind of take all of these various states that are kind of floating around our program loops and wrap them in a single object and then use that single object to kind of, yeah, do the training. So just so you can see that Optex is much more than SGD and Atom. Let me show you this piece of code here. I mentioned here that you can do various schedules, cosine schedules, linear schedules, arbitrary schedules. You can do chaining of multiple optimizers, multiple components. You can do parameter freezing when you want to do some fine tuning. So it's very powerful, much more powerful than using pure JECs. I mean, it's as powerful, but like it's way more concise. Okay. So let's show this code here. This will not run. This is just some simple example of copy pasted from ImageNet examples from the respective repos from Haiku and from Flex. Here you can see this function create learning rate, which basically takes some parameters and then creates this from Optex. You call this linear schedule function. We give it the initial value, the end value, and then the number of steps where we are linearly increasing the learning rate or decreasing. Now here in this case, yeah, you can basically either increase or decrease depending on the initial and end values. Then we have the cosine function here. So the cosine schedule. And finally, we kind of join schedules using Optex join schedules. And we join the warmup, the linear warmup combined with the cosine schedule. And I guess, yeah, and we put the boundaries here to kind of denote the moment when we are switching modes. And then you can use this scheduling function and pass it directly to Optex SGD here with some additional like Nestoro momentum, et cetera. Finally from, you can see here a snippet from Haiku, even more complex, arguably. We have a trace here. We have a chaining operation of this tracing operator. And then we have the scale by schedule, like a bunch of small components which can make arbitrary complex transformations. So that's pretty much it. That was section number two. We got some understanding of how the toy like neural networks, i.e. linear regression in this case, are trained. And now we're going to see how to create custom models in Flex, which is something we care about. So let's start with a simple example of how to build an MLP. So multi-layer perceptron network. And you can see here we are inheriting from NNModule. So we form a class, we inherit from the module. So every, when you want to build a custom model, you always want to inherit from the NNModule. Then you can see that we have this field here. And this is called data field. It's from one of the newer versions of Python. You may be familiar with it. You may be not. But basically NNModule is Python's data class. So if I were to, let me check. I think I have a link here. Yeah, we have a link here. Let me just copy paste this thing here and let me show you what a data class is. So data class, let me kind of zoom here. A data class is basically this. So you have this wrapper called data class, and then you can just kind of put the variable names here. And then a couple of functions are automatically instantiated for you, defined for you. So you can see, let me just find an equivalent representation here. Yeah. So if you were to explicitly do the code op, the data class, this is what you'd have to do for those two fields. You'd have to create this representation function, the quality function. You'd have to call the init. And there is a lot of redundancy here. As you can see, we have rank being called three times here. So all in all, there's just a neat feature of Python more than Flex that, as I said, NNModule is kind of piggybacking on. And yeah, so as a consequence of NNModule being a Python data class, unfortunately, we'll have to deal with a new function name. We're not using init anymore. We're using setup. So setup is old init. And because as I mentioned here, basically, the data class is implicitly using the init function. And we saw that on the tab I just closed. What we do here is we just form a list of layers. So you can see here, we iterate through this sequence of integers, which is a simple list of integers. We create these dense, i.e. feed forward layers. And then during the call function, what we do, we pass the input. We just rename that as activation. And we iterate over the layers. We kind of call, we pass the activation through the layer. And finally, depending on whether or not we are in the last layer, we apply the ReLU, the ReLU nonlinear activation here. And then we finally return the activation. So the output layer, that's the output activation. So yeah, that's your basic definition of an MLP. Fairly, fairly simple. Now let's see how we can kind of do some inference. So again, we split a key. We have the H key. We have the init key. Init is supposed to initialize the parameters. Let me kind of first slowly walk you through this part. So we have MLP. We instantiate it using this number of neurons per layer is going to be 16, 8, 1, which means in the first layer, we have 16 neurons. Then in the second, we have 8. And then we have final, a single neuron. And we instantiate this data set here. We have the four data points that each has four dimensions. And we use the H, so we use the X key here. Then we call the model init here with the init key and we pass X. So we've seen this multiple times already. We have params. We call model apply and we do a forward pass and we get Y here. And that's it. The same as if we were to use a simple, like a dense layer and not a complex MLP. Again, I'm just calling a tree map and we're printing shapes. So let's see how this thing looks like when we run it. And as expected, we have 16, 8, 1, and we have 4 here because remember we have automatic shape inference. And because our data points have four dimensions, basically, that's what we have for here. And this kernel here, and this weight here. By the way, I hate the word kernel. It's so overloaded in machine learning that it's crazy. So yeah, maybe using a weight, like a terminology here is maybe a bit better. I don't know. Cool. I made a note here. Let's try and do the same thing just with this nn-compact wrapper. And what the trick here is, is you don't have to have set up at all. So what you can do is, because we have a fairly simple function here, we can just kind of wrap this call function into nn-compact. And then what we can do here is we can iterate through the number of neurons per layer instead of through layers. And that means we'll have number of neurons here, number of neurons. And then what we do instead of having a layer, we'll just instantiate it on the fly here. So nn-dense, we kind of initialize it using the number of neurons, and then we do the same. Everything else remains the same. And I think this should run unless I've made some mistake. Let me see whether it's going to work. So everything remains the same. It's just a bit more concise way of writing down things. So yeah, I managed to make an error here, and the error is that we should not be using the layers. Layers does not exist anymore. So let me just kind of switch that here. And this should now work. Yeah, it is. Okay, cool. You just saw a second way. So this is a very flex-like way of writing things down using nn-compact. But either way, you get the same semantics. Cool. Let's continue here. The output is, by the way, as expected, we have four times one because single output neuron for data points, remember, and then that's why we have four times one. It's good to kind of check the shapes and make sure those make sense. Now, let's dig even deeper. And there are two very, very crucial concepts you ought to know about. The first one is this thing called a param. And basically, it's a terminology for trainable parameters in neural networks in Flex. So let's implement the simple model that's going to leverage those params and see how this thing works. So here, what we see is if we were to try to implement the dense layer, either feedforward layer, this is how we would approximately do it. So here, again, it's a data class. So we can instantiate these fields here. We have number of neurons. We have the weight and the bias initializer functions. So yeah, bias is going to be initialized as all zeros. And for weights, we're going to use like a normal initializer. Not that important. The details are not that important. The actual implementation of this function is not vital for this explanation. So let's see what we do here. So we call self-param and we give it a name. And here, it's going to be called weight. And it can be arbitrary name. You can choose whatever you want. So that's the name that will appear in that frozen dict under the params collection, if you recall. If not, I'm going to run this. You're going to see it in a second. So now we call the initializer function of our choice. So the random state is going to be implicitly passed into this function through the init and apply function. So you don't have to worry about that here. So the only thing you have to provide is the shape information for this initializer function. And we'll just take the last dimension of our data and then the number of neurons here, which is the desired number of outputs. We do the same thing for bias. We call self.param. So this is, again, this is flex syntax. And we name it bias. We have the bias in it. We pass the number of neurons as the shape because, yeah, that's how bias shape is supposed to look like. And then this is a simple implementation of feed-forward layer. We do a dot product between the inputs and between the kernel, the weights, and then we just add the bias. And that's it. This is how we would implement like a dense, a linear, a feed-forward layer in flex. Again, we have a common pattern here. We just form the random seeds here, the random states here. And then we instantiate the model here. We just form some random data on this line 18. Then we initialize. We get the parameters. We call apply and we get the output. And finally, we can print the shapes. And let me run this and see what we have as the output. And you can see this is like something we're used to seeing. We have params collection. And I'm deliberately saying collection because that's a reserved word in flex. And I'm going to soon explain what exactly it means. And yeah, you can see biases and weights here. And that's it. I've linked a couple of source code files here. You can check those at your own pace. But let me briefly show you, and you may be confused by some of the idiosyncrasies of flex. Like why do we have brackets here and why don't we have brackets here? What type of inconsistency is this? So my OCD was kind of triggered by this. I had to dig into the source code and I found the solution here. Basically what happens, let me zoom in here. So we have a partial. It forms a higher order function. It just kind of plugs in these arguments that you don't have to kind of plug those in later on when you're calling this variance scaling function. So the reason for the inconsistency is that if I were to find the variance scaling implementation, you can see the variance scaling basically returns back this function in it. So that's why we have to call the brackets so that we get back the init function, which we can then call as an initializer. Whereas zeros, if I were to find the zeros, it's already a function. So that's why you don't have to call like the brackets. So minor implementation detail. I just thought explaining the why the inconsistency in case you wondered. Second thing, I want to quickly show you the implementation of the aptly named linear layer, even though they call it dense. So let's see that. Okay. So let me find the class dense, dense general, dense here. Okay. So obviously it's a bit more complicated than the toy implementation, which I just showed you, but like the basics are the same. There is just an additional type conversion depending on whether you're training this thing in a mixed precision mode. You may want to convert the data type of your activations, of your input data, et cetera. So you can see here again, we have the self-param kernel initializer function. We have the shape. Additionally, again, I said we have a data type here and then again, we have some conversion of data type here. So basically, and then we have dot general, which is basically as if we were to call JMP. It's just a more lower level function that says you can see more verbose, but also much more, I guess, less error prone. We formed the bias here and basically if bias exists, we just add it here and then that's, yeah, I mean, that's just a W times X plus B implementation of that. So that's pretty much it. Now let me get back to the notebook and let's continue this journey. Yeah, again, I just have a simple cell here showing that the signature is such that this initializer function, the first argument is actually key, not the shape. So that's why I said that you are implicitly passing in the RNG. So by the way, you usually will not have to care about these things that much because you'll just usually just be using the NN layers, which are already implemented for you. Here I'm just trying to dig a bit deeper into Flex and kind of show you how things work behind the, like in the background. Hopefully you'll appreciate that. Okay, let's continue here. So those were the trainable parameters, but as you may know, when you're training neural networks, you sometimes have parameters which are a part of the model, but they're not exactly trainable. So you're maybe keeping some statistics in the case of batch norm. I guess that's the best example. The most famous example would be the batch norm and batch norm has a trainable parameter as well as those non-trainable batch statistics. And now we're going to see how Flex handles those. And the novel concept I want you to keep in mind here is variable. And I just have a note in terminology, so it's a broader term and it includes both the trainable as well as the non-trainable variables. So you'll see params being called variables as well. They're just trainable variables. Okay, that out of the way, let's see a very simple contrived example that contains both the variables and the parameters alike. So we have a simple class here called very convoluted name bias adder with running mean. Don't try and deduce the semantics of this function from this title, from this name, it's just a contrived example. So what happens in the call function, we're using again this compact notation, we just check whether this variable is already initialized, namely the batch stats collection. And then we take a look whether the EMA variable inside of that collection has been initialized. Now let's see what collection and all of this, what I've just said means. So batch stats is actually not an arbitrary name. You'll see that in, I think I'm gonna show you the implementation of the batch norm, the source code itself. Basically they've hard coded now whether that's a good design or not, I don't know. I guess I don't have enough context, but like in old one, I'd say it's not the best design to hard code like a string like that. But yeah, it is what it is. So we can see here self.variable. So again, that's a flex thing and we form a collection. So this is the collection name, batch stats. And the reason they have different collections is so that you can treat them differently. So that's the point. So we have the params collection which usually contains the trainable parameter, the trainable variables. And then you have batch stats which contains these types of non-trainable variables where you take the batch of images say and you calculate some statistics. So for example, in the case of batch norm, you calculate the mean statistics, you calculate the variance and then you use that to do the logic of the layer. Okay, so it's very similar to self.param. You just have the collection name here, then you have the name, then you have the initializer function here. We are just being explicit here. We just have the zeros function and then we pass in the shape. We have one and then we have a column here because basically the zero dimension is batch dimension. So we're gonna just have the shape of a single data point here. Then we have a bias. So bias is a trainable variable and this will be added. So this thing here is kind of implicitly set to params here. So I noted, I have it here noted. So by default, we'll add this variable to params collection, whereas the above will be added to the batch stats collection. Here because of some of the implementation details, unfortunately, we have to pass the key because that's the expected signature of the function for this self.param, even though we're not using it in this particular case. So yeah. And finally, before I explain this snippet here, EMA stands for Exponentially Moving Average, just for you to understand the semantics there. So what we do here is if we have initialized the EMA variable already, then we can start in the forward pass, we're going to update it using this exponentially moving logic. So what we do is we take the decay. Decay is usually some larger number like 099 or something. And then you keep, you kind of have inertia. You keep the old value with a higher probability and then one minus that and then you do the mean of the current batch of data. So this is kind of damping, has this damping effect. I'm trying to teach you flex here, but like I think that explaining some of the machine learning concepts along the way may help a lot of you. So that's why I'm doing this. Finally we just returned the return is X, which is input data minus the exponentially moving average here plus bias. And you may wonder why do we have dot value and the reason is that self dot variable returns a reference, not the value itself. So that's why you have to kind of do the dot value to get the actual value of the variable. Okay, this will make a bit more sense in a couple of seconds. Again, very, very familiar pattern here. We form the random states. We instantiate the model here. We form the data. So we have 10 data points. Each data point has four dimensions. Again, we initialize the model, we get the variables and then we print them here. So let me just kind of run this and show you already what's going on here. So here is the frozen dict, which we've, we're used to seeing this. And now aside from params, we also have the batch stats where we have this EMA variable. And so yeah, this is the structure, the underlying structure of this, of this contrived model. Only the values are zero because that's the initializer we've been, we've set here, the zeros. And then what we do here is we apply. So this is the novel part of notation. I want you to remember we pass the variables and then we pass the input and then we tell flex, Hey, batch stats are actually mutable, which means that during the forward pass, the value is going to change. Whereas for regular parameters, when you do like a feed forward to a neural network, you will not be changing the weights. But here for these types of variables, they're going to change during the forward pass. And that's the reason why we have to have this distinction and different collections. Okay. So we get the output and we get this time, we get a tuple, not just the output, we get the updated non-drainable parameters and I've printed them here and you can see how all of the zeros now converted into this, which is because we've done, we have data X, we've done the mean across the data. So we've aggregated those 10 vectors into a single vector. And then here using this update rule, we've kind of have a novel, novel value here. So all in all, what I want you to take out of from here is there are multiple collections because we have multiple semantics, depending whether we have trainable parameters or these types of other variables. And again, we have this mutable keyword, which you should remember alongside this self dot variable. So those are some of the important flex idiosyncrasies. Okay. Let's continue here and let's form like an update function, a training function. Let's see how we could train this function. And once you understand this, trust me, it's going to be way easier because we are now digging really deep into all of the nitty gritty details. It's going to be much more easy, much more easy, much easier to actually train flex models than this. Okay. We pass bunch of arguments here because we're still not using the train state, which I'm going to later introduce. So bear with me here. We pass the params into loss function. And what we do here, we pass the apply function. Apply function is basically your model apply. So we have model here and then I'm going to pass model dot apply here into the update step. So that's the apply function. So it's basically how you do the feed forward. Then we pass the variables and here we're being explicit. So I have params collection and I pass the params and we have the non-trainable params, which I just do the double star here to unpack the keys and the values. And this is how you basically form the variables. We will soon see why we have to split these two. And then we pass the input and then we have the mutable. Again here we're just being a bit more general and we say, give me the keys that correspond to those non-trainable parameters. So batch stats would be one of those keys in one of the last examples. So just to give you some idea. Finally here the loss, because this example is super contrived, this doesn't make much sense. Basically once we have the loss function, we do this is a very standard pattern. You've seen this already a couple of times, like only in this notebook and you'll be seeing it even more once you start coding in Flex. Basically you pass the loss function to the value and grad transform function. And we say has auxiliary, we set to true because aside from loss, loss function is returning these updated non-trainable parameters. And remember, once we do a forward pass, we are kind of mutating those parameters. That's why we have to have this mutable in the first place. Okay, cool. So that's what we get here. Once we pass the parameters, we get back the gradients of those parameters with respect to loss function and we return the value because we have value and grad. So we return back the loss and the updated non-trainable parameters. Then we have this common pattern using OptX. We update using the optimizer state and then we apply the updates to the parameters to get novel parameters. And that's it. Then we return a bunch of arguments here because we're still not using the train state and I'm going to introduce that I think in one of the next cells. So yeah, stay tuned. Cool. So let's see how this whole like how everything comes into a single picture. So we have Winstonchi the model here, we have dummy data and so we initialize the model, we get the variables. Now here, because with this time we have both the trainable as well as the non-trainable variables, we split them using this pop function and then because we just have a dictionary, so this is going to work. And then they usually do this in the documentation, they delete the variables to avoid wasting resources. And then we kind of create the optimizer here, we initialize the state and then we have just a simple for loop running a couple of training steps and here is the update step that's being called. So we are correctly updating the parameters as well as treating the non-trainable variables in a correct manner. So this is as nitty gritty as you're going to get in this video. So if you followed me so far, congrats. If not, you're still won't have any problems following the rest of the notebook. So now we're going to go up in the abstraction and see one example where you basically have the most complex syntax you can have in Flex. So I'm kind of going to mix a lot of various layers which are going to introduce both those non-trainable variables as well as some stochasticity in the forward pass via a dropout. So let's see that example here. We have this DDM block and the name DDM because we have dense dropout batch norm which is a fairly regular block you'll see like in also in CNNs or whatnot. It's a fairly regular stack of layers. Again it's a data class so we have number of neurons specified here and the training flag and depending on the value of this training flag, we'll have different behavior from dropout and batch norm here. So what happens is we instantiate the feed forward layer here using the num neurons here, the num neurons variable and then we pass whatever like comes as the output, we pass it into dropout. The dropout has 50% probability rate of kind of dropping the activations and again we have this deterministic. So if we are training, that means not true will be false which means we are not deterministic during the training. We want to have the stochasticity but when we are evaluating then the training flag is set to false so this will be true. We're going to be deterministic and basically dropout is going to be kind of turned off pretty much. Finally we take the output from dropout, we pass it into batch norm, we have use running average again depends on the training flag. If training flag is true, if we are training then we pass false because we don't want to use the running average. We want to use the batch average, that's how batch norm works. So you take a batch, you calculate the statistics and you use those parameters instead of using the parameters that were accumulated throughout the training. So that's it. This is more of a machine learning concept than a flex concept so I'm going to not dwell too much on it. Again familiar pattern, we form a bunch of random states here, we instantiate our block, we get some random, some dummy data here. We initialize the model and here are some additional idiosyncrasies of flex. So this time when we want to initialize and this may be a bit confusing, we have to pass like a random state for both the params collection as well as for the dropout. And it doesn't make much sense semantically because I mean dropout cannot be initialized, you don't have any learnable parameters there but just because of the way how everything is implemented you kind of have to have this additional step. So keep that in mind. Then we pass the X and everything else remains the same. We get the variables back and then we apply and now here comes additional part because now we have dropout and dropout as you saw complicates both the initialization procedure as well as the apply procedure. So here we have to have this RNG's keyword and basically what that means is you set the random state so you seed the dropout procedure so you can again reproduce it later on. So that's why it makes a lot of sense to have it here although it does not make much sense to have dropout here to the best of my knowledge. So yeah, outcome the output, the Y and the non-trainable variables, the non-trainable parameters. I'm kind of mixing the terminology here between hopefully you're getting the gist of this. And if we want to run this in evaluation mode where training flag is set to false, the apply function simplifies by a lot. We now don't have to pass the RNG's because dropout is effectively turned off and we also don't have to care about mutable part because batch norm just won't be updated. Those internal kind of statistics will not be updated so we don't have to include the mutable. So this is as complicated as it gets. Remember a couple of these novel keywords. We have the mutable, we have the inclusion of RNG's for the apply when we have the stochasticity in the forward pass. And we have to during initialization if we have dropout we have to kind of specify keys per collection. And this is why I said Haiku seems to be cleaner in my honest opinion compared to Flex because they don't have the separation. It's kind of cleaner. Like all of these signatures simplified by a lot. Now whether that means Haiku loses some of the expressivity that's I guess worth thinking about but yeah. Okay that was it. Now we are in the last section of this video which is already overly long. Hopefully you're following along. Basically I'm gonna walk you through this fully fledged convolutional neural network example on MNIST. We're gonna code it up in Flex and this is basically now that you have all of the components in your mind this is gonna make hopefully a lot of sense and you're gonna see how Flex loops and training like code actually looks like. Okay let's start by defining the model. We have a convolutional neural network. We're using this MN compact pattern. A lot of hard coding here for the demo purpose I guess it's fine. Basically a lot of convolutional layers followed by relu's and then here and there we have some pooling average pooling then we flatten out the shapes and we have a couple of linear layers here and finally we have as the output the log softmax of the outputs. That's it and we have 10 classes because we are dealing with MNIST as you recall. Okay let's run this thing to define the neural network there and now I'm just going to reuse the the PyTorch data loading functionality from from tutorial number three of Jack series. So I won't be diving deeper here. Basically I have some collate function that's supposed to create non-Py arrays out of out of PyTorch tensors because that's what a Jacks and thus Flex is dealing with. I have some custom transformations to just kind of manipulate the MNIST image and make it a zero one range flow through the two etc. Some need to degree deals. Basically we form the data sets here. We have the data loaders with certain batch size like 128 images per batch. I load the training images and the labels as well as the test images and specifically I load the whole test data set because it's fairly small for MNIST so we can kind of keep this in memory no problem. Cool let's run the cell to load the data and now let's visualize a single image from from this data set. So it's just like a regular MNIST image. You can see number seven here. You can see the ground truth label printed here on the screen as well. So let's now understand a couple of functions that we'll need to to get this CNN training and learning on this simple MNIST data set. So first things first we have this train step and eval step functions. So I guess you're already familiar with most of these patterns but let me quickly walk you through what's going on here. So what's happening is we have the loss function here. What we do we do a forward pass through the CNN. We pass the variables. So again I'm being explicit here with the collections. We have the params collections. We pass the params here. We pass the images which is our data and we get back the logits for those images. Now we convert the ground truth labels using the basically this one hot function. We get the one hot labels and then what this thing here is if you can guess it's a simple cross entropy loss because what we do here is we do element wise multiplication and because both the logits as well as the one hot labels have the shape of batch size times 10 because we have 10 classes in MNIST. What this does is basically like the one hot labels will kind of select the true label from the logits and because logits is log softmax if you remember how we define the CNN here. So let me just go up there. So we have log softmax as the output of the CNN. So what will happen is this will be just log softmax. Let me just find where we stopped. Okay here and then we're going to sum across the minus one. So that's the last dimension i.e. the dimension where we have the 10 labels. So we're going to sum those up. Most of those are going to be zero because as I said this is going to select just a single one. So we end up with a bunch of logs and then we just kind of do the mean across the dimension zero by default. So that's across the batch dimension and we get like a mean of the log sum pretty much which is the cross entropy loss with the minus here. And we pretty much return back the loss as well as the logits. And then we have like a very very familiar pattern already. We have the value and grad function we set as auxiliary to true because we're not only returning loss. We're also returning the logits. And then this is the new thing. So we have this state now and that's the thing I've been mentioning throughout the video. We're going to soon see what exactly and how we construct that. But basically we take out the parameters which is just they form the part of our of our training state. So parameters are basically the weights of our neural network and we pass those and basically we get the gradients back and we get the logits. Then we just call state apply gradients. So now we don't have the two optics line because those are hidden in the apply gradients function. And we just pass the gradients and get the state back which is the updated state. Next up we compute the metrics. We just pass the logits and ground truth labels. We get the metrics which is a simple dictionary containing loss and accuracy. We'll soon see that. Not that important though. Evaluation step again very simple. We pass the test images here. We call we do a forward pass. We get logits back and we just compute the metrics because we want to evaluate the model on the test images and that's it. Okay let's continue here. Now that we have the training step and evaluation step defined we can basically just kind of call those in a loop for the whole epoch. So we iterate here through the PyTorch data loader. We are fetching images and labels. We start calling the train step over and over again and we're just collecting the metrics across all of those batches. And then what we do is we call device get which if you recall from one of the previous videos we'll just take those like the data that's currently on the accelerator whatever that is GPU or TPU or whatnot. It's going to fetch it back to the host memory. So that's why we have the MP suffix here. And then I just do some simple aggregation here across the batches and I return the final state and all of those metrics back from the function. Evaluate model fairly similar. We just call the eval step just once here. We don't have to loop because if you recall I'm loading the whole test data set and so a single line of code here will work. We do the same thing device get to pull on to host and then we have to do this dot item because this is just going to contain a NumPy arrays with a single element. So you have to kind of call the dot item to get the scalar from the NumPy array. And that's it. We return back the metrics. Okay. Finally, these are the two last functions that we need before we see the main training loop. And this one is fairly important. So this one makes things I mean it makes things it's a syntactic sugar but it makes things way cleaner and thus less error prone. So what we do here is we instantiate the CNN. We basically initialize the parameters. We pass the key. We pass the dummy data. So just we pass all ones and then we extract from that frozen dict we just extract the params collection because that's the format that this train state here will expect. And then we just kind of call and form this SGT optimizer without having to optimize it. Without having to initialize it because that's going to happen inside of this train state class. So that's cool. We're kind of delegating a lot of the functionality there. And finally, we also pass the CNN apply so that we can then do some more advanced logic as we saw there in the loop. So here we call the apply gradients. We can do that because we have all of those ingredients inside. Okay. Let me kind of quickly walk you through that class so that it's kind of clear what's going on. Let me open it up. If I open it up and here is the train state function. So we have the apply gradients. We saw that one in the training loop and we have the create function. So basically you can see here the optimizer is initialized internally. So yeah, like as much work as possible was delegated to this function so that we don't have to externally do the initialization. And here you can see what it does when we call the apply gradients. It basically does those two optics functions. So we have the update. We get the updates in the new optimizer state and then we apply the updates. It's just a simple, simple function that makes things a bit cleaner in our code. So that's it. Let me go back here. So let me run this if I haven't already. And we are getting to the final cell of this notebook where we have the training loop. We create the training state. We have it here and you can see how neat the training loop actually looks like. So we just call the train for one epoch. So for every epoch, we just call the train one epoch and we get back from the trains. We convert the train state. So there's going to be modified train state. We're going to get the metrics. We plot the metrics like the loss and accuracy, et cetera. And we do the same thing on the test dataset. So we pass the test images and test labels and that's pretty much it. So I can try and run this, but currently I don't think I have an access to a GPU. So it's going to be a bit slower, but yeah, it's just going to print every single epoch. It's going to print some of this data. And yeah, that's everything you have to know pretty much. And in my opinion, this is probably this example, even though it's way more complex and actually does a real thing, it was way simpler compared to those contrived examples where we had to do the mutable, where we had to pass the random states for the dropout and stuff like that. So as an exercise, how would we go about adding dropout here and also batch norm? So that would like start adding that complexity on the syntactic level. Let's roughly think what would happen if we were to add dropout. So first we'd have to modify this create train state function and we'd have to add not just the key for the params, but also for the, we'd have to create a kind of dictionary here where we'd have like params and then we'd initialize that collection with a certain key and then we'd have the dropout and we'd have to initialize that with another key. And that's on the initialization side. Then in the actual training loop, what we'd have to do is we'd have to pass the key for every single train step, we'd have to have a unique key. So because as you recall, we'd have to add RNGs here and then create and pass the key here. So that means that for every single train step, we'll have to have a key and let me see where we're calling the train step right here. So that means we have to have splitting happening on this line here. So we just have to pass a specific key to the train one epoch function and then kind of keep splitting because we don't want to reuse the key because if you reuse the key, you'll get the same output, the same randomness and thus you don't have this stochasticity. So that's very important to keep in mind. So yeah, it's a fairly simple modification. I just wanted to kind of quickly walk you through that and for batch norm similarly, you'd have to have the mutable keyword here and there and it would complicate things a little bit, but definitely not by a lot. So yeah, let's see whether we have some outputs. Yeah, we have some outputs here and you can see the accuracy is already and we don't have any overfitting and so yeah, we have like a 98% accuracy. That's fairly nice. That's pretty much it. I encourage you to go and check out the Flex documentation A and B. Go and check out the examples they have on their GitHub. So you can see like this ImageNet example is very cool. They have mixed precision training. They have like training over multiple devices, then like all of that additional complexity and I think they can train a model in like an hour or something. Yeah, like wall time, a couple of hours to train these, although you have to have a bunch of TPUs or V100s to do that. So yeah, go check those out at your own pace and I think I mentioned in the beginning, Hugging Face has very nice examples as well and they've organized this community week which you should check out as well. So I'm gonna link all of these resources in the video description, so do check them out as well and yeah, I mean Hugging Face team is doing an amazing job, so just a huge shout out to those guys. Cool, that's pretty much it. Let's do a quick recap. Let's do a short summary of what we've learned, what we've seen in this video about Flex and basically I'm gonna open up this table of content in deep note here to kind of guide my thinking and let me create a code cell here using Control J shortcut. So basically we've seen the, so on top of everything you need to know about Jax already, we've seen the init and apply functions. We've seen that if you wanna create our custom modules, sometimes you'll have to use the self-param and self-variable functions. Then we've seen stuff such as nn-compact which makes the construction of those custom modules a bit more compact obviously. Then we saw some new keywords that we haven't seen in Jax such as mutable if you wanna have the non-trainable variables or RNGs if you wanna have like stochasticity such as dropout like in the forward pass and what else? We've seen the train state object which is pretty much like a wrapper state that makes things a little bit cleaner. So all in all, there is not that much you need to know on top of Jax, so if you already have the basics kinda set up, then you'll be good to go and this video will hopefully give you enough context to get you started. Again, huge shout out to the Deep Note team for sponsoring this video and if you like this one, share it out with your friends. That's the best way you can support the channel. Also check out the reference. I'm gonna link below in the video description. By using the link again, you'll get 20 hours for free using the, you can use these pro machines and that's it. Subscribe to this channel and until next time, bye bye.
[{"start": 0.0, "end": 7.5, "text": " What's cracking you guys? In this video I'm resuming my Jax series of videos and in this one I'm gonna cover Flex."}, {"start": 7.5, "end": 18.8, "text": " Basically a neural network library built on top of Jax and you probably are already aware of frameworks such as TensorFlow, Keras, PyTorch."}, {"start": 18.8, "end": 23.3, "text": " So then like why is there a need for yet another deep learning framework?"}, {"start": 23.3, "end": 32.0, "text": " And the answer would be like twofold probably. One of the most important things I think is the performance component."}, {"start": 32.0, "end": 39.0, "text": " We're gonna later see some numbers on certain benchmarks and Flex kinda was the best, the fastest."}, {"start": 39.0, "end": 43.400000000000006, "text": " I.e. in general like Jax is faster than the other alternatives."}, {"start": 43.4, "end": 53.8, "text": " And B. you get better reproducibility because of the way that Jax and thus Flex and some other frameworks such as Haiku which is also built on top of Jax."}, {"start": 53.8, "end": 61.599999999999994, "text": " How they handle state. And basically because they're using the functional programming paradigm, they handle state externally."}, {"start": 61.599999999999994, "end": 71.4, "text": " They deal with pure functions and so basically even the random number generators handle state externally and all of that leads to better reproducibility."}, {"start": 71.4, "end": 76.2, "text": " So those are some whys behind learning Flex and yeah."}, {"start": 76.2, "end": 86.4, "text": " So before you start watching this video you should have some knowledge of Jax already and if you don't know what Jax is, if you've never done any Jax,"}, {"start": 86.4, "end": 91.7, "text": " I strongly encourage you to check out my Machine Learning with Jax series of videos."}, {"start": 91.7, "end": 96.9, "text": " You can see here the first one from zero to hero which is going to teach you the basics."}, {"start": 96.9, "end": 105.4, "text": " Then we have from hero to hero pro plus which is going to go beyond and teach you like a nitty-gritty details of Jax in general."}, {"start": 105.4, "end": 110.9, "text": " And finally I have a video where I'm covering where I'm coding a neural network from scratch in Pure Jax."}, {"start": 110.9, "end": 117.9, "text": " So yeah, if you don't know anything about Jax, maybe go and check out at least the tutorial number one before you watch this one."}, {"start": 117.9, "end": 122.7, "text": " Okay, now you maybe notice something different in this video."}, {"start": 122.7, "end": 134.7, "text": " I'm not using colab anymore and what this is is deep note and I first want to give up like a huge thank you to the deep note team for deciding to sponsor this video."}, {"start": 134.7, "end": 145.1, "text": " So I want to briefly mention that like I have so far decided to decline like multiple sponsorship opportunities and the reason is I really want to promote something that I truly believe in myself."}, {"start": 145.1, "end": 155.7, "text": " So that means I should either use the product like on a daily or on a weekly basis myself or I think that the product has a lot of value and deep note definitely checks the marks."}, {"start": 155.7, "end": 157.4, "text": " And so yeah, here I am."}, {"start": 157.4, "end": 159.4, "text": " So what's the thing with deep note?"}, {"start": 159.4, "end": 164.0, "text": " If you take a look at it, it looks like a regular Jupyter notebook like your colab notebook."}, {"start": 164.0, "end": 165.0, "text": " So what's the difference?"}, {"start": 165.0, "end": 167.79999999999998, "text": " So you can see here like you've got your markdown cells."}, {"start": 167.79999999999998, "end": 171.6, "text": " You can click shift enter and render the markdown cell."}, {"start": 171.6, "end": 173.29999999999998, "text": " You have the code cells."}, {"start": 173.3, "end": 178.10000000000002, "text": " Well, the one of the most important things about deep note is the collaborative feature."}, {"start": 178.10000000000002, "end": 191.60000000000002, "text": " So that basically means you can share this notebook with multiple collaborators with a team and you can real time collaborate on the notebook and you can basically see other people like typing real time commenting, etc."}, {"start": 191.60000000000002, "end": 193.20000000000002, "text": " It's a very good collaborative tool."}, {"start": 193.20000000000002, "end": 198.9, "text": " And unfortunately, I cannot show you right now those features because I'm so here."}, {"start": 198.9, "end": 204.6, "text": " Secondly, it integrates which bunch of different data sources you can check out those yourself."}, {"start": 204.6, "end": 212.3, "text": " But like drive, Amazon s3, Redshift, BigQuery, a bunch of other applications with GitHub as well."}, {"start": 212.3, "end": 217.9, "text": " They also provide templates which can get you started like you can maybe pick this A B testing."}, {"start": 217.9, "end": 219.3, "text": " Let me show you this one."}, {"start": 219.3, "end": 224.9, "text": " So it's got a very cool widgets, cool dashboards, and you can even create deep note apps."}, {"start": 224.9, "end": 228.20000000000002, "text": " So let me show you briefly how this thing looks like."}, {"start": 228.2, "end": 230.39999999999998, "text": " You can see there is like a lot of these widgets."}, {"start": 230.39999999999998, "end": 233.2, "text": " You can kind of create these interactive plots."}, {"start": 233.2, "end": 238.29999999999998, "text": " And if you go and hit this publishing editor, you can even go and create a dashboard."}, {"start": 238.29999999999998, "end": 253.2, "text": " And so that means no code, which means you can kind of can basically use this as a data science notebook and present the results to your like business partners or whoever is not tech savvy, but wants to see the results and get some insights."}, {"start": 253.2, "end": 254.79999999999998, "text": " So yeah, that's very cool."}, {"start": 254.79999999999998, "end": 258.0, "text": " So let me see what else I think I've covered."}, {"start": 258.0, "end": 260.9, "text": " The main things you can check it out in your own pace."}, {"start": 260.9, "end": 262.2, "text": " You've even got terminals."}, {"start": 262.2, "end": 266.7, "text": " You can kind of edit the environment if for whatever reason you need to do that."}, {"start": 266.7, "end": 269.6, "text": " You can you can track the common history of the whole team."}, {"start": 269.6, "end": 273.1, "text": " In general, the whole history of what happened is fairly nicely tracked."}, {"start": 273.1, "end": 277.0, "text": " So and finally, let me get back to my notebook."}, {"start": 277.0, "end": 279.9, "text": " You can see I've installed some of the packages here."}, {"start": 279.9, "end": 280.8, "text": " You don't have to do that."}, {"start": 280.8, "end": 283.6, "text": " You can create a fully custom Docker container."}, {"start": 283.6, "end": 285.1, "text": " They have a support for that."}, {"start": 285.1, "end": 289.6, "text": " I even created some of the Docker images myself here, but I won't be using that."}, {"start": 289.6, "end": 292.70000000000005, "text": " I'm going to be using this 3.7, which is the same one."}, {"start": 292.70000000000005, "end": 297.5, "text": " I recommend you use if you want to start with this specific notebook."}, {"start": 297.5, "end": 302.1, "text": " And finally, I've got a referral link for you guys down in the video description."}, {"start": 302.1, "end": 312.6, "text": " If you want to support the channel and also if you want to get 20 additional hours of these pro machines on deep node, do click on the referral link and do check out deep node."}, {"start": 312.6, "end": 314.8, "text": " So so yeah, let's get started."}, {"start": 314.8, "end": 333.90000000000003, "text": " The point of this video is to learn to get some basic understanding of flex so that you can build train your own neural networks and then you can take it from there and start reading the documentation, which is probably the best source of information right now because the community is not as strong as like say pie torts."}, {"start": 333.90000000000003, "end": 341.2, "text": " Even though it's growing and it's being adopted by hugging face, they even had a community week where they were coded model in flex."}, {"start": 341.2, "end": 345.5, "text": " I'm going to link some of those in the in the video description so you can check those out."}, {"start": 345.5, "end": 348.2, "text": " But yeah, let's get started here."}, {"start": 348.2, "end": 356.8, "text": " Basically, I'm going to import here some some some like Jax and certain libraries, which is non-pie from Jax."}, {"start": 356.8, "end": 362.2, "text": " So Jax, Numbay, and then we're going to import the flex, which is the library we're going to learn today."}, {"start": 362.2, "end": 376.8, "text": " And it's a it's a library developed by Google Research, in particular, Google Brain, and it's a the name itself is a contraction between like you have flexibility, you have Jax, and then you kind of get flex very creative."}, {"start": 376.8, "end": 378.7, "text": " I'm going to import some functions here."}, {"start": 378.7, "end": 384.4, "text": " They have this leaning API, which is I think the second iteration of their API."}, {"start": 384.4, "end": 393.59999999999997, "text": " And you can see that the convention is to import the API itself as an N, which you may be familiar with if you come from the pie torch background."}, {"start": 393.59999999999997, "end": 396.59999999999997, "text": " It's a fairly common notation."}, {"start": 396.59999999999997, "end": 403.9, "text": " Okay, train state is just a wrapper function that's going to help us make things a bit cleaner."}, {"start": 403.9, "end": 405.4, "text": " We're going to see that in a moment."}, {"start": 405.4, "end": 406.79999999999995, "text": " Yeah, I think I mentioned Haiku."}, {"start": 406.79999999999995, "end": 412.0, "text": " It's it's also like a library built on top of Jax that comes from DeepMind."}, {"start": 412.0, "end": 415.8, "text": " And the only reason I have it here is for comparison purposes."}, {"start": 415.8, "end": 419.6, "text": " So we won't be learning Haiku in this video, although I'll cover it in the next one."}, {"start": 419.6, "end": 421.8, "text": " So stay tuned for that."}, {"start": 421.8, "end": 429.1, "text": " I'm importing Optex here, which is like another library built on top of Jax for Optimizers."}, {"start": 429.1, "end": 432.5, "text": " And the rationale here is the following."}, {"start": 432.5, "end": 439.2, "text": " You want to have in this Jax ecosystem, every one of these libraries takes care of a specific vertical."}, {"start": 439.2, "end": 447.4, "text": " So you have Optex for Optimizers, you've got Jax, you've got Flex and Haiku for like neural network building and construction."}, {"start": 447.4, "end": 449.9, "text": " You've got checks for testing the networks."}, {"start": 449.9, "end": 454.7, "text": " You've got libraries for reinforcement learning, libraries for graph neural networks, etc."}, {"start": 454.7, "end": 460.9, "text": " So that's like kind of the paradigm, the mindset of the Jax ecosystem, so to speak."}, {"start": 460.9, "end": 468.59999999999997, "text": " Finally, I'm going to use a pie torch data loading functionalities here, because if you recall from my Jax tutorials,"}, {"start": 468.6, "end": 471.1, "text": " they did not want to reinvent the wheel here."}, {"start": 471.1, "end": 475.0, "text": " They think that data loading functionality is already good."}, {"start": 475.0, "end": 480.0, "text": " And so they're just kind of utilizing either pie torch or tensorflow data sets, etc."}, {"start": 480.0, "end": 485.40000000000003, "text": " I'm going to use a pie torch because that's what I'm the most familiar with."}, {"start": 485.40000000000003, "end": 490.20000000000005, "text": " Finally, I'm importing some functions here and some common modules."}, {"start": 490.20000000000005, "end": 494.40000000000003, "text": " I already ran these cells because, yeah, just to save some time."}, {"start": 494.4, "end": 500.4, "text": " But now starting from here, we're going to run every single cell and see what is going on."}, {"start": 500.4, "end": 511.29999999999995, "text": " So let's start by building the most simple model, which is like just a simple, this is like a feed forward layer."}, {"start": 511.29999999999995, "end": 516.9, "text": " You can also think about it as a linear regression because we don't have any nonlinear activation functions."}, {"start": 516.9, "end": 519.9, "text": " And so this is just a linear regression model."}, {"start": 519.9, "end": 521.6999999999999, "text": " So let me run this."}, {"start": 521.7, "end": 525.1, "text": " So you just kind of click shift enter the same as in Colab."}, {"start": 525.1, "end": 529.1, "text": " And the notation is dense instead of linear."}, {"start": 529.1, "end": 538.1, "text": " And you can see that the dense class inherits from this module class, which is the same analogous to pie torch."}, {"start": 538.1, "end": 545.5, "text": " And the idea is even though this is a functional programming like environment, like in its core,"}, {"start": 545.5, "end": 551.8, "text": " they tried and built like an objective oriented paradigm wrapper around it. And that's what Flex is."}, {"start": 551.8, "end": 553.9, "text": " And that's what Haiku is."}, {"start": 553.9, "end": 563.7, "text": " So that's you as a programmer are kind of programming in your usual using your usual mental model like the object oriented programming."}, {"start": 563.7, "end": 565.0, "text": " OK, let's continue here."}, {"start": 565.0, "end": 568.5, "text": " So let's see how we can do inference using this model."}, {"start": 568.5, "end": 571.6, "text": " And the answer would be two steps."}, {"start": 571.6, "end": 574.9, "text": " We've got this in that function and we've got the apply function."}, {"start": 574.9, "end": 580.5, "text": " The point of init is to initialize the initial parameters because as you recall,"}, {"start": 580.5, "end": 587.3, "text": " Jax and thus Flex by transitivity property is handling like state externally."}, {"start": 587.3, "end": 589.5, "text": " So that's why we need init and apply."}, {"start": 589.5, "end": 591.1, "text": " And let's see what we do here."}, {"start": 591.1, "end": 594.8, "text": " Basically, hopefully you can see the text reasonably well."}, {"start": 594.8, "end": 598.1, "text": " I'm going to zoom in a little bit more for just in case."}, {"start": 598.1, "end": 601.3, "text": " But the UI is a bit uglier."}, {"start": 601.3, "end": 604.3, "text": " Hopefully you'll see the text, which is the point here."}, {"start": 604.3, "end": 606.5, "text": " OK, we have some random seed here."}, {"start": 606.5, "end": 608.3, "text": " Your regular Jax thing here."}, {"start": 608.3, "end": 614.5, "text": " We are externally handling the random state of the random number generator and we split it into two keys."}, {"start": 614.5, "end": 624.3, "text": " Then we just instantiate a random like data point here using the key one is just a 10 dimensional vector."}, {"start": 624.3, "end": 631.1999999999999, "text": " And then what we do is we take the model, we call this init function and we pass the key."}, {"start": 631.2, "end": 635.9000000000001, "text": " And the key is used obviously because initialization itself is random process."}, {"start": 635.9000000000001, "end": 639.5, "text": " So we want to have key for reproducibility reasons."}, {"start": 639.5, "end": 647.5, "text": " We pass the input X and basically we get back the parameters and the output in the case of using this function here."}, {"start": 647.5, "end": 650.3000000000001, "text": " In it with output, not that it's that important."}, {"start": 650.3000000000001, "end": 651.6, "text": " I'm going to later just using it."}, {"start": 651.6, "end": 653.5, "text": " I'm going to show in a second."}, {"start": 653.5, "end": 658.7, "text": " But yeah, then you can just we can visualize the shape of the parameters and see that it"}, {"start": 658.7, "end": 662.5, "text": " matches our expectation for a simple linear model."}, {"start": 662.5, "end": 668.5, "text": " So let me run the cell and you can see here that the output is exactly what we expect."}, {"start": 668.5, "end": 676.9000000000001, "text": " Bias term has is the dimensionality is five because the as we call here, we have five output neurons"}, {"start": 676.9000000000001, "end": 681.1, "text": " and we have 10 to five for the weight for the kernel."}, {"start": 681.1, "end": 683.7, "text": " And you can always a couple of things here."}, {"start": 683.7, "end": 690.9000000000001, "text": " So first is the automatic shape inference because I haven't specified 10 anywhere here."}, {"start": 690.9000000000001, "end": 693.5, "text": " I've just specified the output dimensionality."}, {"start": 693.5, "end": 698.6, "text": " And so the shape is automatically inferred through this input variable X."}, {"start": 698.6, "end": 700.1, "text": " So that's what happened there."}, {"start": 700.1, "end": 706.1, "text": " The second thing you can notice here is this immutable structure because as you recall again,"}, {"start": 706.1, "end": 708.5, "text": " Jax loves dealing with immutable objects."}, {"start": 708.5, "end": 715.2, "text": " So what flex has done here is just made it easier for Jax to not introduce any box here by"}, {"start": 715.2, "end": 717.6, "text": " returning immutable structures."}, {"start": 717.6, "end": 720.4, "text": " So that's why we have the frozen dict here."}, {"start": 720.4, "end": 724.3, "text": " Finally, I'm using in it with output here, but you usually use just in it."}, {"start": 724.3, "end": 725.9, "text": " You don't care about the output initially."}, {"start": 725.9, "end": 731.2, "text": " So I can just run in it here and then I'm going to delete the Y here."}, {"start": 731.2, "end": 734.1, "text": " I'm going to comment this out and we'll get the same result."}, {"start": 734.1, "end": 738.4, "text": " So yeah, usually you'll see the in it function again."}, {"start": 738.4, "end": 744.6, "text": " To recap here, what this thing does here is it's iterating through the supremes of pie"}, {"start": 744.6, "end": 749.6, "text": " tree and a pie tree object, which is a term from from Jax again."}, {"start": 749.6, "end": 751.1, "text": " And basically you're iterating."}, {"start": 751.1, "end": 755.9, "text": " You're kind of traversing the tree and for every single leaf you're applying this function"}, {"start": 755.9, "end": 762.6999999999999, "text": " here, which means you'll be taking like, let's say you have bias and you have some like a"}, {"start": 762.6999999999999, "end": 764.4, "text": " Jax lump array in there."}, {"start": 764.4, "end": 765.3, "text": " You're going to take it."}, {"start": 765.3, "end": 767.3, "text": " You're going to take its shape and you're going to change."}, {"start": 767.3, "end": 771.4, "text": " You're basically going to map the tree of parameters into tree of shapes."}, {"start": 771.4, "end": 772.9, "text": " So that's what you what you get here."}, {"start": 772.9, "end": 775.5, "text": " So yeah, just a small, small recap there."}, {"start": 775.5, "end": 778.5, "text": " Okay, so that should be fairly familiar."}, {"start": 778.5, "end": 783.4, "text": " So the new thing is the in it function and the second thing is this apply function."}, {"start": 783.4, "end": 784.9, "text": " So let's see what happens here."}, {"start": 784.9, "end": 789.6999999999999, "text": " So we have model we call apply and we pass the parameters and the input X."}, {"start": 789.6999999999999, "end": 795.5999999999999, "text": " So this is the only thing that you have to kind of that you have to adjust your mental"}, {"start": 795.6, "end": 801.6, "text": " model to this new paradigm and it may be very weird initially, but like believe me, I started"}, {"start": 801.6, "end": 806.8000000000001, "text": " playing with Jax in November last year 2021 and I'm already feeling very, very comfortable"}, {"start": 806.8000000000001, "end": 808.6, "text": " with it with this much."}, {"start": 808.6, "end": 814.5, "text": " I feel much better than when I first encountered this this kind of signature."}, {"start": 814.5, "end": 820.2, "text": " So you call apply you get the output and we can run this and get the results as expected"}, {"start": 820.2, "end": 822.0, "text": " five outputs here."}, {"start": 822.0, "end": 826.0, "text": " So again, you cannot call the model like this."}, {"start": 826.0, "end": 830.3, "text": " So this is what you're used to in other frameworks, but here this is not going to work."}, {"start": 830.3, "end": 831.3, "text": " It's going to throw an exception."}, {"start": 831.3, "end": 836.3, "text": " We're going to catch it and print something can call compact methods on unbound modules."}, {"start": 836.3, "end": 839.8, "text": " So the details themselves are not that important here."}, {"start": 839.8, "end": 843.5, "text": " Okay, so that's the very, very basics of flex."}, {"start": 843.5, "end": 849.4, "text": " Now I want to do a small coding exercise here and contrast flex with Haiku and in my personal"}, {"start": 849.4, "end": 853.3, "text": " opinion Haiku is a bit cleaner than flex."}, {"start": 853.3, "end": 858.5, "text": " I like it better me personally, although it may happen the flex is more powerful because"}, {"start": 858.5, "end": 864.56, "text": " of the additional like notation differences you'll see soon see."}, {"start": 864.56, "end": 870.24, "text": " So yeah, I'm not very strongly opinionated about this, but so far for me it seems Haiku"}, {"start": 870.24, "end": 874.4399999999999, "text": " is even nicer than flex even though flex is way more popular in the community."}, {"start": 874.4399999999999, "end": 876.4, "text": " I mentioned hugging face community week."}, {"start": 876.4, "end": 881.48, "text": " I mentioned they have even have like 5,000 plus models on hugging face hub, which is"}, {"start": 881.48, "end": 886.68, "text": " fairly, fairly nice amount of models for a relatively new framework compared to TensorFlow"}, {"start": 886.68, "end": 889.04, "text": " or like PyTorch."}, {"start": 889.04, "end": 892.6, "text": " So let me show you how you would do the same thing in Haiku now."}, {"start": 892.6, "end": 894.16, "text": " It's going to be very brief."}, {"start": 894.16, "end": 896.1999999999999, "text": " So let's create a model."}, {"start": 896.1999999999999, "end": 904.86, "text": " So Haiku uses this HK as a, as a, as a like a notation then linear and then five."}, {"start": 904.86, "end": 908.16, "text": " Now if I were to run this, it's going to fail."}, {"start": 908.16, "end": 909.52, "text": " So let me try and run this."}, {"start": 909.52, "end": 916.1800000000001, "text": " And as you can see here, all of these modules must be initialized inside an HK Haiku transform"}, {"start": 916.1800000000001, "end": 917.3000000000001, "text": " function."}, {"start": 917.3000000000001, "end": 922.16, "text": " So what the trick is Haiku is very explicit about when you want to, you take a function"}, {"start": 922.16, "end": 926.52, "text": " like linear and you, in order to have the init and apply."}, {"start": 926.52, "end": 928.36, "text": " So it's using the same notation as flex."}, {"start": 928.36, "end": 930.6800000000001, "text": " You have to pass it into this HQ transform function."}, {"start": 930.6800000000001, "end": 932.44, "text": " So let me do that."}, {"start": 932.44, "end": 933.72, "text": " Transform."}, {"start": 933.72, "end": 939.8000000000001, "text": " And what we do here, we create a Lambda anonymous function here and we just simply call this,"}, {"start": 939.8000000000001, "end": 942.8000000000001, "text": " this, this, this linear model."}, {"start": 942.8000000000001, "end": 943.8000000000001, "text": " And that's it."}, {"start": 943.8000000000001, "end": 944.84, "text": " Now this will work."}, {"start": 944.84, "end": 947.44, "text": " So if I were to run this, it will not fail."}, {"start": 947.44, "end": 950.1, "text": " But now let me show you what else happens."}, {"start": 950.1, "end": 955.6800000000001, "text": " So let me just copy paste these parts here because they remain the same, the seeds and"}, {"start": 955.6800000000001, "end": 958.76, "text": " the inputs will remain the same."}, {"start": 958.76, "end": 961.6, "text": " Now let me initialize the parameters."}, {"start": 961.6, "end": 968.16, "text": " So we get model init again, we pass the key, we pass the X and we get parameters."}, {"start": 968.16, "end": 969.16, "text": " So so far so good."}, {"start": 969.16, "end": 972.6, "text": " It looks very, very similar to two flex."}, {"start": 972.6, "end": 978.16, "text": " And finally, let's get the output by calling model apply."}, {"start": 978.16, "end": 981.44, "text": " We pass the parameters here."}, {"start": 981.44, "end": 987.2, "text": " We pass none because we don't have any stochasticity in the forward pass."}, {"start": 987.2, "end": 992.6800000000001, "text": " And for now it may be, it may look as high because even more complex and a kind of uglier"}, {"start": 992.6800000000001, "end": 997.0400000000001, "text": " in the notation compared to two, two, two flex, but it's not the case."}, {"start": 997.0400000000001, "end": 999.0200000000001, "text": " You'll see that in the next video."}, {"start": 999.0200000000001, "end": 1000.6, "text": " So you'll do it like this."}, {"start": 1000.6, "end": 1002.4000000000001, "text": " And by the way, I made a small mistake here."}, {"start": 1002.4000000000001, "end": 1003.9200000000001, "text": " You don't want to reuse keys."}, {"start": 1003.9200000000001, "end": 1005.88, "text": " So this should be key two."}, {"start": 1005.88, "end": 1009.76, "text": " So I should be able to run this right now and everything should work."}, {"start": 1009.76, "end": 1012.44, "text": " So this is, this is high Q very similar."}, {"start": 1012.44, "end": 1017.0, "text": " You just have this transform function and, and yeah, so that's cool."}, {"start": 1017.0, "end": 1018.76, "text": " I had the solutions here."}, {"start": 1018.76, "end": 1022.6, "text": " So let me kind of show this code and hopefully this looks the same."}, {"start": 1022.6, "end": 1027.96, "text": " So yeah, we have the model, we form the data here, we call in it, we call apply and we"}, {"start": 1027.96, "end": 1028.96, "text": " print."}, {"start": 1028.96, "end": 1032.0, "text": " And again, we have this, I forgot to do this."}, {"start": 1032.0, "end": 1038.08, "text": " Basically it just shows that high Q follows the same paradigm as flex."}, {"start": 1038.08, "end": 1040.5, "text": " It's inheriting from this module class as well."}, {"start": 1040.5, "end": 1044.1, "text": " So let me just hide this code again."}, {"start": 1044.1, "end": 1046.56, "text": " So hide the code and that's it."}, {"start": 1046.56, "end": 1047.56, "text": " Cool."}, {"start": 1047.56, "end": 1050.36, "text": " That was section number one."}, {"start": 1050.36, "end": 1055.6799999999998, "text": " Hopefully you're not too overwhelmed because as I said, you will have to kind of get accustomed"}, {"start": 1055.6799999999998, "end": 1059.24, "text": " with the, with this functional programming paradigm."}, {"start": 1059.24, "end": 1062.3999999999999, "text": " But yeah, it gets easier after, after, after a quick."}, {"start": 1062.3999999999999, "end": 1069.12, "text": " So as I said, flex and high Q offer a performance boost."}, {"start": 1069.12, "end": 1071.96, "text": " So here's some numbers from the hugging face team."}, {"start": 1071.96, "end": 1078.18, "text": " You can see here using a single GPU and using a flex you get compared to PyTorch, you get"}, {"start": 1078.18, "end": 1079.52, "text": " certain improvements."}, {"start": 1079.52, "end": 1083.6000000000001, "text": " Like in all of these, I think these are, all of these are NLP tasks."}, {"start": 1083.6000000000001, "end": 1086.58, "text": " So you have Cola, you have the natural language inference task."}, {"start": 1086.58, "end": 1091.14, "text": " You have, yeah, semantic, I guess, yeah, whatever."}, {"start": 1091.14, "end": 1095.56, "text": " So in any case, you get a fairly nice improvements here speed wise."}, {"start": 1095.56, "end": 1100.96, "text": " And in total, you can see that here you've kind of saved more than an hour using a single"}, {"start": 1100.96, "end": 1106.48, "text": " GPU, which is not a small thing, especially when you scale these two big language, large"}, {"start": 1106.48, "end": 1107.48, "text": " language models."}, {"start": 1107.48, "end": 1108.48, "text": " Okay."}, {"start": 1108.48, "end": 1113.32, "text": " Let's get going and let's introduce a single, a simple toy example where we're going to"}, {"start": 1113.32, "end": 1119.64, "text": " just form like a two dimensional data set and try and train a simple linear regression"}, {"start": 1119.64, "end": 1126.6000000000001, "text": " model using the NN dense model, which is an implementation and instantiation of the linear"}, {"start": 1126.6000000000001, "end": 1127.88, "text": " regression model."}, {"start": 1127.88, "end": 1131.4, "text": " So I'm going to have 150 data points."}, {"start": 1131.4, "end": 1133.4, "text": " The X dimension is two."}, {"start": 1133.4, "end": 1140.3600000000001, "text": " The output is going to be, we are regressing like a scalar just to make things a plottable,"}, {"start": 1140.3600000000001, "end": 1145.0400000000002, "text": " if that's even a word, because yeah, if we had multiple dimensions here, I wouldn't be"}, {"start": 1145.0400000000002, "end": 1146.0400000000002, "text": " able to plot."}, {"start": 1146.0400000000002, "end": 1147.92, "text": " I'd have to do some dimensionality reduction."}, {"start": 1147.92, "end": 1149.5200000000002, "text": " So I'm going to skip that."}, {"start": 1149.5200000000002, "end": 1156.72, "text": " Again, we're going to form some keys here being explicit about the random state as usual."}, {"start": 1156.72, "end": 1161.76, "text": " We are going to form the ground truth weights and biases here."}, {"start": 1161.76, "end": 1167.4, "text": " So basically we'll be trying to learn this weight and this bias later during the training."}, {"start": 1167.4, "end": 1171.04, "text": " And now what we do here is we form a structure."}, {"start": 1171.04, "end": 1174.52, "text": " That's, that's something that flex is, is used to."}, {"start": 1174.52, "end": 1180.0, "text": " So you have the freeze is supposed to create a frozen dict, a frozen dictionary out of"}, {"start": 1180.0, "end": 1183.2, "text": " the normal Python native dictionary here."}, {"start": 1183.2, "end": 1188.52, "text": " As you can recall from, from this section, we had frozen dict here."}, {"start": 1188.52, "end": 1190.8400000000001, "text": " So let me get back here."}, {"start": 1190.8400000000001, "end": 1195.8600000000001, "text": " And the structure is such that we have this params collection as flex calls it."}, {"start": 1195.8600000000001, "end": 1200.0, "text": " And then we have bias B and kernel W and that's it."}, {"start": 1200.0, "end": 1204.24, "text": " Now we, we've used this ground truth model and generate the data."}, {"start": 1204.24, "end": 1210.44, "text": " So again, we split the key, which we created here, we split it and we get the key that"}, {"start": 1210.44, "end": 1214.16, "text": " we're going to use to generate the axis, the domain of our data."}, {"start": 1214.16, "end": 1219.0800000000002, "text": " So we have 150 samples of, so 150 to two doubles."}, {"start": 1219.0800000000002, "end": 1222.92, "text": " So that means just 150 two dimensional data points."}, {"start": 1222.92, "end": 1229.16, "text": " And then we're going to use it, use the ground truth model and regress the, the YS here."}, {"start": 1229.16, "end": 1233.68, "text": " And then we're just going to do a superposition of noise here using this amplitude."}, {"start": 1233.68, "end": 1234.68, "text": " That's pretty much it."}, {"start": 1234.68, "end": 1236.52, "text": " Let me run this now."}, {"start": 1236.52, "end": 1240.48, "text": " And you can see the shapes, 152 and 151."}, {"start": 1240.48, "end": 1244.16, "text": " And now let's plot the, the, the data set we just created."}, {"start": 1244.16, "end": 1245.24, "text": " And you can see it here."}, {"start": 1245.24, "end": 1250.2, "text": " Now I was lazy and so I didn't try and create an interactive plot here."}, {"start": 1250.2, "end": 1253.78, "text": " And so let's be smart about it."}, {"start": 1253.78, "end": 1256.24, "text": " And let's try and visualize where this data lies in."}, {"start": 1256.24, "end": 1258.08, "text": " And you probably already know the answer."}, {"start": 1258.08, "end": 1261.6399999999999, "text": " If not, let's do a simple exercise and try to approach it."}, {"start": 1261.6399999999999, "end": 1265.68, "text": " So first one is we know how we created, how we generated the data."}, {"start": 1265.68, "end": 1274.8, "text": " So basically what we've done here is we had Y is equal to W X plus some bias."}, {"start": 1274.8, "end": 1284.16, "text": " And because X is two dimensional, this basically turns into W one X one plus W two X two plus"}, {"start": 1284.16, "end": 1291.8400000000001, "text": " B. And you can see that this is basically your, your analytical formulation of a plane,"}, {"start": 1291.84, "end": 1297.08, "text": " of a 2D plane embedded into 3D space."}, {"start": 1297.08, "end": 1299.12, "text": " So that would be one approach."}, {"start": 1299.12, "end": 1301.8799999999999, "text": " The second approach is data driven as a data driven."}, {"start": 1301.8799999999999, "end": 1306.6799999999998, "text": " I mean, what I'm going to do here is increase from 150 to 1500."}, {"start": 1306.6799999999998, "end": 1309.9199999999998, "text": " Okay, let's run this cell now and see what we get."}, {"start": 1309.9199999999998, "end": 1314.76, "text": " So we get 1500 data points and let's run this again."}, {"start": 1314.76, "end": 1319.8, "text": " And hopefully now you can see that it lies, it lies on the, on the 2D plane here."}, {"start": 1319.8, "end": 1323.1599999999999, "text": " And yeah, the reason we added noise, hopefully that's, that's clear enough is to kind of"}, {"start": 1323.1599999999999, "end": 1328.12, "text": " hide the actual, the ground to the, the actual model that, that lies in the background that"}, {"start": 1328.12, "end": 1329.68, "text": " generated this data."}, {"start": 1329.68, "end": 1335.6399999999999, "text": " And now we're going to try next up to, to kind of like a regress and infer the model"}, {"start": 1335.6399999999999, "end": 1337.24, "text": " that generated the data."}, {"start": 1337.24, "end": 1340.68, "text": " So first let's create this, this loss function."}, {"start": 1340.68, "end": 1345.54, "text": " This make MSC is just going to return this MSC loss function, which is a jitted version"}, {"start": 1345.54, "end": 1346.6, "text": " of the MSC."}, {"start": 1346.6, "end": 1349.56, "text": " So just kind of trying to optimize things a little bit here."}, {"start": 1349.56, "end": 1354.1599999999999, "text": " So squared error itself takes a single data point X and Y."}, {"start": 1354.1599999999999, "end": 1356.9199999999998, "text": " We do inference, we just do a prediction."}, {"start": 1356.9199999999998, "end": 1361.6799999999998, "text": " We pass in the single data point, we get a prediction and return the mean squared error"}, {"start": 1361.6799999999998, "end": 1362.6799999999998, "text": " here."}, {"start": 1362.6799999999998, "end": 1369.48, "text": " You may be confused by the inner function here and inner is used just because so as"}, {"start": 1369.48, "end": 1377.36, "text": " to be able to generalize two cases when like Y was a multi-dimensional vector and not a"}, {"start": 1377.36, "end": 1379.04, "text": " scalar like it is right now."}, {"start": 1379.04, "end": 1384.1599999999999, "text": " So this is just a generalization to you of our mean squared error loss."}, {"start": 1384.1599999999999, "end": 1387.24, "text": " Finally we take the square error and we call the V map."}, {"start": 1387.24, "end": 1392.1399999999999, "text": " So that's the transform from Jack's that makes this thing run on a batch of data."}, {"start": 1392.1399999999999, "end": 1397.6399999999999, "text": " So we pass now like a batch of our data, like the, let me, whoops, let me actually reduce"}, {"start": 1397.6399999999999, "end": 1398.8799999999999, "text": " from thousand five hundred."}, {"start": 1398.8799999999999, "end": 1401.1399999999999, "text": " No, I'm going to leave it at thousand five hundred."}, {"start": 1401.1399999999999, "end": 1402.78, "text": " It's still like a small data set."}, {"start": 1402.78, "end": 1404.94, "text": " So we're going to run this on the batch of data."}, {"start": 1404.94, "end": 1410.14, "text": " So this will give us like a thousand five hundred by one vector."}, {"start": 1410.14, "end": 1414.04, "text": " And then we're going to do mean over the X is equal zero."}, {"start": 1414.04, "end": 1418.8, "text": " So that means we're going to kind of average out all of those, of those loss losses for"}, {"start": 1418.8, "end": 1422.6000000000001, "text": " all of these singular data points to get the final loss."}, {"start": 1422.6000000000001, "end": 1423.92, "text": " And that's pretty much it."}, {"start": 1423.92, "end": 1426.28, "text": " That's how we get our MSC loss here."}, {"start": 1426.28, "end": 1431.6000000000001, "text": " And finally, I kind of wrap it into this value and grad transform function from Jack's to"}, {"start": 1431.6000000000001, "end": 1433.72, "text": " get the, yeah."}, {"start": 1433.72, "end": 1437.08, "text": " That's how we calculate gradients and you should be familiar with this already."}, {"start": 1437.08, "end": 1440.72, "text": " So now this will be our model."}, {"start": 1440.72, "end": 1445.48, "text": " So we're kind of cheating here a little bit because this is the exact same model that"}, {"start": 1445.48, "end": 1446.52, "text": " generated the data."}, {"start": 1446.52, "end": 1451.0, "text": " So we, we usually wouldn't have this advantage of knowing the model, but here for the sake"}, {"start": 1451.0, "end": 1456.68, "text": " of just demoing this, I'm going to suppose we, I'm going to assume we know the exact"}, {"start": 1456.68, "end": 1457.68, "text": " model."}, {"start": 1457.68, "end": 1460.96, "text": " So we're going to do the usual stuff here, initialize the model."}, {"start": 1460.96, "end": 1466.52, "text": " We get initial parameters and then we're going to do a simple training here, a training loop."}, {"start": 1466.52, "end": 1470.16, "text": " We're going to have 20 epochs, just an arbitrary number."}, {"start": 1470.16, "end": 1474.0, "text": " You can tweak it, but like it's enough for this toy problem."}, {"start": 1474.0, "end": 1479.48, "text": " We called, we called the, we just passed the params and we get the gradients and we get"}, {"start": 1479.48, "end": 1480.48, "text": " the loss."}, {"start": 1480.48, "end": 1486.22, "text": " So this may be like a dark magic a little bit, but the trick is this value and grad"}, {"start": 1486.22, "end": 1492.96, "text": " function already contains the MSC loss, which already contains our data because we had this"}, {"start": 1492.96, "end": 1500.38, "text": " wrapper function make MSC loss and the data excess and YS is already stored internally."}, {"start": 1500.38, "end": 1503.3, "text": " So that's why we don't have to even pass anything here."}, {"start": 1503.3, "end": 1508.76, "text": " So what this does is for the current set of parameters and for the car, for the dataset,"}, {"start": 1508.76, "end": 1515.1200000000001, "text": " which is static, it figures out the current gradients and the current loss for those parameters."}, {"start": 1515.12, "end": 1518.56, "text": " And then we just do a simple stochastic grid in the send here."}, {"start": 1518.56, "end": 1520.6, "text": " I'm being very explicit here."}, {"start": 1520.6, "end": 1522.56, "text": " You usually will not be doing it like this."}, {"start": 1522.56, "end": 1527.36, "text": " You're going to be using Optex, but here I'm being explicit and doing it in pure Jax way."}, {"start": 1527.36, "end": 1529.6799999999998, "text": " By the way, I'm having some internet connection problems."}, {"start": 1529.6799999999998, "end": 1531.08, "text": " Hopefully that will not interfere."}, {"start": 1531.08, "end": 1532.6399999999999, "text": " I think not, but yeah."}, {"start": 1532.6399999999999, "end": 1536.52, "text": " So you can see here, we do the tree multimap from Jax."}, {"start": 1536.52, "end": 1539.6799999999998, "text": " We basically have two pie trees here, params and grads."}, {"start": 1539.6799999999998, "end": 1544.0, "text": " We are going and traversing the leaves and we do the following simple manipulation from"}, {"start": 1544.0, "end": 1545.28, "text": " the current parameter."}, {"start": 1545.28, "end": 1551.44, "text": " We subtract the learning rate times the gradient, which is a definition of your simplest stochastic"}, {"start": 1551.44, "end": 1552.8, "text": " gradient descent."}, {"start": 1552.8, "end": 1554.12, "text": " And that's it."}, {"start": 1554.12, "end": 1558.36, "text": " So this is actually not stochastic gradient descent because we're using the whole data."}, {"start": 1558.36, "end": 1560.34, "text": " So this will be a batch gradient descent."}, {"start": 1560.34, "end": 1561.96, "text": " So it's exact."}, {"start": 1561.96, "end": 1565.64, "text": " We're not being stochastic here."}, {"start": 1565.64, "end": 1568.64, "text": " And finally, we are logging every now and then."}, {"start": 1568.64, "end": 1572.64, "text": " So every five epochs, I just print the epoch and then print the loss."}, {"start": 1572.64, "end": 1578.2, "text": " And yeah, if I were to run this shift enter, we can get, we can see the initial parameters"}, {"start": 1578.2, "end": 1583.1000000000001, "text": " which were initialized using this key, which is different compared to the way we've initialized"}, {"start": 1583.1000000000001, "end": 1584.1000000000001, "text": " the model here."}, {"start": 1584.1000000000001, "end": 1590.8400000000001, "text": " We use the specialized keys here and then we kind of wrap those into this freeze dictionary,"}, {"start": 1590.8400000000001, "end": 1591.8400000000001, "text": " et cetera."}, {"start": 1591.8400000000001, "end": 1592.8400000000001, "text": " Okay."}, {"start": 1592.8400000000001, "end": 1595.0400000000002, "text": " So those were the initial parameters."}, {"start": 1595.0400000000002, "end": 1598.3200000000002, "text": " Then we had loss going down towards zero."}, {"start": 1598.32, "end": 1604.04, "text": " And finally we get like, as you can see here, the final parameters, the learn parameters"}, {"start": 1604.04, "end": 1607.12, "text": " are here and you can see that they are pretty much the same."}, {"start": 1607.12, "end": 1613.36, "text": " We get two five four here, two five, very similar numbers, which is also indicated by"}, {"start": 1613.36, "end": 1614.96, "text": " the loss."}, {"start": 1614.96, "end": 1619.3999999999999, "text": " So yeah, that's how the training loop in Flex looks like."}, {"start": 1619.3999999999999, "end": 1620.3999999999999, "text": " Almost in Flex."}, {"start": 1620.3999999999999, "end": 1625.28, "text": " As I said, I'm using more, this is more like a puregex type of syntax here."}, {"start": 1625.28, "end": 1632.12, "text": " So let's now use OptX, which is a dedicated library for this purpose in particular, and"}, {"start": 1632.12, "end": 1633.28, "text": " let's initialize it."}, {"start": 1633.28, "end": 1640.48, "text": " So from OptX we call the SGD, we pass in the learning rates and then we have to again initialize"}, {"start": 1640.48, "end": 1644.44, "text": " the OptX because state has to be externally manipulated."}, {"start": 1644.44, "end": 1647.42, "text": " And here we get a state and we kind of print the state."}, {"start": 1647.42, "end": 1652.0, "text": " So let's run this cell and let's see what happens."}, {"start": 1652.0, "end": 1657.8, "text": " And as you can see, we've got a tuple of two empty states and that may not make sense if"}, {"start": 1657.8, "end": 1662.24, "text": " you're not familiar with how optimizers work, but if you are, this makes perfect sense because"}, {"start": 1662.24, "end": 1664.72, "text": " SGD doesn't have any state."}, {"start": 1664.72, "end": 1670.72, "text": " Compare that to Atom, which has to kind of keep the various statistics of gradients,"}, {"start": 1670.72, "end": 1676.2, "text": " like the variance of gradients in memory in state and that's why they're like, we have"}, {"start": 1676.2, "end": 1678.84, "text": " to like kind of explicitly manipulate it."}, {"start": 1678.84, "end": 1685.8, "text": " So if I were to change this to Atom, and yeah, I won't be changing the variable name, I'm"}, {"start": 1685.8, "end": 1687.6799999999998, "text": " just gonna change and use Atom."}, {"start": 1687.6799999999998, "end": 1688.6799999999998, "text": " Let me try and run this."}, {"start": 1688.6799999999998, "end": 1692.76, "text": " You can see that now we have like a meaningful state here and that's something you have to"}, {"start": 1692.76, "end": 1697.8799999999999, "text": " kind of pass externally here in Flex, in Jax in general."}, {"start": 1697.8799999999999, "end": 1704.3999999999999, "text": " So yeah, you can see what happens here is we have two set of parameters here."}, {"start": 1704.3999999999999, "end": 1708.1999999999998, "text": " First is the mu parameters and then we have the nu parameters."}, {"start": 1708.2, "end": 1715.52, "text": " And these are just your Greek alphabet letters that are denoting the same quantities as in"}, {"start": 1715.52, "end": 1716.52, "text": " the original Atom paper."}, {"start": 1716.52, "end": 1723.32, "text": " So if I were to Google here Atom paper, blah, blah, blah, and if I were to open this paper"}, {"start": 1723.32, "end": 1729.3600000000001, "text": " on archive, we're going to see the formulas are using the exact same variable name."}, {"start": 1729.3600000000001, "end": 1732.1200000000001, "text": " So I think it's here, is this the Atom?"}, {"start": 1732.1200000000001, "end": 1733.1200000000001, "text": " Yeah."}, {"start": 1733.12, "end": 1740.3999999999999, "text": " And this is the nu, I don't know whether I'm pronouncing it correctly, probably not, but"}, {"start": 1740.3999999999999, "end": 1745.56, "text": " basically those are the statistics, the gradient statistics that Atom is kind of keeping in"}, {"start": 1745.56, "end": 1747.52, "text": " the background in order to do the update."}, {"start": 1747.52, "end": 1749.3, "text": " So here is the Atom update."}, {"start": 1749.3, "end": 1754.52, "text": " It's way more complicated compared to, let's say, SGD because yeah, you can see here we're"}, {"start": 1754.52, "end": 1760.84, "text": " using the mu and the nu here to form the novel parameter, to form the update."}, {"start": 1760.84, "end": 1762.84, "text": " Okay, so that's it."}, {"start": 1762.84, "end": 1769.6399999999999, "text": " Let me return back to SGD so that we can have comparable results compared to the pure Jax"}, {"start": 1769.6399999999999, "end": 1771.84, "text": " approach we used here."}, {"start": 1771.84, "end": 1773.36, "text": " And so the tree multimap."}, {"start": 1773.36, "end": 1778.36, "text": " So let me go back here and we're going to have the same exact thing replicated."}, {"start": 1778.36, "end": 1783.4399999999998, "text": " The only thing that's different here is we are using instead of explicitly handling the"}, {"start": 1783.4399999999998, "end": 1787.72, "text": " optimization part, we're going to use OptX."}, {"start": 1787.72, "end": 1792.92, "text": " And you can see in this example, it actually takes two lines to implement the code that"}, {"start": 1792.92, "end": 1793.92, "text": " we implemented here."}, {"start": 1793.92, "end": 1797.46, "text": " I think in, yeah, we had just a single line here."}, {"start": 1797.46, "end": 1801.28, "text": " So you may ask yourself why would we use OptX then?"}, {"start": 1801.28, "end": 1802.56, "text": " And then, I mean, it's pretty obvious."}, {"start": 1802.56, "end": 1807.76, "text": " I mean, when you have some more complex optimizers, you'd have in pure Jax, you'd have like 10,"}, {"start": 1807.76, "end": 1809.08, "text": " 15 lines or more."}, {"start": 1809.08, "end": 1811.28, "text": " Here you always have two lines and that's very cool."}, {"start": 1811.28, "end": 1812.5, "text": " So this is how it looks like."}, {"start": 1812.5, "end": 1814.96, "text": " So you have the optimizer, the SGD optimizer here."}, {"start": 1814.96, "end": 1820.72, "text": " We pass the gradients and we pass its state and outcomes, the novel state and outcome,"}, {"start": 1820.72, "end": 1823.28, "text": " the updates, which are the modified gradients."}, {"start": 1823.28, "end": 1828.2, "text": " So in the case of Atom, depending on the internal state, we'd be modifying these updates and"}, {"start": 1828.2, "end": 1832.76, "text": " then we apply, we just call this apply updates on the parameters and we get the novel parameters"}, {"start": 1832.76, "end": 1833.76, "text": " and that's it."}, {"start": 1833.76, "end": 1834.92, "text": " So this should not work."}, {"start": 1834.92, "end": 1838.88, "text": " Let me try and run this and see what we get."}, {"start": 1838.88, "end": 1841.92, "text": " And you can see here that we get like a loss."}, {"start": 1841.92, "end": 1842.92, "text": " I think it's the same loss."}, {"start": 1842.92, "end": 1848.64, "text": " It's 512 here and if we take a look here, it's also 512."}, {"start": 1848.64, "end": 1850.28, "text": " The last two digits are 85."}, {"start": 1850.28, "end": 1851.72, "text": " The last two digits are 85."}, {"start": 1851.72, "end": 1856.76, "text": " We get the same loss because I mean, this implements the SGD perfectly."}, {"start": 1856.76, "end": 1857.88, "text": " Okay."}, {"start": 1857.88, "end": 1858.88, "text": " So that's cool."}, {"start": 1858.88, "end": 1860.6000000000001, "text": " And later, I have a note here."}, {"start": 1860.6000000000001, "end": 1866.1000000000001, "text": " Later we'll see how we can kind of take all of these various states that are kind of floating"}, {"start": 1866.1000000000001, "end": 1871.2, "text": " around our program loops and wrap them in a single object and then use that single object"}, {"start": 1871.2, "end": 1874.76, "text": " to kind of, yeah, do the training."}, {"start": 1874.76, "end": 1879.52, "text": " So just so you can see that Optex is much more than SGD and Atom."}, {"start": 1879.52, "end": 1882.76, "text": " Let me show you this piece of code here."}, {"start": 1882.76, "end": 1888.56, "text": " I mentioned here that you can do various schedules, cosine schedules, linear schedules, arbitrary"}, {"start": 1888.56, "end": 1889.56, "text": " schedules."}, {"start": 1889.56, "end": 1892.8400000000001, "text": " You can do chaining of multiple optimizers, multiple components."}, {"start": 1892.8400000000001, "end": 1896.44, "text": " You can do parameter freezing when you want to do some fine tuning."}, {"start": 1896.44, "end": 1900.52, "text": " So it's very powerful, much more powerful than using pure JECs."}, {"start": 1900.52, "end": 1904.2, "text": " I mean, it's as powerful, but like it's way more concise."}, {"start": 1904.2, "end": 1905.2, "text": " Okay."}, {"start": 1905.2, "end": 1908.24, "text": " So let's show this code here."}, {"start": 1908.24, "end": 1909.24, "text": " This will not run."}, {"start": 1909.24, "end": 1915.2, "text": " This is just some simple example of copy pasted from ImageNet examples from the respective"}, {"start": 1915.2, "end": 1918.2, "text": " repos from Haiku and from Flex."}, {"start": 1918.2, "end": 1922.16, "text": " Here you can see this function create learning rate, which basically takes some parameters"}, {"start": 1922.16, "end": 1925.12, "text": " and then creates this from Optex."}, {"start": 1925.12, "end": 1927.1399999999999, "text": " You call this linear schedule function."}, {"start": 1927.14, "end": 1931.4, "text": " We give it the initial value, the end value, and then the number of steps where we are"}, {"start": 1931.4, "end": 1935.1200000000001, "text": " linearly increasing the learning rate or decreasing."}, {"start": 1935.1200000000001, "end": 1939.8000000000002, "text": " Now here in this case, yeah, you can basically either increase or decrease depending on the"}, {"start": 1939.8000000000002, "end": 1941.6000000000001, "text": " initial and end values."}, {"start": 1941.6000000000001, "end": 1943.88, "text": " Then we have the cosine function here."}, {"start": 1943.88, "end": 1945.24, "text": " So the cosine schedule."}, {"start": 1945.24, "end": 1949.0, "text": " And finally, we kind of join schedules using Optex join schedules."}, {"start": 1949.0, "end": 1953.68, "text": " And we join the warmup, the linear warmup combined with the cosine schedule."}, {"start": 1953.68, "end": 1957.44, "text": " And I guess, yeah, and we put the boundaries here to kind of denote the moment when we"}, {"start": 1957.44, "end": 1959.44, "text": " are switching modes."}, {"start": 1959.44, "end": 1966.48, "text": " And then you can use this scheduling function and pass it directly to Optex SGD here with"}, {"start": 1966.48, "end": 1969.68, "text": " some additional like Nestoro momentum, et cetera."}, {"start": 1969.68, "end": 1974.4, "text": " Finally from, you can see here a snippet from Haiku, even more complex, arguably."}, {"start": 1974.4, "end": 1976.28, "text": " We have a trace here."}, {"start": 1976.28, "end": 1979.8400000000001, "text": " We have a chaining operation of this tracing operator."}, {"start": 1979.84, "end": 1984.32, "text": " And then we have the scale by schedule, like a bunch of small components which can make"}, {"start": 1984.32, "end": 1987.84, "text": " arbitrary complex transformations."}, {"start": 1987.84, "end": 1989.76, "text": " So that's pretty much it."}, {"start": 1989.76, "end": 1991.82, "text": " That was section number two."}, {"start": 1991.82, "end": 1996.56, "text": " We got some understanding of how the toy like neural networks, i.e. linear regression in"}, {"start": 1996.56, "end": 1998.36, "text": " this case, are trained."}, {"start": 1998.36, "end": 2003.32, "text": " And now we're going to see how to create custom models in Flex, which is something we care"}, {"start": 2003.32, "end": 2004.3999999999999, "text": " about."}, {"start": 2004.3999999999999, "end": 2008.1999999999998, "text": " So let's start with a simple example of how to build an MLP."}, {"start": 2008.2, "end": 2010.88, "text": " So multi-layer perceptron network."}, {"start": 2010.88, "end": 2015.46, "text": " And you can see here we are inheriting from NNModule."}, {"start": 2015.46, "end": 2017.56, "text": " So we form a class, we inherit from the module."}, {"start": 2017.56, "end": 2022.48, "text": " So every, when you want to build a custom model, you always want to inherit from the"}, {"start": 2022.48, "end": 2023.48, "text": " NNModule."}, {"start": 2023.48, "end": 2026.22, "text": " Then you can see that we have this field here."}, {"start": 2026.22, "end": 2028.88, "text": " And this is called data field."}, {"start": 2028.88, "end": 2030.94, "text": " It's from one of the newer versions of Python."}, {"start": 2030.94, "end": 2032.5800000000002, "text": " You may be familiar with it."}, {"start": 2032.5800000000002, "end": 2034.44, "text": " You may be not."}, {"start": 2034.44, "end": 2037.96, "text": " But basically NNModule is Python's data class."}, {"start": 2037.96, "end": 2040.08, "text": " So if I were to, let me check."}, {"start": 2040.08, "end": 2041.48, "text": " I think I have a link here."}, {"start": 2041.48, "end": 2042.94, "text": " Yeah, we have a link here."}, {"start": 2042.94, "end": 2047.68, "text": " Let me just copy paste this thing here and let me show you what a data class is."}, {"start": 2047.68, "end": 2049.6, "text": " So data class, let me kind of zoom here."}, {"start": 2049.6, "end": 2052.68, "text": " A data class is basically this."}, {"start": 2052.68, "end": 2058.48, "text": " So you have this wrapper called data class, and then you can just kind of put the variable"}, {"start": 2058.48, "end": 2059.48, "text": " names here."}, {"start": 2059.48, "end": 2064.66, "text": " And then a couple of functions are automatically instantiated for you, defined for you."}, {"start": 2064.66, "end": 2068.6, "text": " So you can see, let me just find an equivalent representation here."}, {"start": 2068.6, "end": 2069.6, "text": " Yeah."}, {"start": 2069.6, "end": 2077.04, "text": " So if you were to explicitly do the code op, the data class, this is what you'd have to"}, {"start": 2077.04, "end": 2078.92, "text": " do for those two fields."}, {"start": 2078.92, "end": 2083.24, "text": " You'd have to create this representation function, the quality function."}, {"start": 2083.24, "end": 2084.24, "text": " You'd have to call the init."}, {"start": 2084.24, "end": 2085.64, "text": " And there is a lot of redundancy here."}, {"start": 2085.64, "end": 2088.8399999999997, "text": " As you can see, we have rank being called three times here."}, {"start": 2088.84, "end": 2095.96, "text": " So all in all, there's just a neat feature of Python more than Flex that, as I said,"}, {"start": 2095.96, "end": 2099.08, "text": " NNModule is kind of piggybacking on."}, {"start": 2099.08, "end": 2107.04, "text": " And yeah, so as a consequence of NNModule being a Python data class, unfortunately,"}, {"start": 2107.04, "end": 2109.8, "text": " we'll have to deal with a new function name."}, {"start": 2109.8, "end": 2111.32, "text": " We're not using init anymore."}, {"start": 2111.32, "end": 2112.36, "text": " We're using setup."}, {"start": 2112.36, "end": 2114.44, "text": " So setup is old init."}, {"start": 2114.44, "end": 2121.28, "text": " And because as I mentioned here, basically, the data class is implicitly using the init"}, {"start": 2121.28, "end": 2122.28, "text": " function."}, {"start": 2122.28, "end": 2124.96, "text": " And we saw that on the tab I just closed."}, {"start": 2124.96, "end": 2128.6, "text": " What we do here is we just form a list of layers."}, {"start": 2128.6, "end": 2132.84, "text": " So you can see here, we iterate through this sequence of integers, which is a simple list"}, {"start": 2132.84, "end": 2134.4, "text": " of integers."}, {"start": 2134.4, "end": 2138.12, "text": " We create these dense, i.e. feed forward layers."}, {"start": 2138.12, "end": 2142.36, "text": " And then during the call function, what we do, we pass the input."}, {"start": 2142.36, "end": 2144.46, "text": " We just rename that as activation."}, {"start": 2144.46, "end": 2146.84, "text": " And we iterate over the layers."}, {"start": 2146.84, "end": 2150.2000000000003, "text": " We kind of call, we pass the activation through the layer."}, {"start": 2150.2000000000003, "end": 2155.2400000000002, "text": " And finally, depending on whether or not we are in the last layer, we apply the ReLU,"}, {"start": 2155.2400000000002, "end": 2157.48, "text": " the ReLU nonlinear activation here."}, {"start": 2157.48, "end": 2159.7200000000003, "text": " And then we finally return the activation."}, {"start": 2159.7200000000003, "end": 2162.84, "text": " So the output layer, that's the output activation."}, {"start": 2162.84, "end": 2165.2000000000003, "text": " So yeah, that's your basic definition of an MLP."}, {"start": 2165.2000000000003, "end": 2166.76, "text": " Fairly, fairly simple."}, {"start": 2166.76, "end": 2169.28, "text": " Now let's see how we can kind of do some inference."}, {"start": 2169.28, "end": 2171.36, "text": " So again, we split a key."}, {"start": 2171.36, "end": 2172.32, "text": " We have the H key."}, {"start": 2172.32, "end": 2174.0800000000004, "text": " We have the init key."}, {"start": 2174.0800000000004, "end": 2176.44, "text": " Init is supposed to initialize the parameters."}, {"start": 2176.44, "end": 2179.46, "text": " Let me kind of first slowly walk you through this part."}, {"start": 2179.46, "end": 2180.7200000000003, "text": " So we have MLP."}, {"start": 2180.7200000000003, "end": 2187.1200000000003, "text": " We instantiate it using this number of neurons per layer is going to be 16, 8, 1, which means"}, {"start": 2187.1200000000003, "end": 2189.88, "text": " in the first layer, we have 16 neurons."}, {"start": 2189.88, "end": 2190.92, "text": " Then in the second, we have 8."}, {"start": 2190.92, "end": 2194.92, "text": " And then we have final, a single neuron."}, {"start": 2194.92, "end": 2200.7200000000003, "text": " And we instantiate this data set here."}, {"start": 2200.72, "end": 2205.56, "text": " We have the four data points that each has four dimensions."}, {"start": 2205.56, "end": 2209.48, "text": " And we use the H, so we use the X key here."}, {"start": 2209.48, "end": 2213.72, "text": " Then we call the model init here with the init key and we pass X."}, {"start": 2213.72, "end": 2216.24, "text": " So we've seen this multiple times already."}, {"start": 2216.24, "end": 2217.2799999999997, "text": " We have params."}, {"start": 2217.2799999999997, "end": 2222.72, "text": " We call model apply and we do a forward pass and we get Y here."}, {"start": 2222.72, "end": 2224.72, "text": " And that's it."}, {"start": 2224.72, "end": 2230.7599999999998, "text": " The same as if we were to use a simple, like a dense layer and not a complex MLP."}, {"start": 2230.7599999999998, "end": 2234.2999999999997, "text": " Again, I'm just calling a tree map and we're printing shapes."}, {"start": 2234.2999999999997, "end": 2239.3399999999997, "text": " So let's see how this thing looks like when we run it."}, {"start": 2239.3399999999997, "end": 2248.2799999999997, "text": " And as expected, we have 16, 8, 1, and we have 4 here because remember we have automatic"}, {"start": 2248.2799999999997, "end": 2249.8399999999997, "text": " shape inference."}, {"start": 2249.8399999999997, "end": 2254.68, "text": " And because our data points have four dimensions, basically, that's what we have for here."}, {"start": 2254.68, "end": 2257.3199999999997, "text": " And this kernel here, and this weight here."}, {"start": 2257.3199999999997, "end": 2259.7599999999998, "text": " By the way, I hate the word kernel."}, {"start": 2259.7599999999998, "end": 2262.2999999999997, "text": " It's so overloaded in machine learning that it's crazy."}, {"start": 2262.2999999999997, "end": 2268.44, "text": " So yeah, maybe using a weight, like a terminology here is maybe a bit better."}, {"start": 2268.44, "end": 2270.24, "text": " I don't know."}, {"start": 2270.24, "end": 2271.3599999999997, "text": " Cool."}, {"start": 2271.3599999999997, "end": 2272.9199999999996, "text": " I made a note here."}, {"start": 2272.9199999999996, "end": 2278.6, "text": " Let's try and do the same thing just with this nn-compact wrapper."}, {"start": 2278.6, "end": 2282.72, "text": " And what the trick here is, is you don't have to have set up at all."}, {"start": 2282.72, "end": 2287.24, "text": " So what you can do is, because we have a fairly simple function here, we can just kind of"}, {"start": 2287.24, "end": 2291.6, "text": " wrap this call function into nn-compact."}, {"start": 2291.6, "end": 2297.54, "text": " And then what we can do here is we can iterate through the number of neurons per layer instead"}, {"start": 2297.54, "end": 2299.16, "text": " of through layers."}, {"start": 2299.16, "end": 2303.12, "text": " And that means we'll have number of neurons here, number of neurons."}, {"start": 2303.12, "end": 2307.4199999999996, "text": " And then what we do instead of having a layer, we'll just instantiate it on the fly here."}, {"start": 2307.42, "end": 2314.7000000000003, "text": " So nn-dense, we kind of initialize it using the number of neurons, and then we do the"}, {"start": 2314.7000000000003, "end": 2315.7000000000003, "text": " same."}, {"start": 2315.7000000000003, "end": 2317.2400000000002, "text": " Everything else remains the same."}, {"start": 2317.2400000000002, "end": 2320.48, "text": " And I think this should run unless I've made some mistake."}, {"start": 2320.48, "end": 2322.88, "text": " Let me see whether it's going to work."}, {"start": 2322.88, "end": 2324.32, "text": " So everything remains the same."}, {"start": 2324.32, "end": 2329.52, "text": " It's just a bit more concise way of writing down things."}, {"start": 2329.52, "end": 2333.64, "text": " So yeah, I managed to make an error here, and the error is that we should not be using"}, {"start": 2333.64, "end": 2334.64, "text": " the layers."}, {"start": 2334.64, "end": 2336.32, "text": " Layers does not exist anymore."}, {"start": 2336.32, "end": 2339.6400000000003, "text": " So let me just kind of switch that here."}, {"start": 2339.6400000000003, "end": 2341.4, "text": " And this should now work."}, {"start": 2341.4, "end": 2342.88, "text": " Yeah, it is."}, {"start": 2342.88, "end": 2343.92, "text": " Okay, cool."}, {"start": 2343.92, "end": 2345.76, "text": " You just saw a second way."}, {"start": 2345.76, "end": 2350.84, "text": " So this is a very flex-like way of writing things down using nn-compact."}, {"start": 2350.84, "end": 2353.7200000000003, "text": " But either way, you get the same semantics."}, {"start": 2353.7200000000003, "end": 2354.96, "text": " Cool."}, {"start": 2354.96, "end": 2356.3, "text": " Let's continue here."}, {"start": 2356.3, "end": 2361.1200000000003, "text": " The output is, by the way, as expected, we have four times one because single output"}, {"start": 2361.1200000000003, "end": 2365.92, "text": " neuron for data points, remember, and then that's why we have four times one."}, {"start": 2365.92, "end": 2369.92, "text": " It's good to kind of check the shapes and make sure those make sense."}, {"start": 2369.92, "end": 2372.56, "text": " Now, let's dig even deeper."}, {"start": 2372.56, "end": 2379.7200000000003, "text": " And there are two very, very crucial concepts you ought to know about."}, {"start": 2379.7200000000003, "end": 2384.16, "text": " The first one is this thing called a param."}, {"start": 2384.16, "end": 2391.0, "text": " And basically, it's a terminology for trainable parameters in neural networks in Flex."}, {"start": 2391.0, "end": 2396.0, "text": " So let's implement the simple model that's going to leverage those params and see how"}, {"start": 2396.0, "end": 2397.34, "text": " this thing works."}, {"start": 2397.34, "end": 2403.24, "text": " So here, what we see is if we were to try to implement the dense layer, either feedforward"}, {"start": 2403.24, "end": 2407.68, "text": " layer, this is how we would approximately do it."}, {"start": 2407.68, "end": 2409.6, "text": " So here, again, it's a data class."}, {"start": 2409.6, "end": 2412.52, "text": " So we can instantiate these fields here."}, {"start": 2412.52, "end": 2414.02, "text": " We have number of neurons."}, {"start": 2414.02, "end": 2418.36, "text": " We have the weight and the bias initializer functions."}, {"start": 2418.36, "end": 2422.56, "text": " So yeah, bias is going to be initialized as all zeros."}, {"start": 2422.56, "end": 2426.6400000000003, "text": " And for weights, we're going to use like a normal initializer."}, {"start": 2426.6400000000003, "end": 2427.6400000000003, "text": " Not that important."}, {"start": 2427.6400000000003, "end": 2430.94, "text": " The details are not that important."}, {"start": 2430.94, "end": 2435.04, "text": " The actual implementation of this function is not vital for this explanation."}, {"start": 2435.04, "end": 2436.6, "text": " So let's see what we do here."}, {"start": 2436.6, "end": 2441.42, "text": " So we call self-param and we give it a name."}, {"start": 2441.42, "end": 2443.6, "text": " And here, it's going to be called weight."}, {"start": 2443.6, "end": 2445.0, "text": " And it can be arbitrary name."}, {"start": 2445.0, "end": 2446.4, "text": " You can choose whatever you want."}, {"start": 2446.4, "end": 2450.6800000000003, "text": " So that's the name that will appear in that frozen dict under the params collection, if"}, {"start": 2450.6800000000003, "end": 2451.6800000000003, "text": " you recall."}, {"start": 2451.6800000000003, "end": 2452.6800000000003, "text": " If not, I'm going to run this."}, {"start": 2452.6800000000003, "end": 2453.82, "text": " You're going to see it in a second."}, {"start": 2453.82, "end": 2456.2400000000002, "text": " So now we call the initializer function of our choice."}, {"start": 2456.2400000000002, "end": 2461.64, "text": " So the random state is going to be implicitly passed into this function through the init"}, {"start": 2461.64, "end": 2462.64, "text": " and apply function."}, {"start": 2462.64, "end": 2465.36, "text": " So you don't have to worry about that here."}, {"start": 2465.36, "end": 2472.1, "text": " So the only thing you have to provide is the shape information for this initializer function."}, {"start": 2472.1, "end": 2477.2799999999997, "text": " And we'll just take the last dimension of our data and then the number of neurons here,"}, {"start": 2477.2799999999997, "end": 2480.48, "text": " which is the desired number of outputs."}, {"start": 2480.48, "end": 2481.72, "text": " We do the same thing for bias."}, {"start": 2481.72, "end": 2482.72, "text": " We call self.param."}, {"start": 2482.72, "end": 2487.36, "text": " So this is, again, this is flex syntax."}, {"start": 2487.36, "end": 2488.96, "text": " And we name it bias."}, {"start": 2488.96, "end": 2490.1, "text": " We have the bias in it."}, {"start": 2490.1, "end": 2495.68, "text": " We pass the number of neurons as the shape because, yeah, that's how bias shape is supposed"}, {"start": 2495.68, "end": 2496.68, "text": " to look like."}, {"start": 2496.68, "end": 2499.7599999999998, "text": " And then this is a simple implementation of feed-forward layer."}, {"start": 2499.76, "end": 2505.5600000000004, "text": " We do a dot product between the inputs and between the kernel, the weights, and then"}, {"start": 2505.5600000000004, "end": 2506.6400000000003, "text": " we just add the bias."}, {"start": 2506.6400000000003, "end": 2507.6400000000003, "text": " And that's it."}, {"start": 2507.6400000000003, "end": 2513.88, "text": " This is how we would implement like a dense, a linear, a feed-forward layer in flex."}, {"start": 2513.88, "end": 2517.5600000000004, "text": " Again, we have a common pattern here."}, {"start": 2517.5600000000004, "end": 2522.28, "text": " We just form the random seeds here, the random states here."}, {"start": 2522.28, "end": 2525.2400000000002, "text": " And then we instantiate the model here."}, {"start": 2525.24, "end": 2530.12, "text": " We just form some random data on this line 18."}, {"start": 2530.12, "end": 2531.12, "text": " Then we initialize."}, {"start": 2531.12, "end": 2532.16, "text": " We get the parameters."}, {"start": 2532.16, "end": 2535.16, "text": " We call apply and we get the output."}, {"start": 2535.16, "end": 2536.9399999999996, "text": " And finally, we can print the shapes."}, {"start": 2536.9399999999996, "end": 2539.8999999999996, "text": " And let me run this and see what we have as the output."}, {"start": 2539.8999999999996, "end": 2543.3799999999997, "text": " And you can see this is like something we're used to seeing."}, {"start": 2543.3799999999997, "end": 2545.2, "text": " We have params collection."}, {"start": 2545.2, "end": 2550.8399999999997, "text": " And I'm deliberately saying collection because that's a reserved word in flex."}, {"start": 2550.8399999999997, "end": 2554.3199999999997, "text": " And I'm going to soon explain what exactly it means."}, {"start": 2554.32, "end": 2556.76, "text": " And yeah, you can see biases and weights here."}, {"start": 2556.76, "end": 2557.76, "text": " And that's it."}, {"start": 2557.76, "end": 2561.7200000000003, "text": " I've linked a couple of source code files here."}, {"start": 2561.7200000000003, "end": 2563.6400000000003, "text": " You can check those at your own pace."}, {"start": 2563.6400000000003, "end": 2568.7200000000003, "text": " But let me briefly show you, and you may be confused by some of the idiosyncrasies of"}, {"start": 2568.7200000000003, "end": 2569.7200000000003, "text": " flex."}, {"start": 2569.7200000000003, "end": 2573.46, "text": " Like why do we have brackets here and why don't we have brackets here?"}, {"start": 2573.46, "end": 2575.28, "text": " What type of inconsistency is this?"}, {"start": 2575.28, "end": 2577.96, "text": " So my OCD was kind of triggered by this."}, {"start": 2577.96, "end": 2583.0800000000004, "text": " I had to dig into the source code and I found the solution here."}, {"start": 2583.0800000000004, "end": 2584.0800000000004, "text": " Basically what happens, let me zoom in here."}, {"start": 2584.08, "end": 2585.08, "text": " So we have a partial."}, {"start": 2585.08, "end": 2586.08, "text": " It forms a higher order function."}, {"start": 2586.08, "end": 2596.84, "text": " It just kind of plugs in these arguments that you don't have to kind of plug those in later"}, {"start": 2596.84, "end": 2599.3199999999997, "text": " on when you're calling this variance scaling function."}, {"start": 2599.3199999999997, "end": 2604.54, "text": " So the reason for the inconsistency is that if I were to find the variance scaling implementation,"}, {"start": 2604.54, "end": 2609.0, "text": " you can see the variance scaling basically returns back this function in it."}, {"start": 2609.0, "end": 2613.56, "text": " So that's why we have to call the brackets so that we get back the init function, which"}, {"start": 2613.56, "end": 2615.96, "text": " we can then call as an initializer."}, {"start": 2615.96, "end": 2620.2799999999997, "text": " Whereas zeros, if I were to find the zeros, it's already a function."}, {"start": 2620.2799999999997, "end": 2622.96, "text": " So that's why you don't have to call like the brackets."}, {"start": 2622.96, "end": 2625.12, "text": " So minor implementation detail."}, {"start": 2625.12, "end": 2630.32, "text": " I just thought explaining the why the inconsistency in case you wondered."}, {"start": 2630.32, "end": 2636.6, "text": " Second thing, I want to quickly show you the implementation of the aptly named linear layer,"}, {"start": 2636.6, "end": 2638.2799999999997, "text": " even though they call it dense."}, {"start": 2638.2799999999997, "end": 2639.7599999999998, "text": " So let's see that."}, {"start": 2639.7599999999998, "end": 2640.7599999999998, "text": " Okay."}, {"start": 2640.76, "end": 2645.6000000000004, "text": " So let me find the class dense, dense general, dense here."}, {"start": 2645.6000000000004, "end": 2646.6000000000004, "text": " Okay."}, {"start": 2646.6000000000004, "end": 2652.1200000000003, "text": " So obviously it's a bit more complicated than the toy implementation, which I just showed"}, {"start": 2652.1200000000003, "end": 2656.1600000000003, "text": " you, but like the basics are the same."}, {"start": 2656.1600000000003, "end": 2659.76, "text": " There is just an additional type conversion depending on whether you're training this"}, {"start": 2659.76, "end": 2662.44, "text": " thing in a mixed precision mode."}, {"start": 2662.44, "end": 2669.28, "text": " You may want to convert the data type of your activations, of your input data, et cetera."}, {"start": 2669.28, "end": 2674.6000000000004, "text": " So you can see here again, we have the self-param kernel initializer function."}, {"start": 2674.6000000000004, "end": 2675.6000000000004, "text": " We have the shape."}, {"start": 2675.6000000000004, "end": 2681.6200000000003, "text": " Additionally, again, I said we have a data type here and then again, we have some conversion"}, {"start": 2681.6200000000003, "end": 2682.6200000000003, "text": " of data type here."}, {"start": 2682.6200000000003, "end": 2688.1200000000003, "text": " So basically, and then we have dot general, which is basically as if we were to call JMP."}, {"start": 2688.1200000000003, "end": 2693.52, "text": " It's just a more lower level function that says you can see more verbose, but also much"}, {"start": 2693.52, "end": 2697.7200000000003, "text": " more, I guess, less error prone."}, {"start": 2697.72, "end": 2704.9199999999996, "text": " We formed the bias here and basically if bias exists, we just add it here and then that's,"}, {"start": 2704.9199999999996, "end": 2709.48, "text": " yeah, I mean, that's just a W times X plus B implementation of that."}, {"start": 2709.48, "end": 2711.2, "text": " So that's pretty much it."}, {"start": 2711.2, "end": 2715.9199999999996, "text": " Now let me get back to the notebook and let's continue this journey."}, {"start": 2715.9199999999996, "end": 2723.9199999999996, "text": " Yeah, again, I just have a simple cell here showing that the signature is such that this"}, {"start": 2723.92, "end": 2728.6, "text": " initializer function, the first argument is actually key, not the shape."}, {"start": 2728.6, "end": 2734.16, "text": " So that's why I said that you are implicitly passing in the RNG."}, {"start": 2734.16, "end": 2738.98, "text": " So by the way, you usually will not have to care about these things that much because"}, {"start": 2738.98, "end": 2743.92, "text": " you'll just usually just be using the NN layers, which are already implemented for you."}, {"start": 2743.92, "end": 2748.4, "text": " Here I'm just trying to dig a bit deeper into Flex and kind of show you how things work"}, {"start": 2748.4, "end": 2751.48, "text": " behind the, like in the background."}, {"start": 2751.48, "end": 2752.64, "text": " Hopefully you'll appreciate that."}, {"start": 2752.64, "end": 2755.72, "text": " Okay, let's continue here."}, {"start": 2755.72, "end": 2763.24, "text": " So those were the trainable parameters, but as you may know, when you're training neural"}, {"start": 2763.24, "end": 2768.2799999999997, "text": " networks, you sometimes have parameters which are a part of the model, but they're not exactly"}, {"start": 2768.2799999999997, "end": 2769.2799999999997, "text": " trainable."}, {"start": 2769.2799999999997, "end": 2772.3199999999997, "text": " So you're maybe keeping some statistics in the case of batch norm."}, {"start": 2772.3199999999997, "end": 2773.64, "text": " I guess that's the best example."}, {"start": 2773.64, "end": 2779.48, "text": " The most famous example would be the batch norm and batch norm has a trainable parameter"}, {"start": 2779.48, "end": 2784.84, "text": " as well as those non-trainable batch statistics."}, {"start": 2784.84, "end": 2789.44, "text": " And now we're going to see how Flex handles those."}, {"start": 2789.44, "end": 2795.0, "text": " And the novel concept I want you to keep in mind here is variable."}, {"start": 2795.0, "end": 2799.64, "text": " And I just have a note in terminology, so it's a broader term and it includes both the"}, {"start": 2799.64, "end": 2803.76, "text": " trainable as well as the non-trainable variables."}, {"start": 2803.76, "end": 2807.56, "text": " So you'll see params being called variables as well."}, {"start": 2807.56, "end": 2809.04, "text": " They're just trainable variables."}, {"start": 2809.04, "end": 2816.48, "text": " Okay, that out of the way, let's see a very simple contrived example that contains both"}, {"start": 2816.48, "end": 2819.32, "text": " the variables and the parameters alike."}, {"start": 2819.32, "end": 2826.52, "text": " So we have a simple class here called very convoluted name bias adder with running mean."}, {"start": 2826.52, "end": 2832.88, "text": " Don't try and deduce the semantics of this function from this title, from this name,"}, {"start": 2832.88, "end": 2835.44, "text": " it's just a contrived example."}, {"start": 2835.44, "end": 2839.44, "text": " So what happens in the call function, we're using again this compact notation, we just"}, {"start": 2839.44, "end": 2847.04, "text": " check whether this variable is already initialized, namely the batch stats collection."}, {"start": 2847.04, "end": 2853.12, "text": " And then we take a look whether the EMA variable inside of that collection has been initialized."}, {"start": 2853.12, "end": 2858.5, "text": " Now let's see what collection and all of this, what I've just said means."}, {"start": 2858.5, "end": 2862.2400000000002, "text": " So batch stats is actually not an arbitrary name."}, {"start": 2862.24, "end": 2866.72, "text": " You'll see that in, I think I'm gonna show you the implementation of the batch norm,"}, {"start": 2866.72, "end": 2869.0, "text": " the source code itself."}, {"start": 2869.0, "end": 2873.52, "text": " Basically they've hard coded now whether that's a good design or not, I don't know."}, {"start": 2873.52, "end": 2877.6, "text": " I guess I don't have enough context, but like in old one, I'd say it's not the best design"}, {"start": 2877.6, "end": 2880.64, "text": " to hard code like a string like that."}, {"start": 2880.64, "end": 2883.24, "text": " But yeah, it is what it is."}, {"start": 2883.24, "end": 2887.08, "text": " So we can see here self.variable."}, {"start": 2887.08, "end": 2890.0, "text": " So again, that's a flex thing and we form a collection."}, {"start": 2890.0, "end": 2892.72, "text": " So this is the collection name, batch stats."}, {"start": 2892.72, "end": 2897.4, "text": " And the reason they have different collections is so that you can treat them differently."}, {"start": 2897.4, "end": 2898.48, "text": " So that's the point."}, {"start": 2898.48, "end": 2902.84, "text": " So we have the params collection which usually contains the trainable parameter, the trainable"}, {"start": 2902.84, "end": 2903.84, "text": " variables."}, {"start": 2903.84, "end": 2910.24, "text": " And then you have batch stats which contains these types of non-trainable variables where"}, {"start": 2910.24, "end": 2913.72, "text": " you take the batch of images say and you calculate some statistics."}, {"start": 2913.72, "end": 2918.08, "text": " So for example, in the case of batch norm, you calculate the mean statistics, you calculate"}, {"start": 2918.08, "end": 2923.36, "text": " the variance and then you use that to do the logic of the layer."}, {"start": 2923.36, "end": 2927.52, "text": " Okay, so it's very similar to self.param."}, {"start": 2927.52, "end": 2932.2, "text": " You just have the collection name here, then you have the name, then you have the initializer"}, {"start": 2932.2, "end": 2934.84, "text": " function here."}, {"start": 2934.84, "end": 2936.6, "text": " We are just being explicit here."}, {"start": 2936.6, "end": 2939.98, "text": " We just have the zeros function and then we pass in the shape."}, {"start": 2939.98, "end": 2946.64, "text": " We have one and then we have a column here because basically the zero dimension is batch"}, {"start": 2946.64, "end": 2947.64, "text": " dimension."}, {"start": 2947.64, "end": 2951.72, "text": " So we're gonna just have the shape of a single data point here."}, {"start": 2951.72, "end": 2953.18, "text": " Then we have a bias."}, {"start": 2953.18, "end": 2958.2799999999997, "text": " So bias is a trainable variable and this will be added."}, {"start": 2958.2799999999997, "end": 2962.9, "text": " So this thing here is kind of implicitly set to params here."}, {"start": 2962.9, "end": 2965.3199999999997, "text": " So I noted, I have it here noted."}, {"start": 2965.3199999999997, "end": 2971.7999999999997, "text": " So by default, we'll add this variable to params collection, whereas the above will"}, {"start": 2971.7999999999997, "end": 2974.7599999999998, "text": " be added to the batch stats collection."}, {"start": 2974.76, "end": 2979.1600000000003, "text": " Here because of some of the implementation details, unfortunately, we have to pass the"}, {"start": 2979.1600000000003, "end": 2983.92, "text": " key because that's the expected signature of the function for this self.param, even"}, {"start": 2983.92, "end": 2988.1200000000003, "text": " though we're not using it in this particular case."}, {"start": 2988.1200000000003, "end": 2989.1200000000003, "text": " So yeah."}, {"start": 2989.1200000000003, "end": 2993.88, "text": " And finally, before I explain this snippet here, EMA stands for Exponentially Moving"}, {"start": 2993.88, "end": 2997.0800000000004, "text": " Average, just for you to understand the semantics there."}, {"start": 2997.08, "end": 3004.7599999999998, "text": " So what we do here is if we have initialized the EMA variable already, then we can start"}, {"start": 3004.7599999999998, "end": 3010.68, "text": " in the forward pass, we're going to update it using this exponentially moving logic."}, {"start": 3010.68, "end": 3012.56, "text": " So what we do is we take the decay."}, {"start": 3012.56, "end": 3016.84, "text": " Decay is usually some larger number like 099 or something."}, {"start": 3016.84, "end": 3019.06, "text": " And then you keep, you kind of have inertia."}, {"start": 3019.06, "end": 3026.36, "text": " You keep the old value with a higher probability and then one minus that and then you do the"}, {"start": 3026.36, "end": 3029.4, "text": " mean of the current batch of data."}, {"start": 3029.4, "end": 3032.84, "text": " So this is kind of damping, has this damping effect."}, {"start": 3032.84, "end": 3037.6400000000003, "text": " I'm trying to teach you flex here, but like I think that explaining some of the machine"}, {"start": 3037.6400000000003, "end": 3039.96, "text": " learning concepts along the way may help a lot of you."}, {"start": 3039.96, "end": 3043.04, "text": " So that's why I'm doing this."}, {"start": 3043.04, "end": 3049.4, "text": " Finally we just returned the return is X, which is input data minus the exponentially"}, {"start": 3049.4, "end": 3052.6, "text": " moving average here plus bias."}, {"start": 3052.6, "end": 3057.92, "text": " And you may wonder why do we have dot value and the reason is that self dot variable returns"}, {"start": 3057.92, "end": 3060.42, "text": " a reference, not the value itself."}, {"start": 3060.42, "end": 3065.04, "text": " So that's why you have to kind of do the dot value to get the actual value of the variable."}, {"start": 3065.04, "end": 3068.72, "text": " Okay, this will make a bit more sense in a couple of seconds."}, {"start": 3068.72, "end": 3072.02, "text": " Again, very, very familiar pattern here."}, {"start": 3072.02, "end": 3074.2799999999997, "text": " We form the random states."}, {"start": 3074.2799999999997, "end": 3076.0, "text": " We instantiate the model here."}, {"start": 3076.0, "end": 3077.8399999999997, "text": " We form the data."}, {"start": 3077.8399999999997, "end": 3079.66, "text": " So we have 10 data points."}, {"start": 3079.66, "end": 3081.72, "text": " Each data point has four dimensions."}, {"start": 3081.72, "end": 3086.56, "text": " Again, we initialize the model, we get the variables and then we print them here."}, {"start": 3086.56, "end": 3090.9599999999996, "text": " So let me just kind of run this and show you already what's going on here."}, {"start": 3090.9599999999996, "end": 3094.8799999999997, "text": " So here is the frozen dict, which we've, we're used to seeing this."}, {"start": 3094.8799999999997, "end": 3101.7999999999997, "text": " And now aside from params, we also have the batch stats where we have this EMA variable."}, {"start": 3101.7999999999997, "end": 3106.9199999999996, "text": " And so yeah, this is the structure, the underlying structure of this, of this contrived model."}, {"start": 3106.92, "end": 3112.16, "text": " Only the values are zero because that's the initializer we've been, we've set here, the"}, {"start": 3112.16, "end": 3113.16, "text": " zeros."}, {"start": 3113.16, "end": 3115.92, "text": " And then what we do here is we apply."}, {"start": 3115.92, "end": 3118.12, "text": " So this is the novel part of notation."}, {"start": 3118.12, "end": 3125.6, "text": " I want you to remember we pass the variables and then we pass the input and then we tell"}, {"start": 3125.6, "end": 3131.56, "text": " flex, Hey, batch stats are actually mutable, which means that during the forward pass,"}, {"start": 3131.56, "end": 3133.32, "text": " the value is going to change."}, {"start": 3133.32, "end": 3137.7200000000003, "text": " Whereas for regular parameters, when you do like a feed forward to a neural network, you"}, {"start": 3137.7200000000003, "end": 3139.7200000000003, "text": " will not be changing the weights."}, {"start": 3139.7200000000003, "end": 3145.56, "text": " But here for these types of variables, they're going to change during the forward pass."}, {"start": 3145.56, "end": 3149.1200000000003, "text": " And that's the reason why we have to have this distinction and different collections."}, {"start": 3149.1200000000003, "end": 3150.1200000000003, "text": " Okay."}, {"start": 3150.1200000000003, "end": 3154.76, "text": " So we get the output and we get this time, we get a tuple, not just the output, we get"}, {"start": 3154.76, "end": 3160.6000000000004, "text": " the updated non-drainable parameters and I've printed them here and you can see how all"}, {"start": 3160.6, "end": 3168.36, "text": " of the zeros now converted into this, which is because we've done, we have data X, we've"}, {"start": 3168.36, "end": 3170.7999999999997, "text": " done the mean across the data."}, {"start": 3170.7999999999997, "end": 3174.24, "text": " So we've aggregated those 10 vectors into a single vector."}, {"start": 3174.24, "end": 3180.56, "text": " And then here using this update rule, we've kind of have a novel, novel value here."}, {"start": 3180.56, "end": 3186.0, "text": " So all in all, what I want you to take out of from here is there are multiple collections"}, {"start": 3186.0, "end": 3189.96, "text": " because we have multiple semantics, depending whether we have trainable parameters or these"}, {"start": 3189.96, "end": 3193.4, "text": " types of other variables."}, {"start": 3193.4, "end": 3198.6, "text": " And again, we have this mutable keyword, which you should remember alongside this self dot"}, {"start": 3198.6, "end": 3199.6, "text": " variable."}, {"start": 3199.6, "end": 3202.68, "text": " So those are some of the important flex idiosyncrasies."}, {"start": 3202.68, "end": 3204.28, "text": " Okay."}, {"start": 3204.28, "end": 3210.12, "text": " Let's continue here and let's form like an update function, a training function."}, {"start": 3210.12, "end": 3211.7200000000003, "text": " Let's see how we could train this function."}, {"start": 3211.7200000000003, "end": 3215.88, "text": " And once you understand this, trust me, it's going to be way easier because we are now"}, {"start": 3215.88, "end": 3219.56, "text": " digging really deep into all of the nitty gritty details."}, {"start": 3219.56, "end": 3227.68, "text": " It's going to be much more easy, much more easy, much easier to actually train flex models"}, {"start": 3227.68, "end": 3228.68, "text": " than this."}, {"start": 3228.68, "end": 3229.68, "text": " Okay."}, {"start": 3229.68, "end": 3233.6, "text": " We pass bunch of arguments here because we're still not using the train state, which I'm"}, {"start": 3233.6, "end": 3235.44, "text": " going to later introduce."}, {"start": 3235.44, "end": 3236.98, "text": " So bear with me here."}, {"start": 3236.98, "end": 3238.96, "text": " We pass the params into loss function."}, {"start": 3238.96, "end": 3241.64, "text": " And what we do here, we pass the apply function."}, {"start": 3241.64, "end": 3245.72, "text": " Apply function is basically your model apply."}, {"start": 3245.72, "end": 3252.0, "text": " So we have model here and then I'm going to pass model dot apply here into the update"}, {"start": 3252.0, "end": 3253.0, "text": " step."}, {"start": 3253.0, "end": 3254.0, "text": " So that's the apply function."}, {"start": 3254.0, "end": 3256.3999999999996, "text": " So it's basically how you do the feed forward."}, {"start": 3256.3999999999996, "end": 3259.8799999999997, "text": " Then we pass the variables and here we're being explicit."}, {"start": 3259.8799999999997, "end": 3266.24, "text": " So I have params collection and I pass the params and we have the non-trainable params,"}, {"start": 3266.24, "end": 3270.9599999999996, "text": " which I just do the double star here to unpack the keys and the values."}, {"start": 3270.9599999999996, "end": 3273.9599999999996, "text": " And this is how you basically form the variables."}, {"start": 3273.96, "end": 3277.12, "text": " We will soon see why we have to split these two."}, {"start": 3277.12, "end": 3280.12, "text": " And then we pass the input and then we have the mutable."}, {"start": 3280.12, "end": 3284.92, "text": " Again here we're just being a bit more general and we say, give me the keys that correspond"}, {"start": 3284.92, "end": 3287.0, "text": " to those non-trainable parameters."}, {"start": 3287.0, "end": 3291.84, "text": " So batch stats would be one of those keys in one of the last examples."}, {"start": 3291.84, "end": 3294.48, "text": " So just to give you some idea."}, {"start": 3294.48, "end": 3300.0, "text": " Finally here the loss, because this example is super contrived, this doesn't make much"}, {"start": 3300.0, "end": 3301.0, "text": " sense."}, {"start": 3301.0, "end": 3305.84, "text": " Basically once we have the loss function, we do this is a very standard pattern."}, {"start": 3305.84, "end": 3309.64, "text": " You've seen this already a couple of times, like only in this notebook and you'll be seeing"}, {"start": 3309.64, "end": 3313.48, "text": " it even more once you start coding in Flex."}, {"start": 3313.48, "end": 3319.68, "text": " Basically you pass the loss function to the value and grad transform function."}, {"start": 3319.68, "end": 3326.4, "text": " And we say has auxiliary, we set to true because aside from loss, loss function is returning"}, {"start": 3326.4, "end": 3329.44, "text": " these updated non-trainable parameters."}, {"start": 3329.44, "end": 3334.92, "text": " And remember, once we do a forward pass, we are kind of mutating those parameters."}, {"start": 3334.92, "end": 3337.44, "text": " That's why we have to have this mutable in the first place."}, {"start": 3337.44, "end": 3338.68, "text": " Okay, cool."}, {"start": 3338.68, "end": 3340.58, "text": " So that's what we get here."}, {"start": 3340.58, "end": 3345.6, "text": " Once we pass the parameters, we get back the gradients of those parameters with respect"}, {"start": 3345.6, "end": 3351.12, "text": " to loss function and we return the value because we have value and grad."}, {"start": 3351.12, "end": 3356.68, "text": " So we return back the loss and the updated non-trainable parameters."}, {"start": 3356.68, "end": 3359.3, "text": " Then we have this common pattern using OptX."}, {"start": 3359.3, "end": 3364.1200000000003, "text": " We update using the optimizer state and then we apply the updates to the parameters to"}, {"start": 3364.1200000000003, "end": 3365.76, "text": " get novel parameters."}, {"start": 3365.76, "end": 3366.76, "text": " And that's it."}, {"start": 3366.76, "end": 3371.2000000000003, "text": " Then we return a bunch of arguments here because we're still not using the train state and"}, {"start": 3371.2000000000003, "end": 3374.0600000000004, "text": " I'm going to introduce that I think in one of the next cells."}, {"start": 3374.0600000000004, "end": 3376.0800000000004, "text": " So yeah, stay tuned."}, {"start": 3376.0800000000004, "end": 3377.2400000000002, "text": " Cool."}, {"start": 3377.2400000000002, "end": 3382.76, "text": " So let's see how this whole like how everything comes into a single picture."}, {"start": 3382.76, "end": 3389.0600000000004, "text": " So we have Winstonchi the model here, we have dummy data and so we initialize the model,"}, {"start": 3389.06, "end": 3390.32, "text": " we get the variables."}, {"start": 3390.32, "end": 3395.0, "text": " Now here, because with this time we have both the trainable as well as the non-trainable"}, {"start": 3395.0, "end": 3400.72, "text": " variables, we split them using this pop function and then because we just have a dictionary,"}, {"start": 3400.72, "end": 3401.94, "text": " so this is going to work."}, {"start": 3401.94, "end": 3407.52, "text": " And then they usually do this in the documentation, they delete the variables to avoid wasting"}, {"start": 3407.52, "end": 3409.2, "text": " resources."}, {"start": 3409.2, "end": 3416.96, "text": " And then we kind of create the optimizer here, we initialize the state and then we have just"}, {"start": 3416.96, "end": 3425.08, "text": " a simple for loop running a couple of training steps and here is the update step that's being"}, {"start": 3425.08, "end": 3426.08, "text": " called."}, {"start": 3426.08, "end": 3432.84, "text": " So we are correctly updating the parameters as well as treating the non-trainable variables"}, {"start": 3432.84, "end": 3434.4, "text": " in a correct manner."}, {"start": 3434.4, "end": 3438.7200000000003, "text": " So this is as nitty gritty as you're going to get in this video."}, {"start": 3438.7200000000003, "end": 3442.12, "text": " So if you followed me so far, congrats."}, {"start": 3442.12, "end": 3446.04, "text": " If not, you're still won't have any problems following the rest of the notebook."}, {"start": 3446.04, "end": 3456.0, "text": " So now we're going to go up in the abstraction and see one example where you basically have"}, {"start": 3456.0, "end": 3459.7, "text": " the most complex syntax you can have in Flex."}, {"start": 3459.7, "end": 3466.2799999999997, "text": " So I'm kind of going to mix a lot of various layers which are going to introduce both those"}, {"start": 3466.2799999999997, "end": 3472.64, "text": " non-trainable variables as well as some stochasticity in the forward pass via a dropout."}, {"start": 3472.64, "end": 3475.12, "text": " So let's see that example here."}, {"start": 3475.12, "end": 3480.24, "text": " We have this DDM block and the name DDM because we have dense dropout batch norm which is"}, {"start": 3480.24, "end": 3485.48, "text": " a fairly regular block you'll see like in also in CNNs or whatnot."}, {"start": 3485.48, "end": 3488.4, "text": " It's a fairly regular stack of layers."}, {"start": 3488.4, "end": 3492.68, "text": " Again it's a data class so we have number of neurons specified here and the training"}, {"start": 3492.68, "end": 3500.52, "text": " flag and depending on the value of this training flag, we'll have different behavior from dropout"}, {"start": 3500.52, "end": 3502.44, "text": " and batch norm here."}, {"start": 3502.44, "end": 3509.04, "text": " So what happens is we instantiate the feed forward layer here using the num neurons here,"}, {"start": 3509.04, "end": 3515.28, "text": " the num neurons variable and then we pass whatever like comes as the output, we pass"}, {"start": 3515.28, "end": 3517.04, "text": " it into dropout."}, {"start": 3517.04, "end": 3524.28, "text": " The dropout has 50% probability rate of kind of dropping the activations and again we have"}, {"start": 3524.28, "end": 3525.28, "text": " this deterministic."}, {"start": 3525.28, "end": 3531.14, "text": " So if we are training, that means not true will be false which means we are not deterministic"}, {"start": 3531.14, "end": 3532.14, "text": " during the training."}, {"start": 3532.14, "end": 3536.52, "text": " We want to have the stochasticity but when we are evaluating then the training flag is"}, {"start": 3536.52, "end": 3538.44, "text": " set to false so this will be true."}, {"start": 3538.44, "end": 3543.68, "text": " We're going to be deterministic and basically dropout is going to be kind of turned off"}, {"start": 3543.68, "end": 3545.44, "text": " pretty much."}, {"start": 3545.44, "end": 3550.72, "text": " Finally we take the output from dropout, we pass it into batch norm, we have use running"}, {"start": 3550.72, "end": 3553.2799999999997, "text": " average again depends on the training flag."}, {"start": 3553.2799999999997, "end": 3558.72, "text": " If training flag is true, if we are training then we pass false because we don't want to"}, {"start": 3558.72, "end": 3560.48, "text": " use the running average."}, {"start": 3560.48, "end": 3563.0, "text": " We want to use the batch average, that's how batch norm works."}, {"start": 3563.0, "end": 3567.68, "text": " So you take a batch, you calculate the statistics and you use those parameters instead of using"}, {"start": 3567.68, "end": 3571.2400000000002, "text": " the parameters that were accumulated throughout the training."}, {"start": 3571.2400000000002, "end": 3573.0, "text": " So that's it."}, {"start": 3573.0, "end": 3576.96, "text": " This is more of a machine learning concept than a flex concept so I'm going to not dwell"}, {"start": 3576.96, "end": 3578.8, "text": " too much on it."}, {"start": 3578.8, "end": 3585.06, "text": " Again familiar pattern, we form a bunch of random states here, we instantiate our block,"}, {"start": 3585.06, "end": 3588.64, "text": " we get some random, some dummy data here."}, {"start": 3588.64, "end": 3593.72, "text": " We initialize the model and here are some additional idiosyncrasies of flex."}, {"start": 3593.72, "end": 3599.8799999999997, "text": " So this time when we want to initialize and this may be a bit confusing, we have to pass"}, {"start": 3599.8799999999997, "end": 3607.56, "text": " like a random state for both the params collection as well as for the dropout."}, {"start": 3607.56, "end": 3614.7799999999997, "text": " And it doesn't make much sense semantically because I mean dropout cannot be initialized,"}, {"start": 3614.7799999999997, "end": 3618.24, "text": " you don't have any learnable parameters there but just because of the way how everything"}, {"start": 3618.24, "end": 3621.52, "text": " is implemented you kind of have to have this additional step."}, {"start": 3621.52, "end": 3623.1, "text": " So keep that in mind."}, {"start": 3623.1, "end": 3625.7999999999997, "text": " Then we pass the X and everything else remains the same."}, {"start": 3625.7999999999997, "end": 3631.12, "text": " We get the variables back and then we apply and now here comes additional part because"}, {"start": 3631.12, "end": 3637.72, "text": " now we have dropout and dropout as you saw complicates both the initialization procedure"}, {"start": 3637.72, "end": 3639.8199999999997, "text": " as well as the apply procedure."}, {"start": 3639.8199999999997, "end": 3646.9599999999996, "text": " So here we have to have this RNG's keyword and basically what that means is you set the"}, {"start": 3646.96, "end": 3652.84, "text": " random state so you seed the dropout procedure so you can again reproduce it later on."}, {"start": 3652.84, "end": 3657.2, "text": " So that's why it makes a lot of sense to have it here although it does not make much sense"}, {"start": 3657.2, "end": 3660.16, "text": " to have dropout here to the best of my knowledge."}, {"start": 3660.16, "end": 3666.52, "text": " So yeah, outcome the output, the Y and the non-trainable variables, the non-trainable"}, {"start": 3666.52, "end": 3667.52, "text": " parameters."}, {"start": 3667.52, "end": 3674.68, "text": " I'm kind of mixing the terminology here between hopefully you're getting the gist of this."}, {"start": 3674.68, "end": 3680.6, "text": " And if we want to run this in evaluation mode where training flag is set to false, the apply"}, {"start": 3680.6, "end": 3683.3999999999996, "text": " function simplifies by a lot."}, {"start": 3683.3999999999996, "end": 3690.3799999999997, "text": " We now don't have to pass the RNG's because dropout is effectively turned off and we also"}, {"start": 3690.3799999999997, "end": 3697.12, "text": " don't have to care about mutable part because batch norm just won't be updated."}, {"start": 3697.12, "end": 3702.96, "text": " Those internal kind of statistics will not be updated so we don't have to include the"}, {"start": 3702.96, "end": 3704.12, "text": " mutable."}, {"start": 3704.12, "end": 3707.8599999999997, "text": " So this is as complicated as it gets."}, {"start": 3707.8599999999997, "end": 3710.12, "text": " Remember a couple of these novel keywords."}, {"start": 3710.12, "end": 3716.64, "text": " We have the mutable, we have the inclusion of RNG's for the apply when we have the stochasticity"}, {"start": 3716.64, "end": 3718.3199999999997, "text": " in the forward pass."}, {"start": 3718.3199999999997, "end": 3723.7999999999997, "text": " And we have to during initialization if we have dropout we have to kind of specify keys"}, {"start": 3723.7999999999997, "end": 3725.56, "text": " per collection."}, {"start": 3725.56, "end": 3733.16, "text": " And this is why I said Haiku seems to be cleaner in my honest opinion compared to Flex because"}, {"start": 3733.16, "end": 3735.2, "text": " they don't have the separation."}, {"start": 3735.2, "end": 3737.44, "text": " It's kind of cleaner."}, {"start": 3737.44, "end": 3740.2, "text": " Like all of these signatures simplified by a lot."}, {"start": 3740.2, "end": 3745.48, "text": " Now whether that means Haiku loses some of the expressivity that's I guess worth thinking"}, {"start": 3745.48, "end": 3747.2799999999997, "text": " about but yeah."}, {"start": 3747.2799999999997, "end": 3748.3999999999996, "text": " Okay that was it."}, {"start": 3748.3999999999996, "end": 3753.56, "text": " Now we are in the last section of this video which is already overly long."}, {"start": 3753.56, "end": 3755.44, "text": " Hopefully you're following along."}, {"start": 3755.44, "end": 3760.04, "text": " Basically I'm gonna walk you through this fully fledged convolutional neural network"}, {"start": 3760.04, "end": 3762.3599999999997, "text": " example on MNIST."}, {"start": 3762.36, "end": 3766.56, "text": " We're gonna code it up in Flex and this is basically now that you have all of the components"}, {"start": 3766.56, "end": 3771.6, "text": " in your mind this is gonna make hopefully a lot of sense and you're gonna see how Flex"}, {"start": 3771.6, "end": 3775.84, "text": " loops and training like code actually looks like."}, {"start": 3775.84, "end": 3777.92, "text": " Okay let's start by defining the model."}, {"start": 3777.92, "end": 3779.6, "text": " We have a convolutional neural network."}, {"start": 3779.6, "end": 3781.84, "text": " We're using this MN compact pattern."}, {"start": 3781.84, "end": 3785.96, "text": " A lot of hard coding here for the demo purpose I guess it's fine."}, {"start": 3785.96, "end": 3791.7200000000003, "text": " Basically a lot of convolutional layers followed by relu's and then here and there we have"}, {"start": 3791.72, "end": 3797.12, "text": " some pooling average pooling then we flatten out the shapes and we have a couple of linear"}, {"start": 3797.12, "end": 3802.2799999999997, "text": " layers here and finally we have as the output the log softmax of the outputs."}, {"start": 3802.2799999999997, "end": 3806.3199999999997, "text": " That's it and we have 10 classes because we are dealing with MNIST as you recall."}, {"start": 3806.3199999999997, "end": 3814.3599999999997, "text": " Okay let's run this thing to define the neural network there and now I'm just going to reuse"}, {"start": 3814.3599999999997, "end": 3820.2799999999997, "text": " the the PyTorch data loading functionality from from tutorial number three of Jack series."}, {"start": 3820.28, "end": 3822.6400000000003, "text": " So I won't be diving deeper here."}, {"start": 3822.6400000000003, "end": 3828.36, "text": " Basically I have some collate function that's supposed to create non-Py arrays out of out"}, {"start": 3828.36, "end": 3833.28, "text": " of PyTorch tensors because that's what a Jacks and thus Flex is dealing with."}, {"start": 3833.28, "end": 3839.5600000000004, "text": " I have some custom transformations to just kind of manipulate the MNIST image and make"}, {"start": 3839.5600000000004, "end": 3842.1200000000003, "text": " it a zero one range flow through the two etc."}, {"start": 3842.1200000000003, "end": 3844.0400000000004, "text": " Some need to degree deals."}, {"start": 3844.0400000000004, "end": 3845.8, "text": " Basically we form the data sets here."}, {"start": 3845.8, "end": 3850.28, "text": " We have the data loaders with certain batch size like 128 images per batch."}, {"start": 3850.28, "end": 3856.44, "text": " I load the training images and the labels as well as the test images and specifically"}, {"start": 3856.44, "end": 3862.0800000000004, "text": " I load the whole test data set because it's fairly small for MNIST so we can kind of keep"}, {"start": 3862.0800000000004, "end": 3864.1200000000003, "text": " this in memory no problem."}, {"start": 3864.1200000000003, "end": 3871.4, "text": " Cool let's run the cell to load the data and now let's visualize a single image from from"}, {"start": 3871.4, "end": 3872.88, "text": " this data set."}, {"start": 3872.88, "end": 3876.04, "text": " So it's just like a regular MNIST image."}, {"start": 3876.04, "end": 3877.6800000000003, "text": " You can see number seven here."}, {"start": 3877.6800000000003, "end": 3882.2000000000003, "text": " You can see the ground truth label printed here on the screen as well."}, {"start": 3882.2000000000003, "end": 3890.28, "text": " So let's now understand a couple of functions that we'll need to to get this CNN training"}, {"start": 3890.28, "end": 3893.12, "text": " and learning on this simple MNIST data set."}, {"start": 3893.12, "end": 3899.58, "text": " So first things first we have this train step and eval step functions."}, {"start": 3899.58, "end": 3904.92, "text": " So I guess you're already familiar with most of these patterns but let me quickly walk"}, {"start": 3904.92, "end": 3907.34, "text": " you through what's going on here."}, {"start": 3907.34, "end": 3911.12, "text": " So what's happening is we have the loss function here."}, {"start": 3911.12, "end": 3913.66, "text": " What we do we do a forward pass through the CNN."}, {"start": 3913.66, "end": 3915.2999999999997, "text": " We pass the variables."}, {"start": 3915.2999999999997, "end": 3917.68, "text": " So again I'm being explicit here with the collections."}, {"start": 3917.68, "end": 3919.56, "text": " We have the params collections."}, {"start": 3919.56, "end": 3921.2999999999997, "text": " We pass the params here."}, {"start": 3921.2999999999997, "end": 3927.12, "text": " We pass the images which is our data and we get back the logits for those images."}, {"start": 3927.12, "end": 3932.16, "text": " Now we convert the ground truth labels using the basically this one hot function."}, {"start": 3932.16, "end": 3938.56, "text": " We get the one hot labels and then what this thing here is if you can guess it's a simple"}, {"start": 3938.56, "end": 3945.3199999999997, "text": " cross entropy loss because what we do here is we do element wise multiplication and because"}, {"start": 3945.3199999999997, "end": 3953.08, "text": " both the logits as well as the one hot labels have the shape of batch size times 10 because"}, {"start": 3953.08, "end": 3955.18, "text": " we have 10 classes in MNIST."}, {"start": 3955.18, "end": 3962.98, "text": " What this does is basically like the one hot labels will kind of select the true label"}, {"start": 3962.98, "end": 3968.2, "text": " from the logits and because logits is log softmax if you remember how we define the"}, {"start": 3968.2, "end": 3969.56, "text": " CNN here."}, {"start": 3969.56, "end": 3971.04, "text": " So let me just go up there."}, {"start": 3971.04, "end": 3975.6, "text": " So we have log softmax as the output of the CNN."}, {"start": 3975.6, "end": 3978.66, "text": " So what will happen is this will be just log softmax."}, {"start": 3978.66, "end": 3980.24, "text": " Let me just find where we stopped."}, {"start": 3980.24, "end": 3984.24, "text": " Okay here and then we're going to sum across the minus one."}, {"start": 3984.24, "end": 3987.72, "text": " So that's the last dimension i.e. the dimension where we have the 10 labels."}, {"start": 3987.72, "end": 3989.8399999999997, "text": " So we're going to sum those up."}, {"start": 3989.8399999999997, "end": 3993.72, "text": " Most of those are going to be zero because as I said this is going to select just a single"}, {"start": 3993.72, "end": 3994.72, "text": " one."}, {"start": 3994.72, "end": 4000.3999999999996, "text": " So we end up with a bunch of logs and then we just kind of do the mean across the dimension"}, {"start": 4000.3999999999996, "end": 4002.3599999999997, "text": " zero by default."}, {"start": 4002.3599999999997, "end": 4008.8599999999997, "text": " So that's across the batch dimension and we get like a mean of the log sum pretty much"}, {"start": 4008.8599999999997, "end": 4012.2, "text": " which is the cross entropy loss with the minus here."}, {"start": 4012.2, "end": 4016.4399999999996, "text": " And we pretty much return back the loss as well as the logits."}, {"start": 4016.4399999999996, "end": 4020.52, "text": " And then we have like a very very familiar pattern already."}, {"start": 4020.52, "end": 4025.68, "text": " We have the value and grad function we set as auxiliary to true because we're not only"}, {"start": 4025.68, "end": 4027.72, "text": " returning loss."}, {"start": 4027.72, "end": 4030.0, "text": " We're also returning the logits."}, {"start": 4030.0, "end": 4031.9199999999996, "text": " And then this is the new thing."}, {"start": 4031.9199999999996, "end": 4036.2, "text": " So we have this state now and that's the thing I've been mentioning throughout the video."}, {"start": 4036.2, "end": 4039.4199999999996, "text": " We're going to soon see what exactly and how we construct that."}, {"start": 4039.42, "end": 4046.08, "text": " But basically we take out the parameters which is just they form the part of our of our training"}, {"start": 4046.08, "end": 4047.6, "text": " state."}, {"start": 4047.6, "end": 4052.88, "text": " So parameters are basically the weights of our neural network and we pass those and basically"}, {"start": 4052.88, "end": 4056.96, "text": " we get the gradients back and we get the logits."}, {"start": 4056.96, "end": 4060.38, "text": " Then we just call state apply gradients."}, {"start": 4060.38, "end": 4067.2400000000002, "text": " So now we don't have the two optics line because those are hidden in the apply gradients function."}, {"start": 4067.24, "end": 4072.7999999999997, "text": " And we just pass the gradients and get the state back which is the updated state."}, {"start": 4072.7999999999997, "end": 4074.6, "text": " Next up we compute the metrics."}, {"start": 4074.6, "end": 4077.7999999999997, "text": " We just pass the logits and ground truth labels."}, {"start": 4077.7999999999997, "end": 4081.56, "text": " We get the metrics which is a simple dictionary containing loss and accuracy."}, {"start": 4081.56, "end": 4082.56, "text": " We'll soon see that."}, {"start": 4082.56, "end": 4084.7599999999998, "text": " Not that important though."}, {"start": 4084.7599999999998, "end": 4086.58, "text": " Evaluation step again very simple."}, {"start": 4086.58, "end": 4088.6, "text": " We pass the test images here."}, {"start": 4088.6, "end": 4090.56, "text": " We call we do a forward pass."}, {"start": 4090.56, "end": 4094.72, "text": " We get logits back and we just compute the metrics because we want to evaluate the model"}, {"start": 4094.72, "end": 4098.32, "text": " on the test images and that's it."}, {"start": 4098.32, "end": 4100.28, "text": " Okay let's continue here."}, {"start": 4100.28, "end": 4106.719999999999, "text": " Now that we have the training step and evaluation step defined we can basically just kind of"}, {"start": 4106.719999999999, "end": 4110.24, "text": " call those in a loop for the whole epoch."}, {"start": 4110.24, "end": 4113.92, "text": " So we iterate here through the PyTorch data loader."}, {"start": 4113.92, "end": 4116.04, "text": " We are fetching images and labels."}, {"start": 4116.04, "end": 4122.92, "text": " We start calling the train step over and over again and we're just collecting the metrics"}, {"start": 4122.92, "end": 4125.92, "text": " across all of those batches."}, {"start": 4125.92, "end": 4130.26, "text": " And then what we do is we call device get which if you recall from one of the previous"}, {"start": 4130.26, "end": 4135.6, "text": " videos we'll just take those like the data that's currently on the accelerator whatever"}, {"start": 4135.6, "end": 4138.36, "text": " that is GPU or TPU or whatnot."}, {"start": 4138.36, "end": 4141.2, "text": " It's going to fetch it back to the host memory."}, {"start": 4141.2, "end": 4143.76, "text": " So that's why we have the MP suffix here."}, {"start": 4143.76, "end": 4150.02, "text": " And then I just do some simple aggregation here across the batches and I return the final"}, {"start": 4150.02, "end": 4154.240000000001, "text": " state and all of those metrics back from the function."}, {"start": 4154.240000000001, "end": 4156.080000000001, "text": " Evaluate model fairly similar."}, {"start": 4156.080000000001, "end": 4158.68, "text": " We just call the eval step just once here."}, {"start": 4158.68, "end": 4164.46, "text": " We don't have to loop because if you recall I'm loading the whole test data set and so"}, {"start": 4164.46, "end": 4167.240000000001, "text": " a single line of code here will work."}, {"start": 4167.240000000001, "end": 4173.72, "text": " We do the same thing device get to pull on to host and then we have to do this dot item"}, {"start": 4173.72, "end": 4179.02, "text": " because this is just going to contain a NumPy arrays with a single element."}, {"start": 4179.02, "end": 4183.84, "text": " So you have to kind of call the dot item to get the scalar from the NumPy array."}, {"start": 4183.84, "end": 4184.84, "text": " And that's it."}, {"start": 4184.84, "end": 4186.56, "text": " We return back the metrics."}, {"start": 4186.56, "end": 4187.56, "text": " Okay."}, {"start": 4187.56, "end": 4194.240000000001, "text": " Finally, these are the two last functions that we need before we see the main training loop."}, {"start": 4194.240000000001, "end": 4196.1, "text": " And this one is fairly important."}, {"start": 4196.1, "end": 4200.72, "text": " So this one makes things I mean it makes things it's a syntactic sugar but it makes things"}, {"start": 4200.72, "end": 4204.96, "text": " way cleaner and thus less error prone."}, {"start": 4204.96, "end": 4208.240000000001, "text": " So what we do here is we instantiate the CNN."}, {"start": 4208.24, "end": 4212.719999999999, "text": " We basically initialize the parameters."}, {"start": 4212.719999999999, "end": 4213.719999999999, "text": " We pass the key."}, {"start": 4213.719999999999, "end": 4215.0199999999995, "text": " We pass the dummy data."}, {"start": 4215.0199999999995, "end": 4220.48, "text": " So just we pass all ones and then we extract from that frozen dict we just extract the"}, {"start": 4220.48, "end": 4228.76, "text": " params collection because that's the format that this train state here will expect."}, {"start": 4228.76, "end": 4234.92, "text": " And then we just kind of call and form this SGT optimizer without having to optimize it."}, {"start": 4234.92, "end": 4240.26, "text": " Without having to initialize it because that's going to happen inside of this train state"}, {"start": 4240.26, "end": 4241.26, "text": " class."}, {"start": 4241.26, "end": 4242.56, "text": " So that's cool."}, {"start": 4242.56, "end": 4246.28, "text": " We're kind of delegating a lot of the functionality there."}, {"start": 4246.28, "end": 4252.8, "text": " And finally, we also pass the CNN apply so that we can then do some more advanced logic"}, {"start": 4252.8, "end": 4254.76, "text": " as we saw there in the loop."}, {"start": 4254.76, "end": 4257.28, "text": " So here we call the apply gradients."}, {"start": 4257.28, "end": 4260.68, "text": " We can do that because we have all of those ingredients inside."}, {"start": 4260.68, "end": 4262.12, "text": " Okay."}, {"start": 4262.12, "end": 4266.8, "text": " Let me kind of quickly walk you through that class so that it's kind of clear what's going"}, {"start": 4266.8, "end": 4267.8, "text": " on."}, {"start": 4267.8, "end": 4269.76, "text": " Let me open it up."}, {"start": 4269.76, "end": 4273.5599999999995, "text": " If I open it up and here is the train state function."}, {"start": 4273.5599999999995, "end": 4275.16, "text": " So we have the apply gradients."}, {"start": 4275.16, "end": 4279.74, "text": " We saw that one in the training loop and we have the create function."}, {"start": 4279.74, "end": 4285.44, "text": " So basically you can see here the optimizer is initialized internally."}, {"start": 4285.44, "end": 4291.08, "text": " So yeah, like as much work as possible was delegated to this function so that we don't"}, {"start": 4291.08, "end": 4294.2, "text": " have to externally do the initialization."}, {"start": 4294.2, "end": 4298.24, "text": " And here you can see what it does when we call the apply gradients."}, {"start": 4298.24, "end": 4301.0199999999995, "text": " It basically does those two optics functions."}, {"start": 4301.0199999999995, "end": 4302.58, "text": " So we have the update."}, {"start": 4302.58, "end": 4305.96, "text": " We get the updates in the new optimizer state and then we apply the updates."}, {"start": 4305.96, "end": 4311.4, "text": " It's just a simple, simple function that makes things a bit cleaner in our code."}, {"start": 4311.4, "end": 4312.4, "text": " So that's it."}, {"start": 4312.4, "end": 4313.84, "text": " Let me go back here."}, {"start": 4313.84, "end": 4316.4, "text": " So let me run this if I haven't already."}, {"start": 4316.4, "end": 4321.96, "text": " And we are getting to the final cell of this notebook where we have the training loop."}, {"start": 4321.96, "end": 4324.16, "text": " We create the training state."}, {"start": 4324.16, "end": 4329.48, "text": " We have it here and you can see how neat the training loop actually looks like."}, {"start": 4329.48, "end": 4331.679999999999, "text": " So we just call the train for one epoch."}, {"start": 4331.679999999999, "end": 4336.639999999999, "text": " So for every epoch, we just call the train one epoch and we get back from the trains."}, {"start": 4336.639999999999, "end": 4338.12, "text": " We convert the train state."}, {"start": 4338.12, "end": 4340.04, "text": " So there's going to be modified train state."}, {"start": 4340.04, "end": 4341.32, "text": " We're going to get the metrics."}, {"start": 4341.32, "end": 4344.2, "text": " We plot the metrics like the loss and accuracy, et cetera."}, {"start": 4344.2, "end": 4346.76, "text": " And we do the same thing on the test dataset."}, {"start": 4346.76, "end": 4350.92, "text": " So we pass the test images and test labels and that's pretty much it."}, {"start": 4350.92, "end": 4354.84, "text": " So I can try and run this, but currently I don't think I have an access to a GPU."}, {"start": 4354.84, "end": 4359.84, "text": " So it's going to be a bit slower, but yeah, it's just going to print every single epoch."}, {"start": 4359.84, "end": 4362.66, "text": " It's going to print some of this data."}, {"start": 4362.66, "end": 4366.36, "text": " And yeah, that's everything you have to know pretty much."}, {"start": 4366.36, "end": 4371.48, "text": " And in my opinion, this is probably this example, even though it's way more complex and actually"}, {"start": 4371.48, "end": 4377.679999999999, "text": " does a real thing, it was way simpler compared to those contrived examples where we had to"}, {"start": 4377.679999999999, "end": 4383.959999999999, "text": " do the mutable, where we had to pass the random states for the dropout and stuff like that."}, {"start": 4383.959999999999, "end": 4392.0, "text": " So as an exercise, how would we go about adding dropout here and also batch norm?"}, {"start": 4392.0, "end": 4398.599999999999, "text": " So that would like start adding that complexity on the syntactic level."}, {"start": 4398.6, "end": 4401.4800000000005, "text": " Let's roughly think what would happen if we were to add dropout."}, {"start": 4401.4800000000005, "end": 4409.4400000000005, "text": " So first we'd have to modify this create train state function and we'd have to add not just"}, {"start": 4409.4400000000005, "end": 4413.72, "text": " the key for the params, but also for the, we'd have to create a kind of dictionary here"}, {"start": 4413.72, "end": 4418.52, "text": " where we'd have like params and then we'd initialize that collection with a certain"}, {"start": 4418.52, "end": 4425.08, "text": " key and then we'd have the dropout and we'd have to initialize that with another key."}, {"start": 4425.08, "end": 4428.56, "text": " And that's on the initialization side."}, {"start": 4428.56, "end": 4435.080000000001, "text": " Then in the actual training loop, what we'd have to do is we'd have to pass the key for"}, {"start": 4435.080000000001, "end": 4438.4400000000005, "text": " every single train step, we'd have to have a unique key."}, {"start": 4438.4400000000005, "end": 4445.04, "text": " So because as you recall, we'd have to add RNGs here and then create and pass the key"}, {"start": 4445.04, "end": 4446.120000000001, "text": " here."}, {"start": 4446.120000000001, "end": 4451.400000000001, "text": " So that means that for every single train step, we'll have to have a key and let me"}, {"start": 4451.400000000001, "end": 4456.240000000001, "text": " see where we're calling the train step right here."}, {"start": 4456.24, "end": 4461.8, "text": " So that means we have to have splitting happening on this line here."}, {"start": 4461.8, "end": 4465.76, "text": " So we just have to pass a specific key to the train one epoch function and then kind"}, {"start": 4465.76, "end": 4470.76, "text": " of keep splitting because we don't want to reuse the key because if you reuse the key,"}, {"start": 4470.76, "end": 4474.96, "text": " you'll get the same output, the same randomness and thus you don't have this stochasticity."}, {"start": 4474.96, "end": 4476.96, "text": " So that's very important to keep in mind."}, {"start": 4476.96, "end": 4479.48, "text": " So yeah, it's a fairly simple modification."}, {"start": 4479.48, "end": 4483.719999999999, "text": " I just wanted to kind of quickly walk you through that and for batch norm similarly,"}, {"start": 4483.72, "end": 4489.360000000001, "text": " you'd have to have the mutable keyword here and there and it would complicate things a"}, {"start": 4489.360000000001, "end": 4492.04, "text": " little bit, but definitely not by a lot."}, {"start": 4492.04, "end": 4494.6, "text": " So yeah, let's see whether we have some outputs."}, {"start": 4494.6, "end": 4500.2, "text": " Yeah, we have some outputs here and you can see the accuracy is already and we don't have"}, {"start": 4500.2, "end": 4504.360000000001, "text": " any overfitting and so yeah, we have like a 98% accuracy."}, {"start": 4504.360000000001, "end": 4505.68, "text": " That's fairly nice."}, {"start": 4505.68, "end": 4506.68, "text": " That's pretty much it."}, {"start": 4506.68, "end": 4513.12, "text": " I encourage you to go and check out the Flex documentation A and B. Go and check out the"}, {"start": 4513.12, "end": 4515.86, "text": " examples they have on their GitHub."}, {"start": 4515.86, "end": 4519.5199999999995, "text": " So you can see like this ImageNet example is very cool."}, {"start": 4519.5199999999995, "end": 4521.8, "text": " They have mixed precision training."}, {"start": 4521.8, "end": 4528.92, "text": " They have like training over multiple devices, then like all of that additional complexity"}, {"start": 4528.92, "end": 4531.32, "text": " and I think they can train a model in like an hour or something."}, {"start": 4531.32, "end": 4535.5599999999995, "text": " Yeah, like wall time, a couple of hours to train these, although you have to have a bunch"}, {"start": 4535.5599999999995, "end": 4539.04, "text": " of TPUs or V100s to do that."}, {"start": 4539.04, "end": 4545.72, "text": " So yeah, go check those out at your own pace and I think I mentioned in the beginning,"}, {"start": 4545.72, "end": 4552.0, "text": " Hugging Face has very nice examples as well and they've organized this community week"}, {"start": 4552.0, "end": 4553.24, "text": " which you should check out as well."}, {"start": 4553.24, "end": 4560.08, "text": " So I'm gonna link all of these resources in the video description, so do check them out"}, {"start": 4560.08, "end": 4565.8, "text": " as well and yeah, I mean Hugging Face team is doing an amazing job, so just a huge shout"}, {"start": 4565.8, "end": 4566.8, "text": " out to those guys."}, {"start": 4566.8, "end": 4568.4, "text": " Cool, that's pretty much it."}, {"start": 4568.4, "end": 4569.4, "text": " Let's do a quick recap."}, {"start": 4569.4, "end": 4576.799999999999, "text": " Let's do a short summary of what we've learned, what we've seen in this video about Flex and"}, {"start": 4576.799999999999, "end": 4582.4, "text": " basically I'm gonna open up this table of content in deep note here to kind of guide"}, {"start": 4582.4, "end": 4589.7, "text": " my thinking and let me create a code cell here using Control J shortcut."}, {"start": 4589.7, "end": 4596.12, "text": " So basically we've seen the, so on top of everything you need to know about Jax already,"}, {"start": 4596.12, "end": 4599.88, "text": " we've seen the init and apply functions."}, {"start": 4599.88, "end": 4604.44, "text": " We've seen that if you wanna create our custom modules, sometimes you'll have to use the"}, {"start": 4604.44, "end": 4610.76, "text": " self-param and self-variable functions."}, {"start": 4610.76, "end": 4617.36, "text": " Then we've seen stuff such as nn-compact which makes the construction of those custom modules"}, {"start": 4617.36, "end": 4620.12, "text": " a bit more compact obviously."}, {"start": 4620.12, "end": 4626.72, "text": " Then we saw some new keywords that we haven't seen in Jax such as mutable if you wanna have"}, {"start": 4626.72, "end": 4633.92, "text": " the non-trainable variables or RNGs if you wanna have like stochasticity such as dropout"}, {"start": 4633.92, "end": 4638.12, "text": " like in the forward pass and what else?"}, {"start": 4638.12, "end": 4643.5199999999995, "text": " We've seen the train state object which is pretty much like a wrapper state that makes"}, {"start": 4643.5199999999995, "end": 4645.0599999999995, "text": " things a little bit cleaner."}, {"start": 4645.0599999999995, "end": 4649.24, "text": " So all in all, there is not that much you need to know on top of Jax, so if you already"}, {"start": 4649.24, "end": 4655.88, "text": " have the basics kinda set up, then you'll be good to go and this video will hopefully"}, {"start": 4655.88, "end": 4658.08, "text": " give you enough context to get you started."}, {"start": 4658.08, "end": 4663.2, "text": " Again, huge shout out to the Deep Note team for sponsoring this video and if you like"}, {"start": 4663.2, "end": 4665.5599999999995, "text": " this one, share it out with your friends."}, {"start": 4665.5599999999995, "end": 4667.5, "text": " That's the best way you can support the channel."}, {"start": 4667.5, "end": 4669.5199999999995, "text": " Also check out the reference."}, {"start": 4669.5199999999995, "end": 4672.24, "text": " I'm gonna link below in the video description."}, {"start": 4672.24, "end": 4676.84, "text": " By using the link again, you'll get 20 hours for free using the, you can use these pro"}, {"start": 4676.84, "end": 4678.88, "text": " machines and that's it."}, {"start": 4678.88, "end": 4681.8, "text": " Subscribe to this channel and until next time, bye bye."}]