Notebook: github.com/gordicaleksa/get-started-with-JAX/blob/main/Tutorial_4_Flax_Zero2Hero_Colab.ipynb My referral link: deepnote.com/referral?token=823d18856ad5 (you get 20h of "Pro machine" for free!). And you support the channel. ❤️ Again, I'll never promote anything I don't find valuable myself - so I strongly suggest you check them out. ❤️
@@TheAIEpiphany hey. first of all thank you for all the Jax and Flax videos. Still basically the only good ones on the whole web. Is Haiku video still planned?
At 01:10:58 I thought you said "ingradients" and thought it was a great pun, but then I listened again and it was just "ingredients" (which did make more sense, after all). I think I will save this video to come back to reference as there was a lot of info given, all of it great. I'm in awe and wish I could create the kind of content you do.
@@TheAIEpiphany Yes, please! I've run into so many problems with Haiku even though I'm more of a fan of its design principles. It doesn't help that the community seems much smaller. Could I request for something along the lines of an ActorCritic class where - the "forward" outputs a tuple of the actor and critics outputs - the actor and critic share a layer I ran into massive problems trying to get that off the ground
Let me try and clarify. First of all - It doesn't matter that much, the optimal point in the param space is invariant to loss scaling by a positive const (or in general, in theory, you could apply an arbitrary strictly monotonic function - like log), and the optimization process should be fairly robust to "whatever" positive constant we put in there. Div by 2 is there simply because of a convention (d/dx(x^2) = 2*x and that's why you'll see div by 2 in the literature the analytical expression looks cleaner, you end up with only x instead of 2*x). Even in the multi dim case where inner would give you L = x1^2 + x2^2 + ... when you do dL/dx1 you end up with 2*x1, and similarly for other dimensions (x2, etc.).
You're welcome! Well I use OneNote, Keep for notes. I set alarms on my phone not to forget stuff (although I just started doing that recently). Other than that, I just minimize distractions and try to have deep work sessions very often. :)
Nice tutorial ! But I still have a little question: what if we have multiple dropouts in a model? I noticed that the params are organized in dicts to be passed into init or apply, then what are the names of the multiple dropouts? I dont feel like they are using the same rngs. How to distinguish them?
Of course! DeepMind and Google Research are building everything in Haiku/Flax - that should answer your question. 😅 (TFDS is usually used for data loading itself)
Could anyone help out/point to some help on using JAX for timeseries modeling. Very difficult to use the common sliding window method when generating features. Is there a fix or does it need rethought entirely?
Currently (19 Mar 2022) your notebook fails at "import flax" stage (on colab, didn't try other options). I could only run it by replacing stuff with the CPU-only !pip install jax jaxlib flax
Thanks a lot. I must say, at the first glance, this FLAX/JAX thing seems to be even a bigger crap than pytorch is (which is a big thing to me). It's hard to understand why people cannot do things such simply like with keras. Why they need to do all these dance with a simple training? With keras, you just throw the layers on the top of each other, do a model.compile() + model.fit(), and ready. If you want to do tricks, you can, but you don't need it. With FLAX, generate this variable, add that param. Param has one syntax, variable has the other. Keep in mind that this layer needs this variable that needs another. Just doing a synchronized batch normalization seems to be rocket science... Anyway thx. again, hopefully, I will get used to this thing too.
Because the people who design these frameworks are purist Computer Scientists and Mathematicians with little to no consumer product experience. So they design what they want and your job is to praise them for doing so because they are very very very easily triggered and really ought to resolve their academic PTSD, but, I digress. That said, Jax is the least crappy just for treating accelerators like GPUs and TPUs as first class citizens instead of a sad afterthought like everything else. Both Andrej Karpathy and Francois Chollet have demonstrated how to simplify as you desire. And yet they don't get the backing to build a new framework that way. Maybe John Carmack can do a thing or two about this.
Keras is pretty good for rapid prototyping, but its not that good when it comes to configuring you model. Tbh, I have seen people don't even know how 3d matrix multiplication is done and using keras and building models. It works fine til you get nans, or get stuck at a local minima loss value or get other errors. That time you don't know what to do in Keras cause you never visualized your model or cared about hyper-parameters and just created a model with any random architecture you liked. pytorch on the other hands is very close to mathematical side so you realize what is actually happening, compare to keras model which are just `feed the data and get the output whatever it is`. Also another part, in research papers every tiny bit of edge matters in performance, this is where things like JAX shines.
@@shubhampatel6908 "Also another part, in research papers every tiny bit of edge matters in performance, this is where things like JAX shines." you said. That is quite a bold statement, more papers use torch than JAX.
@@MrEnyecz your srgument doesn't make much sense. When torch was new, more papers used tensorflow than torch lol. Also I didn't say every research, I said the research where every bit of edge matters.
