Coding a Neural Network from Scratch in Pure JAX | Machine Learning with JAX | Tutorial #3
❤️ Become The AI Epiphany Patreon ❤️
https://www.patreon.com/theaiepiphany

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

Watch me code a Neural Network from Scratch! 🥳 In this 3rd video of the JAX tutorials series.

In this video, I create an MLP (multi-layer perceptron) and train it as a classifier on MNIST (although it's trivial to use a more complex dataset) - all this in pure JAX (no Flax/Haiku/Optax).

I then add cool visualizations such as:
* Visualizing MLP's learned weights
* Visualizing embeddings of a batch of images in t-SNE
* Finally, we analyze the dead neurons

Credit: 
Got the inspiration from the official advanced JAX tutorial here: https://jax.readthedocs.io/en/latest/notebooks/Neural_Network_and_Data_Loading.html

▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬ 
✅ Get started with JAX GitHub: https://github.com/gordicaleksa/get-started-with-JAX
✅ Dead neuron article: https://pytorch-lightning.readthedocs.io/en/latest/notebooks/course_UvA-DL/02-activation-functions.html
▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬ 

⌚️ Timetable: 
00:00:00 Intro, structuring the code
00:03:10 MLP initialization function
00:13:30 Prediction function
00:24:10 PyTorch MNIST dataset
00:31:40 PyTorch data loaders
00:39:55 Training loop
00:49:15 Adding the accuracy metric
01:01:45 Visualize the image and prediction
01:04:40 Small code refactoring
01:09:25 Visualizing MLP weights
01:11:30 Visualizing embeddings using t-SNE
01:17:55 Analyzing dead neurons
01:24:35 Outro

▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬ 
💰 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ć
Bartłomiej Danek
Zvonimir Sabljic

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

💼 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/

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

#jax #neuralnetwork #coding
 What's cracking guys? This is the third video in the Machine Learning with Jack's series and in this one We're going to build a neural network from scratch I'm going to be using a multilayer perceptron and I'm gonna train a classification model on top of the MNIST dataset By the way, I'm going to be using MNIST just for convenience Conceptually, even if we were to use ImageNet, the code would be pretty much the same Just for the sake of me coding this real-time and this video not taking three hours I'm going to be using this small dataset such as MNIST Finally, we're going to be using a PyTorch data loaders to load the data from MNIST and the reason is Jack's designers by design decided not to develop loaders because other libraries such as frameworks such as TensorFlow and PyTorch have already done a decent job at Adding this functionality so they did not want to reinvent the wheel So that's in a nutshell what this video will be about and then as a bonus points I'm gonna do some visualizations and finally, we're hopefully I'll be able to parallelize the training in the case that Colab allows me to use the TPU if not Previous videos already explained how to do that. So that's just a bonus point Anyways, let's let's start with that and structure this this whole video. So first thing we're going to do MLP training on on MNIST dataset, right? That's the first thing the second thing we're going to do is we're going to do some visualizations and Finally, the third thing is going to be basically adding the Parallelization so yeah, hopefully I spelled that correctly So now let me break these tasks into smaller sub tasks first thing where we need to create a init function in it MLP we need to and add the Predict function. So those are going to be two things we need to do here Whoops to do whoa. I don't know how to write today. Okay, that's the first to do the second one will be basically add data loading in pytorch and the third to do is going to be add the training loop loss function and whatnot maybe accuracy or something like that So that's the the breakdown of the first section of this video now here. I'm gonna just do something like maybe visualize visualize the MLP weights so we're just gonna See whether we can notice some patterns especially in the shallower layers whether MLP learn to extract certain features Which are salient, which are also like understandable and intuitive to us humans Then I'm going to do something like I'm gonna visualize embeddings using t-sne and finally, I'm gonna do something like visualize dead neurons if I have the time I Won't be breaking this parallelization part because that's that's pending and we're gonna see whether we have time A and B whether I'll have TPUs on disposal. So let's start by breaking down Solving one task at a time. So first let me start with initial init and MLP Initial init MLP function. So that function should given some configuration of the MLP network It's gonna return the weights and biases of the network. So that's a high-level description of what I want to achieve here Now let's go and implement the function. So we're gonna call it Something descriptive like init MLP, I guess and I'm gonna pass layer widths So that's gonna be the configuration and then we're going to start Implementing the function here. So params is just going to be some list basically, we're gonna iterate through this like layer widths and we're gonna take the In width and out width will be the names. So that's the number of Neurons that comes into the layer and number of neurons that goes from the layer And basically I'm going to iterate through the zip version of this so layer widths and I'm gonna index it like this So basically we want to start here with the 0th one and go and just exclude the last element and here I'm gonna Oops, let me paste it. I'm gonna start from the first one and go all the way including the last element So hopefully you'll understand a second why I'm doing this And let me maybe first create the high-level let me call the function and see how that will work So I'm gonna have ML params here. I'm gonna call the init MLP I'm gonna call the init MLP. I'm gonna pass the list here. So First thing that we're gonna do is we're gonna add like the number of Elements of pixels in the MNIST data set then I'm gonna arbitrarily create some configuration So we're gonna have a hidden layer with 512 neurons then 256 and finally 10 So what this function is doing is as you can see here It started from it starts from the 0th one and this one starts from the first one So that means we're gonna be fetching tuples like this this one and then this one and then this one because that will enable me to Create the weight and biases we need for each of the layers. So let's do that So I'm gonna append to params the following thing. So basically We're gonna create a numpy Random we're just gonna randomly Sample data from the normal distribution That is gonna have a shape like this so out with and this is just convention We usually just first put the out Dimension here and then the input dimension. So there's gonna be for example 512 and then 784 here you could it vice versa, but it's gonna complicate things later on. So there's just a convention to keep in mind So next thing we'll need We'll need a bias term. So there's gonna be something like this random And and we're gonna need just the out with because that's the bias part, right? So this should now work Let me close this list and let me just return back the parameters here. Let's return the params and Let's see whether this is going to work. So In order to understand whether this is correct. We can maybe use the tree map functionality from Jax So I'm gonna just print Jax tree map I'm gonna create a lambda function here Which is going to just return the shape of the leaf and the leaves are going to be fetched from the MLP params Which means we're going to be taking weight and biases these matrices here and we're gonna be printing their shapes To understand whether this is implemented correctly. So obviously did not import Jax neither Neither did I import non-pice is gonna fail. So I'm gonna first import those libraries import numpy as MP import Jax and I'm gonna run this and I'm going to run this as well Let's see whether this is going to work. So it's invalid syntax. Let's see what I got wrong here NumPy random. Oh, yeah, I added a Closing bracket here so that failed. Let me rerun it Okay, so this looks correct. I think so we have 512 here and then we have 256 here and then we have 10 I'd say this is this is correct. So there's a couple of modifications We need to add here first thing is a however to just use the standard deviation one as the initialization This thing is gonna explode. We need to scale these parameters. I'm not adding any fancy initialization method here I'm just gonna scale these like this just gonna like basically Make the standard deviation smaller By a factor of hundred and that's it. The second thing we need to add here is now that we have implemented numpy's random Number generator so we'll actually be using Jax random number generator So that's why I'm gonna just do that right now So it was easier for me conceptually to first start with numpy now. I'm gonna just quickly Change this so I guess checks random PRNG key or something is that the functionality I'm gonna add a seed here. Whoops I'm just gonna add some constants here. So C is gonna be zero I'm gonna be using seed here and I'm gonna pass RNG inside of the initmlp function we need to add it here as the argument and now the thing we need to do is create the keys out of the this generator So this is actually also a key. So I'm gonna rename it because it will be misleading I'm gonna pass the key here and from the key, I'm gonna create keys using the split method So I'm gonna create keys using Jax random split method and I'm gonna pass in the key and I'm gonna set the number of keys to Whatever the number of layers is so that's length of this thing minus one, right? Because the number of layers is by one smaller than the length of this configuration file because we're taking tuples Remember so this this and then finally 256 and 10 as the last layer So that's gonna be the number of keys and now we are going to additionally add to this zip function I'm gonna add the keys and I'm gonna then add the key part here. And finally, I'm going to split it additionally into the weight key so we need to split it into weight key and Bias key so something like this. So I'm gonna do again Jax random split and I'm gonna pass in the key and that's gonna split Let me see whether this is going to create a problem because this thing I'm gonna call this the parent key just for the sake of Avoiding some bugs here. So parent key here and creating keys and I'm iterating through the keys here I'm splitting that into the weight key and the bias key and finally we're going to now Input the weight key and the bias key and create the correct random function from Jax Not completely sure about this syntax here I'm gonna just try something that makes sense and then I'm gonna search if it doesn't work So Jax random something like that and then I'm gonna pass Like the key so I'm gonna pass the weight key And then I'm gonna paste Pass the shape here. So this should be the syntax and then here I'm gonna again Do Jax random and I'm gonna pass the bias key here and the shape is going to be something like this So this should now work. Let me kind of split this into two lines. It's more readable Let me let me kind of try and refactor this a little bit something like this and then I'm gonna indent this a little bit more Fuck my life. Okay, and I'm gonna do it like this Okay, this looks a bit better. So hopefully this works. Let's see whether it makes sense. We have the parent key We split it into multiple keys here. We use those keys So we'll have a key per layer every layer We have a specific key and then we do an additional split because every layer has weights and bias We pass the weights we paste the bias and that's it. This should now work as expected. Let me try and run it hopefully it's gonna work and It's not working because a module object is not callable. Let me see. What have I done here? Aha I called random is the module not the function. So let's see how the the actual syntax looks like Jax random random. I don't know. Let's see Whether this is going to Yeah, this is the wrong file, let me try and see whether we have some interesting information here So we have random uniform here, but do we have random something? That's not uniform generator because uniform distribution is not as good for parameter initialization as Gaussian because Gaussian tends to because of the mean zero We're gonna have an expectation smaller weights than using uniform I guess and also it's gonna be symmetric around zero. That's important. So I really need to find the Actual Gaussian Random, let me just use Something like this Gaussian Jax random Gaussian Will this work? Let me open up a couple files here I Think we saw this one Okay, let me let me see whether I can just Grab a number from here Nope random seed random uniform They're just using uniform numbers all the time here. What what what's this? Oh Random normal. Okay. Okay. This is the syntax random normal. Whoops that that took some time. Okay Jax random and then dot normal and that's gonna work random dot normal and hopefully now this will work as expected and It seems it does everything looks correct here, I think we are we are good to go to the next section so We have the correct scale. We have the correct Structure three layers every layer has weights and every layer has bias We've used Jax's pseudo run random number generator and we had some primitive initialization scheme here using the scale parameter Now let's do the second part and that's the add the predict function. So now that we have this Let me just add another cell here Let me add a function called MLP predict we're gonna pass in the parameters of the MLP and we're going to What else we need some data? So let's pass in the input data here. This is going to be a single image And a flatten one because we're using an MLP. We're not using a convolutional neural network here. So How we're gonna split this is we're gonna take the hidden layers here So maybe let's call this Yeah, I'm just gonna do it like this so hidden layers I'm gonna take the The params all the params except for the last one That's gonna be the hidden layers the two one the first two layers and then I'm gonna iterate through those layers So I'm gonna just iterate like this We're iterating through the through the hidden layers and taking the weight and bias matrices Then we need to Set activation something like this activation is gonna be initially equal to whatever the input is And then we're just gonna reuse it here. So activation activity Reasion, I don't know how to spell today Basically, this is going to be What we need to do activation function and we need to create the the the mapping so mapping will be basically Jax I guess stop product No, I'm gonna add the jack jacks non-pay API here. So I'm gonna add Jacks non-pay as JMP because that's gonna be useful. I'm gonna rerun the cell and Finally here so JMP dot should work. I'm gonna pass in the weight and then the example and then I'm gonna add the bias term so this thing should be This is the functionality of the linear layer. Finally, we need the relu So I think and then contains the rally unit. Yep, it does so something like this should now function and this should not be Exit should be activation. Okay So this show work Let me see whether it makes sense So we're doing a dot product. We are adding the bias. We are doing the activation function and we get the activation out This looks reasonable. Let me now do the final step. The final step is we get the last and The last layers weight and bias so something like this params Of the last element of the last layer and we do apply this again. So I'm just gonna copy paste this We're gonna copy paste this I'm gonna use the last one The reason I'm doing it like this is because the last layer does not have the relu activation so these are going to be just logits and We could alternatively do it in a single for loop but then I'd have to have a switch statement like if Last layer then apply relu otherwise don't apply relu. So at the end is gonna be even Less readable. So I'm just gonna do it like this and finally we're going to return What we're going to return logits But not quite because Jax does not have cross entropy loss implemented. We'll have to be resourceful here and Create something that's going to be Maybe confusing at the first point. We're going to be implement importing from the Jax SciPy, I think special we're gonna import something called Like X log some or something X log some I don't know. Let me try and Search for this not completely sure So, yeah log some X log some X log some X so this thing is gonna allow us to Implement a numerically stable version of cross entropy So what this what this is is the following so imagine you we have ten outputs, right? We have ten outputs from the this amnest MLP model And let's take a single one. Maybe the single output will be called a one So this is the raw output from from from the logits and we can Equivalently represent this without changing anything like this log X of of one because These two are just inverse of each other. So this thing is equivalent to this thing here. So now What this will do is it's gonna add a log sum So it's gonna do something like log and it's gonna have a sum of Exponentials so something like this Oh one and Then it's gonna have obviously because it's a sum. It's gonna have as a second argument The second output and then the third output and you get the point all the way down So if you just look at this expression now, let me just kind of copy paste it here So that it's more Visible and now we're just gonna apply the simple rule. Basically if you try and subtract the logarithms That's equivalent as if you were to take the arguments and just divide them So we're just gonna take the this argument here. So The left argument and we're just gonna take the left argument And we're just gonna take the right argument and we're just gonna take the left argument And we're just gonna divide it by this sum here of all these basically exponents and you can see that this is Pretty much we just implemented softmax The only difference is that we have log softmax and that's it So this is what this thing implements and that's gonna be super convenient a bit later We're gonna see why in a second, but let me now test this function as well So as you can see here every time I write a certain module like this init MLP, I test it here so this is a Not not that rigorous testing but some testing nonetheless so let me try and Pass some data here. So I'm gonna just create a simple dummy Like image here. So dummy image flat. I'm gonna call it that way and this is going to be ran then Basically again is gonna be MP prod of the M NIST Image size this should work correctly Let me just print the shape to make sure this is this has the correct shape. It should be a flat image so 784 Should be the shape and we're going to pass that into the MLP predict and see what we get as the prediction So prediction is going to be MLP predict We pass the MLP parameters we generated up there. We pass it here. We pass in the Image here So something like this and this should now work Hopefully, so I'm gonna just print out the prediction Shape here. So let me try and run this thing So name log some XP is not defined because I have not rerun this thing Whoops, okay. Let's see what's wrong here. Let me reopen the tab here Jacks SciPy special Jacks SciPy special aha From Jacks SciPy special import. Yeah I'm really bad today. Okay. Let me try and rerun this again And see whether it works. It does not because unsupported operand types unsupported operand types, blah blah blah the device array and function Okay, something went wrong here. Let me see what so I'm passing the MP prod should give me so how you debug this is you just open up another cell here I'm just gonna do it like this I'm gonna just print out the MP prod of this thing and see whether this gives me the number it and it does so this thing returns back the number and For some reason this is not working Okay, I got it all wrong. This thing actually worked. The problem is here and even have the red squiggly line I'm stupid basically. I just I just subtracted a function from from from this I need to pass in the logits here and now this should work. Hopefully, okay, and it does work. So the shapes do work as expected so now because this is Jax without Reimplementing this function for the batch of data. We're just gonna wrap it up into V map and let's see whether that's gonna work So I'm gonna create a like a batched version so batched Whoops Let me zoom in again. So batched MLP predict is going to be equal V map and then I'm gonna wrap the MLP predict here and Then we're gonna specify the in axis Let me import V map so that I don't forget so from Jax import We're gonna need JIT. We're gonna need V map and we are we're gonna need P map and grad as well So let me rerun this and now let me continue here So in axis is gonna be none because we don't we just want to broadcast the parameters and it's gonna be zero for the Because we have the the batch dimension will be on the zero on the zero dimension So this should now work I think that's it and let me try and test this function again. So again Let's test it Small test here Basically, I'm gonna create this same thing Except that this time I'm gonna add a batch dimension of 16 here Let me check the shape here So this is gonna be images flat and I'm gonna pass that into the batched version and see what the predictions are So predictions predictions Batched version we pass in the again MLP parameters. We pass in the images flat And if I were to run this I'm gonna just print the predictions Shape and I'm gonna comment out this first test here So this should now work as expected hopefully and it does so we get the oh, I forgot to actually use this this version of the images so this should be 16 Yeah, so this seems to to work correctly and that means we have done the first part here And now let's jump to adding the data loading functionality. So I'm just gonna annotate this like this So this was a small test and this thing These things here These things here were also tests and Finally, okay, let's add another cell here We need to add the data loading so First thing we need to add the MNIST dataset and then we'll need to add the data loader Which is going to load the images that we're gonna pass into this into this MLP So let's start with the syntax PyTorch dataset MNIST I'm gonna search for it and let's see how the syntax goes like So basically TorchVision datasets MNIST something Let's let's see some example whether they have MNIST here. They do not have MNIST Let me see whether I can type something in like TorchVision datasets MNIST example Okay, let me return back here I think this is going to work I'm just gonna copy Copy this thing. Let me put this tab here. I'm gonna need it probably a bit later So something like this should work. I'm gonna import here so import So from these datasets, I'm gonna import the MNIST dataset And now I'm just gonna go ahead and add the data loader and import the MNIST dataset And now I'm just gonna be using MNIST. So I'm gonna run this Hopefully this import will work Let me see Okay, it worked. So now I'm gonna use MNIST here Basically, let me delete this and I'm gonna create a train dataset By calling MNIST and passing in let's see what the arguments here are So we have root that's where the data will be downloaded. Then we have the split I'm just gonna leave this as the default one Train, download and transform So it sucks that it doesn't have an example here that can just copy paste So let me just do it like this. We need roots. We need basically split We need to train, download and transform. Those are the four we'll need So root is just gonna be basically The structure, how it looks like, is the following So if I was to add an additional cell here Import OS package and print the OS get current working directory It's gonna print content. So this thing here is called the content directory So I'm gonna, inside of this directory, I'm gonna into the I guess train MNIST I'm gonna download the train dataset And we're gonna need, what else? We need the train equals to true So that's gonna make sure that this is the training dataset And finally we need download and transform So download, I guess download is by default true here, right? So download, if true, downloads the dataset from the internet and puts it into root directory If dataset is already downloaded, it is not downloaded again So yeah, we wanna download it and then transform is gonna be needed as well We're gonna see why in a second So for now let me just omit it I'm just gonna put transform to none And that's it Let me just print the type of this thing Train dataset just to have something as a placeholder And let me rerun this, let me close this, let me rerun this and see whether this works And it does It's downloading the data, I think it's finished We have train MNIST dataset right here in this Google Colabs storage space So let's now analyze what this dataset is So basically if I were to fetch like a something from the train dataset So I'm fairly sure it has the index functionality So I can just index and take the 0th one I'm gonna print the type of this something And that's going to be basically, as you can see here, pill image and we have a label So that's the structure of our dataset So because we're working with Jax, we're not going to be using PyTorch tensors So that means we need to convert this into a NumPy array and not pill image So how do we do that? And the answer is using this transform functions We're just gonna add a simple transform function So I'm just gonna call it custom transform And what the function is going to do is It's gonna take the image and it's going to convert the image into NumPy So something like this, so NumPy array of this x And I'm just gonna set the data type to MPFloat32 And I'm gonna return this thing I don't even need a temporary variable here So just return this And now if I were to run this, let me see whether it's gonna be converted into NumPy And it's not Because I forgot to add the custom transform here So let me run this now And yep, I'm getting unveiled the output here, which is a good sign actually So let me, instead of printing something, let me print the something 0 And then the type, this should be, I guess, NumPy And it is Okay, so we got this sorted Basically we now have a NumPy Float32 image And the second, let me just check what the second part of this something tuple is So let me just check whether that's integer or whatnot So it is, it's an integer So we have, because it's a label So I think that should be alright Now, we need to additionally create the test dataset So I'm just gonna copy paste this part here I'm gonna copy paste it here This is going to be a test Train is false because we want a test dataset Download equals true is okay And custom transform, that's all okay So let me just run this again and see whether it's alright So it's downloading the test data And as you can see here, we have the test MNIST in the storage space Okay, cool Second thing we need to do is add the data loader So we need to be able to load these images in I think I just understood as I made a mistake here We need to flatten out images If I were to, again, if I was to take the image label tuple Because I know what this is now I'm just gonna take the train dataset I'm gonna take the 0th one And I'm just gonna take the, actually just Let me just grab the image By adding 0 here This is an image, but the thing is The shape of this image is, I guess, 28 by 28, yep So we need to add the Just flatten the function, flatten the thing here And if I were to run this, we get the correct shape now So this is now correct Basically the last thing we need to add here is just data loader So if I was to search it PyTorch data loader Let's see what I get here Data loader Okay, we have the signature here I'm just gonna copy paste it And let's see whether that's gonna work for us So, I'm gonna obviously need to import the data loader functionality So let's see where can I import it from Torch vision or what not, I don't know Let's see here Import PyTorch data loader Let's open up this data tutorial So basically data loader Oh my god, I hate this Okay, from Torch, utils data, import data loader Okay, let me go back here Let me import that part Let me rerun the cell And now we have the data loader Okay, so this will work I'm gonna use the train data set I'm gonna create the train loader And I'm gonna add some random parameters here So batch size will be set to, for example, 128 And so just passing the batch size Shuffle is gonna be equal to true Because this is a train data set We wanna shuffle We don't need the custom sampler We don't need a batch sampler Number of workers is not important We'll need a collate function Because now we have a Basically we're not using PyTorch sensors These are all just optimization parameters Such as pin memory So we don't need none of those So let me just do it like this So we additionally need to create a collate function Because let me see what happens if I were to just Use it without a custom collate function So batch of data is equal to next I'm just gonna convert the train loader into iterator And then just fetch the next batch using the next function Let me see what this thing is Batch of data And whether it's gonna work So it's a list Okay, let's see what this list consists of So if I were to take the zeroth element And run it again It's a tensor So as you can see PyTorch by default creates a tensor If we do not add the custom collate function So that's why I'm gonna add a custom collate function here So def custom collate function We pass in the batch of data Basically we want to have all of the images Batched into a single lump array And we want to have all of the labels in a separate data structure Let's see what we currently get So I'm gonna just grab this custom collate I'm gonna paste it here So this batch input is just going to be a list of tuples And each tuple contains the numpy image and contains the label So we need to convert that into a bit different shape So instead of having a list of tuples We need a list with two elements One is all of our images batched into a numpy array And the second one is all of our labels basically Inside of a single numpy array So just to make sure I'm correct here Let me print the type of batch that should be a list Then let me print the type of the batch of zero that should be a tuple Finally let me print the type of the batch of zero And then zero this should be numpy And if I were to copy paste this This should be I guess integer And yeah This is gonna fail obviously because we're not returning anything Hopefully the print functionality will work nonetheless Yeah we have a list, then we have a tuple, then we have numpy array, we have integers So I already knew the structure obviously If you did not know the structure you can just kind of play with it Just analyze the shapes, analyze the types, etc. Okay so what we need to do here is we need to transpose this data So I'm gonna do it something like this Transposed data is gonna be I'm just going to unpack the batch of data here and then zip it So what this will do is we'll now have So this is a list of tuples After we do this we'll just have a bunch of tuples and zip We'll take the first element from every tuple that means all of our images And it's gonna put it into a single list And then it's gonna take all of our labels from all the tuples And it's gonna combine it again into a list So we can make sure that's correct So this again should be a list This thing here And basically the zeroth element should be a list as well And the first element should be a list as well So this should be All of these should be lists Let me check whether that's the case Zip is not subscriptable Let me see what I have done So basically let me check Oh yeah, I have to convert this into a list explicitly here So let me now run it And yeah, so we have list tuple tuple So yeah, instead of a list, they're actually using tuple in the background If I were to try and So the length of this thing should be 2 I guess Transpose data length should be 2 because now we have Yeah, it's 2 Because now we have all of the images and all of the labels in a single tuple So the next thing we need to do is basically Convert the tuples into a NumPy array So I'm gonna do something with this Labels is a NumPy array of this transposed data Of 1 because the second tuple is the labels And then we need images And here it's a bit more complex We need to do MPStack of transposed data of 0 And this will hopefully work as expected Now I'm gonna return back the images and the labels And this collate function should now properly work Let me run this and see whether we have the correct result It seems nothing is crashing So I managed to get the batch of data We're printing the type and it seems Okay, it's a NumPy array Let's see what the shape is So if I was to take the Let me do it like this So batch of 0 are images So that's like this, images And oops And finally, let's see what's the shape of the images Let me delete this thing And run this And yeah, we have a batch of 128 images Which are flattened out And if I was to take the labels Same as this here I'm just gonna add Instead of 0, I'm gonna put 1 for labels And let's see the shape of these labels This looks correct And finally, we need to check the data type So images Let's take the 0th image and print the data type And also let's take the 0th label and print the data type So it's always a smart idea to understand your shapes and your data types To avoid some bugs down the road So flow 32 looks good and in 64 I think that's gonna be totally fine Okay, we have our data loading functionality Let's continue on and see what else do we need to do And that's adding the looping function So the training loop functionality Let me delete this because this is done And let me now finally, let me also delete this And let me add a final cell adding a looping function A training loop function Okay, this should be fairly simple What we need to do is define the number of epochs So number of epochs is gonna be, I don't know, for now let's put something like 10 So we're gonna iterate through the epochs So in range num of epochs We're gonna do it 10 times We're going to do what? So we have this pattern where we take the MLP parameters So we have the MLP params We're gonna do some update function To which we're gonna pass the MLP params We're gonna pass the data So images and GT labels And this should just update our parameters And that should be the high level functionality So now we need to write this update function Let's see how we can do that So def update We pass in the params We pass in the images We pass in the labels Okay, obviously I forgot to add data loading functionality here So we'll have basically for what? We'll have images We'll have labels in the train loader And now we can pass the images and the GT labels But there is a small catch here So images Images are currently on a CPU And that's gonna be handled correctly by the Jack Slump API So we don't need to worry about it The thing we need to do here though Is to convert the labels into one hot representation We'll soon see why So let's first start developing this update function And then we'll understand what else we need to develop Okay, we obviously need to call the batched MLP predict function Okay, this one We're gonna pass in the parameters That's this argument here We're gonna pass in the images And this will return the expected shape of this Will be whatever the batch dimension is And then 10 because we have 10 outputs, remember And these are log softmax outputs So let's call those predictions So those are predictions here Finally, this is the trick that I told you about That will help us achieve the numerical stability of cross entropy loss So we'll need to just do something like this We need to multiply predictions with these labels So these are just GT labels So ground truth labels If we multiply them and if they are one hot That means we'll fetch the log softmax component From the predictions wherever the true value is And we'll want to push that to one So we'll want to have a probability one there Because GT label basically says that That's where the true class should lie in So now we just need to basically add minus JMP mean here And this is going to work as expected Because this is going to fetch the correct component, this part here And then we just do the mean This is gonna give us the cross entropy implementation pretty much Later, if I have enough time, I'm gonna do a more intuitive implementation We're gonna see it's not working because of numerical problems So the final component, as you can see here We need to convert this into one hot representation So that means that GT labels are actually JECs Let me see whether it's in... yeah, it's here So basically we just convert labels into one hot I'm gonna make sure this shape here is correct So labels shape And let's see whether that's gonna work or not Okay, this is obviously a dummy mistake This should be in a loss function, not in the update function So in this loss function, we pass in the parameters We pass in the, I guess, images and the GT labels And here we need to do this thing All of this goes into the loss function And now in the actual update function What we'll do is we'll take the gradients So we'll just find the grads So we'll call the grad function And we'll wrap in the loss function We're gonna pass in whatever is required So that's params, images, and GT labels And this is going to give us basically derivatives Of every single parameter in this pyTree here So we get the grads here And finally we need to return the JECsTreeMultiMap Is the function I need here We're gonna just do a simple function here So we get the weights, we get the gradients And what we do here is Okay, I don't need this here Let me just do the stochastic gradient descent update here So basically from the weight we subtract learning rate times the gradient Whereas I'm gonna just add some learning rate here by default Something like, I don't know, 0.01 And finally let me pass in the pyTree So that's params and grads So this should now work as expected But we don't have any metrics here Or any logging of whatever kind So neither the loss nor the accuracy So I'm just gonna do a small break here And then a break here just to make sure that this thing works So one hot required one poly... Okay, so I forgot to add the number of classes of MNIST So MNIST classes I'm just gonna pass in the length of that thing here And see whether this works now So learning rate is not defined Why not? I just added it here So let's see Oh my god, I actually placed this here I don't know why I'm making so many mistakes Learning rate here, now let's try and run it Yeah, I'm gonna expect this is gonna last a little bit Okay, so this thing actually worked Now let me try and grab the loss Not just the grads here So well, I think the function is called val and grad So let me just check that Val and grad, jacks Let's see whether that's the name of the function Jacks package Value and grad is the name of the function, okay? So we have value and grad here Value and grad So it's gonna return the loss and the grad And here I'm gonna return the loss and the parameters And so let me just do loss here And let me print the value of the loss And let's see what happens here Okay, I forgot to obviously import this thing here Ummm I guess it's from jacks, right? Yep, it seems that's correct Let me run this And we get some reasonable value of the loss function here So that's cool Let's now try and remove the break And let's run this for a single epoch And see whether the loss goes down So let me run this and see what we get So it seems that the loss is going down And that's very, very reassuring The only problem is we're getting too many output here So let me just interrupt the execution here And let me clear this output Wow, there's so many Okay, let me just clear this I'm gonna add a conditional logging here So if, I'm just gonna place a numerate here So that we know the batch ID So something like this And if the counter So every, because every batch has 128 images So let's say we can log every 50 batches or something So if this equals zero Finally, let's move the loss inside of this if statement So let me indent it here I'm gonna also just print this thing once And now this looks like it should work Let me run it And see, so 128, 10, this looks to be a good format So I'm just gonna interrupt this execution I'm basically going to just leave the loss printing here And the final thing before I'll remove this break Is to add the accuracy metrics So let's add the accuracy logging as well And then we'll be certain that MLP is learning something So accuracy Again, I'm not sure whether Jax Let me see whether Jax has accuracy metric Yeah, I don't see anything So I'm just gonna implement it myself So basically we're gonna pass the parameters again We're gonna pass the loader So either test or train loader And that should be it I think So the thing we need to do is find the number of correctly classified images So I'm just gonna basically iterate through images and labels Here in the loader I'm gonna again run the So I'm gonna call the batch So something like this I mean this is going to be inefficient Because I'm gonna run it in general if we had some arbitrary data set But this implantation is gonna work for MNIST pretty well So I'm just gonna go through the whole training set and through the whole test set And find the accuracy and report that somewhere here So after we finish the epoch I'm gonna print something like this So epoch And then I'm gonna print the ID of the epoch And then I'm gonna print the train accuracy equals to this And test accuracy equals to this I'm just gonna call the basically accuracy I'm gonna pass the MLP params after the epoch has been finished I'm gonna pass the train loader So this thing And I just forgot to So I basically just realized that I forgot to add the test loader So I'm gonna quickly add that as well So something like this Let me paste it here So now I have to add the test loader Somewhere above Let me see where I've defined Yeah, here So I'm gonna add the test loader It's gonna be the same as the train data loader I'm just gonna swap this for test data set Batch size, I can leave it the same size Shuffle, false No need to shuffle that data set And custom collate function is also fine So I'm gonna rerun the cell And now we have the test loader So this should now all be correctly set up We just need to implement the actual accuracy metric here Let's see, we need to do the same thing as here We need some predictions And then I'm gonna find the basically Instead of finding the raw predictions here So the log softmax I'm gonna do immediately call the soft argmax function And the axis should be 1 So that's gonna find the highest probability class For the particular prediction And we're gonna get basically prediction classes here And we're gonna contrast that with labels pretty much So we need to do something like this We'll need some sum temporary variables So sum, I'm gonna call it accumulator Basically accumulator is gonna add here the MP sum Between, which is gonna do predict classes And we're gonna see where those predicted classes are the same as the ground truth classes So I can maybe do something like this just for consistency sake So this, and we're gonna sum them up Find where we had true, that means we had the same class So this is going to work correctly And finally in order to return the metric We need to just divide the accumulated value With the length of the actual data loader I guess Is that correct? No it's not because we have to also add We care about the number of images So I guess we need to multiply that with the batch size So let me see whether loader has batch size variable I'm not sure whether this is going to work Whether Colab autocomplete is just tricking me here But this should work And this may be misleading because this is not the accuracy This is just the accumulated value So I'm just gonna replace this with something like Xsum And yeah this is my find and replace method for Google Colab Okay, this should not work I think Aside from this thing maybe not being correct Let me try and run this Maybe instead of waiting for the whole epoch I'm just going to run it right at the beginning here once Let's just see whether this is going to work Okay this actually failed And the reason is already quite familiar to me from PyTorch So what I forgot to do is basically we have to set the drop last to true here Otherwise we'll have the... So potentially the last batch will not have 128 images And that's where the discrepancy comes to be So if I were to rerun the cell here and then start this thing again Now hopefully we'll get the correct results So basically what happened here is at one point the last batch was smaller than 128 And there was a mismatch doing a comparison here and that's why it failed Here you can see it's not working but the thing is it's very very slow So what I'm going to do instead is I'm going to interrupt this And I'm going to load the whole MNIST into memory because it's a fairly small dataset with only 50k, 60k images So I'm going to do the following I'm going to change my implementation of accuracy So instead of doing the loader thing and then iterating through the whole dataset I'm just going to pass the dataset itself And that means this is going to change a lot So I'm going to just do the following Pass the whole dataset all at once into the MLP predict So I'm going to do something like this params the whole dataset So actually the images So these will be dataset images I'm going to pass them here We're going to have the again the JMP argmax alongside xis equals one So this will give us the predicted classes And I'm also going to pass obviously the dataset labels So and I'm going to compare those two and return the mean So basically JMP mean of this comparison And this is neater implementation so it's more concise And it's also way more fast compared to the last one So let me see where this makes sense So I'm doing inference on the whole dataset And I'm then converting my raw output So the log softmax outputs into the actual class indices here And I'm comparing them to the ground truth and finding the mean Okay so this looks like it should work So now I have to do the following thing So I have to load the whole thing into the memory So I'm going to do the following I'm going to take this thing So train images There is this internal member called data I'm going to fetch it that's going to get me the images I'm going to convert this into like a Jack's NumPy array And I'm going to reshape this into Wait let me check what this will return So this will going to be print train images shape Let me see what this is going to return So yeah as you can see the problem here is it's avoiding So this transform is not being called when I just fetched the data like this So we still need to flatten out the images So I'm going to do that I'm going to just reshape this into basically Whatever the length of the train data set is And I'm going to put minus one here So that's going to hopefully flatten out the whole thing And now yeah now we have the correct basically shape Let me just check the type train images It's going to be on the device So that means whatever accelerator I'm using is going to be on that accelerator I'm currently set to using none So yeah because MLP is a fairly simple model and MNIST is a super small data set Let me now do the same thing for labels I need train labels I'm going to do the same thing here I think so at least So I'm going to this time I'm going to access target so that's labels I'm going to convert them into Jack's array I think I need the reshaping so let me check the train label shape And let me check the type of the train labels So we have device array 60K this looks correct I think this is going to work So we need to do the same thing for test data set Let me just copy paste this Replace this with test Test images test labels Test data set here Test data set here Finally test data set here targets this should be working now Okay cool Now I'm going to instead of passing the I'm going to remove this line here And here I'm going to pass the data set and the label So I'm going to pass the train data set The train images And I'm going to pass the train labels And here I'm going to pass basically the same thing just test images And test labels Again I should have waited until I checked this works let me just run it like this See whether this is going to work And it's way faster as you can see and the accuracy is for some reason very very very high already I guess that's because we were assuming using the same MLP parameters so I have to I know what I have done So I'm going to interrupt this The thing is I'm using I'm constantly updating the MLP params so the state is saved here I have to call the init function here So we need to paste the let me find it let me find the init function we need to paste this here So that we start fresh every time we start the training so this now looks like it should work let me run it We have the results very very fast the accuracy is now low that seems to be working Let me interrupt the execution again and I'm going to delete this I'm going to yeah this now should work We're logging the loss every 50 batches we're logging the accuracy every epoch and we have how many epochs we have 10 epochs So let's start with like five and I'm going to run this and yeah let's see what the results are So the loss is going down that's cool that seems to be promising Okay after a single epoch we have 91% accuracy and we do not see any kind any sign of overfitting because the accuracy on both training and test data sets are pretty much the same 91% And now we have 93 I'm going to just skip until the training is done Okay the training is finally done we have 96% accuracy with a couple more epochs we would get to I guess 98 at least or something It could be further improved but the thing is this thing is already working So let me just kind of clear the output and let me try and run a prediction on a single image and I should have done this previously So let me fetch a single image visualize it and see whether the ground truth label like aligns with the predicted label So I'm going to do the following thing I'm going to fetch the single image so images let me take the like test loader I'm going to wrap it into the iterator here and take a single batch I'm going to take a single batch and I'm going to take the only the images and then I'm going to take a single image And make let's make sure that this makes sense so image shape should be 28 by 28 Oh yeah I have to just reshape it into MNIST image size let me see where this is going to work it is Now let me import ad hoc here I'm going to import just map plot live pyplot as plt Plt image show the image and let me just plot that thing So this looks correct we have number seven We'll also need labels so I have I'll have to modify this I'm going to grab labels as well so labels And let me check the the the ground truth label here so ground truth label is going to be labels of zero And let me print the the actual label here print the GT label it should be seven and it is OK So now let's see what our network is giving us if I was to do the MLP MLP predict on top of a image But I'll have to flatten it out so I'm going to pass the MLP params and I'm going to pass in the what I'm going to pass MP Revel of the image will this work let me check so let me check whether this is going to work These are just predictions and yeah we get some results let's now just do the argmax instead And see what we have so yeah we had we have the predicted value of seven so let's just add some string here And this is GT like this I'm going to paste it here let me run this and see whether it makes sense So predicted seven and GT seven so I have I should have done this image analysis while I was developing the Loaders here so basically while I was developing these data sets and loaders but yeah anyways we have it now We see everything is working correctly so that was pretty much it when it comes to the model training So was the first section it took around an hour to do this so now I want to do the following I want to quickly refactor The whole code code base and then I'm going to do some visualizations I was basically cutting some corners Because I'm filming this and I'm really I'm coding real time so it's hard to kind of be very very pedantic So now I want to improve some of those some of the things so let me see what first so this thing this this function Looks fairly nice we have parameters here we have keys I think the structure of this function was fairly nice We have this small testing functionality here so I'm just going to call it test we have a key we created We just called the function to see whether the results make sense and I think this cell was was neat I was still I was still very very fresh I guess this was some initial like this cell was used for debugging so I'm just gonna basically delete it Let me see okay so I had some explanations here etc etc again the design of this function is is fairly neat I'd say You can obviously always add some like doc strings comments etc whatnot okay let me see here we had some tests here So this was this first part was testing basically whether the MLP predict is working as expected And so I just printed the shape of the image I printed the shape of the output and the second part was actually testing whether So this test test the single example function and this was testing the batch test the batched Batched function basically here okay I can uncomment this part as well So yeah I mean it would obviously take a lot of time to to create this to make it perfect I'm just I want to make sure that I didn't That I don't have any any red flags like a huge huge red flags so I defined these custom transforming custom custom colloid functions These look neat as well so some consistency with my naming of the variables would be appreciated I'm sometimes using labels the full word and sometimes I'm abbreviating into LBLS so that sucks but yeah this looks nice We have the data sets here I did some debugging here I can I can I'm just going to comment it out and I'm basically just going to delete this part Now we have the data loaders then I did some testing here so this can be considered like a testing section again testing section What I did here I made sure that we have images labels of correct types and shapes and finally This thing here is just the optimization part where I was loading the whole data set load the whole data set into memory Okay I'm going to refactor this notebook a little bit after after I finish this video and I'm going to upload it to my to my github Hopefully you saw this one already so github let me let me open it up so I have this get started with Jax like repo do check it out I'm going to be posting all of these tutorials there as well as some other useful resources and videos here so that's it Let's continue here okay this to do is done we have the loss function I'm going to just delete this here we have the update function We have the the the accuracy maybe the metric and the loss should be one close to each other here that's it and finally we we create the network here Create create an MLP I guess a MLP and then we have the training loop here finally I could have probably just replace this with a product Function so to make this a little bit more robust but this is not that robust in any case because I'm hard coding names here but yeah You get the point in general you could kind of you'd have a section where you're defining these variables depending so when when the data set changes All of these variables change and basically MLP is out of the box is going to have a compatible shape this this was the point of doing this here And yeah you don't want to hard code stuff around your code base so yeah all in all for for a code that was written very very very ad hoc I think this this looks pretty decent So we have the loading here we convert the representation into the one hot representation then we call the update with this some logging of the loss We do some logging of the accuracy and I think it's fairly fairly cool now let's do a couple of visualizations first let me visualize the MLP weights So I'm gonna copy paste this into a novel cell here and what I mean by that is the following so I'm going to take the MLP weights So I'm going to take the MLP parameters I'm gonna fetch the zero player I'm gonna fetch the weight matrix and I'm gonna just take the that weight matrix and see what the shape is So the shape of that matrix is let me see 512 784 so what I want to do here is for a single output neuron I'm gonna I'm gonna fetch this all the weights for a single neuron and see whether we can see some pattern some interesting pattern I'm gonna do something like this weights single equals weights I'm gonna fetch a single row I'm gonna fetch all the columns and I'm gonna reshape that I'm gonna reshape that into the amnest image size so that's 28 by 28 and this should now be of correct shape with W single shape equals 28 by 28 I'm gonna just plot this and let's see whether we have some meaningful results here so I'm gonna plot it and we do not have any meaningful results here so if I was to randomly pick some other like some other weights for some other output neuron like 4 maybe I don't know let's see whether we can find some meaningful representation here and I do not see anything meaningful here I'm gonna try one more nope I guess not that interpretable but we still have 96% accuracy so that's cool and it was worth a try to just do this initial visualization so this task is done whoops let me just delete this to do and let's jump to the second one so visualize embeddings using t-sneak let me do that one I'm just gonna add a new cell here we're gonna import basically let me see side pie t-sneak okay yeah from scikit-learn manifold t-sneak so this will probably come in handy this documentation here I'm just gonna put the tab right here what the fuck okay let me put it here import from scikit-learn sklearn we import manifold module and we import t-sneak so from this package from this module import the function t-sneak my god okay okay this should now work so t-sneak is going to do a dimensionality reduction and for that I'll have to fetch the activations from the MLP just before we apply the linear layer that means we'll have to modify this predict function so let me find the MLP predict function I'm gonna copy paste it to this cell here I'm gonna modify it so first things first let's call it something else like fetch activations I guess we're gonna pass the params nothing changes there the only difference is we do not have to apply this this last part so let me let me think this will this will be yeah we basically do not want this part we just want to return back the activations the last ones just before we apply the linear layer so activation something like this and this should now work so now let me just create a batch version of this thing so batched fetch activations is gonna be remap we're gonna paste this function in and finally in axis will be again none and zero and this should now work so I'm just gonna fetch a batch of images and I'm gonna get those activations project them into 2d space and plot those images just to see whether the MLP the network learned how to separate different classes so let's see how that will look like so batch of activations is going to be equal to basically this and we need to pass MLP params so that's our train network and we also need to pass a batch of data so that means let me fetch a batch of images so batch I'm gonna take take like this images labels is gonna be next error and I'm gonna I'm going to again take from the test loader so test loader and yeah this should be it now I'm gonna take the images and I'm just gonna test them here and this should now work let me see what's the shape of these batch activations if this thing works so shape let me run this and it's 128 256 and that sounds about right because remember we had let me find a configuration so it was yeah 512 256 256 is the dimensionality of the activations just before we project them into the logits okay so let's now do the following we need to pass this into the t-sne method in order to get 128 comma 2 and this time I'm just gonna use my github because I have this pre-existing code already and I don't want to search through the documentation this video is already too long so I have this annotated get like method and it got pretty popular and anyways I had I used this new here so I'm just gonna copy paste a snippet let me find where I've been using it so here I'm just gonna take this part and I'm just gonna copy paste it into right into our notebook so let me do it like this okay so t-sne is here we'll take two components means basically we want to down project into two dimensions perplexity 30 is just a number that works we don't need this special method here and now finally I need to pass the batch activations here and yeah so this is how the usual normal workflow of a guy who does machine learning looks like you take somebody else's code or your own code even better and you just kind of tweak it so I'll also we won't need this part as well so number of classes is going to be equal to 10 because this is amnest and finally let's see I'm just gonna erase these comments for now so the only thing we need to change here is to put labels instead of node labels here because those are our ground truth labels so basically what this does is for class 0 we're going to extract only those embeddings which we know correspond to images which are of class 0 I have the number 0 in it because this is amnest so let me change this to labels this should all work now we don't need these special I'm not sure whether this is going to work in colab so the only thing we are left to do is to find this dictionary so let me find the dictionary and it's here so it's just a mapping from from a number to a color and let me just define it somewhere here and that's it this should work now hopefully and it does okay so we can see that this is just some stupid error from scatterplot I've already seen this error but like it does not change the semantics so everything works correctly here so we can see that different classes are clustered together in a particular parts of the of the studio space we see some outliers like there is this yellow dot here if you can see it on your screen but yeah in any case we can see that the MLP did learn something so that's cool now the final thing I thought doing is basically seeing how many dead neurons we have so I'm going to delete this cell here I'm gonna add a new cell just below this one so what we need to do here is to find those neurons that do not activate on all of the input images I'm going to feed into the network and those are considered dead neurons because their gradients are going to be zero which means the weights that come before that neuron output can be changed and thus yeah the neuron is pretty much dead I'm going to link a cool article down in the description so you can check it out it has a like a very thorough analysis of various activation functions and there is a dead neuron analysis section as well so you can check it out there so what we need to do is again we need to modify the MLP predict function so we'll have to paste this again here so let me take it and we'll need to fetch every single one of these activations so let me create a collector variable here and I'm just going to collect every single activations here so I'm going to append activations and that's going to be it we again do not need this part because we don't have rally activations here we just care about the outputs of the rally unit so finally I'm going to return back the collector here and that's it everything else remains the same yeah combining this one with this one is probably possible but yeah I just decided to do it ad hoc like this at the moment I'm just going to name this fetch activations two because I'm being super creative right now okay so what we need to do is basically again batch this one we're going to create a batched version of the function so a batched fetch activation is two function with we map we map here and I'm going to pass the function in axis again none and zero and that's it now yeah I guess I can just copy paste this all of this and let me do it here and now I'm just going to need to swap this for the activation two version everything else remains the same and we should expect to have something like let's see where this is going to work at all yeah it's not we're actually just getting a list here and that makes sense because remember that's how we are collecting the variables here the activations here so I actually need to do something like this in order to get the correct result I guess 128 512 that makes sense and let's see we need to have 128 256 in this one and yep we do okay so now that we have these activations we just need to find those neurons where each of these 128 images produced are like a value zero that means this neuron was dead for all these 128 images and that's fairly easy I guess let me form a list here that neurons so this is going to be a flag data structure in the sense that if there is a one that means that the neuron is dead otherwise it's not dead so let's assume for now that we have that all of the neurons are dead so I'm going to just pre populate this with a shape something like this so for for act in bad activations I'm going to so this is going to be hundred so remember the shape of this act will be something like 128 and then 512 or 256 so we need to fetch the shape and we take everything except for the zero basically dimension and now we have for every single neuron we have a corresponding flag here so this is just a set of flags for every single neuron that we care about and those are the ones that basically have a rally unit after them now let's iterate through this through the activations so I'm going to have a layer ID here and I'm going to have the actual activations in enumerate and enumerate which activations so now I'm going to do the following dead neurons from layer ID is going to fetch the corresponding structure equals non-pi I guess it's logical and or something the illogical and we're going to end whatever is here already with basically activations equals zero so we're asking whether the the the neuron is dead and if all of the neurons are dead alongside the batch dimension then that neuron corresponding your own remains that otherwise it's considered alive your own I guess and I think all X is equal zero should do the job here let me see what this is going to work so finally we're going to print the so we're going to iterate through the layers so we read through the layers of dead neurons and we print the how many of them are actually dead something like this let's let's see what's where this is going to work and it does so we have zero dead neurons in the first layer we have four dead neurons in the in the last layer now let's try and do something to make sure this works so let's initialize novel network that's not trained and let's see whether we have different results so now here I'm going to create a novel LP which I'm going to call MLP parameters parameters to and now let's use that one instead to fetch the activations everything else remains the same as far as I can see here and yeah so let me just search MLP params yeah this should be it let's see whether we have different results here now now a useful thing would be to actually analyze this on the whole training data set because if there is a debt if there is a lot of dead neurons in the when you evaluate on the whole training data set that may be problematic and especially problematic neurons are the ones in the shallower layers to play with this code at your own pace let me try I'm fairly sure that I won't have a TPU at disposal currently so if I were to try to use a TPU and click save here let's see whether we can grab a TPU let me click reconnect here and no back end with TPU is available so in any case I'm fairly sure that this was already super enough because one of the previous videos I already explained how to paralyze the computation so it's fairly trivial to make this code parallel just a couple of lines of code and hopefully you got something out of this video this is the first time I'm doing a real time coding session so let me know if you have any feedback down below and whether you like this type of a content so I'm gonna start making more of this if you like it in any case consider subscribing share out this video and join our discord community and until next time bye bye
[{"start": 0.0, "end": 5.22, "text": " What's cracking guys? This is the third video in the Machine Learning with Jack's series and in this one"}, {"start": 5.22, "end": 7.7, "text": " We're going to build a neural network from scratch"}, {"start": 8.14, "end": 14.68, "text": " I'm going to be using a multilayer perceptron and I'm gonna train a classification model on top of the MNIST dataset"}, {"start": 15.0, "end": 18.16, "text": " By the way, I'm going to be using MNIST just for convenience"}, {"start": 18.78, "end": 23.0, "text": " Conceptually, even if we were to use ImageNet, the code would be pretty much the same"}, {"start": 23.28, "end": 28.04, "text": " Just for the sake of me coding this real-time and this video not taking three hours"}, {"start": 28.04, "end": 30.84, "text": " I'm going to be using this small dataset such as MNIST"}, {"start": 31.34, "end": 39.0, "text": " Finally, we're going to be using a PyTorch data loaders to load the data from MNIST and the reason is Jack's designers"}, {"start": 39.48, "end": 46.56, "text": " by design decided not to develop loaders because other libraries such as frameworks such as TensorFlow"}, {"start": 47.04, "end": 50.28, "text": " and PyTorch have already done a decent job at"}, {"start": 51.0, "end": 53.980000000000004, "text": " Adding this functionality so they did not want to reinvent the wheel"}, {"start": 53.98, "end": 60.739999999999995, "text": " So that's in a nutshell what this video will be about and then as a bonus points"}, {"start": 60.739999999999995, "end": 63.18, "text": " I'm gonna do some visualizations and finally, we're hopefully"}, {"start": 63.699999999999996, "end": 69.96, "text": " I'll be able to parallelize the training in the case that Colab allows me to use the TPU if not"}, {"start": 70.42, "end": 73.62, "text": " Previous videos already explained how to do that. So that's just a bonus point"}, {"start": 74.02, "end": 82.38, "text": " Anyways, let's let's start with that and structure this this whole video. So first thing we're going to do MLP training on"}, {"start": 82.38, "end": 87.46, "text": " on MNIST dataset, right? That's the first thing the second thing we're going to do is"}, {"start": 88.22, "end": 90.5, "text": " we're going to do some visualizations and"}, {"start": 91.25999999999999, "end": 93.5, "text": " Finally, the third thing is going to be"}, {"start": 94.5, "end": 96.5, "text": " basically adding the"}, {"start": 97.25999999999999, "end": 100.3, "text": " Parallelization so yeah, hopefully I spelled that correctly"}, {"start": 100.3, "end": 106.97999999999999, "text": " So now let me break these tasks into smaller sub tasks first thing where we need to create a init function in it MLP"}, {"start": 107.53999999999999, "end": 108.97999999999999, "text": " we need to"}, {"start": 108.97999999999999, "end": 110.74, "text": " and add the"}, {"start": 110.74, "end": 114.86, "text": " Predict function. So those are going to be two things we need to do here"}, {"start": 115.46, "end": 120.66, "text": " Whoops to do whoa. I don't know how to write today. Okay, that's the first to do the second one will be"}, {"start": 121.22, "end": 122.61999999999999, "text": " basically add"}, {"start": 122.61999999999999, "end": 124.3, "text": " data loading"}, {"start": 124.3, "end": 130.14, "text": " in pytorch and the third to do is going to be add the training loop"}, {"start": 130.94, "end": 134.74, "text": " loss function and whatnot maybe accuracy or something like that"}, {"start": 134.74, "end": 140.9, "text": " So that's the the breakdown of the first section of this video now here. I'm gonna just do something like"}, {"start": 141.62, "end": 143.62, "text": " maybe visualize"}, {"start": 143.62, "end": 144.66, "text": " visualize"}, {"start": 144.66, "end": 145.66, "text": " the"}, {"start": 145.66, "end": 147.18, "text": " MLP weights"}, {"start": 147.18, "end": 149.18, "text": " so we're just gonna"}, {"start": 149.18, "end": 154.70000000000002, "text": " See whether we can notice some patterns especially in the shallower layers whether MLP learn to extract certain features"}, {"start": 154.70000000000002, "end": 158.94, "text": " Which are salient, which are also like understandable and intuitive to us humans"}, {"start": 158.94, "end": 164.26, "text": " Then I'm going to do something like I'm gonna visualize embeddings"}, {"start": 164.82, "end": 166.82, "text": " using t-sne and"}, {"start": 167.02, "end": 169.46, "text": " finally, I'm gonna do something like"}, {"start": 170.34, "end": 172.94, "text": " visualize dead neurons if I have the time I"}, {"start": 173.46, "end": 179.42, "text": " Won't be breaking this parallelization part because that's that's pending and we're gonna see whether we have time A and B whether I'll have"}, {"start": 179.42, "end": 181.94, "text": " TPUs on disposal. So let's start by breaking down"}, {"start": 182.46, "end": 187.82, "text": " Solving one task at a time. So first let me start with initial init and MLP"}, {"start": 187.82, "end": 193.26, "text": " Initial init MLP function. So that function should given some"}, {"start": 194.06, "end": 196.06, "text": " configuration of the MLP network"}, {"start": 196.42, "end": 202.42, "text": " It's gonna return the weights and biases of the network. So that's a high-level description of what I want to achieve here"}, {"start": 202.5, "end": 205.06, "text": " Now let's go and implement the function. So we're gonna call it"}, {"start": 205.57999999999998, "end": 210.74, "text": " Something descriptive like init MLP, I guess and I'm gonna pass layer widths"}, {"start": 210.98, "end": 214.38, "text": " So that's gonna be the configuration and then we're going to start"}, {"start": 214.38, "end": 217.22, "text": " Implementing the function here. So params is just going to be some list"}, {"start": 217.74, "end": 223.34, "text": " basically, we're gonna iterate through this like layer widths and we're gonna take the"}, {"start": 224.54, "end": 228.85999999999999, "text": " In width and out width will be the names. So that's the number of"}, {"start": 229.38, "end": 233.9, "text": " Neurons that comes into the layer and number of neurons that goes from the layer"}, {"start": 234.38, "end": 241.42, "text": " And basically I'm going to iterate through the zip version of this so layer widths and I'm gonna index it like this"}, {"start": 241.42, "end": 250.14, "text": " So basically we want to start here with the 0th one and go and just exclude the last element and here I'm gonna"}, {"start": 250.7, "end": 256.62, "text": " Oops, let me paste it. I'm gonna start from the first one and go all the way including the last element"}, {"start": 256.62, "end": 259.74, "text": " So hopefully you'll understand a second why I'm doing this"}, {"start": 260.21999999999997, "end": 265.5, "text": " And let me maybe first create the high-level let me call the function and see how that will work"}, {"start": 265.5, "end": 269.34, "text": " So I'm gonna have ML params here. I'm gonna call the init MLP"}, {"start": 269.34, "end": 273.82, "text": " I'm gonna call the init MLP. I'm gonna pass the list here. So"}, {"start": 274.61999999999995, "end": 278.61999999999995, "text": " First thing that we're gonna do is we're gonna add like the number of"}, {"start": 279.38, "end": 284.41999999999996, "text": " Elements of pixels in the MNIST data set then I'm gonna arbitrarily create some configuration"}, {"start": 284.41999999999996, "end": 288.88, "text": " So we're gonna have a hidden layer with 512 neurons then 256 and finally 10"}, {"start": 289.09999999999997, "end": 292.05999999999995, "text": " So what this function is doing is as you can see here"}, {"start": 292.05999999999995, "end": 295.34, "text": " It started from it starts from the 0th one and this one starts from the first one"}, {"start": 295.34, "end": 301.97999999999996, "text": " So that means we're gonna be fetching tuples like this this one and then this one and then this one because that will enable me to"}, {"start": 301.97999999999996, "end": 306.06, "text": " Create the weight and biases we need for each of the layers. So let's do that"}, {"start": 306.06, "end": 310.38, "text": " So I'm gonna append to params the following thing. So basically"}, {"start": 311.34, "end": 313.34, "text": " We're gonna create a numpy"}, {"start": 313.85999999999996, "end": 315.85999999999996, "text": " Random we're just gonna randomly"}, {"start": 316.46, "end": 319.5, "text": " Sample data from the normal distribution"}, {"start": 319.94, "end": 324.62, "text": " That is gonna have a shape like this so out with and this is just convention"}, {"start": 324.62, "end": 326.68, "text": " We usually just first put the out"}, {"start": 327.86, "end": 332.1, "text": " Dimension here and then the input dimension. So there's gonna be for example 512 and then"}, {"start": 332.66, "end": 338.98, "text": " 784 here you could it vice versa, but it's gonna complicate things later on. So there's just a convention to keep in mind"}, {"start": 339.66, "end": 341.66, "text": " So next thing we'll need"}, {"start": 342.14, "end": 345.84000000000003, "text": " We'll need a bias term. So there's gonna be something like this random"}, {"start": 346.86, "end": 353.4, "text": " And and we're gonna need just the out with because that's the bias part, right? So this should now work"}, {"start": 353.4, "end": 359.62, "text": " Let me close this list and let me just return back the parameters here. Let's return the params and"}, {"start": 360.38, "end": 363.09999999999997, "text": " Let's see whether this is going to work. So"}, {"start": 364.02, "end": 370.02, "text": " In order to understand whether this is correct. We can maybe use the tree map functionality from Jax"}, {"start": 370.29999999999995, "end": 373.78, "text": " So I'm gonna just print Jax tree map"}, {"start": 374.38, "end": 376.59999999999997, "text": " I'm gonna create a lambda function here"}, {"start": 376.6, "end": 383.52000000000004, "text": " Which is going to just return the shape of the leaf and the leaves are going to be fetched from the MLP params"}, {"start": 383.62, "end": 389.58000000000004, "text": " Which means we're going to be taking weight and biases these matrices here and we're gonna be printing their shapes"}, {"start": 389.74, "end": 395.5, "text": " To understand whether this is implemented correctly. So obviously did not import Jax neither"}, {"start": 395.58000000000004, "end": 401.66, "text": " Neither did I import non-pice is gonna fail. So I'm gonna first import those libraries import numpy as"}, {"start": 401.66, "end": 408.36, "text": " MP import Jax and I'm gonna run this and I'm going to run this as well"}, {"start": 408.36, "end": 413.74, "text": " Let's see whether this is going to work. So it's invalid syntax. Let's see what I got wrong here"}, {"start": 415.16, "end": 417.92, "text": " NumPy random. Oh, yeah, I added a"}, {"start": 418.6, "end": 422.06, "text": " Closing bracket here so that failed. Let me rerun it"}, {"start": 422.64000000000004, "end": 429.36, "text": " Okay, so this looks correct. I think so we have 512 here and then we have 256 here and then we have 10"}, {"start": 429.36, "end": 433.36, "text": " I'd say this is this is correct. So there's a couple of modifications"}, {"start": 433.36, "end": 440.44, "text": " We need to add here first thing is a however to just use the standard deviation one as the initialization"}, {"start": 441.0, "end": 447.66, "text": " This thing is gonna explode. We need to scale these parameters. I'm not adding any fancy initialization method here"}, {"start": 447.66, "end": 452.52000000000004, "text": " I'm just gonna scale these like this just gonna like basically"}, {"start": 453.04, "end": 455.04, "text": " Make the standard deviation smaller"}, {"start": 455.04, "end": 463.08000000000004, "text": " By a factor of hundred and that's it. The second thing we need to add here is now that we have implemented numpy's random"}, {"start": 464.24, "end": 468.44, "text": " Number generator so we'll actually be using Jax random number generator"}, {"start": 468.44, "end": 471.64000000000004, "text": " So that's why I'm gonna just do that right now"}, {"start": 471.84000000000003, "end": 475.76, "text": " So it was easier for me conceptually to first start with numpy now. I'm gonna just quickly"}, {"start": 475.76, "end": 484.2, "text": " Change this so I guess checks random PRNG key or something is that the functionality I'm gonna add a seed here. Whoops"}, {"start": 484.88, "end": 488.12, "text": " I'm just gonna add some constants here. So C is gonna be zero"}, {"start": 488.12, "end": 494.56, "text": " I'm gonna be using seed here and I'm gonna pass RNG inside of the initmlp function"}, {"start": 494.84, "end": 501.44, "text": " we need to add it here as the argument and now the thing we need to do is create the keys out of the"}, {"start": 501.76, "end": 502.88, "text": " this"}, {"start": 502.88, "end": 504.44, "text": " generator"}, {"start": 504.44, "end": 508.88, "text": " So this is actually also a key. So I'm gonna rename it because it will be misleading"}, {"start": 508.88, "end": 514.12, "text": " I'm gonna pass the key here and from the key, I'm gonna create keys using the split method"}, {"start": 514.12, "end": 517.48, "text": " So I'm gonna create keys using Jax random"}, {"start": 518.24, "end": 522.88, "text": " split method and I'm gonna pass in the key and I'm gonna set the number of keys to"}, {"start": 523.88, "end": 528.48, "text": " Whatever the number of layers is so that's length of this thing minus one, right?"}, {"start": 528.72, "end": 534.24, "text": " Because the number of layers is by one smaller than the length of this configuration file because we're taking tuples"}, {"start": 534.24, "end": 538.62, "text": " Remember so this this and then finally 256 and 10 as the last layer"}, {"start": 539.46, "end": 544.84, "text": " So that's gonna be the number of keys and now we are going to additionally add to this zip function"}, {"start": 544.84, "end": 546.64, "text": " I'm gonna add the keys and"}, {"start": 546.64, "end": 553.6, "text": " I'm gonna then add the key part here. And finally, I'm going to split it additionally into the weight key"}, {"start": 553.6, "end": 555.96, "text": " so we need to split it into weight key and"}, {"start": 557.36, "end": 562.32, "text": " Bias key so something like this. So I'm gonna do again Jax random"}, {"start": 562.32, "end": 564.24, "text": " split"}, {"start": 564.24, "end": 569.0400000000001, "text": " and I'm gonna pass in the key and that's gonna split"}, {"start": 569.88, "end": 576.84, "text": " Let me see whether this is going to create a problem because this thing I'm gonna call this the parent key just for the sake of"}, {"start": 578.2800000000001, "end": 583.96, "text": " Avoiding some bugs here. So parent key here and creating keys and I'm iterating through the keys here"}, {"start": 583.96, "end": 587.72, "text": " I'm splitting that into the weight key and the bias key and finally we're going to now"}, {"start": 587.72, "end": 592.12, "text": " Input the weight key and the bias key and create the correct random function from Jax"}, {"start": 593.32, "end": 595.6, "text": " Not completely sure about this syntax here"}, {"start": 595.6, "end": 600.0, "text": " I'm gonna just try something that makes sense and then I'm gonna search if it doesn't work"}, {"start": 600.0, "end": 603.32, "text": " So Jax random something like that and then I'm gonna pass"}, {"start": 604.28, "end": 606.72, "text": " Like the key so I'm gonna pass the weight key"}, {"start": 607.6800000000001, "end": 609.6800000000001, "text": " And then I'm gonna paste"}, {"start": 610.48, "end": 616.48, "text": " Pass the shape here. So this should be the syntax and then here I'm gonna again"}, {"start": 616.48, "end": 624.04, "text": " Do Jax random and I'm gonna pass the bias key here and the shape is going to be"}, {"start": 624.6, "end": 626.12, "text": " something like this"}, {"start": 626.12, "end": 630.96, "text": " So this should now work. Let me kind of split this into two lines. It's more readable"}, {"start": 631.72, "end": 638.4, "text": " Let me let me kind of try and refactor this a little bit something like this and then I'm gonna indent this a little bit more"}, {"start": 639.2, "end": 642.48, "text": " Fuck my life. Okay, and I'm gonna do it like this"}, {"start": 642.48, "end": 649.84, "text": " Okay, this looks a bit better. So hopefully this works. Let's see whether it makes sense. We have the parent key"}, {"start": 650.9200000000001, "end": 654.86, "text": " We split it into multiple keys here. We use those keys"}, {"start": 654.86, "end": 657.74, "text": " So we'll have a key per layer every layer"}, {"start": 657.74, "end": 662.24, "text": " We have a specific key and then we do an additional split because every layer has weights and bias"}, {"start": 662.48, "end": 669.34, "text": " We pass the weights we paste the bias and that's it. This should now work as expected. Let me try and run it"}, {"start": 669.34, "end": 672.34, "text": " hopefully it's gonna work and"}, {"start": 673.46, "end": 681.86, "text": " It's not working because a module object is not callable. Let me see. What have I done here?"}, {"start": 684.34, "end": 690.34, "text": " Aha I called random is the module not the function. So let's see how the the actual syntax looks like"}, {"start": 691.46, "end": 694.7800000000001, "text": " Jax random random. I don't know. Let's see"}, {"start": 694.78, "end": 698.18, "text": " Whether this is going to"}, {"start": 700.02, "end": 705.9, "text": " Yeah, this is the wrong file, let me try and see whether we have some interesting information here"}, {"start": 706.38, "end": 710.3199999999999, "text": " So we have random uniform here, but do we have random something?"}, {"start": 711.06, "end": 717.66, "text": " That's not uniform generator because uniform distribution is not as good for"}, {"start": 718.06, "end": 721.98, "text": " parameter initialization as Gaussian because Gaussian tends to because of the mean zero"}, {"start": 721.98, "end": 726.98, "text": " We're gonna have an expectation smaller weights than using uniform"}, {"start": 726.98, "end": 732.02, "text": " I guess and also it's gonna be symmetric around zero. That's important. So I really need to find the"}, {"start": 733.86, "end": 735.1800000000001, "text": " Actual"}, {"start": 735.1800000000001, "end": 736.3000000000001, "text": " Gaussian"}, {"start": 736.3000000000001, "end": 738.3000000000001, "text": " Random, let me just use"}, {"start": 738.54, "end": 740.54, "text": " Something like this Gaussian"}, {"start": 740.54, "end": 742.54, "text": " Jax random Gaussian"}, {"start": 743.58, "end": 745.58, "text": " Will this work?"}, {"start": 747.54, "end": 749.54, "text": " Let me open up a couple files here"}, {"start": 749.54, "end": 751.38, "text": " I"}, {"start": 751.38, "end": 753.38, "text": " Think we saw this one"}, {"start": 753.74, "end": 755.9399999999999, "text": " Okay, let me let me see whether I can just"}, {"start": 756.6999999999999, "end": 758.6999999999999, "text": " Grab a number from here"}, {"start": 760.02, "end": 762.4599999999999, "text": " Nope random seed random uniform"}, {"start": 763.3, "end": 767.26, "text": " They're just using uniform numbers all the time here. What what what's this? Oh"}, {"start": 768.5, "end": 774.4599999999999, "text": " Random normal. Okay. Okay. This is the syntax random normal. Whoops that that took some time. Okay"}, {"start": 774.46, "end": 782.46, "text": " Jax random and then dot normal and that's gonna work random dot normal and hopefully now this will work as"}, {"start": 783.14, "end": 785.14, "text": " expected and"}, {"start": 786.0600000000001, "end": 792.46, "text": " It seems it does everything looks correct here, I think we are we are good to go to the next section so"}, {"start": 793.34, "end": 795.58, "text": " We have the correct scale. We have the correct"}, {"start": 796.4200000000001, "end": 800.5, "text": " Structure three layers every layer has weights and every layer has bias"}, {"start": 800.5, "end": 808.42, "text": " We've used Jax's pseudo run random number generator and we had some primitive initialization scheme here using the scale parameter"}, {"start": 808.74, "end": 813.1, "text": " Now let's do the second part and that's the add the predict function. So now that we have this"}, {"start": 813.1, "end": 815.1, "text": " Let me just add another cell here"}, {"start": 816.1, "end": 819.46, "text": " Let me add a function called MLP predict"}, {"start": 820.18, "end": 823.14, "text": " we're gonna pass in the parameters of the MLP and"}, {"start": 824.06, "end": 825.86, "text": " we're going to"}, {"start": 825.86, "end": 832.22, "text": " What else we need some data? So let's pass in the input data here. This is going to be a single image"}, {"start": 832.5, "end": 837.74, "text": " And a flatten one because we're using an MLP. We're not using a convolutional neural network here. So"}, {"start": 838.78, "end": 843.2, "text": " How we're gonna split this is we're gonna take the hidden layers here"}, {"start": 844.22, "end": 846.22, "text": " So maybe let's call this"}, {"start": 848.62, "end": 851.66, "text": " Yeah, I'm just gonna do it like this so hidden layers I'm gonna take the"}, {"start": 851.66, "end": 856.38, "text": " The params all the params except for the last one"}, {"start": 856.38, "end": 863.1, "text": " That's gonna be the hidden layers the two one the first two layers and then I'm gonna iterate through those layers"}, {"start": 863.14, "end": 865.2199999999999, "text": " So I'm gonna just iterate like this"}, {"start": 867.1, "end": 871.7199999999999, "text": " We're iterating through the through the hidden layers and taking the weight and bias matrices"}, {"start": 872.98, "end": 874.98, "text": " Then we need to"}, {"start": 875.06, "end": 880.3399999999999, "text": " Set activation something like this activation is gonna be initially equal to whatever the input is"}, {"start": 880.34, "end": 883.9, "text": " And then we're just gonna reuse it here. So activation"}, {"start": 885.0600000000001, "end": 886.46, "text": " activity"}, {"start": 886.46, "end": 888.46, "text": " Reasion, I don't know how to spell today"}, {"start": 889.7, "end": 891.7, "text": " Basically, this is going to be"}, {"start": 893.12, "end": 900.3000000000001, "text": " What we need to do activation function and we need to create the the the mapping so mapping will be basically Jax"}, {"start": 900.3000000000001, "end": 902.38, "text": " I guess stop product"}, {"start": 903.82, "end": 908.5, "text": " No, I'm gonna add the jack jacks non-pay API here. So I'm gonna add"}, {"start": 908.5, "end": 915.22, "text": " Jacks non-pay as JMP because that's gonna be useful. I'm gonna rerun the cell and"}, {"start": 915.86, "end": 920.94, "text": " Finally here so JMP dot should work. I'm gonna pass in the"}, {"start": 921.46, "end": 928.74, "text": " weight and then the example and then I'm gonna add the bias term so this thing should be"}, {"start": 929.18, "end": 932.66, "text": " This is the functionality of the linear layer. Finally, we need the relu"}, {"start": 932.66, "end": 940.18, "text": " So I think and then contains the rally unit. Yep, it does so something like this should now function and this should not be"}, {"start": 940.4599999999999, "end": 943.02, "text": " Exit should be activation. Okay"}, {"start": 943.6999999999999, "end": 945.6999999999999, "text": " So this show work"}, {"start": 945.9399999999999, "end": 947.9399999999999, "text": " Let me see whether it makes sense"}, {"start": 948.06, "end": 954.4599999999999, "text": " So we're doing a dot product. We are adding the bias. We are doing the activation function and we get the activation out"}, {"start": 954.46, "end": 961.9000000000001, "text": " This looks reasonable. Let me now do the final step. The final step is we get the last and"}, {"start": 962.5400000000001, "end": 965.9000000000001, "text": " The last layers weight and bias so something like this"}, {"start": 966.7800000000001, "end": 968.38, "text": " params"}, {"start": 968.38, "end": 975.58, "text": " Of the last element of the last layer and we do apply this again. So I'm just gonna copy paste this"}, {"start": 977.22, "end": 980.3000000000001, "text": " We're gonna copy paste this I'm gonna use the last one"}, {"start": 980.3, "end": 985.8599999999999, "text": " The reason I'm doing it like this is because the last layer does not have the relu activation"}, {"start": 986.14, "end": 988.54, "text": " so these are going to be just logits and"}, {"start": 989.14, "end": 992.5, "text": " We could alternatively do it in a single for loop"}, {"start": 993.26, "end": 995.9399999999999, "text": " but then I'd have to have a switch statement like if"}, {"start": 996.4599999999999, "end": 1001.9799999999999, "text": " Last layer then apply relu otherwise don't apply relu. So at the end is gonna be even"}, {"start": 1002.5799999999999, "end": 1007.42, "text": " Less readable. So I'm just gonna do it like this and finally we're going to return"}, {"start": 1007.42, "end": 1009.9, "text": " What we're going to return logits"}, {"start": 1010.42, "end": 1017.2199999999999, "text": " But not quite because Jax does not have cross entropy loss implemented. We'll have to be resourceful here and"}, {"start": 1017.78, "end": 1020.42, "text": " Create something that's going to be"}, {"start": 1021.18, "end": 1025.98, "text": " Maybe confusing at the first point. We're going to be implement importing from the"}, {"start": 1026.62, "end": 1028.06, "text": " Jax"}, {"start": 1028.06, "end": 1031.86, "text": " SciPy, I think special we're gonna import something called"}, {"start": 1031.86, "end": 1037.5, "text": " Like X log some or something X log some I don't know. Let me try and"}, {"start": 1038.06, "end": 1040.62, "text": " Search for this not completely sure"}, {"start": 1043.82, "end": 1049.02, "text": " So, yeah log some X log some X log some X"}, {"start": 1049.54, "end": 1053.06, "text": " so this thing is gonna allow us to"}, {"start": 1054.34, "end": 1057.06, "text": " Implement a numerically stable version of cross entropy"}, {"start": 1057.06, "end": 1062.1799999999998, "text": " So what this what this is is the following so imagine you we have ten outputs, right?"}, {"start": 1062.1799999999998, "end": 1065.1799999999998, "text": " We have ten outputs from the this amnest MLP"}, {"start": 1065.8999999999999, "end": 1066.8999999999999, "text": " model"}, {"start": 1066.8999999999999, "end": 1071.1799999999998, "text": " And let's take a single one. Maybe the single output will be called a one"}, {"start": 1071.1799999999998, "end": 1074.98, "text": " So this is the raw output from from from the logits and we can"}, {"start": 1075.3, "end": 1083.06, "text": " Equivalently represent this without changing anything like this log X of of one because"}, {"start": 1083.06, "end": 1088.6599999999999, "text": " These two are just inverse of each other. So this thing is equivalent to this thing here. So now"}, {"start": 1089.22, "end": 1092.3799999999999, "text": " What this will do is it's gonna add a log sum"}, {"start": 1092.3799999999999, "end": 1096.34, "text": " So it's gonna do something like log and it's gonna have a sum of"}, {"start": 1097.1, "end": 1100.26, "text": " Exponentials so something like this Oh one and"}, {"start": 1100.78, "end": 1104.82, "text": " Then it's gonna have obviously because it's a sum. It's gonna have as a second argument"}, {"start": 1105.34, "end": 1109.8999999999999, "text": " The second output and then the third output and you get the point all the way down"}, {"start": 1109.9, "end": 1115.38, "text": " So if you just look at this expression now, let me just kind of copy paste it here"}, {"start": 1116.7, "end": 1118.7, "text": " So that it's more"}, {"start": 1119.02, "end": 1123.5400000000002, "text": " Visible and now we're just gonna apply the simple rule. Basically if you try and subtract the logarithms"}, {"start": 1123.5400000000002, "end": 1127.5, "text": " That's equivalent as if you were to take the arguments and just divide them"}, {"start": 1127.5, "end": 1130.8600000000001, "text": " So we're just gonna take the this argument here. So"}, {"start": 1131.5, "end": 1134.5800000000002, "text": " The left argument and we're just gonna take the left argument"}, {"start": 1134.5800000000002, "end": 1138.5, "text": " And we're just gonna take the right argument and we're just gonna take the left argument"}, {"start": 1138.5, "end": 1142.62, "text": " And we're just gonna divide it by this sum here"}, {"start": 1143.46, "end": 1145.46, "text": " of all these"}, {"start": 1145.74, "end": 1149.02, "text": " basically exponents and you can see that this is"}, {"start": 1149.86, "end": 1151.98, "text": " Pretty much we just implemented softmax"}, {"start": 1152.58, "end": 1156.58, "text": " The only difference is that we have log softmax and that's it"}, {"start": 1156.58, "end": 1160.58, "text": " So this is what this thing implements and that's gonna be super convenient a bit later"}, {"start": 1160.58, "end": 1164.98, "text": " We're gonna see why in a second, but let me now test this function as well"}, {"start": 1164.98, "end": 1170.18, "text": " So as you can see here every time I write a certain module like this init MLP, I test it here"}, {"start": 1170.26, "end": 1172.26, "text": " so this is a"}, {"start": 1172.42, "end": 1175.94, "text": " Not not that rigorous testing but some testing nonetheless"}, {"start": 1176.34, "end": 1178.34, "text": " so let me try and"}, {"start": 1179.58, "end": 1183.8600000000001, "text": " Pass some data here. So I'm gonna just create a simple"}, {"start": 1185.26, "end": 1186.82, "text": " dummy"}, {"start": 1186.82, "end": 1193.26, "text": " Like image here. So dummy image flat. I'm gonna call it that way and this is going to be ran then"}, {"start": 1193.26, "end": 1198.14, "text": " Basically again is gonna be MP prod of the M NIST"}, {"start": 1199.78, "end": 1202.62, "text": " Image size this should work correctly"}, {"start": 1203.34, "end": 1209.02, "text": " Let me just print the shape to make sure this is this has the correct shape. It should be a flat image"}, {"start": 1210.3, "end": 1211.3, "text": " so"}, {"start": 1211.3, "end": 1212.46, "text": " 784"}, {"start": 1212.46, "end": 1218.86, "text": " Should be the shape and we're going to pass that into the MLP predict and see what we get as the prediction"}, {"start": 1218.86, "end": 1220.86, "text": " So prediction is going to be MLP predict"}, {"start": 1220.86, "end": 1227.1799999999998, "text": " We pass the MLP parameters we generated up there. We pass it here. We pass in the"}, {"start": 1228.8999999999999, "end": 1230.58, "text": " Image here"}, {"start": 1230.58, "end": 1233.5, "text": " So something like this and this should now work"}, {"start": 1234.1399999999999, "end": 1237.26, "text": " Hopefully, so I'm gonna just print out the prediction"}, {"start": 1237.82, "end": 1240.4199999999998, "text": " Shape here. So let me try and run this thing"}, {"start": 1242.1799999999998, "end": 1246.4599999999998, "text": " So name log some XP is not defined because I have not rerun this thing"}, {"start": 1246.46, "end": 1251.46, "text": " Whoops, okay. Let's see what's wrong here. Let me reopen the tab here"}, {"start": 1252.18, "end": 1257.22, "text": " Jacks SciPy special Jacks SciPy special aha"}, {"start": 1258.1000000000001, "end": 1261.3, "text": " From Jacks SciPy special import. Yeah"}, {"start": 1262.3, "end": 1266.3400000000001, "text": " I'm really bad today. Okay. Let me try and rerun this again"}, {"start": 1267.14, "end": 1272.46, "text": " And see whether it works. It does not because unsupported operand types"}, {"start": 1272.46, "end": 1276.66, "text": " unsupported operand types, blah blah blah the"}, {"start": 1277.22, "end": 1279.22, "text": " device array and function"}, {"start": 1279.98, "end": 1284.54, "text": " Okay, something went wrong here. Let me see what so I'm passing the"}, {"start": 1285.94, "end": 1290.98, "text": " MP prod should give me so how you debug this is you just open up another cell here"}, {"start": 1291.3400000000001, "end": 1292.66, "text": " I'm just gonna do it like this"}, {"start": 1292.66, "end": 1299.26, "text": " I'm gonna just print out the MP prod of this thing and see whether this gives me the number it and it does"}, {"start": 1299.26, "end": 1302.58, "text": " so this thing returns back the number and"}, {"start": 1304.54, "end": 1306.82, "text": " For some reason this is not working"}, {"start": 1307.66, "end": 1312.86, "text": " Okay, I got it all wrong. This thing actually worked. The problem is here and even have the red squiggly line"}, {"start": 1312.86, "end": 1318.3799999999999, "text": " I'm stupid basically. I just I just subtracted a function from from from this"}, {"start": 1318.3799999999999, "end": 1325.3, "text": " I need to pass in the logits here and now this should work. Hopefully, okay, and it does work. So the shapes do"}, {"start": 1325.3, "end": 1328.58, "text": " work as expected so now because this is Jax without"}, {"start": 1329.1399999999999, "end": 1334.74, "text": " Reimplementing this function for the batch of data. We're just gonna wrap it up into V map and let's see whether that's gonna work"}, {"start": 1334.74, "end": 1338.82, "text": " So I'm gonna create a like a batched version so batched"}, {"start": 1339.62, "end": 1340.6599999999999, "text": " Whoops"}, {"start": 1340.6599999999999, "end": 1345.7, "text": " Let me zoom in again. So batched MLP predict is going to be equal"}, {"start": 1346.18, "end": 1350.1, "text": " V map and then I'm gonna wrap the MLP predict here and"}, {"start": 1350.6599999999999, "end": 1352.6599999999999, "text": " Then we're gonna specify the in axis"}, {"start": 1352.66, "end": 1356.8200000000002, "text": " Let me import V map so that I don't forget so from Jax import"}, {"start": 1357.22, "end": 1363.22, "text": " We're gonna need JIT. We're gonna need V map and we are we're gonna need P map and grad as well"}, {"start": 1363.22, "end": 1366.02, "text": " So let me rerun this and now let me continue here"}, {"start": 1366.02, "end": 1374.1000000000001, "text": " So in axis is gonna be none because we don't we just want to broadcast the parameters and it's gonna be zero for the"}, {"start": 1374.8200000000002, "end": 1380.1000000000001, "text": " Because we have the the batch dimension will be on the zero on the zero dimension"}, {"start": 1380.1, "end": 1382.1, "text": " So this should now work"}, {"start": 1382.1799999999998, "end": 1386.74, "text": " I think that's it and let me try and test this function again. So again"}, {"start": 1387.3799999999999, "end": 1389.3799999999999, "text": " Let's test it"}, {"start": 1389.3799999999999, "end": 1391.3799999999999, "text": " Small test here"}, {"start": 1391.3799999999999, "end": 1394.5, "text": " Basically, I'm gonna create this same thing"}, {"start": 1395.9399999999998, "end": 1399.62, "text": " Except that this time I'm gonna add a batch dimension of 16 here"}, {"start": 1400.1799999999998, "end": 1401.54, "text": " Let me check the shape here"}, {"start": 1401.54, "end": 1407.78, "text": " So this is gonna be images flat and I'm gonna pass that into the batched version and see what the predictions are"}, {"start": 1407.78, "end": 1410.58, "text": " So predictions predictions"}, {"start": 1411.86, "end": 1417.3, "text": " Batched version we pass in the again MLP parameters. We pass in the"}, {"start": 1417.94, "end": 1419.94, "text": " images flat"}, {"start": 1419.94, "end": 1423.78, "text": " And if I were to run this I'm gonna just print the predictions"}, {"start": 1425.06, "end": 1429.22, "text": " Shape and I'm gonna comment out this first test here"}, {"start": 1430.02, "end": 1434.02, "text": " So this should now work as expected hopefully and it does"}, {"start": 1434.02, "end": 1441.7, "text": " so we get the oh, I forgot to actually use this this version of the images so this should be 16"}, {"start": 1441.7, "end": 1447.1399999999999, "text": " Yeah, so this seems to to work correctly and that means we have done the first part here"}, {"start": 1447.1399999999999, "end": 1453.54, "text": " And now let's jump to adding the data loading functionality. So I'm just gonna annotate this like this"}, {"start": 1453.54, "end": 1455.54, "text": " So this was a small test"}, {"start": 1455.54, "end": 1457.54, "text": " and this"}, {"start": 1457.86, "end": 1459.86, "text": " thing"}, {"start": 1460.18, "end": 1462.18, "text": " These things here"}, {"start": 1462.18, "end": 1464.18, "text": " These things here were also tests"}, {"start": 1465.22, "end": 1467.22, "text": " and"}, {"start": 1467.22, "end": 1469.22, "text": " Finally, okay, let's add another cell here"}, {"start": 1470.1000000000001, "end": 1472.1000000000001, "text": " We need to add the data loading so"}, {"start": 1473.22, "end": 1477.22, "text": " First thing we need to add the MNIST dataset and then we'll need to add the data loader"}, {"start": 1477.22, "end": 1481.54, "text": " Which is going to load the images that we're gonna pass into this into this"}, {"start": 1482.1000000000001, "end": 1483.14, "text": " MLP"}, {"start": 1483.14, "end": 1488.42, "text": " So let's start with the syntax PyTorch dataset MNIST"}, {"start": 1488.42, "end": 1490.42, "text": " I'm gonna search for it"}, {"start": 1490.42, "end": 1492.42, "text": " and let's see how"}, {"start": 1492.42, "end": 1494.42, "text": " the syntax goes like"}, {"start": 1496.42, "end": 1500.42, "text": " So basically TorchVision datasets MNIST something"}, {"start": 1500.42, "end": 1504.42, "text": " Let's let's see some example whether they have MNIST here. They do not have MNIST"}, {"start": 1506.42, "end": 1512.42, "text": " Let me see whether I can type something in like"}, {"start": 1512.42, "end": 1518.42, "text": " TorchVision datasets MNIST example"}, {"start": 1518.42, "end": 1522.42, "text": " Okay, let me return back here I think this is going to work"}, {"start": 1522.42, "end": 1524.42, "text": " I'm just gonna copy"}, {"start": 1524.42, "end": 1528.42, "text": " Copy this thing. Let me put this tab here. I'm gonna need it probably a bit later"}, {"start": 1528.42, "end": 1534.42, "text": " So something like this should work. I'm gonna import here so import"}, {"start": 1534.42, "end": 1538.42, "text": " So from these datasets, I'm gonna import"}, {"start": 1538.42, "end": 1540.42, "text": " the MNIST dataset"}, {"start": 1540.42, "end": 1544.42, "text": " And now I'm just gonna go ahead and add the data loader"}, {"start": 1544.42, "end": 1548.42, "text": " and import the MNIST dataset"}, {"start": 1548.42, "end": 1552.42, "text": " And now I'm just gonna be using MNIST. So I'm gonna run this"}, {"start": 1552.42, "end": 1554.42, "text": " Hopefully this import will work"}, {"start": 1554.42, "end": 1556.42, "text": " Let me see"}, {"start": 1556.42, "end": 1560.42, "text": " Okay, it worked. So now I'm gonna use MNIST here"}, {"start": 1560.42, "end": 1566.42, "text": " Basically, let me delete this and I'm gonna create a train dataset"}, {"start": 1566.42, "end": 1570.42, "text": " By calling MNIST and passing in let's see what the arguments here are"}, {"start": 1570.42, "end": 1576.42, "text": " So we have root that's where the data will be downloaded. Then we have the split"}, {"start": 1576.42, "end": 1580.42, "text": " I'm just gonna leave this as the default one"}, {"start": 1580.42, "end": 1584.42, "text": " Train, download and transform"}, {"start": 1584.42, "end": 1588.42, "text": " So it sucks that it doesn't have an example here that can just copy paste"}, {"start": 1588.42, "end": 1592.42, "text": " So let me just do it like this. We need roots. We need basically split"}, {"start": 1592.42, "end": 1596.42, "text": " We need to train, download and transform. Those are the four we'll need"}, {"start": 1596.42, "end": 1600.42, "text": " So root is just gonna be basically"}, {"start": 1600.42, "end": 1602.42, "text": " The structure, how it looks like, is the following"}, {"start": 1602.42, "end": 1606.42, "text": " So if I was to add an additional cell here"}, {"start": 1606.42, "end": 1612.42, "text": " Import OS package and print the OS get current working directory"}, {"start": 1612.42, "end": 1616.42, "text": " It's gonna print content. So this thing here is called the content directory"}, {"start": 1616.42, "end": 1620.42, "text": " So I'm gonna, inside of this directory, I'm gonna into the"}, {"start": 1620.42, "end": 1624.42, "text": " I guess train MNIST"}, {"start": 1624.42, "end": 1628.42, "text": " I'm gonna download the train dataset"}, {"start": 1628.42, "end": 1633.42, "text": " And we're gonna need, what else? We need the train equals to true"}, {"start": 1633.42, "end": 1636.42, "text": " So that's gonna make sure that this is the training dataset"}, {"start": 1636.42, "end": 1639.42, "text": " And finally we need download and transform"}, {"start": 1639.42, "end": 1646.42, "text": " So download, I guess download is by default true here, right?"}, {"start": 1646.42, "end": 1650.42, "text": " So download, if true, downloads the dataset from the internet and puts it into root directory"}, {"start": 1650.42, "end": 1653.42, "text": " If dataset is already downloaded, it is not downloaded again"}, {"start": 1653.42, "end": 1657.42, "text": " So yeah, we wanna download it and then transform is gonna be needed as well"}, {"start": 1657.42, "end": 1659.42, "text": " We're gonna see why in a second"}, {"start": 1659.42, "end": 1661.42, "text": " So for now let me just omit it"}, {"start": 1661.42, "end": 1665.42, "text": " I'm just gonna put transform to none"}, {"start": 1665.42, "end": 1666.42, "text": " And that's it"}, {"start": 1666.42, "end": 1669.42, "text": " Let me just print the type of this thing"}, {"start": 1669.42, "end": 1673.42, "text": " Train dataset just to have something as a placeholder"}, {"start": 1673.42, "end": 1679.42, "text": " And let me rerun this, let me close this, let me rerun this and see whether this works"}, {"start": 1679.42, "end": 1680.42, "text": " And it does"}, {"start": 1680.42, "end": 1684.42, "text": " It's downloading the data, I think it's finished"}, {"start": 1684.42, "end": 1691.42, "text": " We have train MNIST dataset right here in this Google Colabs storage space"}, {"start": 1691.42, "end": 1694.42, "text": " So let's now analyze what this dataset is"}, {"start": 1694.42, "end": 1701.42, "text": " So basically if I were to fetch like a something from the train dataset"}, {"start": 1701.42, "end": 1704.42, "text": " So I'm fairly sure it has the index functionality"}, {"start": 1704.42, "end": 1706.42, "text": " So I can just index and take the 0th one"}, {"start": 1706.42, "end": 1709.42, "text": " I'm gonna print the type of this something"}, {"start": 1709.42, "end": 1716.42, "text": " And that's going to be basically, as you can see here, pill image and we have a label"}, {"start": 1716.42, "end": 1718.42, "text": " So that's the structure of our dataset"}, {"start": 1718.42, "end": 1723.42, "text": " So because we're working with Jax, we're not going to be using PyTorch tensors"}, {"start": 1723.42, "end": 1728.42, "text": " So that means we need to convert this into a NumPy array and not pill image"}, {"start": 1728.42, "end": 1730.42, "text": " So how do we do that?"}, {"start": 1730.42, "end": 1732.42, "text": " And the answer is using this transform functions"}, {"start": 1732.42, "end": 1735.42, "text": " We're just gonna add a simple transform function"}, {"start": 1735.42, "end": 1739.42, "text": " So I'm just gonna call it custom transform"}, {"start": 1739.42, "end": 1743.42, "text": " And what the function is going to do is"}, {"start": 1743.42, "end": 1749.42, "text": " It's gonna take the image and it's going to convert the image into NumPy"}, {"start": 1749.42, "end": 1753.42, "text": " So something like this, so NumPy array of this x"}, {"start": 1753.42, "end": 1759.42, "text": " And I'm just gonna set the data type to MPFloat32"}, {"start": 1759.42, "end": 1761.42, "text": " And I'm gonna return this thing"}, {"start": 1761.42, "end": 1763.42, "text": " I don't even need a temporary variable here"}, {"start": 1763.42, "end": 1765.42, "text": " So just return this"}, {"start": 1765.42, "end": 1771.42, "text": " And now if I were to run this, let me see whether it's gonna be converted into NumPy"}, {"start": 1771.42, "end": 1773.42, "text": " And it's not"}, {"start": 1774.42, "end": 1778.42, "text": " Because I forgot to add the custom transform here"}, {"start": 1778.42, "end": 1780.42, "text": " So let me run this now"}, {"start": 1780.42, "end": 1785.42, "text": " And yep, I'm getting unveiled the output here, which is a good sign actually"}, {"start": 1785.42, "end": 1789.42, "text": " So let me, instead of printing something, let me print the something 0"}, {"start": 1789.42, "end": 1795.42, "text": " And then the type, this should be, I guess, NumPy"}, {"start": 1796.42, "end": 1797.42, "text": " And it is"}, {"start": 1797.42, "end": 1800.42, "text": " Okay, so we got this sorted"}, {"start": 1800.42, "end": 1804.42, "text": " Basically we now have a NumPy Float32 image"}, {"start": 1804.42, "end": 1810.42, "text": " And the second, let me just check what the second part of this something tuple is"}, {"start": 1810.42, "end": 1813.42, "text": " So let me just check whether that's integer or whatnot"}, {"start": 1813.42, "end": 1815.42, "text": " So it is, it's an integer"}, {"start": 1815.42, "end": 1817.42, "text": " So we have, because it's a label"}, {"start": 1817.42, "end": 1819.42, "text": " So I think that should be alright"}, {"start": 1819.42, "end": 1824.42, "text": " Now, we need to additionally create the test dataset"}, {"start": 1824.42, "end": 1827.42, "text": " So I'm just gonna copy paste this part here"}, {"start": 1828.42, "end": 1830.42, "text": " I'm gonna copy paste it here"}, {"start": 1830.42, "end": 1832.42, "text": " This is going to be a test"}, {"start": 1832.42, "end": 1836.42, "text": " Train is false because we want a test dataset"}, {"start": 1836.42, "end": 1838.42, "text": " Download equals true is okay"}, {"start": 1838.42, "end": 1840.42, "text": " And custom transform, that's all okay"}, {"start": 1840.42, "end": 1847.42, "text": " So let me just run this again and see whether it's alright"}, {"start": 1847.42, "end": 1849.42, "text": " So it's downloading the test data"}, {"start": 1849.42, "end": 1854.42, "text": " And as you can see here, we have the test MNIST in the storage space"}, {"start": 1854.42, "end": 1855.42, "text": " Okay, cool"}, {"start": 1855.42, "end": 1858.42, "text": " Second thing we need to do is add the data loader"}, {"start": 1858.42, "end": 1862.42, "text": " So we need to be able to load these images in"}, {"start": 1863.42, "end": 1866.42, "text": " I think I just understood as I made a mistake here"}, {"start": 1866.42, "end": 1868.42, "text": " We need to flatten out images"}, {"start": 1868.42, "end": 1873.42, "text": " If I were to, again, if I was to take the image label tuple"}, {"start": 1873.42, "end": 1875.42, "text": " Because I know what this is now"}, {"start": 1875.42, "end": 1877.42, "text": " I'm just gonna take the train dataset"}, {"start": 1877.42, "end": 1879.42, "text": " I'm gonna take the 0th one"}, {"start": 1879.42, "end": 1881.42, "text": " And I'm just gonna take the, actually just"}, {"start": 1881.42, "end": 1883.42, "text": " Let me just grab the image"}, {"start": 1883.42, "end": 1885.42, "text": " By adding 0 here"}, {"start": 1885.42, "end": 1887.42, "text": " This is an image, but the thing is"}, {"start": 1887.42, "end": 1890.42, "text": " The shape of this image is, I guess, 28 by 28, yep"}, {"start": 1890.42, "end": 1892.42, "text": " So we need to add the"}, {"start": 1893.42, "end": 1896.42, "text": " Just flatten the function, flatten the thing here"}, {"start": 1896.42, "end": 1900.42, "text": " And if I were to run this, we get the correct shape now"}, {"start": 1900.42, "end": 1902.42, "text": " So this is now correct"}, {"start": 1902.42, "end": 1905.42, "text": " Basically the last thing we need to add here is just data loader"}, {"start": 1905.42, "end": 1907.42, "text": " So if I was to search it"}, {"start": 1907.42, "end": 1909.42, "text": " PyTorch data loader"}, {"start": 1910.42, "end": 1912.42, "text": " Let's see what I get here"}, {"start": 1913.42, "end": 1914.42, "text": " Data loader"}, {"start": 1914.42, "end": 1917.42, "text": " Okay, we have the signature here"}, {"start": 1917.42, "end": 1919.42, "text": " I'm just gonna copy paste it"}, {"start": 1920.42, "end": 1923.42, "text": " And let's see whether that's gonna work for us"}, {"start": 1923.42, "end": 1928.42, "text": " So, I'm gonna obviously need to import the data loader functionality"}, {"start": 1928.42, "end": 1931.42, "text": " So let's see where can I import it from"}, {"start": 1934.42, "end": 1937.42, "text": " Torch vision or what not, I don't know"}, {"start": 1937.42, "end": 1939.42, "text": " Let's see here"}, {"start": 1942.42, "end": 1944.42, "text": " Import"}, {"start": 1944.42, "end": 1946.42, "text": " PyTorch data loader"}, {"start": 1948.42, "end": 1951.42, "text": " Let's open up this data tutorial"}, {"start": 1951.42, "end": 1955.42, "text": " So basically data loader"}, {"start": 1955.42, "end": 1957.42, "text": " Oh my god, I hate this"}, {"start": 1959.42, "end": 1963.42, "text": " Okay, from Torch, utils data, import data loader"}, {"start": 1963.42, "end": 1965.42, "text": " Okay, let me go back here"}, {"start": 1965.42, "end": 1968.42, "text": " Let me import that part"}, {"start": 1968.42, "end": 1970.42, "text": " Let me rerun the cell"}, {"start": 1970.42, "end": 1972.42, "text": " And now we have the data loader"}, {"start": 1972.42, "end": 1974.42, "text": " Okay, so this will work"}, {"start": 1974.42, "end": 1977.42, "text": " I'm gonna use the train data set"}, {"start": 1977.42, "end": 1979.42, "text": " I'm gonna create the train loader"}, {"start": 1979.42, "end": 1985.42, "text": " And I'm gonna add some random parameters here"}, {"start": 1985.42, "end": 1990.42, "text": " So batch size will be set to, for example, 128"}, {"start": 1990.42, "end": 1993.42, "text": " And so just passing the batch size"}, {"start": 1993.42, "end": 1995.42, "text": " Shuffle is gonna be equal to true"}, {"start": 1995.42, "end": 1997.42, "text": " Because this is a train data set"}, {"start": 1997.42, "end": 1999.42, "text": " We wanna shuffle"}, {"start": 1999.42, "end": 2001.42, "text": " We don't need the custom sampler"}, {"start": 2001.42, "end": 2003.42, "text": " We don't need a batch sampler"}, {"start": 2003.42, "end": 2005.42, "text": " Number of workers is not important"}, {"start": 2005.42, "end": 2007.42, "text": " We'll need a collate function"}, {"start": 2007.42, "end": 2009.42, "text": " Because now we have a"}, {"start": 2009.42, "end": 2012.42, "text": " Basically we're not using PyTorch sensors"}, {"start": 2012.42, "end": 2015.42, "text": " These are all just optimization parameters"}, {"start": 2015.42, "end": 2016.42, "text": " Such as pin memory"}, {"start": 2016.42, "end": 2019.42, "text": " So we don't need none of those"}, {"start": 2019.42, "end": 2021.42, "text": " So let me just do it like this"}, {"start": 2021.42, "end": 2023.42, "text": " So we additionally need to create a collate function"}, {"start": 2023.42, "end": 2025.42, "text": " Because let me see what happens if I were to just"}, {"start": 2025.42, "end": 2027.42, "text": " Use it without a custom collate function"}, {"start": 2027.42, "end": 2032.42, "text": " So batch of data is equal to next"}, {"start": 2032.42, "end": 2035.42, "text": " I'm just gonna convert the train loader into iterator"}, {"start": 2035.42, "end": 2042.42, "text": " And then just fetch the next batch using the next function"}, {"start": 2042.42, "end": 2044.42, "text": " Let me see what this thing is"}, {"start": 2044.42, "end": 2046.42, "text": " Batch of data"}, {"start": 2046.42, "end": 2048.42, "text": " And whether it's gonna work"}, {"start": 2048.42, "end": 2049.42, "text": " So it's a list"}, {"start": 2049.42, "end": 2051.42, "text": " Okay, let's see what this list consists of"}, {"start": 2051.42, "end": 2054.42, "text": " So if I were to take the zeroth element"}, {"start": 2054.42, "end": 2056.42, "text": " And run it again"}, {"start": 2056.42, "end": 2057.42, "text": " It's a tensor"}, {"start": 2057.42, "end": 2058.42, "text": " So as you can see"}, {"start": 2058.42, "end": 2060.42, "text": " PyTorch by default creates a tensor"}, {"start": 2060.42, "end": 2062.42, "text": " If we do not add the custom collate function"}, {"start": 2062.42, "end": 2066.42, "text": " So that's why I'm gonna add a custom collate function here"}, {"start": 2066.42, "end": 2070.42, "text": " So def custom collate function"}, {"start": 2070.42, "end": 2074.42, "text": " We pass in the batch of data"}, {"start": 2074.42, "end": 2078.42, "text": " Basically we want to have all of the images"}, {"start": 2078.42, "end": 2080.42, "text": " Batched into a single lump array"}, {"start": 2080.42, "end": 2084.42, "text": " And we want to have all of the labels in a separate data structure"}, {"start": 2084.42, "end": 2088.42, "text": " Let's see what we currently get"}, {"start": 2088.42, "end": 2090.42, "text": " So I'm gonna just grab this custom collate"}, {"start": 2090.42, "end": 2092.42, "text": " I'm gonna paste it here"}, {"start": 2092.42, "end": 2096.42, "text": " So this batch input is just going to be a list of tuples"}, {"start": 2096.42, "end": 2101.42, "text": " And each tuple contains the numpy image and contains the label"}, {"start": 2101.42, "end": 2104.42, "text": " So we need to convert that into a bit different shape"}, {"start": 2104.42, "end": 2107.42, "text": " So instead of having a list of tuples"}, {"start": 2107.42, "end": 2109.42, "text": " We need a list with two elements"}, {"start": 2109.42, "end": 2112.42, "text": " One is all of our images batched into a numpy array"}, {"start": 2112.42, "end": 2117.42, "text": " And the second one is all of our labels basically"}, {"start": 2117.42, "end": 2120.42, "text": " Inside of a single numpy array"}, {"start": 2120.42, "end": 2123.42, "text": " So just to make sure I'm correct here"}, {"start": 2123.42, "end": 2126.42, "text": " Let me print the type of batch that should be a list"}, {"start": 2126.42, "end": 2131.42, "text": " Then let me print the type of the batch of zero that should be a tuple"}, {"start": 2131.42, "end": 2138.42, "text": " Finally let me print the type of the batch of zero"}, {"start": 2138.42, "end": 2140.42, "text": " And then zero this should be numpy"}, {"start": 2140.42, "end": 2142.42, "text": " And if I were to copy paste this"}, {"start": 2142.42, "end": 2144.42, "text": " This should be I guess integer"}, {"start": 2144.42, "end": 2147.42, "text": " And yeah"}, {"start": 2147.42, "end": 2150.42, "text": " This is gonna fail obviously because we're not returning anything"}, {"start": 2150.42, "end": 2153.42, "text": " Hopefully the print functionality will work nonetheless"}, {"start": 2153.42, "end": 2157.42, "text": " Yeah we have a list, then we have a tuple, then we have numpy array, we have integers"}, {"start": 2157.42, "end": 2159.42, "text": " So I already knew the structure obviously"}, {"start": 2159.42, "end": 2162.42, "text": " If you did not know the structure you can just kind of play with it"}, {"start": 2162.42, "end": 2165.42, "text": " Just analyze the shapes, analyze the types, etc."}, {"start": 2165.42, "end": 2169.42, "text": " Okay so what we need to do here is we need to transpose this data"}, {"start": 2169.42, "end": 2171.42, "text": " So I'm gonna do it something like this"}, {"start": 2171.42, "end": 2178.42, "text": " Transposed data is gonna be"}, {"start": 2178.42, "end": 2183.42, "text": " I'm just going to unpack the batch of data here and then zip it"}, {"start": 2183.42, "end": 2187.42, "text": " So what this will do is we'll now have"}, {"start": 2187.42, "end": 2188.42, "text": " So this is a list of tuples"}, {"start": 2188.42, "end": 2191.42, "text": " After we do this we'll just have a bunch of tuples and zip"}, {"start": 2191.42, "end": 2194.42, "text": " We'll take the first element from every tuple that means all of our images"}, {"start": 2194.42, "end": 2198.42, "text": " And it's gonna put it into a single list"}, {"start": 2198.42, "end": 2202.42, "text": " And then it's gonna take all of our labels from all the tuples"}, {"start": 2202.42, "end": 2204.42, "text": " And it's gonna combine it again into a list"}, {"start": 2204.42, "end": 2206.42, "text": " So we can make sure that's correct"}, {"start": 2206.42, "end": 2209.42, "text": " So this again should be a list"}, {"start": 2209.42, "end": 2211.42, "text": " This thing here"}, {"start": 2211.42, "end": 2216.42, "text": " And basically the zeroth element should be a list as well"}, {"start": 2216.42, "end": 2219.42, "text": " And the first element should be a list as well"}, {"start": 2219.42, "end": 2222.42, "text": " So this should be"}, {"start": 2222.42, "end": 2224.42, "text": " All of these should be lists"}, {"start": 2224.42, "end": 2227.42, "text": " Let me check whether that's the case"}, {"start": 2227.42, "end": 2230.42, "text": " Zip is not subscriptable"}, {"start": 2230.42, "end": 2232.42, "text": " Let me see what I have done"}, {"start": 2232.42, "end": 2238.42, "text": " So basically let me check"}, {"start": 2238.42, "end": 2242.42, "text": " Oh yeah, I have to convert this into a list explicitly here"}, {"start": 2242.42, "end": 2244.42, "text": " So let me now run it"}, {"start": 2244.42, "end": 2248.42, "text": " And yeah, so we have list tuple tuple"}, {"start": 2248.42, "end": 2251.42, "text": " So yeah, instead of a list, they're actually using tuple in the background"}, {"start": 2251.42, "end": 2254.42, "text": " If I were to try and"}, {"start": 2254.42, "end": 2257.42, "text": " So the length of this thing should be 2 I guess"}, {"start": 2257.42, "end": 2260.42, "text": " Transpose data length should be 2 because now we have"}, {"start": 2260.42, "end": 2262.42, "text": " Yeah, it's 2"}, {"start": 2262.42, "end": 2265.42, "text": " Because now we have all of the images and all of the labels in a single tuple"}, {"start": 2265.42, "end": 2270.42, "text": " So the next thing we need to do is basically"}, {"start": 2270.42, "end": 2273.42, "text": " Convert the tuples into a NumPy array"}, {"start": 2273.42, "end": 2275.42, "text": " So I'm gonna do something with this"}, {"start": 2275.42, "end": 2281.42, "text": " Labels is a NumPy array of this transposed data"}, {"start": 2281.42, "end": 2285.42, "text": " Of 1 because the second tuple is the labels"}, {"start": 2285.42, "end": 2287.42, "text": " And then we need images"}, {"start": 2287.42, "end": 2289.42, "text": " And here it's a bit more complex"}, {"start": 2289.42, "end": 2295.42, "text": " We need to do MPStack of transposed data of 0"}, {"start": 2295.42, "end": 2298.42, "text": " And this will hopefully work as expected"}, {"start": 2298.42, "end": 2301.42, "text": " Now I'm gonna return back the images and the labels"}, {"start": 2301.42, "end": 2304.42, "text": " And this collate function should now properly work"}, {"start": 2304.42, "end": 2307.42, "text": " Let me run this and see whether we have the correct result"}, {"start": 2307.42, "end": 2309.42, "text": " It seems nothing is crashing"}, {"start": 2309.42, "end": 2313.42, "text": " So I managed to get the batch of data"}, {"start": 2313.42, "end": 2316.42, "text": " We're printing the type and it seems"}, {"start": 2316.42, "end": 2319.42, "text": " Okay, it's a NumPy array"}, {"start": 2319.42, "end": 2321.42, "text": " Let's see what the shape is"}, {"start": 2321.42, "end": 2323.42, "text": " So if I was to take the"}, {"start": 2323.42, "end": 2324.42, "text": " Let me do it like this"}, {"start": 2324.42, "end": 2327.42, "text": " So batch of 0 are images"}, {"start": 2327.42, "end": 2331.42, "text": " So that's like this, images"}, {"start": 2331.42, "end": 2333.42, "text": " And oops"}, {"start": 2333.42, "end": 2339.42, "text": " And finally, let's see what's the shape of the images"}, {"start": 2339.42, "end": 2341.42, "text": " Let me delete this thing"}, {"start": 2341.42, "end": 2342.42, "text": " And run this"}, {"start": 2342.42, "end": 2346.42, "text": " And yeah, we have a batch of 128 images"}, {"start": 2346.42, "end": 2348.42, "text": " Which are flattened out"}, {"start": 2348.42, "end": 2350.42, "text": " And if I was to take the labels"}, {"start": 2350.42, "end": 2352.42, "text": " Same as this here"}, {"start": 2352.42, "end": 2354.42, "text": " I'm just gonna add"}, {"start": 2354.42, "end": 2358.42, "text": " Instead of 0, I'm gonna put 1 for labels"}, {"start": 2358.42, "end": 2364.42, "text": " And let's see the shape of these labels"}, {"start": 2364.42, "end": 2365.42, "text": " This looks correct"}, {"start": 2365.42, "end": 2368.42, "text": " And finally, we need to check the data type"}, {"start": 2368.42, "end": 2370.42, "text": " So images"}, {"start": 2370.42, "end": 2373.42, "text": " Let's take the 0th image and print the data type"}, {"start": 2373.42, "end": 2378.42, "text": " And also let's take the 0th label and print the data type"}, {"start": 2378.42, "end": 2383.42, "text": " So it's always a smart idea to understand your shapes and your data types"}, {"start": 2383.42, "end": 2385.42, "text": " To avoid some bugs down the road"}, {"start": 2385.42, "end": 2391.42, "text": " So flow 32 looks good and in 64 I think that's gonna be totally fine"}, {"start": 2391.42, "end": 2394.42, "text": " Okay, we have our data loading functionality"}, {"start": 2394.42, "end": 2397.42, "text": " Let's continue on and see what else do we need to do"}, {"start": 2397.42, "end": 2399.42, "text": " And that's adding the looping function"}, {"start": 2399.42, "end": 2403.42, "text": " So the training loop functionality"}, {"start": 2403.42, "end": 2407.42, "text": " Let me delete this because this is done"}, {"start": 2407.42, "end": 2411.42, "text": " And let me now finally, let me also delete this"}, {"start": 2411.42, "end": 2415.42, "text": " And let me add a final cell adding a looping function"}, {"start": 2415.42, "end": 2417.42, "text": " A training loop function"}, {"start": 2417.42, "end": 2419.42, "text": " Okay, this should be fairly simple"}, {"start": 2419.42, "end": 2422.42, "text": " What we need to do is define the number of epochs"}, {"start": 2422.42, "end": 2426.42, "text": " So number of epochs is gonna be, I don't know, for now let's put something like 10"}, {"start": 2426.42, "end": 2429.42, "text": " So we're gonna iterate through the epochs"}, {"start": 2429.42, "end": 2432.42, "text": " So in range num of epochs"}, {"start": 2432.42, "end": 2434.42, "text": " We're gonna do it 10 times"}, {"start": 2434.42, "end": 2436.42, "text": " We're going to do what?"}, {"start": 2436.42, "end": 2440.42, "text": " So we have this pattern where we take the MLP parameters"}, {"start": 2440.42, "end": 2442.42, "text": " So we have the MLP params"}, {"start": 2442.42, "end": 2444.42, "text": " We're gonna do some update function"}, {"start": 2444.42, "end": 2447.42, "text": " To which we're gonna pass the MLP params"}, {"start": 2447.42, "end": 2449.42, "text": " We're gonna pass the data"}, {"start": 2449.42, "end": 2452.42, "text": " So images and GT labels"}, {"start": 2452.42, "end": 2454.42, "text": " And this should just update our parameters"}, {"start": 2454.42, "end": 2457.42, "text": " And that should be the high level functionality"}, {"start": 2457.42, "end": 2460.42, "text": " So now we need to write this update function"}, {"start": 2460.42, "end": 2461.42, "text": " Let's see how we can do that"}, {"start": 2461.42, "end": 2463.42, "text": " So def update"}, {"start": 2463.42, "end": 2465.42, "text": " We pass in the params"}, {"start": 2465.42, "end": 2466.42, "text": " We pass in the images"}, {"start": 2466.42, "end": 2468.42, "text": " We pass in the labels"}, {"start": 2468.42, "end": 2471.42, "text": " Okay, obviously I forgot to add data loading functionality here"}, {"start": 2471.42, "end": 2475.42, "text": " So we'll have basically for what?"}, {"start": 2475.42, "end": 2476.42, "text": " We'll have images"}, {"start": 2476.42, "end": 2483.42, "text": " We'll have labels in the train loader"}, {"start": 2483.42, "end": 2488.42, "text": " And now we can pass the images and the GT labels"}, {"start": 2488.42, "end": 2489.42, "text": " But there is a small catch here"}, {"start": 2489.42, "end": 2493.42, "text": " So images"}, {"start": 2493.42, "end": 2496.42, "text": " Images are currently on a CPU"}, {"start": 2496.42, "end": 2501.42, "text": " And that's gonna be handled correctly by the Jack Slump API"}, {"start": 2501.42, "end": 2503.42, "text": " So we don't need to worry about it"}, {"start": 2503.42, "end": 2505.42, "text": " The thing we need to do here though"}, {"start": 2505.42, "end": 2509.42, "text": " Is to convert the labels into one hot representation"}, {"start": 2509.42, "end": 2511.42, "text": " We'll soon see why"}, {"start": 2511.42, "end": 2517.42, "text": " So let's first start developing this update function"}, {"start": 2517.42, "end": 2520.42, "text": " And then we'll understand what else we need to develop"}, {"start": 2520.42, "end": 2527.42, "text": " Okay, we obviously need to call the batched MLP predict function"}, {"start": 2527.42, "end": 2529.42, "text": " Okay, this one"}, {"start": 2529.42, "end": 2532.42, "text": " We're gonna pass in the parameters"}, {"start": 2532.42, "end": 2534.42, "text": " That's this argument here"}, {"start": 2534.42, "end": 2536.42, "text": " We're gonna pass in the images"}, {"start": 2536.42, "end": 2541.42, "text": " And this will return the expected shape of this"}, {"start": 2541.42, "end": 2543.42, "text": " Will be whatever the batch dimension is"}, {"start": 2543.42, "end": 2546.42, "text": " And then 10 because we have 10 outputs, remember"}, {"start": 2546.42, "end": 2550.42, "text": " And these are log softmax outputs"}, {"start": 2550.42, "end": 2552.42, "text": " So let's call those predictions"}, {"start": 2552.42, "end": 2554.42, "text": " So those are predictions here"}, {"start": 2554.42, "end": 2560.42, "text": " Finally, this is the trick that I told you about"}, {"start": 2560.42, "end": 2564.42, "text": " That will help us achieve the numerical stability of cross entropy loss"}, {"start": 2564.42, "end": 2568.42, "text": " So we'll need to just do something like this"}, {"start": 2568.42, "end": 2574.42, "text": " We need to multiply predictions with these labels"}, {"start": 2574.42, "end": 2576.42, "text": " So these are just GT labels"}, {"start": 2576.42, "end": 2578.42, "text": " So ground truth labels"}, {"start": 2578.42, "end": 2581.42, "text": " If we multiply them and if they are one hot"}, {"start": 2581.42, "end": 2585.42, "text": " That means we'll fetch the log softmax component"}, {"start": 2585.42, "end": 2589.42, "text": " From the predictions wherever the true value is"}, {"start": 2589.42, "end": 2591.42, "text": " And we'll want to push that to one"}, {"start": 2591.42, "end": 2594.42, "text": " So we'll want to have a probability one there"}, {"start": 2594.42, "end": 2597.42, "text": " Because GT label basically says that"}, {"start": 2597.42, "end": 2600.42, "text": " That's where the true class should lie in"}, {"start": 2600.42, "end": 2605.42, "text": " So now we just need to basically add minus JMP mean here"}, {"start": 2605.42, "end": 2608.42, "text": " And this is going to work as expected"}, {"start": 2608.42, "end": 2612.42, "text": " Because this is going to fetch the correct component, this part here"}, {"start": 2612.42, "end": 2614.42, "text": " And then we just do the mean"}, {"start": 2614.42, "end": 2617.42, "text": " This is gonna give us the cross entropy implementation pretty much"}, {"start": 2617.42, "end": 2621.42, "text": " Later, if I have enough time, I'm gonna do a more intuitive implementation"}, {"start": 2621.42, "end": 2625.42, "text": " We're gonna see it's not working because of numerical problems"}, {"start": 2625.42, "end": 2627.42, "text": " So the final component, as you can see here"}, {"start": 2627.42, "end": 2629.42, "text": " We need to convert this into one hot representation"}, {"start": 2629.42, "end": 2636.42, "text": " So that means that GT labels are actually JECs"}, {"start": 2636.42, "end": 2639.42, "text": " Let me see whether it's in... yeah, it's here"}, {"start": 2639.42, "end": 2642.42, "text": " So basically we just convert labels into one hot"}, {"start": 2642.42, "end": 2645.42, "text": " I'm gonna make sure this shape here is correct"}, {"start": 2645.42, "end": 2648.42, "text": " So labels shape"}, {"start": 2648.42, "end": 2652.42, "text": " And let's see whether that's gonna work or not"}, {"start": 2652.42, "end": 2654.42, "text": " Okay, this is obviously a dummy mistake"}, {"start": 2654.42, "end": 2657.42, "text": " This should be in a loss function, not in the update function"}, {"start": 2657.42, "end": 2660.42, "text": " So in this loss function, we pass in the parameters"}, {"start": 2660.42, "end": 2665.42, "text": " We pass in the, I guess, images and the GT labels"}, {"start": 2665.42, "end": 2668.42, "text": " And here we need to do this thing"}, {"start": 2668.42, "end": 2672.42, "text": " All of this goes into the loss function"}, {"start": 2672.42, "end": 2675.42, "text": " And now in the actual update function"}, {"start": 2675.42, "end": 2678.42, "text": " What we'll do is we'll take the gradients"}, {"start": 2678.42, "end": 2680.42, "text": " So we'll just find the grads"}, {"start": 2680.42, "end": 2682.42, "text": " So we'll call the grad function"}, {"start": 2682.42, "end": 2685.42, "text": " And we'll wrap in the loss function"}, {"start": 2685.42, "end": 2687.42, "text": " We're gonna pass in whatever is required"}, {"start": 2687.42, "end": 2694.42, "text": " So that's params, images, and GT labels"}, {"start": 2694.42, "end": 2698.42, "text": " And this is going to give us basically derivatives"}, {"start": 2698.42, "end": 2701.42, "text": " Of every single parameter in this pyTree here"}, {"start": 2701.42, "end": 2703.42, "text": " So we get the grads here"}, {"start": 2703.42, "end": 2708.42, "text": " And finally we need to return the JECsTreeMultiMap"}, {"start": 2708.42, "end": 2710.42, "text": " Is the function I need here"}, {"start": 2710.42, "end": 2712.42, "text": " We're gonna just do a simple function here"}, {"start": 2712.42, "end": 2715.42, "text": " So we get the weights, we get the gradients"}, {"start": 2715.42, "end": 2717.42, "text": " And what we do here is"}, {"start": 2717.42, "end": 2719.42, "text": " Okay, I don't need this here"}, {"start": 2719.42, "end": 2722.42, "text": " Let me just do the stochastic gradient descent update here"}, {"start": 2722.42, "end": 2728.42, "text": " So basically from the weight we subtract learning rate times the gradient"}, {"start": 2728.42, "end": 2732.42, "text": " Whereas I'm gonna just add some learning rate here by default"}, {"start": 2732.42, "end": 2735.42, "text": " Something like, I don't know, 0.01"}, {"start": 2735.42, "end": 2738.42, "text": " And finally let me pass in the pyTree"}, {"start": 2738.42, "end": 2742.42, "text": " So that's params and grads"}, {"start": 2742.42, "end": 2745.42, "text": " So this should now work as expected"}, {"start": 2745.42, "end": 2747.42, "text": " But we don't have any metrics here"}, {"start": 2747.42, "end": 2750.42, "text": " Or any logging of whatever kind"}, {"start": 2750.42, "end": 2755.42, "text": " So neither the loss nor the accuracy"}, {"start": 2755.42, "end": 2757.42, "text": " So I'm just gonna do a small break here"}, {"start": 2757.42, "end": 2761.42, "text": " And then a break here just to make sure that this thing works"}, {"start": 2761.42, "end": 2763.42, "text": " So one hot required one poly..."}, {"start": 2763.42, "end": 2767.42, "text": " Okay, so I forgot to add the number of classes of MNIST"}, {"start": 2767.42, "end": 2770.42, "text": " So MNIST classes"}, {"start": 2770.42, "end": 2773.42, "text": " I'm just gonna pass in the length of that thing here"}, {"start": 2773.42, "end": 2776.42, "text": " And see whether this works now"}, {"start": 2776.42, "end": 2780.42, "text": " So learning rate is not defined"}, {"start": 2780.42, "end": 2782.42, "text": " Why not? I just added it here"}, {"start": 2782.42, "end": 2784.42, "text": " So let's see"}, {"start": 2784.42, "end": 2786.42, "text": " Oh my god, I actually placed this here"}, {"start": 2786.42, "end": 2789.42, "text": " I don't know why I'm making so many mistakes"}, {"start": 2789.42, "end": 2792.42, "text": " Learning rate here, now let's try and run it"}, {"start": 2792.42, "end": 2795.42, "text": " Yeah, I'm gonna expect this is gonna last a little bit"}, {"start": 2795.42, "end": 2798.42, "text": " Okay, so this thing actually worked"}, {"start": 2798.42, "end": 2804.42, "text": " Now let me try and grab the loss"}, {"start": 2804.42, "end": 2805.42, "text": " Not just the grads here"}, {"start": 2805.42, "end": 2808.42, "text": " So well, I think the function is called val and grad"}, {"start": 2808.42, "end": 2810.42, "text": " So let me just check that"}, {"start": 2810.42, "end": 2812.42, "text": " Val and grad, jacks"}, {"start": 2812.42, "end": 2814.42, "text": " Let's see whether that's the name of the function"}, {"start": 2814.42, "end": 2817.42, "text": " Jacks package"}, {"start": 2817.42, "end": 2820.42, "text": " Value and grad is the name of the function, okay?"}, {"start": 2820.42, "end": 2822.42, "text": " So we have value and grad here"}, {"start": 2822.42, "end": 2823.42, "text": " Value and grad"}, {"start": 2823.42, "end": 2826.42, "text": " So it's gonna return the loss and the grad"}, {"start": 2826.42, "end": 2829.42, "text": " And here I'm gonna return the loss and the parameters"}, {"start": 2829.42, "end": 2832.42, "text": " And so let me just do loss here"}, {"start": 2832.42, "end": 2834.42, "text": " And let me print the value of the loss"}, {"start": 2834.42, "end": 2837.42, "text": " And let's see what happens here"}, {"start": 2837.42, "end": 2842.42, "text": " Okay, I forgot to obviously import this thing here"}, {"start": 2842.42, "end": 2845.42, "text": " Ummm"}, {"start": 2845.42, "end": 2848.42, "text": " I guess it's from jacks, right?"}, {"start": 2848.42, "end": 2850.42, "text": " Yep, it seems that's correct"}, {"start": 2850.42, "end": 2852.42, "text": " Let me run this"}, {"start": 2852.42, "end": 2855.42, "text": " And we get some reasonable value of the loss function here"}, {"start": 2855.42, "end": 2857.42, "text": " So that's cool"}, {"start": 2857.42, "end": 2862.42, "text": " Let's now try and remove the break"}, {"start": 2862.42, "end": 2865.42, "text": " And let's run this for a single epoch"}, {"start": 2865.42, "end": 2867.42, "text": " And see whether the loss goes down"}, {"start": 2867.42, "end": 2870.42, "text": " So let me run this and see what we get"}, {"start": 2870.42, "end": 2872.42, "text": " So it seems that the loss is going down"}, {"start": 2872.42, "end": 2875.42, "text": " And that's very, very reassuring"}, {"start": 2875.42, "end": 2877.42, "text": " The only problem is we're getting too many output here"}, {"start": 2877.42, "end": 2880.42, "text": " So let me just interrupt the execution here"}, {"start": 2880.42, "end": 2883.42, "text": " And let me clear this output"}, {"start": 2883.42, "end": 2884.42, "text": " Wow, there's so many"}, {"start": 2884.42, "end": 2886.42, "text": " Okay, let me just clear this"}, {"start": 2886.42, "end": 2890.42, "text": " I'm gonna add a conditional logging here"}, {"start": 2890.42, "end": 2895.42, "text": " So if, I'm just gonna place a numerate here"}, {"start": 2895.42, "end": 2898.42, "text": " So that we know the batch ID"}, {"start": 2898.42, "end": 2900.42, "text": " So something like this"}, {"start": 2900.42, "end": 2902.42, "text": " And if the counter"}, {"start": 2902.42, "end": 2905.42, "text": " So every, because every batch has 128 images"}, {"start": 2905.42, "end": 2911.42, "text": " So let's say we can log every 50 batches or something"}, {"start": 2911.42, "end": 2913.42, "text": " So if this equals zero"}, {"start": 2913.42, "end": 2917.42, "text": " Finally, let's move the loss inside of this if statement"}, {"start": 2917.42, "end": 2920.42, "text": " So let me indent it here"}, {"start": 2920.42, "end": 2924.42, "text": " I'm gonna also just print this thing once"}, {"start": 2924.42, "end": 2928.42, "text": " And now this looks like it should work"}, {"start": 2928.42, "end": 2930.42, "text": " Let me run it"}, {"start": 2930.42, "end": 2934.42, "text": " And see, so 128, 10, this looks to be a good format"}, {"start": 2934.42, "end": 2937.42, "text": " So I'm just gonna interrupt this execution"}, {"start": 2937.42, "end": 2943.42, "text": " I'm basically going to just leave the loss printing here"}, {"start": 2943.42, "end": 2948.42, "text": " And the final thing before I'll remove this break"}, {"start": 2948.42, "end": 2951.42, "text": " Is to add the accuracy metrics"}, {"start": 2951.42, "end": 2954.42, "text": " So let's add the accuracy logging as well"}, {"start": 2954.42, "end": 2958.42, "text": " And then we'll be certain that MLP is learning something"}, {"start": 2958.42, "end": 2961.42, "text": " So accuracy"}, {"start": 2961.42, "end": 2962.42, "text": " Again, I'm not sure whether Jax"}, {"start": 2962.42, "end": 2967.42, "text": " Let me see whether Jax has accuracy metric"}, {"start": 2967.42, "end": 2969.42, "text": " Yeah, I don't see anything"}, {"start": 2969.42, "end": 2970.42, "text": " So I'm just gonna implement it myself"}, {"start": 2970.42, "end": 2974.42, "text": " So basically we're gonna pass the parameters again"}, {"start": 2974.42, "end": 2976.42, "text": " We're gonna pass the loader"}, {"start": 2976.42, "end": 2978.42, "text": " So either test or train loader"}, {"start": 2978.42, "end": 2980.42, "text": " And that should be it I think"}, {"start": 2980.42, "end": 2987.42, "text": " So the thing we need to do is find the number of correctly classified images"}, {"start": 2987.42, "end": 2992.42, "text": " So I'm just gonna basically iterate through images and labels"}, {"start": 2992.42, "end": 2995.42, "text": " Here in the loader"}, {"start": 2995.42, "end": 2999.42, "text": " I'm gonna again run the"}, {"start": 2999.42, "end": 3001.42, "text": " So I'm gonna call the batch"}, {"start": 3001.42, "end": 3003.42, "text": " So something like this"}, {"start": 3003.42, "end": 3006.42, "text": " I mean this is going to be inefficient"}, {"start": 3006.42, "end": 3010.42, "text": " Because I'm gonna run it in general if we had some arbitrary data set"}, {"start": 3010.42, "end": 3013.42, "text": " But this implantation is gonna work for MNIST pretty well"}, {"start": 3013.42, "end": 3017.42, "text": " So I'm just gonna go through the whole training set and through the whole test set"}, {"start": 3017.42, "end": 3020.42, "text": " And find the accuracy and report that somewhere here"}, {"start": 3020.42, "end": 3024.42, "text": " So after we finish the epoch"}, {"start": 3024.42, "end": 3027.42, "text": " I'm gonna print something like this"}, {"start": 3027.42, "end": 3029.42, "text": " So epoch"}, {"start": 3029.42, "end": 3032.42, "text": " And then I'm gonna print the ID of the epoch"}, {"start": 3032.42, "end": 3038.42, "text": " And then I'm gonna print the train accuracy equals to this"}, {"start": 3038.42, "end": 3041.42, "text": " And test accuracy equals to this"}, {"start": 3041.42, "end": 3045.42, "text": " I'm just gonna call the basically accuracy"}, {"start": 3045.42, "end": 3051.42, "text": " I'm gonna pass the MLP params after the epoch has been finished"}, {"start": 3051.42, "end": 3053.42, "text": " I'm gonna pass the train loader"}, {"start": 3053.42, "end": 3055.42, "text": " So this thing"}, {"start": 3055.42, "end": 3057.42, "text": " And I just forgot to"}, {"start": 3057.42, "end": 3061.42, "text": " So I basically just realized that I forgot to add the test loader"}, {"start": 3061.42, "end": 3063.42, "text": " So I'm gonna quickly add that as well"}, {"start": 3063.42, "end": 3065.42, "text": " So something like this"}, {"start": 3065.42, "end": 3067.42, "text": " Let me paste it here"}, {"start": 3067.42, "end": 3069.42, "text": " So now I have to add the test loader"}, {"start": 3069.42, "end": 3071.42, "text": " Somewhere above"}, {"start": 3071.42, "end": 3073.42, "text": " Let me see where I've defined"}, {"start": 3073.42, "end": 3074.42, "text": " Yeah, here"}, {"start": 3074.42, "end": 3076.42, "text": " So I'm gonna add the test loader"}, {"start": 3076.42, "end": 3079.42, "text": " It's gonna be the same as the train data loader"}, {"start": 3079.42, "end": 3083.42, "text": " I'm just gonna swap this for test data set"}, {"start": 3083.42, "end": 3086.42, "text": " Batch size, I can leave it the same size"}, {"start": 3086.42, "end": 3088.42, "text": " Shuffle, false"}, {"start": 3088.42, "end": 3090.42, "text": " No need to shuffle that data set"}, {"start": 3090.42, "end": 3093.42, "text": " And custom collate function is also fine"}, {"start": 3093.42, "end": 3096.42, "text": " So I'm gonna rerun the cell"}, {"start": 3096.42, "end": 3098.42, "text": " And now we have the test loader"}, {"start": 3098.42, "end": 3101.42, "text": " So this should now all be correctly set up"}, {"start": 3101.42, "end": 3104.42, "text": " We just need to implement the actual accuracy metric here"}, {"start": 3104.42, "end": 3107.42, "text": " Let's see, we need to do the same thing as here"}, {"start": 3107.42, "end": 3109.42, "text": " We need some predictions"}, {"start": 3109.42, "end": 3114.42, "text": " And then I'm gonna find the basically"}, {"start": 3114.42, "end": 3116.42, "text": " Instead of finding the raw predictions here"}, {"start": 3116.42, "end": 3118.42, "text": " So the log softmax"}, {"start": 3118.42, "end": 3123.42, "text": " I'm gonna do immediately call the soft argmax function"}, {"start": 3123.42, "end": 3126.42, "text": " And the axis should be 1"}, {"start": 3126.42, "end": 3130.42, "text": " So that's gonna find the highest probability class"}, {"start": 3130.42, "end": 3132.42, "text": " For the particular prediction"}, {"start": 3132.42, "end": 3136.42, "text": " And we're gonna get basically prediction classes here"}, {"start": 3136.42, "end": 3139.42, "text": " And we're gonna contrast that with labels pretty much"}, {"start": 3139.42, "end": 3142.42, "text": " So we need to do something like this"}, {"start": 3142.42, "end": 3145.42, "text": " We'll need some sum temporary variables"}, {"start": 3145.42, "end": 3148.42, "text": " So sum, I'm gonna call it accumulator"}, {"start": 3148.42, "end": 3153.42, "text": " Basically accumulator is gonna add here the MP sum"}, {"start": 3153.42, "end": 3156.42, "text": " Between, which is gonna do predict classes"}, {"start": 3156.42, "end": 3162.42, "text": " And we're gonna see where those predicted classes are the same as the ground truth classes"}, {"start": 3162.42, "end": 3168.42, "text": " So I can maybe do something like this just for consistency sake"}, {"start": 3168.42, "end": 3171.42, "text": " So this, and we're gonna sum them up"}, {"start": 3171.42, "end": 3174.42, "text": " Find where we had true, that means we had the same class"}, {"start": 3174.42, "end": 3177.42, "text": " So this is going to work correctly"}, {"start": 3177.42, "end": 3179.42, "text": " And finally in order to return the metric"}, {"start": 3179.42, "end": 3184.42, "text": " We need to just divide the accumulated value"}, {"start": 3184.42, "end": 3188.42, "text": " With the length of the actual data loader I guess"}, {"start": 3188.42, "end": 3191.42, "text": " Is that correct?"}, {"start": 3191.42, "end": 3195.42, "text": " No it's not because we have to also add"}, {"start": 3195.42, "end": 3197.42, "text": " We care about the number of images"}, {"start": 3197.42, "end": 3202.42, "text": " So I guess we need to multiply that with the batch size"}, {"start": 3202.42, "end": 3206.42, "text": " So let me see whether loader has batch size variable"}, {"start": 3206.42, "end": 3208.42, "text": " I'm not sure whether this is going to work"}, {"start": 3208.42, "end": 3213.42, "text": " Whether Colab autocomplete is just tricking me here"}, {"start": 3213.42, "end": 3214.42, "text": " But this should work"}, {"start": 3214.42, "end": 3217.42, "text": " And this may be misleading because this is not the accuracy"}, {"start": 3217.42, "end": 3219.42, "text": " This is just the accumulated value"}, {"start": 3219.42, "end": 3222.42, "text": " So I'm just gonna replace this with something like Xsum"}, {"start": 3222.42, "end": 3226.42, "text": " And yeah this is my find and replace method for Google Colab"}, {"start": 3226.42, "end": 3229.42, "text": " Okay, this should not work I think"}, {"start": 3229.42, "end": 3233.42, "text": " Aside from this thing maybe not being correct"}, {"start": 3233.42, "end": 3236.42, "text": " Let me try and run this"}, {"start": 3236.42, "end": 3240.42, "text": " Maybe instead of waiting for the whole epoch"}, {"start": 3240.42, "end": 3244.42, "text": " I'm just going to run it right at the beginning here once"}, {"start": 3244.42, "end": 3248.42, "text": " Let's just see whether this is going to work"}, {"start": 3248.42, "end": 3250.42, "text": " Okay this actually failed"}, {"start": 3250.42, "end": 3254.42, "text": " And the reason is already quite familiar to me from PyTorch"}, {"start": 3254.42, "end": 3263.42, "text": " So what I forgot to do is basically we have to set the drop last to true here"}, {"start": 3263.42, "end": 3266.42, "text": " Otherwise we'll have the..."}, {"start": 3266.42, "end": 3269.42, "text": " So potentially the last batch will not have 128 images"}, {"start": 3269.42, "end": 3272.42, "text": " And that's where the discrepancy comes to be"}, {"start": 3272.42, "end": 3278.42, "text": " So if I were to rerun the cell here and then start this thing again"}, {"start": 3278.42, "end": 3282.42, "text": " Now hopefully we'll get the correct results"}, {"start": 3282.42, "end": 3287.42, "text": " So basically what happened here is at one point the last batch was smaller than 128"}, {"start": 3287.42, "end": 3291.42, "text": " And there was a mismatch doing a comparison here and that's why it failed"}, {"start": 3291.42, "end": 3295.42, "text": " Here you can see it's not working but the thing is it's very very slow"}, {"start": 3295.42, "end": 3299.42, "text": " So what I'm going to do instead is I'm going to interrupt this"}, {"start": 3299.42, "end": 3306.42, "text": " And I'm going to load the whole MNIST into memory because it's a fairly small dataset with only 50k, 60k images"}, {"start": 3306.42, "end": 3308.42, "text": " So I'm going to do the following"}, {"start": 3308.42, "end": 3311.42, "text": " I'm going to change my implementation of accuracy"}, {"start": 3311.42, "end": 3316.42, "text": " So instead of doing the loader thing and then iterating through the whole dataset"}, {"start": 3316.42, "end": 3319.42, "text": " I'm just going to pass the dataset itself"}, {"start": 3319.42, "end": 3322.42, "text": " And that means this is going to change a lot"}, {"start": 3322.42, "end": 3325.42, "text": " So I'm going to just do the following"}, {"start": 3325.42, "end": 3330.42, "text": " Pass the whole dataset all at once into the MLP predict"}, {"start": 3330.42, "end": 3335.42, "text": " So I'm going to do something like this params the whole dataset"}, {"start": 3335.42, "end": 3337.42, "text": " So actually the images"}, {"start": 3337.42, "end": 3340.42, "text": " So these will be dataset images"}, {"start": 3340.42, "end": 3342.42, "text": " I'm going to pass them here"}, {"start": 3342.42, "end": 3350.42, "text": " We're going to have the again the JMP argmax alongside xis equals one"}, {"start": 3350.42, "end": 3355.42, "text": " So this will give us the predicted classes"}, {"start": 3355.42, "end": 3360.42, "text": " And I'm also going to pass obviously the dataset labels"}, {"start": 3360.42, "end": 3365.42, "text": " So and I'm going to compare those two and return the mean"}, {"start": 3365.42, "end": 3369.42, "text": " So basically JMP mean of this comparison"}, {"start": 3369.42, "end": 3373.42, "text": " And this is neater implementation so it's more concise"}, {"start": 3373.42, "end": 3377.42, "text": " And it's also way more fast compared to the last one"}, {"start": 3377.42, "end": 3379.42, "text": " So let me see where this makes sense"}, {"start": 3379.42, "end": 3382.42, "text": " So I'm doing inference on the whole dataset"}, {"start": 3382.42, "end": 3385.42, "text": " And I'm then converting my raw output"}, {"start": 3385.42, "end": 3392.42, "text": " So the log softmax outputs into the actual class indices here"}, {"start": 3392.42, "end": 3396.42, "text": " And I'm comparing them to the ground truth and finding the mean"}, {"start": 3396.42, "end": 3399.42, "text": " Okay so this looks like it should work"}, {"start": 3399.42, "end": 3401.42, "text": " So now I have to do the following thing"}, {"start": 3401.42, "end": 3404.42, "text": " So I have to load the whole thing into the memory"}, {"start": 3404.42, "end": 3405.42, "text": " So I'm going to do the following"}, {"start": 3405.42, "end": 3408.42, "text": " I'm going to take this thing"}, {"start": 3408.42, "end": 3410.42, "text": " So train images"}, {"start": 3410.42, "end": 3414.42, "text": " There is this internal member called data"}, {"start": 3414.42, "end": 3417.42, "text": " I'm going to fetch it that's going to get me the images"}, {"start": 3417.42, "end": 3423.42, "text": " I'm going to convert this into like a Jack's NumPy array"}, {"start": 3423.42, "end": 3428.42, "text": " And I'm going to reshape this into"}, {"start": 3428.42, "end": 3430.42, "text": " Wait let me check what this will return"}, {"start": 3430.42, "end": 3436.42, "text": " So this will going to be print train images shape"}, {"start": 3436.42, "end": 3438.42, "text": " Let me see what this is going to return"}, {"start": 3438.42, "end": 3441.42, "text": " So yeah as you can see the problem here is it's avoiding"}, {"start": 3441.42, "end": 3446.42, "text": " So this transform is not being called when I just fetched the data like this"}, {"start": 3446.42, "end": 3450.42, "text": " So we still need to flatten out the images"}, {"start": 3450.42, "end": 3455.42, "text": " So I'm going to do that I'm going to just reshape this into basically"}, {"start": 3455.42, "end": 3459.42, "text": " Whatever the length of the train data set is"}, {"start": 3459.42, "end": 3461.42, "text": " And I'm going to put minus one here"}, {"start": 3461.42, "end": 3464.42, "text": " So that's going to hopefully flatten out the whole thing"}, {"start": 3464.42, "end": 3469.42, "text": " And now yeah now we have the correct basically shape"}, {"start": 3469.42, "end": 3473.42, "text": " Let me just check the type train images"}, {"start": 3473.42, "end": 3475.42, "text": " It's going to be on the device"}, {"start": 3475.42, "end": 3479.42, "text": " So that means whatever accelerator I'm using is going to be on that accelerator"}, {"start": 3479.42, "end": 3482.42, "text": " I'm currently set to using none"}, {"start": 3482.42, "end": 3488.42, "text": " So yeah because MLP is a fairly simple model and MNIST is a super small data set"}, {"start": 3488.42, "end": 3491.42, "text": " Let me now do the same thing for labels I need train labels"}, {"start": 3491.42, "end": 3495.42, "text": " I'm going to do the same thing here"}, {"start": 3495.42, "end": 3497.42, "text": " I think so at least"}, {"start": 3497.42, "end": 3501.42, "text": " So I'm going to this time I'm going to access target so that's labels"}, {"start": 3501.42, "end": 3505.42, "text": " I'm going to convert them into Jack's array"}, {"start": 3505.42, "end": 3510.42, "text": " I think I need the reshaping so let me check the train label shape"}, {"start": 3510.42, "end": 3515.42, "text": " And let me check the type of the train labels"}, {"start": 3515.42, "end": 3519.42, "text": " So we have device array 60K this looks correct"}, {"start": 3519.42, "end": 3521.42, "text": " I think this is going to work"}, {"start": 3521.42, "end": 3524.42, "text": " So we need to do the same thing for test data set"}, {"start": 3524.42, "end": 3526.42, "text": " Let me just copy paste this"}, {"start": 3526.42, "end": 3529.42, "text": " Replace this with test"}, {"start": 3529.42, "end": 3533.42, "text": " Test images test labels"}, {"start": 3533.42, "end": 3536.42, "text": " Test data set here"}, {"start": 3536.42, "end": 3539.42, "text": " Test data set here"}, {"start": 3539.42, "end": 3545.42, "text": " Finally test data set here targets this should be working now"}, {"start": 3545.42, "end": 3547.42, "text": " Okay cool"}, {"start": 3547.42, "end": 3553.42, "text": " Now I'm going to instead of passing the I'm going to remove this line here"}, {"start": 3553.42, "end": 3556.42, "text": " And here I'm going to pass the data set and the label"}, {"start": 3556.42, "end": 3560.42, "text": " So I'm going to pass the train data set"}, {"start": 3560.42, "end": 3562.42, "text": " The train images"}, {"start": 3562.42, "end": 3567.42, "text": " And I'm going to pass the train labels"}, {"start": 3567.42, "end": 3574.42, "text": " And here I'm going to pass basically the same thing just test images"}, {"start": 3574.42, "end": 3578.42, "text": " And test labels"}, {"start": 3578.42, "end": 3584.42, "text": " Again I should have waited until I checked this works let me just run it like this"}, {"start": 3584.42, "end": 3588.42, "text": " See whether this is going to work"}, {"start": 3588.42, "end": 3596.42, "text": " And it's way faster as you can see and the accuracy is for some reason very very very high already"}, {"start": 3596.42, "end": 3602.42, "text": " I guess that's because we were assuming using the same MLP parameters so I have to I know what I have done"}, {"start": 3602.42, "end": 3604.42, "text": " So I'm going to interrupt this"}, {"start": 3604.42, "end": 3612.42, "text": " The thing is I'm using I'm constantly updating the MLP params so the state is saved here I have to call the init function here"}, {"start": 3612.42, "end": 3622.42, "text": " So we need to paste the let me find it let me find the init function we need to paste this here"}, {"start": 3622.42, "end": 3632.42, "text": " So that we start fresh every time we start the training so this now looks like it should work let me run it"}, {"start": 3632.42, "end": 3637.42, "text": " We have the results very very fast the accuracy is now low that seems to be working"}, {"start": 3637.42, "end": 3646.42, "text": " Let me interrupt the execution again and I'm going to delete this I'm going to yeah this now should work"}, {"start": 3646.42, "end": 3656.42, "text": " We're logging the loss every 50 batches we're logging the accuracy every epoch and we have how many epochs we have 10 epochs"}, {"start": 3656.42, "end": 3663.42, "text": " So let's start with like five and I'm going to run this and yeah let's see what the results are"}, {"start": 3663.42, "end": 3669.42, "text": " So the loss is going down that's cool that seems to be promising"}, {"start": 3669.42, "end": 3688.42, "text": " Okay after a single epoch we have 91% accuracy and we do not see any kind any sign of overfitting because the accuracy on both training and test data sets are pretty much the same 91%"}, {"start": 3688.42, "end": 3694.42, "text": " And now we have 93 I'm going to just skip until the training is done"}, {"start": 3694.42, "end": 3702.42, "text": " Okay the training is finally done we have 96% accuracy with a couple more epochs we would get to I guess 98 at least or something"}, {"start": 3702.42, "end": 3706.42, "text": " It could be further improved but the thing is this thing is already working"}, {"start": 3706.42, "end": 3713.42, "text": " So let me just kind of clear the output and let me try and run a prediction on a single image and I should have done this previously"}, {"start": 3713.42, "end": 3721.42, "text": " So let me fetch a single image visualize it and see whether the ground truth label like aligns with the predicted label"}, {"start": 3721.42, "end": 3731.42, "text": " So I'm going to do the following thing I'm going to fetch the single image so images let me take the like test loader"}, {"start": 3731.42, "end": 3737.42, "text": " I'm going to wrap it into the iterator here and take a single batch"}, {"start": 3737.42, "end": 3744.42, "text": " I'm going to take a single batch and I'm going to take the only the images and then I'm going to take a single image"}, {"start": 3744.42, "end": 3751.42, "text": " And make let's make sure that this makes sense so image shape should be 28 by 28"}, {"start": 3751.42, "end": 3760.42, "text": " Oh yeah I have to just reshape it into MNIST image size let me see where this is going to work it is"}, {"start": 3760.42, "end": 3769.42, "text": " Now let me import ad hoc here I'm going to import just map plot live pyplot as plt"}, {"start": 3769.42, "end": 3776.42, "text": " Plt image show the image and let me just plot that thing"}, {"start": 3776.42, "end": 3781.42, "text": " So this looks correct we have number seven"}, {"start": 3781.42, "end": 3788.42, "text": " We'll also need labels so I have I'll have to modify this I'm going to grab labels as well so labels"}, {"start": 3788.42, "end": 3797.42, "text": " And let me check the the the ground truth label here so ground truth label is going to be labels of zero"}, {"start": 3797.42, "end": 3806.42, "text": " And let me print the the actual label here print the GT label it should be seven and it is OK"}, {"start": 3806.42, "end": 3817.42, "text": " So now let's see what our network is giving us if I was to do the MLP MLP predict on top of a image"}, {"start": 3817.42, "end": 3827.42, "text": " But I'll have to flatten it out so I'm going to pass the MLP params and I'm going to pass in the what I'm going to pass MP"}, {"start": 3827.42, "end": 3835.42, "text": " Revel of the image will this work let me check so let me check whether this is going to work"}, {"start": 3835.42, "end": 3845.42, "text": " These are just predictions and yeah we get some results let's now just do the argmax instead"}, {"start": 3845.42, "end": 3855.42, "text": " And see what we have so yeah we had we have the predicted value of seven so let's just add some string here"}, {"start": 3855.42, "end": 3864.42, "text": " And this is GT like this I'm going to paste it here let me run this and see whether it makes sense"}, {"start": 3864.42, "end": 3873.42, "text": " So predicted seven and GT seven so I have I should have done this image analysis while I was developing the"}, {"start": 3873.42, "end": 3880.42, "text": " Loaders here so basically while I was developing these data sets and loaders but yeah anyways we have it now"}, {"start": 3880.42, "end": 3885.42, "text": " We see everything is working correctly so that was pretty much it when it comes to the model training"}, {"start": 3885.42, "end": 3891.42, "text": " So was the first section it took around an hour to do this so now I want to do the following I want to quickly refactor"}, {"start": 3891.42, "end": 3896.42, "text": " The whole code code base and then I'm going to do some visualizations I was basically cutting some corners"}, {"start": 3896.42, "end": 3903.42, "text": " Because I'm filming this and I'm really I'm coding real time so it's hard to kind of be very very pedantic"}, {"start": 3903.42, "end": 3911.42, "text": " So now I want to improve some of those some of the things so let me see what first so this thing this this function"}, {"start": 3911.42, "end": 3918.42, "text": " Looks fairly nice we have parameters here we have keys I think the structure of this function was fairly nice"}, {"start": 3918.42, "end": 3925.42, "text": " We have this small testing functionality here so I'm just going to call it test we have a key we created"}, {"start": 3925.42, "end": 3933.42, "text": " We just called the function to see whether the results make sense and I think this cell was was neat I was still I was still very very fresh"}, {"start": 3933.42, "end": 3941.42, "text": " I guess this was some initial like this cell was used for debugging so I'm just gonna basically delete it"}, {"start": 3941.42, "end": 3949.42, "text": " Let me see okay so I had some explanations here etc etc again the design of this function is is fairly neat I'd say"}, {"start": 3949.42, "end": 3957.42, "text": " You can obviously always add some like doc strings comments etc whatnot okay let me see here we had some tests here"}, {"start": 3957.42, "end": 3965.42, "text": " So this was this first part was testing basically whether the MLP predict is working as expected"}, {"start": 3965.42, "end": 3972.42, "text": " And so I just printed the shape of the image I printed the shape of the output and the second part was actually testing whether"}, {"start": 3972.42, "end": 3981.42, "text": " So this test test the single example function and this was testing the batch test the batched"}, {"start": 3981.42, "end": 3988.42, "text": " Batched function basically here okay I can uncomment this part as well"}, {"start": 3988.42, "end": 3996.42, "text": " So yeah I mean it would obviously take a lot of time to to create this to make it perfect I'm just I want to make sure that I didn't"}, {"start": 3996.42, "end": 4005.42, "text": " That I don't have any any red flags like a huge huge red flags so I defined these custom transforming custom custom colloid functions"}, {"start": 4005.42, "end": 4013.42, "text": " These look neat as well so some consistency with my naming of the variables would be appreciated"}, {"start": 4013.42, "end": 4022.42, "text": " I'm sometimes using labels the full word and sometimes I'm abbreviating into LBLS so that sucks but yeah this looks nice"}, {"start": 4022.42, "end": 4033.42, "text": " We have the data sets here I did some debugging here I can I can I'm just going to comment it out and I'm basically just going to delete this part"}, {"start": 4033.42, "end": 4042.42, "text": " Now we have the data loaders then I did some testing here so this can be considered like a testing section again testing section"}, {"start": 4042.42, "end": 4050.42, "text": " What I did here I made sure that we have images labels of correct types and shapes and finally"}, {"start": 4050.42, "end": 4061.42, "text": " This thing here is just the optimization part where I was loading the whole data set load the whole data set into memory"}, {"start": 4061.42, "end": 4069.42, "text": " Okay I'm going to refactor this notebook a little bit after after I finish this video and I'm going to upload it to my to my github"}, {"start": 4069.42, "end": 4079.42, "text": " Hopefully you saw this one already so github let me let me open it up so I have this get started with Jax like repo do check it out"}, {"start": 4079.42, "end": 4086.42, "text": " I'm going to be posting all of these tutorials there as well as some other useful resources and videos here so that's it"}, {"start": 4086.42, "end": 4096.42, "text": " Let's continue here okay this to do is done we have the loss function I'm going to just delete this here we have the update function"}, {"start": 4096.42, "end": 4107.42, "text": " We have the the the accuracy maybe the metric and the loss should be one close to each other here that's it and finally we we create the network here"}, {"start": 4107.42, "end": 4120.42, "text": " Create create an MLP I guess a MLP and then we have the training loop here finally I could have probably just replace this with a product"}, {"start": 4120.42, "end": 4128.42, "text": " Function so to make this a little bit more robust but this is not that robust in any case because I'm hard coding names here but yeah"}, {"start": 4128.42, "end": 4135.42, "text": " You get the point in general you could kind of you'd have a section where you're defining these variables depending so when when the data set changes"}, {"start": 4135.42, "end": 4145.42, "text": " All of these variables change and basically MLP is out of the box is going to have a compatible shape this this was the point of doing this here"}, {"start": 4145.42, "end": 4154.42, "text": " And yeah you don't want to hard code stuff around your code base so yeah all in all for for a code that was written very very very ad hoc I think this this looks pretty decent"}, {"start": 4154.42, "end": 4163.42, "text": " So we have the loading here we convert the representation into the one hot representation then we call the update with this some logging of the loss"}, {"start": 4163.42, "end": 4173.42, "text": " We do some logging of the accuracy and I think it's fairly fairly cool now let's do a couple of visualizations first let me visualize the MLP weights"}, {"start": 4173.42, "end": 4183.42, "text": " So I'm gonna copy paste this into a novel cell here and what I mean by that is the following so I'm going to take the MLP weights"}, {"start": 4183.42, "end": 4196.42, "text": " So I'm going to take the MLP parameters I'm gonna fetch the zero player I'm gonna fetch the weight matrix and I'm gonna just take the that weight matrix and see what the shape is"}, {"start": 4196.42, "end": 4210.42, "text": " So the shape of that matrix is let me see 512 784 so what I want to do here is for a single output neuron I'm gonna I'm gonna fetch this all the weights for a single neuron and see whether we can see some pattern some interesting pattern"}, {"start": 4210.42, "end": 4234.42, "text": " I'm gonna do something like this weights single equals weights I'm gonna fetch a single row I'm gonna fetch all the columns and I'm gonna reshape that I'm gonna reshape that into the amnest image size so that's 28 by 28 and this should now be of correct shape with W single"}, {"start": 4234.42, "end": 4262.42, "text": " shape equals 28 by 28 I'm gonna just plot this and let's see whether we have some meaningful results here so I'm gonna plot it and we do not have any meaningful results here so if I was to randomly pick some other like some other weights for some other output neuron like 4 maybe I don't know"}, {"start": 4262.42, "end": 4282.42, "text": " let's see whether we can find some meaningful representation here and I do not see anything meaningful here I'm gonna try one more nope I guess not that interpretable but we still have 96% accuracy so that's cool and it was worth a try to just do this initial visualization"}, {"start": 4282.42, "end": 4311.42, "text": " so this task is done whoops let me just delete this to do and let's jump to the second one so visualize embeddings using t-sneak let me do that one I'm just gonna add a new cell here we're gonna import basically let me see side pie t-sneak okay yeah from scikit-learn manifold t-sneak so this will probably"}, {"start": 4311.42, "end": 4334.42, "text": " come in handy this documentation here I'm just gonna put the tab right here what the fuck okay let me put it here import from scikit-learn sklearn we import manifold module and we import t-sneak so from this package from this module import the function"}, {"start": 4334.42, "end": 4355.42, "text": " t-sneak my god okay okay this should now work so t-sneak is going to do a dimensionality reduction and for that I'll have to fetch the activations from the MLP just before we apply the linear layer that means we'll have to modify this predict"}, {"start": 4355.42, "end": 4374.42, "text": " function so let me find the MLP predict function I'm gonna copy paste it to this cell here I'm gonna modify it so first things first let's call it something else like fetch activations I guess we're gonna pass the params nothing changes there the only"}, {"start": 4374.42, "end": 4393.42, "text": " difference is we do not have to apply this this last part so let me let me think this will this will be yeah we basically do not want this part we just want to return back the activations the last ones just before we apply the linear layer so activation"}, {"start": 4393.42, "end": 4419.42, "text": " something like this and this should now work so now let me just create a batch version of this thing so batched fetch activations is gonna be remap we're gonna paste this function in and finally in axis will be again none and zero and this should now work so I'm just gonna fetch a batch of images"}, {"start": 4419.42, "end": 4440.42, "text": " and I'm gonna get those activations project them into 2d space and plot those images just to see whether the MLP the network learned how to separate different classes so let's see how that will look like so batch of activations is going to be equal to basically"}, {"start": 4440.42, "end": 4455.42, "text": " this and we need to pass MLP params so that's our train network and we also need to pass a batch of data so that means let me fetch a batch of images so batch"}, {"start": 4455.42, "end": 4475.42, "text": " I'm gonna take take like this images labels is gonna be next error and I'm gonna I'm going to again take from the test loader so test loader and yeah this should be it now I'm gonna take the images and I'm just gonna"}, {"start": 4475.42, "end": 4498.42, "text": " test them here and this should now work let me see what's the shape of these batch activations if this thing works so shape let me run this and it's 128 256 and that sounds about right because remember we had"}, {"start": 4498.42, "end": 4522.42, "text": " let me find a configuration so it was yeah 512 256 256 is the dimensionality of the activations just before we project them into the logits okay so let's now do the following we need to pass this into the t-sne method in order to get 128 comma 2"}, {"start": 4522.42, "end": 4540.42, "text": " and this time I'm just gonna use my github because I have this pre-existing code already and I don't want to search through the documentation this video is already too long so I have this annotated get like method and it got pretty popular and anyways I had"}, {"start": 4540.42, "end": 4559.42, "text": " I used this new here so I'm just gonna copy paste a snippet let me find where I've been using it so here I'm just gonna take this part and I'm just gonna copy paste it into right into our notebook so let me do it like this okay so t-sne is here we'll take two"}, {"start": 4559.42, "end": 4576.42, "text": " components means basically we want to down project into two dimensions perplexity 30 is just a number that works we don't need this special method here and now finally I need to pass the batch activations here and yeah so this is how the usual normal"}, {"start": 4576.42, "end": 4592.42, "text": " workflow of a guy who does machine learning looks like you take somebody else's code or your own code even better and you just kind of tweak it so I'll also we won't need this part as well so number of classes is going to be equal to 10 because this is"}, {"start": 4592.42, "end": 4612.42, "text": " amnest and finally let's see I'm just gonna erase these comments for now so the only thing we need to change here is to put labels instead of node labels here because those are our ground truth labels so basically what this does is for class 0 we're going to extract"}, {"start": 4612.42, "end": 4634.42, "text": " only those embeddings which we know correspond to images which are of class 0 I have the number 0 in it because this is amnest so let me change this to labels this should all work now we don't need these special I'm not sure whether this is going to work in colab so the only thing we are left to do is to find this dictionary"}, {"start": 4634.42, "end": 4660.42, "text": " so let me find the dictionary and it's here so it's just a mapping from from a number to a color and let me just define it somewhere here and that's it this should work now hopefully and it does okay so we can see that this is just some stupid error from scatterplot I've already seen this error but like it does not change the semantics so everything works correctly here"}, {"start": 4660.42, "end": 4686.42, "text": " so we can see that different classes are clustered together in a particular parts of the of the studio space we see some outliers like there is this yellow dot here if you can see it on your screen but yeah in any case we can see that the MLP did learn something so that's cool now the final thing I thought doing is basically seeing how many dead neurons we have so I'm going to delete this cell here"}, {"start": 4686.42, "end": 4706.42, "text": " I'm gonna add a new cell just below this one so what we need to do here is to find those neurons that do not activate on all of the input images I'm going to feed into the network and those are considered dead neurons because their gradients are going to be zero which means the weights that come before that neuron output"}, {"start": 4706.42, "end": 4723.42, "text": " can be changed and thus yeah the neuron is pretty much dead I'm going to link a cool article down in the description so you can check it out it has a like a very thorough analysis of various activation functions and there is a dead neuron analysis section as well so you can check it out there"}, {"start": 4723.42, "end": 4746.42, "text": " so what we need to do is again we need to modify the MLP predict function so we'll have to paste this again here so let me take it and we'll need to fetch every single one of these activations so let me create a collector variable here and I'm just going to collect every single activations here"}, {"start": 4746.42, "end": 4763.42, "text": " so I'm going to append activations and that's going to be it we again do not need this part because we don't have rally activations here we just care about the outputs of the rally unit so finally I'm going to return back the collector here and that's it everything else remains the same"}, {"start": 4763.42, "end": 4787.42, "text": " yeah combining this one with this one is probably possible but yeah I just decided to do it ad hoc like this at the moment I'm just going to name this fetch activations two because I'm being super creative right now okay so what we need to do is basically again batch this one we're going to create a batched version of the function"}, {"start": 4787.42, "end": 4814.42, "text": " so a batched fetch activation is two function with we map we map here and I'm going to pass the function in axis again none and zero and that's it now yeah I guess I can just copy paste this all of this and let me do it here and now I'm just going to need to swap this for the activation two version"}, {"start": 4814.42, "end": 4830.42, "text": " everything else remains the same and we should expect to have something like let's see where this is going to work at all yeah it's not we're actually just getting a list here and that makes sense because remember that's how we are collecting the variables here the activations here"}, {"start": 4830.42, "end": 4854.42, "text": " so I actually need to do something like this in order to get the correct result I guess 128 512 that makes sense and let's see we need to have 128 256 in this one and yep we do okay so now that we have these activations we just need to find those neurons where each of these 128 images produced are like a value zero"}, {"start": 4854.42, "end": 4882.42, "text": " that means this neuron was dead for all these 128 images and that's fairly easy I guess let me form a list here that neurons so this is going to be a flag data structure in the sense that if there is a one that means that the neuron is dead otherwise it's not dead so let's assume for now that we have that all of the neurons are dead so I'm going to just pre populate this with a shape something like this"}, {"start": 4882.42, "end": 4911.42, "text": " so for for act in bad activations I'm going to so this is going to be hundred so remember the shape of this act will be something like 128 and then 512 or 256 so we need to fetch the shape and we take everything except for the zero basically dimension and now we have for every"}, {"start": 4911.42, "end": 4939.42, "text": " single neuron we have a corresponding flag here so this is just a set of flags for every single neuron that we care about and those are the ones that basically have a rally unit after them now let's iterate through this through the activations so I'm going to have a layer ID here and I'm going to have the actual activations in enumerate"}, {"start": 4939.42, "end": 4943.42, "text": " and enumerate"}, {"start": 4943.42, "end": 4969.42, "text": " which activations so now I'm going to do the following dead neurons from layer ID is going to fetch the corresponding structure equals non-pi I guess it's logical and or something the illogical and we're going to end whatever is here already with basically"}, {"start": 4969.42, "end": 4981.42, "text": " activations equals zero so we're asking whether the the the neuron is dead and if all of the neurons are dead alongside the"}, {"start": 4981.42, "end": 5008.42, "text": " batch dimension then that neuron corresponding your own remains that otherwise it's considered alive your own I guess and I think all X is equal zero should do the job here let me see what this is going to work so finally we're going to print the so we're going to iterate through the layers so we read through the layers of dead neurons and we print the"}, {"start": 5008.42, "end": 5022.42, "text": " how many of them are actually dead something like this let's let's see what's where this is going to work and it does so we have zero dead neurons in the first layer we have four dead neurons in the in the last layer"}, {"start": 5022.42, "end": 5038.42, "text": " now let's try and do something to make sure this works so let's initialize novel network that's not trained and let's see whether we have different results so now"}, {"start": 5038.42, "end": 5064.42, "text": " here I'm going to create a novel LP which I'm going to call MLP parameters parameters to and now let's use that one instead to fetch the activations everything else remains the same as far as I can see here and yeah so let me just search MLP params yeah this should be it let's see whether we have different results here"}, {"start": 5064.42, "end": 5080.42, "text": " now now a useful thing would be to actually analyze this on the whole training data set because if there is a debt if there is a lot of dead neurons in the when you evaluate on the whole training data set that may be problematic and especially problematic neurons are the ones in the shallower layers"}, {"start": 5080.42, "end": 5096.42, "text": " to play with this code at your own pace let me try I'm fairly sure that I won't have a TPU at disposal currently so if I were to try to use a TPU and click save here let's see whether we can grab a TPU"}, {"start": 5096.42, "end": 5119.42, "text": " let me click reconnect here and no back end with TPU is available so in any case I'm fairly sure that this was already super enough because one of the previous videos I already explained how to paralyze the computation so it's fairly trivial to make this code parallel just a couple of lines of code and hopefully you got something out of this video"}, {"start": 5119.42, "end": 5138.42, "text": " this is the first time I'm doing a real time coding session so let me know if you have any feedback down below and whether you like this type of a content so I'm gonna start making more of this if you like it in any case consider subscribing share out this video and join our discord community and until next time bye bye"}]