Train Self-Supervised Models with LightlyTrain + DINOv2

Key Takeaways

In this video, we're going to be using self-supervised learning and specifically DNOV2 and we're going to be fine-tuning it uh on a large amount of unlabeled meme data. So, this video is made in a collaboration with Lightly Train, which is a library that makes it very simple in this case to use state-of-the-art self-supervised learning approaches on computer vision models on completely unlabeled data. So, yeah, the main thing here is that we basically need no labels. uh and you're going to see in the code as well that it's going to be surprisingly simple to be able to fine-tune this. So let's jump into the code and see exactly how we can do this. So the structure for this project is basically we have a folder with images. It can be any kind of images. You can imagine that this is your proprietary data. Um but in this case we've scraped you know a huge amount of meme data and uh basically we put all of that into one folder. And then the structure of the code here is going to be that we have one script for fine-tuning. We have another script for generating the embeddings uh from that fine-tune model and then one final script just to see how it looks like before and after. But before we get uh started with the actual code, let's create an environment to see you know how we can install this uh lightly AI package. So for this I'll just kind of create [Music] lightly and then specifically to get this package we need to do pip install lightly train and uh they come with a huge amount of uh models that you're able to use. One of them that I find particularly uh relevant is the TIM library which has a huge amount of computer vision models. And as mentioned, we also want to be using in this case uh models specifically from TIM library. So we'll also do pip install TIM. And that's all we need to do for the installation. So now if we jump into the code, I'm going to show you how we can fine-tune in this case. uh and we're going to start with import lightly train and then we'll just create uh first of all like an initial uh model so basically I want to do this with no fine-tuning um just lightly train we'll specify the output directory that we want let's just call it um out and then we'll call it initial for this initial model and data we'll specify to be images in this case we're using a tiny vit so basically I went to Tim's there's a leaderboard and I just looked at which one is performant for the size and uh and tiny vit in that case shows up as one of the top models but you know there's a huge amount of models you can choose from here method we're going to be choosing distillation and this one uses DOV2 in the background in this case we want to set epoch to zero because I want to be able to create a comparison so we'll just create epoch zero this will just save the model directly ly and then we'll do override true and then we'll create another function for fine-tuning the model. So we'll do the same thing out finetune uh we'll use images again same model distillation but we'll set epoch to 500 uh overwrite true resume true just in case we stop and continue at future points and that is actually the entire code needed to be fine-tuning. So obviously there's a huge amount of arguments that we can take a look at. So if we go to definition of of lightly train for example. So you can choose from a lot of arguments here. Uh basically it's built on top of pytor lightning. So you'll recognize a lot of the um arguments here. For example, precision um it's default to 32 which you know we could optimize by changing to float 16. Uh you can choose the devices if you're going to use multiple GPUs. And so since it's built on on top of Pyrus Lightning, it has all of these inbuilt um really powerful optimization things that we can do. In this case though, we're just going to use the bare like the central part very uh to keep things very simple. All right. So if we now just run this pi python a light likely denov2 and then it's going to start the training process. Uh in this case an epoch takes about 4 minutes. Um on my data set here uh we can see things like data weight. Uh we can also see the actual model. In this case you can see it's a very very uh small model. It only takes about for me uh 2 or 3 GB in VRAM and it's running on a 3090 GPU in this case. So assuming you've run that uh which I've already done. I ran for 500 epochs. It took about I believe 12 hours or so on my data set. Uh we now want to generate the embeddings. uh so that we can actually see you know what's been done on this uh model in this case and the results of it. Uh and you can imagine you know like this is we're doing it on a toy data set of memes but you can imagine doing this on proprietary data that you know you that you want to find tune that is in in house of your company etc. right um and so this to do this quite simply is actually quite interesting. Uh so import lightly train we can do another function generate initial embeddings uh and we'll just output it as initial embeddings. PTH in this case we'll use the checkpoint so as I said I already did this uh previously it's uh in t tiny vit initial is what the folder I named it to and then in checkpoints we'll just use the last checkpoint uh we'll use data images or write true and we'll do the same thing for the fine-tuned except now we'll use instead of using tiny vit initial we'll set it tiny vit finetune uh and that's it so now we're basically waiting for it to generate the embeddings for both the initial model that wasn't fine-tuned on this data and the finetune model where we ran for about um I believe 300 epochs or so. Uh so basically the idea is uh this is completely unlabeled. So now we're basically generating the embeddings for both the initial model that wasn't fine-tuned on the memes data and the fine-tuned uh embeddings. And the basic idea is then that we will be able to generate um a cosine similarity between let's say we input an image we'll generate an embedding for just that one and then we'll do a cosine similarity with our uh data set and then we'll do a quick lookup to see you know which are the most similar items and that's exactly what I've done in the C in the compare result script. Uh I don't think this is too relevant to go through because it's mostly just uh matlib and create the plot. uh but I'll show you the result of how it looks like after we actually run this and all the code will be on a GitHub repository in case you want to check the details of this. All right. So then if we run this compare results uh we need to also install pip install math.lib. So it's basically taking a random set set of examples and it's generating. So for example, the above here is the initial uh which you know still does something that makes sense I would say because in this case like we can see the colors are sometimes pretty similar. Um but if we look at kind of the fine-tuned model then we can see that okay these are statues the the top results now are statues. Uh this one it has you know two statues as well. There's a building but then there's a statue. It's finding that kind of similarity. Additionally, like this one is pretty similar as well with a sky in the background. Uh and this is actually like statues as well. Um and two of them. Um so like in this one we can see that it's you know it's obviously done something towards the right uh direction. Here's another one. Uh basically here we can actually find very similar uh directly like meme templates. But here we can see that there were more basically that we didn't find previously uh that had now been replaced kind of with uh more more of those meme templates. Uh here's another one for example. We can see this uh kind of meme template figure. I don't know exactly what this is to be honest, but uh uh we can see obviously that the two next to it are the same figure, but then here it's like kind of same kind of style and and and theme basically of the memes here. I think it's also pretty interesting like it's kind of the same kind of animated kind of template. Uh previously it was you know a bit off I would say compared to what we're getting now. Um and it's quite interesting that that this is actually possible to do in a completely unsupervised way uh without any labels. Right? Similarly here we can see these are kind of similar type of structures to the to the meme. They have some text at the top and they have some really chaotic scene uh which we can see in all of these uh which wasn't really obvious from the previous. So I actually think this is a very interesting library. Uh it's built on top of Python's lightning uh and it makes it very simple. They have a bunch of other libraries as well. I mostly took a look at the lightly train one which makes it um you know easy to fine-tune using self-supervised learning. Uh my idea is basically that I could uh also add this to the nani meme project that I've been building. Currently you can query for memes um for if you search. So for example, we can do Trump tariff and we'll get memes, you know, relevant to what's going on right now to that query, right? Um and basically though this is only text based. So, what would be really interesting is if we take this model that we now fine-tune and we can um input uh the meme and we'll get similar type of meme templates um you know memes to that template that we uploaded. Um and so that's one thing I wanted to extend this kind of project that I've done with the result of u what we just fine-tuned. All right, that's it for this video. I hope you found it interesting. Uh I know it like doing this what we're actually doing is very complex. If we would build it from scratch, it would take us several hours. This video is probably going to be 15 minutes. Um, so yeah, that's I guess the power of using these kinds of wrappers on top and and fine-tuning in this way. All right, that's it for this video. I hope to see you in the next one. [Music] [Applause]