Summary

The web content provides a tutorial on how to save and load a machine learning model in Google Colab using Google Drive to prevent data loss from runtime restarts.

Abstract

The article outlines a method for preserving machine learning model progress in Google Colab by saving model checkpoints to Google Drive. It begins by explaining the necessity of saving models due to potential runtime restarts in Colab, which could otherwise result in the loss of all training progress. The author then guides the reader through the process of mounting Google Drive to Colab, saving the model checkpoint to the mounted drive, and subsequently loading it back into Colab for continued training or evaluation. The tutorial emphasizes the importance of regularly saving the model to avoid starting from scratch and suggests saving the model along with its accuracy to facilitate the selection of the best performing model. The author also offers personal assistance via Slack for any further help.

Opinions

The author believes that losing a trained model due to runtime restarts is suboptimal and should be avoided by saving model checkpoints.
The author suggests that accessing Google Drive from Google Colab is essential for preserving work, indicating a strong endorsement for integrating the two services.
The author provides a subjective recommendation to save the model's performance metrics alongside the model itself, implying that this practice is beneficial for model selection.
By offering personal assistance, the author conveys a willingness to help and support the reader, reflecting a community-oriented approach to problem-solving.

How to save our model to Google Drive and reuse it

If you are using Google Colab and if the runtime restarts during training, you will lose your trained model. Then you have to start again from the scratch, which is not optimal. Instead, you should save your model checkpoint to Google Drive at every epoch and reload it next time when you start. In this short tutorial, I will explain how to mount Google Drive, save our model checkpoint to Google Drive and reload it back.

How to mount Google Drive

Google Colab runs isolated from Google Drive, so you cannot access it directly. To access it, you need to authenticate, give permissions to Colab so that it can access it and mount the drive.

Add the following code to a cell:

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

Once run, it will print an URL, asking for an authorization code:

Visit that URL, authorize with your Google account and copy the code which it displays. Paste the code back in the colab notebook. Once authorized, it displays that our drive has been mounted:

You should able to see your drive in sidebar, under Files tab. Notice that our drive is mounted under path `/content/gdrive/My Drive`. If you do ls you should be able to see your drive contents:

!ls /content/gdrive/My Drive

How to save your model in Google Drive

Make sure you have mounted your Google Drive. Now, to save our model checkpoint (or any file), we need to save it at the drive’s mounted path.

To save our model, we just use torch.savemethod:

model_save_name = 'classifier.pt'
path = F"/content/gdrive/My Drive/{model_save_name}" 
torch.save(model.state_dict(), path)

Now, if you visit your google drive at https://drive.google.com/drive/my-drive you will be able to see classifier.pt file saved!

How to load the model from Google Drive

Make sure you have mounted your Google Drive. Now, we will access our saved model checkpoint from the Google Drive and use it. We know the path and we will use that in torch.load :

model_save_name = 'classifier.pt'
path = F"/content/gdrive/My Drive/{model_save_name}"
model.load_state_dict(torch.load(path))

That’s it! Make sure you load the state dict just before your training starts.

Pro tip: Save the model name along with its accuracy. So that you can pick the best model available. Hope this helps!

If you need any help, feel free to ping me on Slack. My id is avinashss