Facebook AI's DINO | PyTorch Code Explained

Key Takeaways

what's cracking guys in this video we're gonna do something very different so i'm gonna walk you through the code of the dyno model which i previously covered in one of my videos uh from facebook a research so we're gonna step through the code and understand how exactly the training works line by line in pi torch so if you find that useful let me know and also leave a feedback of what you think about this video so i can improve it over the for the next for the next time so this is super experimental before i dig into the code i'll do a like a super short overview of the paper just one minute overview for those of you who haven't watched the video go ahead and watch it i'm going to link it somewhere here and so let me do a quick recap of what happens so that you have some context before we dig into the code okay so basically in that video what what i've explained was i explained how exactly this attention mechanism works but i explained it on a very high level using these diagrams uh and so from from this image we get this this this attention map and we're gonna see how exactly that works by the way if you if you want me to share the squiggles of the paper with you let me know down in the comments i'm gonna leave a comment and if i get enough thumbs up i'm gonna leave i'm gonna think about how i can share these with you okay so uh next up we we saw this high level diagram of how dyno works so basically we have the student branch we have the teacher branch the teacher branch is updated using this exponentially moving average from the student uh like weights and we also have a stop gradient here which means we won't be updating the teacher uh finally after applying the soft max we get a distrib output distributions here and the whole idea is to make these two distributions be the same uh and we do that across entropy as you can see here so the p2 times log of p1 so the idea is to extract the essence of the image we input images which had different augmentations as the input to student and teacher and hopefully we want to make sure that the distributions are the same because that's the same image and that's how we extract the kind of the essence of the image so that would be a high level explanation uh so we saw the pseudocode now we're going to see the actual code we saw that we are using uh so i mentioned augmentation so we're using these uh global crops and we're using these local crops and we are again trying to um make the model output the same distributions for both of these uh finally let me see what else i had so explain how k nearest neighbor algorithm works and uh probably in the next video i'm going to show you how exactly this looks in code because i think the things i just explained will take like at least an hour to to kind of cover so let's get back to the code and see how this looks like okay so you've watched the video you have some context from this short recap now let's start digging into actual code okay so let's put some breakpoint here as i'm going to step through this code and show you exactly how it works so let's start with this get rx parser function okay so i'm gonna just quickly skim over these uh arguments so that you have some context an idea of which variables are gonna be circulating around the code so first things first we have the architecture we're going to be using the vision transformer small they have many different options but we're going to focus on that one batch size i had to increase it to 32 because it was 16 by default and it's for some reason crashing my machine i still haven't been debugging the amount of memory it's eating up on my gpu so yeah so outputting dimension is just 65 65 000 and that's the actual output distributions we just saw in the paper overview which we're going to compare and make sure they are the same for different crops okay uh then we have just normalized last layer i'm gonna skip this momentum teacher so 0.996 that's the at the actual coefficient that's being used in the exponentially moving average when we are forming the teacher weights from the student weights uh again veteranization this is not that important uh warm-up temperature teacher temperature um number of epochs all of these are not that important there's just a schedule of how the temperature of the teacher is going to change throughout the training using mixed precision training is going to be on so i obviously won't be focusing on the optimization stuff in this video so weight decay okay that's again we have some schedules so they have the end and initial value uh clipping gradients so that's also basically when you're updating the the weights uh they're going to use this to kind of clip the norm of the of all of the weights before doing an update we're going to see that in a couple of minutes so number of epochs 100 by default where it or not to freeze the last layer this is just some trick that makes the training a bit easier uh learning rate um so they basically have a linear warm up first and then they have some kind of uh i think cosine schedule uh until the end of the training so warm up epochs just tells you across uh how many epochs you're you're you're taking to do that linear warm-up thing uh minimum learning rate that's the pretty much learning rate at the end of the of the training optimizer they're using adam uh white adam so that's because they have distributed training this dyno model is trained on 8 or 16 gpus originally but i have only a single gpu here so that's that's it drop path rate is just an argument that is going to modify the the vision transformer architecture then we have the global crop scale so it's going to be from zero four to one and contrast that to so i'm gonna quickly get back to that one but we have eight local crops and we have um and we have so local crop scale is going to be from zero zero five to zero point four and that's the same number as here so that means um so that means that uh this this for for the small crops we're going to take uh five percent to 40 of the image size and then we're going to resize that back but like for the big crops we're going to take from 40 to 100 of the image and then we're gonna do the resizing to two to four as for the small crops they're going to be resized to 96 96 pixels uh again some details but like uh yeah this video is going to be about a lot of details we're going to use imagenet the thing i did is i actually downloaded a small subset of imagenet because this thing is huge more than like 150 gigabytes i think and so from fastai basically it's called imagenet or i don't know how to pronounce it with two t's and aside from me having to download image that small subset of imagenet i also had to do a couple of tricks to make this work on windows because i'm running on windows machine so if you want to uh know about those nitty-gritty details and small bugs and problems i had to solve i do joined the discord server actually had a thread where as i was in analyzing this code over the past two days i was kind of writing down the the steps i took and the problems i encountered okay so finally uh as i said imagenet uh where i'm going to save the checkpoints and logs um seeds c basically serves to make this code reproducible because otherwise there is a lot of randomness and every time you run this code it's going to have different output and different results and the images are going to be loaded in different order it makes it easy for me to have some consistency between different runs and it's easier to debug the code so number of workers not that important all of this is just some distributed settings and that's not that vital okay uh we parse the arguments we create the output directory where we're going to dump the logs uh the log files and the checkpoints and finally let's go to the main function then that's train dino um so again this is the init distributed mode it's just a function that's uh setting up certain settings uh when you're doing distributed training on 8 or 16 gpus again i'm using only one gpu i had to do some modification because of windows uh but yeah uh do check out this discord server for that so we're fixing the random seats uh they're just printing i'm gonna rise up the console here so they're basically just printing the the uh shot code uh the the hash of this commit with of the code i'm currently running uh then they're just printing some arguments uh coding m benchmark set to true is just a certain optimization again that makes the code run faster under certain assumptions again not that important okay this is fairly important data augmentation dino i'm gonna quickly walk you through this code and then i'm going to step over it in the execution loop but basically as you can see here they have flip and color jitter uh transform so that's just a compose that means we're composing a couple of different transforms and these are organized so that they do the flipping and the color uh like photometric augmentations so we have as you can see here random horizontal flip with fifty percent uh with fifty percent chance we're gonna flip it uh either flip it or not flip it then we have the uh we're gonna randomly apply with certain probability like eighty percent uh the score jitter which which is going to change the brightness the contrast the situation the hue so basically your photometric augmentations and finally random grayscale is with 20 chance we're gonna uh kind of uh apply the the grayscale transformation uh the then we have normalization this is your standard stuff we are converting uh the the images to pi torch tensor uh we are normalizing them and don't get confused by these magic numbers even though it'd be really nice if the authors actually extracted this into your constant because they've been using this all around the place so this is just your image imagenet statistics so the mean and the standard deviation of the imagenet images okay so for the fun part this random resize crop is the thing that does the magic of dyno so global crop scale is again the thing i just mentioned so from we're going to take 40 to 100 of the image we're gonna cut it there and then we're gonna resize it back to 224 by 224 uh using bicubic interpolation method uh then they're going to apply some photometric augmentations and the flipping the gaussian blur and we're gonna normalize the images uh we have a second global uh crop because uh if you remember the details from the paper we have two global crops and we have eight local crops although obviously that's you can you can experiment with that uh so again a random resize crop everything is the same they just apply solarization here if you're not familiar with solarization i'm just going to put the image on the screen so you can see a visual example i guess the image is going to explain that much faster than i will uh finally we have the the local crops so that means uh we are applying the local crop scale so that means we are now using that 1.5 percent to 40 of the image and then we resize that back into 96 by 96 images so these crops are going to these images that came from the smaller crop of the image are going to be of smaller resolution and 96 by 96 again flipping and color tutoring gaussian blur normalization standard stuff so that's it so basically what's gonna happen is this uh data augmentation dyno uh is gonna be passed to the data set in pytorch and then every time you try to fetch an image from the data set it's gonna apply it's going to append one crop here one global crop the second global crop and then we're going to have iterate in a for loop so eight times and do these small local crops and then we're going to return like a list of 10 crops in total uh this is just some code i wrote i'm gonna experiment with that a bit later but let me get back to the actual line where we stopped okay so that was the that was the uh data augmentation of dyno i think that was a pretty important detail by the way i'll be adding timestamps to all of these logical sections so feel free to skip if you understand some part you can you can just skip to the part you you want to understand better uh okay continuing on we are creating this data set uh so it's going to link to the as you can see here imagenet train subdirectory and we're going to pass the transform i already mentioned which means every time you fetch an image from this data set it's going to apply those 10 transformations and we're going to get those 10 crops back uh distributed sampler so just ignore the distributed part this is just going to shuffle the images once we input that into the data loader so as we can see we have a data set here we have the sampler we have i think 64 images per gpu and since i have a single gpu that's going to be 64 images in total a number of workers just optimization stuff pin memory as well that just kind of helps you basically transfer the images from the cpu to the gpu faster because there's this dedicated slot of the memory so i'm currently going through every single detail uh because i'm experimenting but like uh i'm later on gonna abstract many details so let me know whether you actually want everything explained or you just want me to cover the main idea uh this is experimental video so please give me that feedback and yeah droplets means because we have a data set of i think uh like i have around nine thousand something images because every batch has 64 images as we saw here that means that probably the last batch won't have 64 because unless the number of images in our training data set is divisible by 64 we're going to have the last batch being less than 64 and this drop last just means we're going to drop the last batch because that can cause some issues probably because of the shape being different uh so that's the reason you sometimes see this droplet set to true okay data loaded there are uh as you can see here 9460 69 images okay this is just some bureaucracy stuff uh they did because they used to to name name these models date now uh they're just kind of replacing them with vit but the architecture is already v80 small so nothing's going to happen there so basically this says that so vit is just a like a short hand notation for the all of the vision transformer models that exist in this vision transformer pi file so we're gonna because we are vision transformer we're gonna enter this if statement and we're going to as you can see here we're gonna form so this statement here may confuse you so this is just a list of all of the different uh like uh functions from this vision transformer file and this is just going to instantiate the model and we're going to instantiate it with the 32 by 32 so the patch will be 32 by 32 pixels and we're going to use this drop path rate of 0.1 so the quickest way you can understand how this thing uh what this thing does is just put a print statement in there do this uh like just print the dictionary uh set a breakpoint there stop the code re-execute again and we're going gonna see what happens there so i'm gonna just skip these skip these i'm gonna get here i'm gonna open the console and after i print this you're gonna see uh whoops bunch of stuff and as you can see here it's actually printing this whole vision transformer pi file and the thing is uh it's some of the some of the functions are organized into dictionary so this is just your python stuff so that means we're going to have somewhere in here we're going to have uh like these functions let me let me jump to that file so we're going to have somewhere in there we're going to have this v80 small uh function and then we're going to pass the patch size and the additional we are passing that drop path uh coefficient so that's the the thing that happens there so basically we're instantiating the vit small model here uh then we're doing the same thing for the teacher networks there was a student network we have the teacher network and since we're using uh vision transformer small the embedding dimension is going to be 384 uh that's just the hyper parameter they they kind of decided to use okay so they have in these other branches they're using excite and some other stuff we don't care about so um now this is another important part basically uh what they do is they they wrap the student network so this is the vision transformer and they put this dino head on top of it so dyno hat is just a simple mlp uh with some additional tweaks and i'm gonna jump into those details right now so the main point is they are forming uh this this mlp and this number of layers is three by default that means the mlp is gonna have three linear layers that means we're gonna uh execute the else branch and we're gonna form the mlp that just consists out of as you can see here we we have some potentially some batch norms potentially some uh like jelly uh like uh activation functions uh but most of all most importantly we have this linear function uh so basically it's an mlp and then we just initialize the way it's using this self-apply that's an inbuilt uh function that's a part of this nn module uh class and then we're going to form the last layer by creating yet another linear layer and the linear layer is going to have so the output dimension here is that 65 000 something so that's the actual output of the whole pipeline and then we're going to compare those outputs by cross entropy later on so that's the last layer uh the interesting part here is this weight norm is just a hook function uh what it does is uh once you put this weight g to all of ones that means they're going to normalize the actual linear layer so it's just going to make sure that the l2 norm of the weights of that linear layer are equal to one that's basically it and if they decide to do this normalization of the last layer if this is set to true then as you can see here the the those gradients are going to be set to false which means the following during the training those uh those weights will not will not be updated which means that the l2 norm is going to be kept at one which means we're normalizing the last layer hence the normal last layer i think it's pretty self-explanatory so again simple mlp on top of it we have the last layer which has the 65 000 something number of neurons as the output and that's it optionally we have these uh these uh like normalization stuff okay getting back here so that's the dyno head we stick it on top of the student and uh one additional detail which we need to go through this is this multi-crop wrapper again i have timestamps just jump if you don't if you're not interested into this section uh so the multi-crop wrapper what it does is the following um now i'm actually just gonna put a break point here and then after we we get to this part we will understand what multi-crop wrapper does but basically on a high level for now it's basically just going to deal with those crops past and through the the pipeline and then out comes the result from the student the final the final output and the final output from the teacher we'll see that in a once we hit that part of the code okay uh getting back here so that was how we formed the student part so again it's a multi-crop wrapper around student which is vision transformer and dyno head which is mlp same thing goes for for the teacher we have teacher we have dino hat and that's it some minor differences like whether we are not we are normalizing the last layer or not but that's it and we are not using the bad normalization so this is uh there is no batch normalization in this whole pipeline which is something we we strive for um okay uh we're just pushing the the student and teacher to the gpu that's this dot cuda that that's what it means and finally because as i said we don't have any any batch norms uh that means we're going to enter the else branch and that means that the teacher network will not be using this distributed data parallel which is just a wrapper that makes it easy to train these models in a distributed fashion uh okay continuing on we wrap the student network uh into this distributed data parallel but i only have a single gpu here so the index zero means the default gpu you have on your system i have a rtx 2080 so that's it then we are loading the as you can see here we are loading the state dictionary directly from the student module so and they say here nicely trend student start with the same weights so initially we're going to start with the same weights later on we're going to use this exponentially moving average to update the teacher weights uh finally as i said there is a stop gradient uh if you saw the diagram there is a stop gradient and this is how you implement the stop gradient in pi torch you basically just iterate through all of the uh like uh parameters and you set their gradients to false so that means it won't be trainable i'm going to skip all of this just f9 and that's it so we we have an output here the student and teacher are built they're both vision transformers small networks okay this part is fairly important and i'm going to probably uh step through it a bit later but let me just give you a glimpse so i'm going to put like a like a break point here in the forward function and as you can see here it's just accumulating the parameters that we passed in the internet function uh it's just registering this uh center vector so if register buffer what it does it makes sure that the this this this vector here uh is a part of the state dictionary even though it's not trainable so this center vector will not be trainable we're not going to update it using gradients we're going to update it using the exponentially moving average again so that's the reason they're using register buffer so it's not trainable so it's not a model parameter in the pi torch like uh terminology and but we still need it in the actual state dictionary of the model um okay finally they have some teacher temperature schedule so with the current parameters this is actually going to return 0.04 uh throughout all of this all of the epochs but what it's it's kind of envisioned to have a linear ramp up here although that's not that important let's get back to the dyno loss okay i'm going to step over this i'm going to later once we are actually using the dynalos i'm going to step through the forward function um this just not important where we are regularizing or not regulars in certain parameters of the student network we're going to create the atom uh wide optimizer and finally because we're using the mixed precision training we're just going to wrap this create this fb16 so that means we're using 16 bits training these models okay here they're just creating schedulers for the learning rate so it's going to be a cosign scheduler then we have weight decay schedule against cosine scheduler and finally momentum scheduler uh basically that's the thing that is used to update the teacher weights using the student weights so that's that parameter uh and again it's just a cosign scheduler okay so that's it uh i think we're we're good to go we're getting closer to the actual training loop uh this already took a lot of time so uh we're just going to this this restart from checkpoint i think we're going to skip it because we don't have any checkpoints for now so just put a breakpoint there we're going to enter it okay let me step over step over okay and we're going to just return from that function since we have don't have any checkpoints okay so that's just uh dummy for now um start epoch it's gonna start with zero we're just uh logging the starting time just some bureaucracy stuff and finally here is the main loop okay we are here so we are starting from the epoch zero we're going to run for 100 epochs and uh ignore this line this is just something that has to do with this distributed sampler otherwise you wouldn't have this line so this is the actual meat of the whole whole script so it's training one epoch is going to train the the dyno for one epoch and then later on which is gonna have some some logging nothing nothing very fancy so i'm gonna step into this f7 and let's see how this function looks like okay so metric logger not that important again okay this is where the fun starts so we have the data loader and we're going to fetch like this images so this images is going to be a list of 10 elements and each element is going to be a batch of images so each element is going to be 64 images so that means so because 10 elements because as you remember we have two global crops we have eight local crops so that means if we take like the zeroth element of this images list that's those are the global crops and we have 64 images and if we take uh like the second index that means that's the first local crop like set of images a batch of local crops and so as you can see here i just pre-wrote some code here to make this a bit faster i'm going to stop the code execution and i'm going to rerun it again i'll i'll need to make another tweak to make this work so that's that means i'm gonna have to add additional crop uh so what this does is we are just going to preserve the image as it is without doing any any crops so that's just the reference image and that was that's going to be prepended as the first crop so now we have 11 crops and hopefully now this is going to work let me let me let me run this uh let me go back to the code where we were okay here we are so what i'm doing here is as you remember i prepended uh this original image so now we have 11 elements instead of 10 and this 0 will just be indexing the original images and i'm just going to fetch like one image out of those 64 images i'm going to convert those to numpy do some move axis because the channel is actually channel first and now we want to have channel last for for the plot function of matplotlib so second line what it does is it will extract the global crop so remember now the the the index one and zero um the index one and two are actually referencing the global crops and again i'm going to take just one single image finally we have the local crop so index three through ten are all going to contain local crops so i'm just going to take one arbitrary local crop and finally let's let's plot those results so um here you can see the the original image just some dude playing golf so i'm gonna now exit that remember that image and now we're gonna see a global crop as you can see it's uh there is a crop but there is also severe solarization and other photometric effects so you can't really see it perfectly but let's see the the local crop so the local crop is as you can see now it's totally clear we're just taking a single small crop and we're resizing it into this image of 96 by 96 and we also apply some augmentations here so that's it that's the now you have an understanding of how images look like and i encourage you when you're debugging and understanding the code do make these visualizations as much as possible because that helps you understand what's going on in your code that plus print statement i guess that's pretty much how i debug the code although there are sometimes more elegant ways to do it okay so um let's let's continue here so okay this line is just forming a global step depending on the epoch we are in depending on the number of images in the dataset depending on the current iteration inside of this uh inside of this loop okay so that's just global step and uh then what they're going to do here is depending on the schedule current schedule we're going to update the weight decay and learning rate depending of these param groups it's not that important i'm just going to skip that part so we're just adapting the weight decay and learning rate depending on the current step okay so skipping that we're putting the images on the gpu so that's this step let me just oops let me just skip over that uh and finally this part uh is just a context context manager because we have this mixed precision training and it's not none so we're gonna enter this but now it's not going to work because i prepended all of this so i'm just gonna quickly comment this out uh and comment this out and i'm gonna restart the training and get back to you okay here we are uh so this is as i said very important step we're taking the global crops and we're passing only the global crops through the teacher to get the teacher output and let's see so analyzing the shape of the output is super important i can't stress that enough so let's now see whoops i have a okay so we have a multi-crop wrapper okay this is the part we we want to understand a little bit better um so x is just a list of two elements again uh each contains a batch of crops these are global crops so because it's a list we're gonna skip this part so this is just a fancy way of separating the global from the local crops and it doesn't make much sense for the teacher network because it's only using global crop so we're just gonna have number two here so as you can see here uh index crops just has number two uh this is going to make more sense for the for the student network so now what we do is uh we iterate here and we're gonna take as you can see from zero to two so that means we're gonna concatenate all of those uh global crops we have 64 images in one in one in one element we have 64 images in the second one we're just going to concatenate them and now we have 128 images in a single tensor and then we're going to pass that through the backbone which is basically the vision transformer of the teach network so that's going to form what so let's try and figure out the shape so 128 and because we are outputting uh only this cls token so that will be the output from the vision transformer so it has 384 uh like uh neurons right so we're gonna have hopefully um what so 128 times 384 okay let's let's see whether the prediction was right so it's kind of good to have these predictions and then uh tweak your brain depending on whether you made made a mistake or not i guess that's a similar way how these neural networks work so the out shape here as you can see the shape is 128 comma 384 so that's it i'm going to skip this part and now we're going this is going to accumulate the the the output and this will again make more sense for the student network now it's just gonna uh we're gonna just copy paste pretty much this output from the vision transformer here and that's it we're gonna exit the loop and now this is the actual dyno head so this is the mlp and now we after passing it here we're gonna end up with those 65 000 something so we're going to have 128 times 65 000 something let's let's see whether that that's true okay uh whoops let me get out of here teacher output let's see whether the teacher output is as i predicted so teacher output let me kind of uh so 128 65 000 that's it so the predictions were right so now this is the interesting part we have a student network and we are passing both the global crops as well as the local crops so let's enter the loop again again it's a list so we're going to skip this and now it's a bit more interesting so it's going to make return the ids where the resolution changes so we have 20 remember the global crops are 224 by 224 and the local crops are 96 by 96. so what this line of code actually does is it finds a separation and returns the indices of the place where the resolution changes so it's going to return as it will soon see 2 and 10 i think or something like that yeah 2 and 10. so that's because after 2 we drop to small resolution and then we keep the small resolution until the end and that's 10 crops so now we are iterating again as you can see here we're first going to pass zero through two so that's global crops uh and we're gonna pass that to the the the network here so that's gonna be again 128 so that was that is going to be as you can see here 128 384 uh and we're just going to concatenate that with this output which is just a placeholder variable you can see here is stored torch empty and we're just concatenating the results to it and that's it now we're going to put this the starting x is now going to be 2 and the ending index is going to be 10. so now we're passing 2 through 10 which means we are fetching the local crops we're concatenating them so we're gonna form what like let me let me do it like calculator we have eight crops all of these have what uh 64 64 images so that's 512. so let's see whether that's true so we have 512 and we're going to expect the output 512 uh 384 so let's see whether that's true so we have out output variable 512 384 that's it so everything works as we expect so let's continue here again we're just concatenating these so now we're going to concatenate these with the with the outputs from the from the global crops and that means we're going to have what like uh 512 for the global crops plus 128. so 512 was for the for the local crops 128 is for the global crop so that's 640. so if we do that we're gonna see that the output is let me see here so the output is 640 384 okay that's it let's continue we're gonna exit the loop we're gonna pass all of these through the mlp and that means we're going to end up with instead of 384 we're going to have those 65 000 something activations so that's it that's the whole trick let me again exit this thing okay exit come on come on come on this success okay so this is the last important detail let's just analyze the dyno loss and that's pretty much it you know should be able to understand the training so here what hap what happens we have we take the student outputs and remember that's again uh 600 something times 65 000 something so 640 times 65 000 and we're just going to scale it with this temperature coefficient so remember that's remember from the diagram from the paper that's what they are doing uh now they're going to chunk this into 10 pieces because remember we have 10 different crops so this is just a way to take these 640 outputs and kind of cluster them into uh semantically meaningful categories so that means global crops and local crops so we're going to see that the student out let's find that variable so that's student out here let me expand it we have 640 and now we're going to have a list of 10 elements and every element will have 64 images as it should be so as you can see here 10 elements in the list and every single one contains 64 images that's it we do a similar thing for the teacher network we have some temperature stuff we have the centering of the output and then we divide by the temperature across the future dimensions so that means across the that vector that has 65 000 activations we're going to do the soft mix across it and then we do the detaching because again we're not training teacher network and we trunk it into two because remember we have uh two global crops being passed through the teacher network so teacher out let me just expand that variable so teacher out has 128 it's going to be split into two elements each will have 64 images so let's step over it and as you can see here two elements 64 images each so that's it so now we have those different views and now we're gonna do the actual uh like loss so what is happening here is we are taking uh a single so this is going to be a single like a global crop and it's going to contain a batch of images for that global crop and then so this just make sure we are not operating on the same view so that's what we're gonna skip in the first iteration because both are zero so that means both are indexing the same global crop and now in the second iteration this is going to be one we're gonna skip that and what happens now is we're gonna take a batch of images so 64 images from the from the from the second global crop and we're going to contrast it with q which is the the first global crop and that means and we're going to do simple cross entropy here so that's q times the log of the student view soft max is here because we already we already done soft max for the teacher network so we have to do soft max and then we have to apply log and then we have just cross entropy and that's it so focusing on dimensions here on the shape again so if we take expand the q variable you can see it's 64 times 65 000. the same thing is going to be for the student view so we have 64 images here and here and we have their outputs and we're just going to make sure that the distributions are the same for all of these 64 images and then what we're going to do is uh basically uh sum it up so we're going to sum up across the that dimension the dimension number one and then this mean will do the averaging across the second dimension and that's gonna end up with a scalar so after we do this and this we end up with a scalar and that's it that's pretty much it now we iterate through every single combination of views so we have we trade through different global crops here we trade through both the global and local crops here we take those views we do the cross entropy and that's how we train this whole thing okay and the final last detail is we are updating this uh this center vector which is used as remember here so the self center is going to be updated using the teacher outputs uh in this function here but going to skip it you can analyze it yourself it's fairly simple uh it's basically following the formula from the paper we're doing the every uh exponentially moving average update of the of the center vector that's it okay uh let's exit this thing let's exit this thing and that's it we now have just the bureaucracy stuff so we clear the gradients of the of the model uh we're going to now do uh backward which is going to calculate the gradients throughout the whole network of all those parameters that are obviously a trainable uh we're gonna do some uh clipping of gradients basically here this line is just going to freeze uh and not train the last layer if you remember there was this parameter called uh what's the name of it freeze so since this uh freezeless layer is only one and current epoch is uh zero that means we're actually going to freeze the last layer of the model but that's again a small optimization trick uh to make this uh thing more stable and trainable uh finally we do the step that means we're gonna actually update the the weights and that's it just some details around mixed precision training nothing that spectacular uh finally we take the the the uh moment value depending on the on the epoch we're in on the global step we're in uh and this is the part where we are updating the teacher networks uh using the uh student weights so we are basically doing exponential moving average uh here so as you can see the teacher weights are going to be multiplied by m which is the this momentum parameter and we're going to add up 1 minus m times the weights from the student network and you can see it on the formula on the screen and that's that's pretty much it okay that was it that was pretty much the the core parts of the training everything else is just logging dumping those log files and checkpoints this was the the very meat of the dyno training and how the training looks like again any feedback is appreciated the more feedback i receive the better i'll know for the next time how to make these videos okay i want to end this video by showing you this visualize attention function let's quickly run through how this thing works and i think it's fairly fairly much simpler compared to the training so let's kind of dig into it i'm gonna put a breakpoint here i'm gonna start debugging and let's see what happens okay again we're using uh v80 small patch size eight uh no pre-trained weights this is not that important we have some image path i already uh inputted the path to this image so we're going to be using this image here this this bird image and uh so that's the image path image size is going to be 480 not that important again we're just going to dump in the current directory uh threshold is also not that important for now let me see what i am using currently so my current threshold is 0.1 uh i'm actually going to delete this so i'm going to just input the image path here as you can see and now let's apply this and restart the code again jumping to here okay i'm going to jump here let's see here so device if i have if you have gpu it's going to be using gpu so it's cuda here for me i have a gpu so we're going to form we're going to create a model so this is the same line as in dyno training code so we're going to basically instantiate the vision transformer small because of this string with patch size eight and zero classes means it won't have the output head it's just going to have the the final layer it's going to be just a bunch of tokens the cls tokens and all of the other tokens but no no actual linear classifier had okay let's do this so these lines we're just gonna freeze the the model parameters that means we are not we're gonna set them to not trainable so that means this model is just going to be used as is no training involved eval evaluation mode uh that basically sets up depending whether you have dropout or batch norm this thing can change the the those layers and that's why it's important when you're using a model in an inference mode to uh call this eval function uh we're gonna push the device to the gpu to make things a bit faster and finally because i don't have any pre-trained weights we're just gonna fetch from a certain url as you can see here we form an url we fetch those from the torch hub we're gonna just concatenate the url to this uh repository and we're gonna fetch the state dict and we're gonna load the state dictionary so the download actually already happened on my local machine so it's just fetching it from the local storage instead of having to go to the cloud okay my image the image name was uh like incorrect so i had to tweak that in the settings and now i rerun the code and here we are again so because we have an image path we're gonna open it load the image we're gonna load the image we're gonna do some transformations so we're gonna resize the image into 480 by 480 we're gonna convert it into tensor and again here is here these uh numbers we already saw that much the magic numbers they just basically represent the the mean and standard deviation of the imagenet data set um so okay so after that we're gonna transform the image so now we have a tensor that's normalized and resized as you can see here finally we make sure that it's divisible by the patch size so the patch size is currently what eight so we're gonna make sure that the width and height are are are actually uh divisible by eight and now we're gonna crop the image here if necessary and unsqueeze zero means we're just gonna add another dimension uh in the front so that's just some pie torch uh like uh stuff and if we take a look at the image dimension here so let's take a look at the shape so the shape is just 3 480 480 and after we execute this line is going to be 1 3 480 480. so that's the unsqueezed part then we're going to divide the resolution by the patch size so this will tell you how many tokens you have across the width dimension and we have 60 and here again we have 60 because remember it's 480 times 480 image uh finally what this does it's gonna affect from the vision transformer it's gonna get the the the self-attention of the image so we're gonna take the image we're gonna like send it to the gpu because the model is also in gpu so you don't wanna have a discrepancy there we're gonna feed the image into the model and it's going to take from the last layer it's to return not the tokens it's going to return the attention coefficients so that means the following so remember we have 60 times 60 uh tokens as we just saw up there so we have 60 here and 60 here so that means we have 300 3 600 tokens in in total plus one for the cls token so that's 3 601 so the attention dimension we should expect should be basically three thousand six hundred one times three hundred six thousand three three thousand six hundred and one and we also have uh six attention heads so we're gonna have something like six times this number here so that's the the the numbers i'm the shape i'm expecting uh so if i step over and after we now let's find the attentions variable so the shape is as you can see here six three thousand six six hundred blah blah the number we expected basically plus the uh this one as the first dimension again uh those are just some idiosyncrasies of pie twerk quick recap how those performed is basically we have so we have as i said three thousand six hundred one tokens you take a single token and you basically form its query and then you do a dot product a scale dot product with all of the other keys and that's why you have three thousand 3601 attention coefficients for that single token but since we have 3601 tokens we're going to have this this this shape here so hopefully that's clear so as you can see here we extract the number of heads the number of heads is six and now what we do here is we take we just extract this is just because of one we take this zero here uh and then we take all of the heads but we take only the zeroth token which means cls token and we take one through the end which means we find the attention coefficients that the cls has across all of the other let's call them spatial tokens so that's the whole idea there so that's why we have these numbers here so let's see what what is the shape we're expecting we're expecting something like let me write down here maybe it's easier to see so we're expecting like six uh and then we expect because we're taking just the cls token uh so we expect six three hundred three thousand six hundred so that's the shape we're actually expecting after after this uh this line of code so let's see whether that's true so attentions as you can see here that's that's the actual number six three thousand 3600 okay i'm just going to ignore this threshold part and we're just going to going to reshape uh the those attentions into uh into different shapes so instead of having it flattened out we're going to return it back to 6060 and finally this is just a nearest neighbor interpolation which means you're going to take that token whatever the coefficient is we're just going to replicate the coefficient into the 8 times 8 pixels because that was the original image shape so that we can visualize the actual attention so that's why we need the interpolation part finally just some making some directories saving the original image and finally here in this in this section here we're going to as you can see here we're going to take uh the attentions for every single head and save them in an image so let me just kind of do f9 there and we're gonna end up having six images which i'm going to show you in a second and finally i'm gonna ignore this thresholding part again and that's that's pretty much it and here are the attention uh maps we got so here you can see the bird image every single head has different attention pattern and that's it and these are the images you you saw in the actual paper now you you know how they are calculated and visualized awesome this was a long video hopefully liked it if you did share it out with a friend subscribe hit that bell icon and uh join the discrete community until next time bye bye [Music] you