Colab comes "batteries included" with many popular Python packages installed, making it a choice tool for easy model experimentation. For this reason, the Roboflow Model Library includes many free, open source computer vision models available on Google Colab.
Colab does come with limitations. The compute resources allocated are limited to a 12-hour span. If the browser session senses inactivity (e.g. the user is not active in that tab) for an hour, the kernel disconnects.
Because of this, it's important to save a model's weights and reload them in a future session. Saving a model's weights means saving the fit of a model after training. Reloading the model weights means using those saved weights in a future experiment – even if that is a new session in Colab.
In this post, we walkthrough how to save and reload model weights from YOLOv5 in the notebook from the Roboflow Model Library, but the steps we follow are applicable to any Colab session (with file path modification).
Saving Model Weights
To save model weights, we must first have weights we want to save and a destination where we seek to save those weights.
Identify the Weights File Path
After training a model, the weights of that model are stored as a file in the Colab session. In our example YOLOv5 notebook, these weights are saved in the
runs folder. Specifically, the file path is:
where the weights are called
(PRO TIP: To identify the file path of any file in Colab, you can use the navigation pane on the left, right click a given file, and select, "Copy file path.")
Download the Weights from Colab (Locally or to Google Drive)
Once we have the file path of our weights file, we can save this file locally or to our Google Drive. We recommend saving weights to your Google Drive.
How do you download a file from Google Colab? It's this simple:
example.txt should be the file we want to download – in our case, we need to replace this with the path to our weights file.
We can also download and save files to our Google Drive, which may be a little more convenient as we may want to have our weights file available with our Google account.
How do you download a file from Google Colab to Google Drive? There's two steps.
First, connect your Google Drive to your Google Colab session by running the below. This will prompt you to visit a separate page and copy/paste an authorization code.
Second, copy the file from your Google Colab notebook to your Google Drive. This requires us to specify (1) the path of the file we want to copy (our weights, in this case) and (2) the location of where we're saving the weights in our Google Drive.
In the snippet below, we're presuming we're using the YOLOv5 notebook model weights location and that we're saving them to the root directory of our Google Drive.
After running this cell, the weights will be available in our Google Drive as a file called
Loading Model Weights in Colab
Once we have our weights saved, we may want to reload them in a later Google Colab session. Loading model weights requires that we have already built our model architecture and environment.
But, instead of using that architecture to train a new model, we will use that architecture to run inference on a new set of data.
Uploading Model Weights into a Colab Session
To add weights to our Colab session, we can either load weights from our local computer or from our Google Drive. We need to know the existing location of our model weights and, critically, where in our Colab notebook we are adding them. We'll walkthrough loading files from your local machine and loading files from Google Drive.
How do you upload files to Google Colab?
For files on your local machine, we'll use the user interface in Google Colab. In the upper left hand corner, select the folder icon to see the list of available folders in a given directory.
Here, there is an icon for uploading a file. Select this icon, and then select wherever the weights are on your local machine. (Note that this upload icon loads the file from your local directory to wherever you selected the upload icon.)
How do you upload a file to Google Colab from Google Drive?
Here, we need to first link our Google Drive and our Google Colab notebook using this code snippet in a cell (which will require an authorization code):
Now, we can load our weights from our Google Drive into our Colab notebook. This requires knowing the file path of the weights in our Google Drive and where we want to copy the weights into our Colab session.
We'll assume the weights are at the base directory of our Google Drive (which is where we saved them in our prior example in this blog post). We will also assume we're loading the model weights into a directory in the YOLOv5 folder structure called
To copy our weights (called
best.pt) from Google Drive to our Colab notebook weights directory, we'll run this in a cell:
Using Our Uploaded Model Weights
Now that we have our model weights in our Colab Session, we can use them with our model. This requires running all cells to set up the model (in the YOLOv5 notebook, this is everything before running the
We need to know the file path of our weights file, and pass that file path to our model running inference. In our YOLOv5 example above where we linked our Google Drive, we loaded our weights called
best.pt into the
content/yolov5/weights directory. Thus, the script where we use our model weights to call inference should reference this weights file path. It'd look like this:
!python detect.py --weights weights/best.pt --img 416 --conf 0.4 --source ../test/images
Voila! We now can save and load model weights from Google Colab to our Google Drive for faster experimentation.