Skip to main content
Loading Explore guides…
Training YOLO in Colab · Acedra
Home
Explore
Computer Science
Training YOLO in Colab
Sign in to save
Chat
If YouTube asks you to sign in, open the video on YouTube.
Watch on YouTube
Reload player
Transcript
42 segments
Copy
Hide video
0:13
Hey everyone, today we're going to walk through training a YOLO model inside Google Colab. I'll go through the process of gathering a data set training the model in Colab and then running it locally on your PC with a custom Python script. We'll show how to train a candy detection model that can detect and identify different popular types of candies. You can follow along and use my candy data set as an example or you can bring your own data set. We'll work inside
0:39
a Google Colab notebook I wrote for training YOLO models. It works with YOLO11, YOLOv8, and YOLOv5. if you want to train your models on a local PC instead, I wrote an article that shows how to do so. You can find it in the video description below. All right let's get started! The first step to training a custom model is to create a data set. A good starting point for a proof of concept model is 100 to 200 images. The training
1:06
images should show the objects in a variety of backgrounds, perspectives, and lighting conditions that are similar to what the model will see in the actual application. The best way to build a data set is to manually take pictures of your objects in a variety of conditions. For my candy detection model, I took pictures of candy in various locations around my house using my phone and a webcam.
1:30
You can also use online sources like Roboflow Universe, Kaggle, or Google's Open Images V7 data set to find pre-labeled images of your objects. For tips and best practices on building an effective data set see my video on how to capture and label training data. Once you finished taking pictures or downloading them from the internet, move all of the images to a single folder on your computer. I'll create a folder called YOLO in my Documents folder, and another folder inside
2:02
that called candy-images. Then I'll move all the photos from my phone into the candy images folder. Now that you've got your training images gathered and saved on your PC, it's time to label them. We'll use an open-source tool called Label Studio to label our images. I like Label Studio because it offers a simple workflow while providing advanced features like automated labeling. Best of all it's free and open source. Because I like Label Studio, I reached out to them and asked if
2:34
they would sponsor this video and they agreed to it. So, this video is sponsored by HumanSignal, the team that developed Label Studio. If you want to learn more about their tool, click the first link in the video description below. To install Label Studio we first need to install Anaconda, which is a tool for creating and managing Python environments. Go to anaconda.com/download and click the skip registration button and then download
2:59
download the package for your OS. When it's finished downloading, run the installer and click through the installation steps you can use the default options for installation. Once it's installed, run Anaconda Prompt by searching for it on the Start bar and clicking on it. This opens a command window that we'll run our commands from. Let's create a new Python environment for installing Label Studio and the other libraries we need for running YOLO models.
3:27
Create a new environment by issuing "conda create --name yolo-env1 python=3.12" . When it's finished creating the environment activate it by issuing "conda activate yolo-env1". Now we can install label Studio by issuing "pip install label-studio". Depending on your internet connection, it'll take a few minutes to install.
4:03
Once it's installed start Label Studio by issuing "label-studio start". After a few seconds, Label Studio will start in a new web browser window. Label Studio runs on a server that's hosted locally on your computer, and you interact with it through your web browser. Even though it's in a web browser, all the accounts files and label data you create are stored entirely on your computer and not on the web (unless you want to set it up that way). To
4:36
start with, you'll need to make an account for your local Label Studio server. Again, this is all hosted on your local computer so you aren't actually creating a web account. Click the sign up link and then enter a fake email address and password. I'll use fake@fake.com and fakepassword. Click create account to continue. You'll enter the project page for Label Studio click the Create Project button to start a new project. Enter a project name like candy images and then click
5:08
the data import tab. On this screen, you can drag and drop images from your folder on your computer into the web browser to import them. If you try and import more than a 100 images at once, you'll get an error about exceeding a max number of files. But, you can get around it by just making sure you only import 100 images at a time, and then going to the next 100 and so on. If you have more than a thousand images to label, it's strongly recommended to link a
5:38
local storage folder or cloud account to the project. See the Label Studio documentation on adding project storage for more information on how to do so. Finally, click the Labeling Setup tab and then click the object detection option. This opens a configuration page for the labeling interface. To start with, delete the airplane and car labels. In the Add Label Names window, type your own class names. I'm labeling pictures of candy, so I'll add each of the
6:07
candy brands like AirHeads, Nerds, Skittles, and so on with a enter after each line. Once all the classes are typed out, click Add. Click Save to finish creating the project. Okay, we're ready to start labeling! Click the first image in the list to open the labeling interface. Create a new bounding box by clicking a class at the bottom of the window and clicking a dragging a box around the object in the image,
6:40
like this Starburst for example. You can click the box and adjust it as needed. You can also use class hotkeys to create a new box. So for example I can press one to create a bounding box for the Airheads in this in this image, and I can press eight to label the peanut M&M's.
7:01
Once you've labeled all the objects in the image, click Submit and then click the next image in the list. On this image, let's zoom in a little to draw tighter bounding boxes around the objects. We'll start with these gummy worms and get a nice tight bounding box. Okay Three Musketeers, Snickers, and Milky Way. Okay, and then when you're done with the image, click Submit to save the labels. Continue working through each image
7:31
if you need tips on best practices for labeling images, check out my object detection labeling video. Labeling takes a while, so grab a drink and put on some music while you're working. Once you've labeled all your images, it's time to export them. Click the project name at the top of the window and then click the Export button. Scroll down and select the YOLO export option and then click Export. Label Studio will package
8:01
the images and label data into a zip file and download it through your web browser. The files are organized in a directory structure required for training YOLO models, with the images and an images folder and the labels in a labels folder. It also generates a classes.txt labelmap file which lists all the classes in your data set. Let's move that downloaded file back into the YOLO folder in Documents and rename it to data.zip. We'll upload this file to the
8:37
Google Colab training notebook later. Now, we're ready to start training our YOLO model. We'll train our model in Google Colab, a service from Google that allows you to write and run Python code in your web browser on a free GPU. Click the Colab link in the video description below to get started. This is the Colab notebook where we'll upload our data set and train our custom YOLO model. First, go to Runtime and Change
9:04
Runtime Type and make sure that a T4 GPU is selected. Click Save then click Connect. Once it's connected, scroll down and click play on this !nvidia-smi block. You can just click Run Anyway on that window that appears. Verify that it reports a Tesla T4 or similar GPU is available. We've already gathered our images so, we can scroll past Step 1. Next, we'll upload our data.zip file into this Colab instance. Click the folder icon on the left
9:36
and then drag and drop your data.zip file in into the sidebar. This little blue circle will show the upload progress. You can also upload data to Google Drive and copy it from there, which goes a little faster. See these instructions on how to do so. If you got your data set from a different source like Roboflow Universe, make sure it has the correct folder structure shown here. Again, all the images are in the images folder, the labels are in the labels folder, and there should
10:05
be a classes.txt file that lists every class. You can just ignore this notes file. You may have to move the files around or create the classes file manually to match this data structure. If you just want to test the training process on a pre-made data set, you can use my candy data set or my coin data set. Click this block here to download it.
10:30
Next unzip the data set by clicking this code block. Now that the files are set up, we're ready to split them into train and validation folders. The train folder will contain the actual images that are used to train the model, while the validation folder has images that are used to check the model's performance after each training epoch. I wrote a Python script that automatically creates the folder structure
10:54
required for training and will randomly move 90% of the data set to the train folder and 10% to the validation folder. Run this code block to download and execute the script. When script finishes, all the images will be in a data folder where the train and validation images and labels are each in their respective folders. Next, we'll install the Ultralytics library, which is the Python library we'll be using to train the YOLO model. Click play to run the pip install command.
11:29
Before we can run training, we need to create a training configuration file. This file sets the location of the train and validation folders and defines the model's classes. Here's an example of what it looks like. Run the next code block to automatically generate the data.yaml config file. When it finishes running, it'll display the contents of the file. Double-check that everything looks correct.
11:55
Okay, we're finally ready to run training! But there's a few important parameters we need to decide on first. We need to pick which YOLO model we want to use. There's a variety of YOLO architectures and model sizes. YOLO11 is the latest architecture and has the best performance, but YOLOv8 and YOLOv5 are more mature and have broader support on hardware platforms. Model sizes range from nano to extra large. I made a quick YouTube video comparing YOLO model performance
12:24
on a Raspberry Pi 5 and a laptop equipped with an RTX 4050 GPU. You can watch this video to get a sense of the frame rates and accuracy you'll get with different YOLO models. Generally, you want to choose a model that's small enough to meet the speed requirements for your application, while still being accurate enough to detect your objects. If you aren't sure which model to use, the YOLO11s model is a good starting point. For my candy model, I'll use YOLO11s. Next, the number
12:54
of epochs sets how long the model will train for. If your data set has less than 200 images, a good starting point is 60 epochs. If your data set has more than 200 images, a good starting point is 40 epochs. Resolution is another parameter we can tweak, but generally it's safe to use the default resolution of 640x 640. However, if you want your model to run faster or if you know you'll be working with low resolution images, try using a lower resolution like 480x480.
13:24
With the parameters decided, it's time to run training. I've got the notebook set up to train a yolo11s model for 60 epochs at a 640x 640 resolution. You can change the parameters of this command if you want to try something different. Click play to start training. The training algorithm will parse the images in the train and validation directories and then start training the model. At the end of each training epoch, the script runs the
13:50
model on the validation data set and reports the resulting mAP, precision, and recall. mAP is a general measure of model accuracy. As training continues, the mAP should generally increase with each epoch. Training will end once it goes through the specified number of epochs. Depending on how big your data set is and the number of epochs used, training can take anywhere from 5 minutes to a few hours. Be careful though, because the Colab instance will time out and disconnect
14:20
after about 5 hours. If you want more time for training, consider upgrading to Colab Pro.
14:32
When training finishes, the best model weights will be saved in runs/detect/train/weights/best.pt. The results.png file shows how metrics like loss, precision, recall, and mAP progressed over each training epoch. Let's make that a little bigger. Generally, mAP should increase and then begin to flatten out by the end of training.
15:03
Now that the model's finished training, let's test it! Run this first code block to perform inference on all the images in the validation folder. Then, run the next code block to display the results for the first 10 images. The model should draw a box around each object of interest in each image. My model is doing a pretty good job of identifying each of the types of candy in these images. Looks like it's getting them mostly
15:31
right. If your model isn't very accurate, there's a couple things you can try. First, double-check your data set to make sure there's no label errors. You can also try increasing the number of training epochs or using a larger model. Okay, now we've got a trained model, so what can we do with it? Drawing boxes on images is great, but it isn't very useful in itself. It's also not very helpful
15:55
to just run the model inside a Colab notebook, so let's download and run it locally on our computer with a custom Python script. First zip and download the trained model by running this code block. It will rename the model to my_model.pt, move it to a folder named my_model, and zip it. It also zips the training results in case you want to reference them later. Then you can download the my_model.zip file from the sidebar. Now we'll walk through instructions for setting it up on
16:26
your PC. Once the download is finished, move the downloaded zip file back into the YOLO folder that we were working in and extract it. We'll work inside this extracted folder to run our model. Go back to the Anaconda Prompt window. If it's still running Label Studio, press CTRL+C a couple times to stop it. Or, if you've already closed the window, you can re-open it, re-activate the environment by issuing "conda activate yolo-env1" and then copy the path to the my_model folder and
17:10
do "cd" and enter the path to the folder to move into it. We'll use the Ultralytics library to run inference with our model. Install it by issuing "pip install ultralytics". This also installs other important libraries like OpenCV-Python, Numpy, and PyTorch, so it may take a few minutes. If you have an NVIDIA GPU, you can install the GPU enabled version of PyTorch by going to the PyTorch Get Started page at pytorch.or/get-started-locally and scrolling down, selecting your operating
17:50
system, and the latest version of CUDA. Copy the installation command and paste it in your Anaconda Prompt window. Make sure to add "--upgrade" after install so that it upgrades the existing installation. Hit enter to run the command. This automatically installs the CUDA and cuDNN drivers needed for running inference on GPU. They're pretty big drivers, so it may take a while. Now we're ready to run the model! I wrote a custom Python script that shows how to load the model,
18:29
run inference on an image, and parse the inference results. The intent for this script is to show you how to work with Ultralytics YOLO models in Python, and it can be used as a starting point for other applications. Download the script by issuing "curl -o yolo_detect.py https://www.ejtech.io/yolo_detect.py" .
18:58
You can run the script by entering the path to your model and an image source. For example, to run inference on a USB camera at 1280x720 resolution, I can issue "python yolo_detect.py --model my_model.pt --source usb0 --resolution 1280x720". All right, so your camera view will appear and the model will draw boxes around the different objects of interest in each frame. So as you can see my, candy model doesn't do too bad
19:35
at detecting different types of candy here. And since it's a YOLO11s model, it runs pretty fast on my NVIDIA 4050 graphics card. So, not bad! You can press Q to quit the program. You can also run the model on a video file image or folder of images. For example, I'll move a video file into the folder and run the script on it by issuing "python yolo_detect.py --model my_model.pt --source test1.mov" . Again, a window will appear showing detection results.
20:23
To see a full description of how to use yolo_detect.py, see the readme file in the Train-YOLO-Models GitHub repository linked in the video description below. Okay, congratulations! You successfully trained a YOLO model. Next, see if you can extend the application beyond just drawing boxes on the screen and counting the number of objects. As an example, I created a candy calorie counter app that
20:50
will tell you the number of calories and the amount of sugar and a handful of candy. You can check out the source code for it in the video description below. I'll be adding more examples to the GitHub repository that show cool ways you can use these YOLO models. Also, keep an eye on my channel for more guides, like a guide showing how to run YOLO models on the Raspberry Pi. Thanks so much for watching, and as always good luck with your projects! [Music]
Here to help you learn
Sign in to keep a private conversation with your course tutor.
Ask me anything about the material…
Sign in to chat