Transcription
In this video, I'm going to show you how to train your very first PyTorch model. If you're interested in deep learning, you really have a few options: PyTorch or TensorFlow. But in recent years, PyTorch has really exploded as being the front runner, and the main reason is because of its flexibility. But with that comes a lot of added complexity.
I've competed a lot on the website Kaggle, which hosts competitions for deep learning and machine learning problems, and over the years, I've trained thousands, if not tens of thousands, of PyTorch models. Now, there are already a bunch of PyTorch tutorials out there, but I want this one to be different. Instead of focusing on the details of PyTorch, we're going to actually start by trying to solve a problem and create a model, and we're going to focus on the three main paradigms that you'll need to know in order to train a PyTorch model. And this knowledge will transfer to whatever type of PyTorch model you're looking to create.
Training most deep learning models requires that you have a GPU, so you can run these large computations very fast. But everything I'm going to show you in this video, I've done in a Kaggle notebook. The nice thing about that is you'll be able to go copy that notebook and run it yourself on a hosted virtual cloud instance that has a GPU, and everything you need already installed.
Okay, I hope you're excited. I know I am. Let's train our very first PyTorch model. So here, I have a Kaggle notebook that we'll be working in. I'm going to link this notebook in the description, so you should be able to load it up and, like I said, copy it and run it in the exact same environment. But everything we're going to learn here, you could also run on your own environment, even a Colab notebook, if you'd like. You just need to make sure that you'd have the GPU with CUDA installed and PyTorch in the correct PyTorch version in your Python environment.
Now, a few things I want to show you in the Kaggle notebook environment. If I click over here to the right, we can see some information about this environment, and the main thing I want to show you is that we are running with a GPU enabled. And the GPU is going to be critical in speeding up this training process. And if I scroll up here, I want to show you that I have a dataset linked to this notebook, and this is an image dataset for classifying playing cards. And we can see the way the data is set up. There's a training folder, and in each of these, we have a folder for different cards in a playing card deck. And if I click on each of these images, you can see the playing card images show here at the bottom. So this is the type of data that we're working with.
In this tutorial, we're going to learn through doing, and what we're trying to achieve is create a model that is capable of taking in images and then detecting which playing card is in that image. And we're going to break this down into three parts: and that's understanding the PyTorch datasets and dataloaders and how that works, learning how we set up the PyTorch model module, and then understanding the PyTorch training loop. Now, these three steps are critical to training almost any PyTorch model. It doesn't have to be an image classifier; it could be for object detection, audio classification. They all mainly will follow this idea.
All right, so the first thing we need to do is import all the packages we're going to be using for this tutorial. So obviously, we're going to be importing PyTorch, which we'll import as torch. We're going to import the NN module from torch. This is all the specific functions related to neural networks. We're also going to import the optim library, which will define our optimizer later. Datasets and dataloader will be the very first step in the process. We need to import them, and we're also going to import a few things from TorchVision, which will make working with these image files a lot easier. Timm is a great library for loading in architectures specific for image classifications. We'll import that, and then, like always, I usually import matplotlib for visualization and pandas and numpy in case we need to use those for data exploration.
I like printing all the versions out here, so if you're running this later, you can compare it to what versions that you're running, and there may be some slight changes. We are running PyTorch 2, and we're running Python version 3.10.
The first thing you're going to need to do to train a PyTorch model is to set up the dataset, because that is how PyTorch is going to be loading in the data as it trains the model and evaluates the model. Now, the nice thing about PyTorch datasets is they're extremely flexible. You can really modify them to load the data in whatever way you want, and then you can take that dataset and just wrap it with a PyTorch dataloader to make it parallelized and load in the data in batches to the model.
And we create this dataset by simply creating a Python class. In our class, we're going to call `PlayingCardDataset`. This will inherit from PyTorch's main `Dataset`, which we imported above. Now, if you're not that familiar with using classes in Python, you might want to do a refresher, but it's not as complicated as you might think. Each class takes an `__init__` method, which basically tells this class what to do when it's created. And then in PyTorch datasets, we'll need two additional methods: one is the `__len__`, this is important because the dataloader will need to know how many examples we have in a dataset once we create it. And then this method called `__getitem__`, and this method takes in an index location in our dataset and we'll return one item.
So let's go and make this actually work for our dataset as we have it situated. So our `__init__` is actually going to take the data directory where our data is sitting. Then we'll also initialize with a transform. This is what will apply to each item in the dataset, and we're just going to keep it super simple here and only provide it a transform that will resize all the images to the same size. We're going to keep this as simple as possible and use from the TorchVision package this `ImageFolder` class, which will take our data directory and whatever transform we provide it. The nice thing about this `ImageFolder` class is it'll assume that any subfolders in that folder have the class name for the image, and it'll handle creating all the labels for us. We could do this manually if we wanted to, but we'll do this for simplicity.
Now, our `__len__` we just want this to return the length of our data, and for `__getitem__` we'll want to return the data item of whatever index that our `__getitem__` is called with. This `ImageFolder` dataset when called on an item will give the image and the class. And I'm going to show you how that works now.
I want to add one additional method that will just make it easier for us to find the class names, and that's going to be called `classes`, and it'll return the data classes from this `ImageFolder`, and that's it. Now we have our `PlayingCardDataset`, and we can create an instance of this dataset. Now, we do need to give it our data directory, and in this Kaggle environment, that's going to be this location where we have the playing card dataset loaded.
Now we can test out our methods to make sure it works, right? We can run a length on the dataset, and we can see here our dataset is 7,600. And we can test the `__getitem__` method by calling this dataset on one of these index locations. It should work anywhere between zero and the max length of this. We could see when we call this at various indexes, we're returned a tuple which contains the image and then the class, which is a number between 0 and 53 here in this dataset, because we have 53 different classes of the playing cards, and then also a joker card. We can also see that these classes are in order. So if I go later in this dataset, 6,000, that will get this different class returned.
Now, we also could call it like this, where we would pull out from this tuple the image, which I can display here because it's a PIL image, and then the label, which we could print: 41. Now, our labels are represented by numbers, and they need, in order for PyTorch to work with them, but we don't know which number is associated with each label. So that's why we added that `classes` method. Don't worry too much about these lines of code, but basically, what we're doing here is creating a dictionary which allows us to link up each of these numbers with the correct label. So running our example here above, we see this is a 10 of clubs, the number of the label is 41, and this dictionary just allows us to see that 41 is, in fact, the 10 of clubs in our label lookup.
Now, there's one more thing we need to do, and that's ensure that our dataset outputs all of these images in the same size, because the model will expect the input to always be consistent. And that's where we're going to use these transforms that we've imported from TorchVision. And we're keeping this transform as simple as possible, and all we're doing is taking the image and making sure that they're always 128 by 128 in size, and then we're converting it to a PyTorch tensor.
So now let's recreate our dataset with this transform added. And if we call one of the values in the index, we can see we have our image and our label. But now our image is a PyTorch tensor, and we can run a shape on this tensor to get an idea of the size of the 3D tensor, and it's in 128 by 128 because we've resized it. And the three here is the number of channels, that's the red, green, and blue channels of the image. Now, I do have a whole video on working with image data that you should check out if you haven't already, where I explain this in more detail.
Now, one other thing to keep in mind about the dataset is we can iterate over this dataset, and we'll just have it break out of this after the first time we run it. But since the dataset is an iterable, we can call this loop and run over each image and label in our dataset, and you can see the image here is the same as before, it's a tensor, and the label is the number associated with each label.
So creating the PyTorch dataset is the main thing that you want to do. And the nice part about the datasets is then we can take it and wrap it with a PyTorch dataloader, which will automatically handle the processing to parallelize reading in each of these images.
So creating our dataloader is as simple as calling the PyTorch dataloader and providing it our dataset. There are a bunch of other settings we can set here, but the main ones we're going to focus on are batch size and the shuffle parameter. So batch size tells us how many of these examples we should pull each time we iterate over the dataloader, and a nice one to start with, we'll say, is 32. And then shuffle allows us to tell the dataloader that every time we load a new example from our dataset, if we want that to be pulled randomly from the dataset or it to be pulled in order. For reasons I won't go into too much detail, shuffling is typically done when training the data, but you don't need to shuffle the data when you're running it on a test set or a validation set.
So we've created our dataloader, and now we can iterate over our dataloader just like we could our dataset, but we'll notice one main thing has changed. So let's iterate over this by calling `for images, labels in dataloader`, and then we'll break. Now, if we look at the images, it's a tensor just like before, but the shape of it is 32 by 3 by 128 by 128. That's because our dataset has been batched into 32 examples. And our labels is a PyTorch tensor with one dimension that has 32 different labels. Also, note because we put shuffle here, and we know that the labels were in order of the original dataset, if we just look at these labels now, we can see that they're random, and each time that we run this, they should be a random order of the labels.
So just to recap, we need to first create a PyTorch dataset before we train our model. We create this by inheriting from the dataset class and creating a few methods like the `__len__` and `__getitem__` method. This creates an iterable object which we can loop over. Then we can create a PyTorch dataloader, which will handle batching these into. And then we can wrap our PyTorch dataset with the dataloader class, which will automatically batch this data when feeding it into our model. Main reason we want to do this is because the model trains much faster when it takes in batches of examples instead of one example at a time.
Now we're on to step two, which is actually creating the PyTorch model which we want to train on this dataset. PyTorch models are extremely configurable, and you could go in and build these models from the ground up, but I want to keep it simple here so that we're really learning about these paradigms and PyTorch, and we're going to use a predefined architecture that exists, and we're going to import that from the library Timm. The library Timm has pre-packaged a lot of the state-of-the-art architectures that are already created for image classification. This allows us to get some really good results without having to define much from scratch at all. The one thing I'll say about creating a PyTorch model is it's all about understanding the shapes in each layer of the model, and we're going to do a little bit of that when we create our classifier.
So let's create our PyTorch model. We're going to create a class that's called `SimpleCardClassifier`, and this is going to inherit from the neural network module of PyTorch, the base `Module`. Now, when creating a PyTorch model, you can customize this a lot, but there really are two main methods that we're going to create: the `__init__` method and the `forward` method. I like to think about it like this, where in `__init__` we define all the different parts of the model, the `forward` method we take in an example or a batch of examples and we connect all these parts which we've defined up here and return our output.
So to make this actually work, let's start building the structure of our model. When we initialize this, we want to give it the number of classes that the model has, which is 53. And we're going to use the EfficientNet B0 model architecture, which is a really nice architecture and it trains fast. The B0 indicates the size of the model; this is one of the smaller ones. By setting `pretrained=True`, we're saying the model weights within this model have already been trained on the ImageNet dataset. Now, this depends on the model, but the EfficientNet B0 by default will output a feature size of 1280, and we're going to need to take this output size and resize it to our number of classes, which is 53. So we're going to create a classifier, which will simply be a linear layer from our EfficientNet output size to our number of classes. Because this base model actually has an additional layer that we want to remove, we're just going to remove that last layer with this line of code. Don't be too confused about it; just think about this as if we're cutting off the very end of this Timm model so that we can instead have it output at our classifier level.
Now, in our `forward` method, we are going to take in an example or a batch of examples, and we want it to return the output. We're going to call this `features` on our input data, and then we're going to call our `classifier` on the output of these features, and let's just call this `output` so it's obvious what the last layer is and what we're returning. I try to make this as simple as possible so you can see how the basic structure and paradigm of creating PyTorch models looks like. Almost forgot, but we need to add this line of code to our `__init__`, which will initialize this object with everything from the parent class.
And let's go ahead and create an example of this. We could give it the number of classes. Now we have our model, and if I print this, it'll even show the structure of this model in detail with everything from the EfficientNet. Since that's really long, I'll just curtail it in our print statement so that we can just see the beginning part of it.
Now, the first thing I like to do after creating the model, just to check and make sure it works correctly. If you remember before, we tested out our dataloader by iterating over it and testing to make sure it returned the images and labels. We can test our model just to make sure the input and output are what is expected by providing it now those images from one example batch. And we see this outputs a tensor, and this is a good sign because it did run correctly, which tells us that the structure of the model can accept the input data that we provided. And if I call this `example_output` and look at the shape, we can see that the `example_output` is 32 by 53, and this is our batch size by our number of classes. So this is exactly how we would expect it: for each of these examples, we have probabilities for each class.
Recapping the model that we've created: we make sure that we import from PyTorch's neural network module, we create an `__init__` and a `forward` method. `__init__` is where we've created the structure of this model, and our `forward` method defines how these parts of the model are connected and will produce our output when called on some input data. And we tested this `forward` method by calling the model on some example images from our dataloader.
Now we're on to step number three, which is creating the training loop that we'll use to train this model. Now, we're going to write this training loop from scratch because it will help us to really understand how the model is trained. So essentially, the idea here is we're going to feed this data into our model many times, and then we're going to apply some sort of loss function to the outputs that we receive when we send in these training examples. This is how the model learns. We do this in batches, that's why we create a PyTorch dataloader, and then we call one epoch when we run through all the batches in one training dataset. And we're keeping this training loop as simple as possible, but there are really two things that we'll need to select before we move forward, and that's the optimizer we want to use and the loss function we want to use.
Let's start with the loss function. We're going to call this the `criterion`, and PyTorch has a bunch of built-in loss functions that you can use for common tasks. We have a classification task here, and cross-entropy loss is a common loss function used for this type of case. We also want to define our optimizer. We'll pull this from PyTorch's `optim` library, and we're just going to use Adam, which is the best place to start when using an optimizer. It works really well in most cases. This optimizer needs to take in our model parameters, and then we need to give it a learning rate. There's a lot more advanced techniques that we can learn about in future videos, like learning rate schedulers, but here we're just going to keep the learning rate constant as we train.
Similar to how we tested out our model and our datasets before to make sure everything was working right, let's test out our loss function. We're going to give it the example output from the model that we ran before here and the labels for that data, and you can see this worked correctly because it imputed our loss on that batch of data. Now, the model hasn't learned anything yet, so this loss is going to be pretty high, but we're just checking to make sure this loss function does work with the input data that we have and the output from the model.
Now, we also want to create datasets like we did before, but these are the actual datasets we'll use for training and evaluating our model. Our dataset actually has a folder for training examples and a folder with validation examples. That makes our job easy because we don't have to define the validation split. I do have a whole video on how to properly create cross-validation for your dataset that you should check out if you haven't already.
So here I've set up our datasets. Since there is a test folder, I'm including that here too, but we don't necessarily need this for our training loop. But you can see we have this one transform, which just resizes our images. We've created our train dataset and our validation dataset based on these folders, and then we've created dataloaders for each with a batch size of 32. And for our train loader, we're shuffling the data as it's training, but for our validation or test, we are not having the dataloader shuffle.
Now that we've defined a loss function, an optimizer, we have datasets and a model, let's go ahead and create this training loop. As I mentioned before, an epoch is one run through the entire training dataset, and we're going to set to train for five epochs. Then we also need to set up some lists which we will store our training and validation losses. Then we'll create the model and we'll start looping through five different epochs.
So `for epoch in range(num_epochs)` and we're going to loop through our training dataloader to train the model. But before we do that, we need to set the model to train and we'll set our running loss, and then we're going to keep track of this loss as we train. Now we're going to loop through our training dataloader. We're going to say `for images, labels in train_loader`. We need to set our optimizer to zero grad, and then like before, we'll create our outputs by calling the forward method on these images. Then we'll take that output and we'll calculate our loss for this batch using our criterion, which we defined above. We'll provide this our outputs and the labels, and this is the part where we run backpropagation on the model, which will update the model weights in each of these steps using `loss.backward()`, and then we need to run a step on our optimizer. And we're going to keep track of our running loss by just adding this loss like this.
Once this inner loop has completed, one epoch is done, and we want to store our training loss. So let's first compute that by taking our running loss and divide it by the length of our train loader. And mainly for tracking, we're going to append this `train_losses` list with our training loss. And this is pretty much it for training. But because we want to see how well the model is doing on the validation set as we train, let's add a validation portion to this training loop. We need to make sure we change the model from being in training mode to evaluation mode, and then we're going to track running loss, but this time it's on our validation dataset, and we'll essentially calculate loss similar to how we did on the training set, but this time on the validation set. One thing we need to do just to make sure the model weights are not touched is to do this under `torch.no_grad()`.
And then let's write our loop. So we've run through the entire validation phase of this model training loop, and we're just going to print epoch stats each epoch, and this will allow us to track how well we're doing by showing the training and validation loss. Need to fix a few names here. And if I run this cell, it'll start running through our training stage.
Now, I purposely left something out, and that's that our training is currently being done on this instance's CPU. If I look here, we can see that the instance's CPU is being used, but the GPU is not being used at all. If we were to train this way, it's going to train very slowly, and we want to use the GPU. So let's go ahead and stop this. And with some slight modifications, we can modify this loop to instead run on the GPU. And to do this, we're going to define the device that we want the model to run on. Now, this could either be CUDA for a GPU or CPU. And this one-liner will automatically use the GPU if it exists. And we can see that now the device is set to CUDA. And then let's go and modify our loop slightly by making sure we move both the model and the data over to this GPU when we're running these loops. And we do this by taking the model and telling it where we want it to go to, which is our device. And similarly, with our images and labels, we are just going to put these on the GPU: `images.to(device)`, `labels.to(device)`. It'll give you an error if the device is not the same between tensors or the model when you're trying to run operations between them.
Now, while I'm running this cell, I can click here and see that the GPU is actually being used as the model is training. Looks like I had this spelled wrong, so let's run this again. And now what we need to do is just wait and let the model train.
Now, if you're impatient like me, you'd like to see a progress bar as the training's going on. So let's go ahead and from `tqdm.notebook import tqdm`. By wrapping tqdm around these loaders, it allows us to get a progress bar as we're training. But as this runs, we do have a progress bar as it runs each training loop and validation loop, and we're printing out the losses, train and validation, as we train. Now, what we'd hope to see here as we're training is the training loss would decrease, and the validation loss will also decrease, or else we'd see something known as overfitting to the training set. But let's let this finish and see what our results look like.
Our training loop is done. It's run five epochs, and we see that our training and validation loss have both decreased. But let's go ahead and visualize. But since we were tracking these losses as we trained, let's just quickly visualize what the loss looked like while it trained. We did have a strange spike here in the validation loss, but overall, our training and validation loss was going down, which is a good sign. And we're training for not that many epochs here, so typically you'd train for a little bit longer or maybe adjust the learning rate.
So just to summarize what we learned about the training loop: we had to pick a criterion or a loss function, we had to set an optimizer. At that point, we create our datasets, both training and validation, then we wrote our training loop from scratch, where we loop over our train dataloader, backpropagating our losses, and then at the end of each epoch, calculating our validation loss.
Now, the whole point of this exercise was to get familiar with PyTorch datasets, creating models, and creating the training loop. But let's go ahead and evaluate our results just to see how well it did visually. I'm not going to spend too much time on this code explaining what it does, but essentially, what it will do is let us visualize what the predictions look like from the model when we feed it a test image. So I'm going to run this code on one of the example test images for a five of diamonds. We're going to see what the results look like. So on the left here, we see the image that was fed into the model. We can see that it is very confident in the fact that this is a five of diamonds.
So let's just print out 10 test examples randomly to see what they look like. For this, I'm going to take all of our test images, which I found using glob, and I'm going to use NumPy to sample 10 of them, and we'll run this code over these 10 examples. So, four of diamonds, it looks like it's right. Ace of clubs, king of clubs, and you can see here most of these results, it's predicting the correct class. And we could just apply an argmax, which is a post-processing technique to identify the highest class predicted, and that would be our ultimate prediction. This maybe looks like the one where it was the most confused. It still got the correct one right, but you can see how the image is a little bit different than the others, so the model struggled a little bit.
Now, if you've made it this far in the video, I first just want to thank you for spending the time watching this. Make sure you like and subscribe. I want to challenge you to do something with these results. You should be able to follow the link in the description to this notebook, and I want you to then here calculate the accuracy of our model on the test set. And you should be able to do that, given everything that I've provided you in this notebook, by just finding out which of these predictions were correct. Thanks for watching. Put some comments below and let me know what you thought of this video, and I can take a deeper dive into some of the more advanced features with PyTorch in a future video if you guys like this one. See you next time.