How to Save and Load Model Weights in Google Colab

Google Colab is Google's hosted Jupyter Notebook product that provides a free compute environment, including GPU and TPU.

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:

/content/yolov5/runs/exp0_yolov5s_results/weights/best.pt

where the weights are called best.pt.

(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:

from google.colab import files
files.download('example.txt') 

where 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.

from google.colab import drive
drive.mount('/content/gdrive')

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.

%cp /content/yolov5/runs/exp0_yolov5s_results/weights/best.pt /content/gdrive/My\ Drive

After running this cell, the weights will be available in our Google Drive as a file called best.pt.

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.

In our case, we will again be basing our example on the YOLOv5 notebook for image recognition and object detection. (In the YOLOv5 notebook, we call inference when we call the detect.py script.)

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):

from google.colab import drive
drive.mount('/content/gdrive')

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 weights.

To copy our weights (called best.pt) from Google Drive to our Colab notebook weights directory, we'll run this in a cell:

%cp /content/gdrive/My\ Drive/best.pt /content/yolov5/weights

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 train.py cell.)

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.

Happy training!