Setting your working directory to Google Drive in a Colab notebook

June 13, 2020
(coding)

There are plenty of articles detailing how to mount your Google Drive to a Google Colab notebook. It’s simple, you just run this in your notebook:

from google.colab import drivedrive.mount('/content/drive', force_remount=True)

But what if you want your notebook to use a folder in your Google Drive as the working directory?

That way, any files that are created within your project will automatically be saved in your Google Drive. If your project runs 12 hours and then times out due to inactivity, any files or checkpoints created will be in your Google Drive.

I wrote a quick script to do this. Just change the project_folder to the folder you want to use. The script will create the project_folder from scratch if it doesn’t exist and then set it to the current working directory.

It also creates a sample text file in that directory to verify that it’s working.

You can run !pwd in a notebook cell at any time to verify your current working directory.

import os 

# Set your working directory to a folder in your Google Drive. This way, if your notebook times out,
# your files will be saved in your Google Drive!

# the base Google Drive directory
root_dir = "/content/drive/My Drive/"

# choose where you want your project files to be saved
project_folder = "Colab Notebooks/My Project Folder/"

def create_and_set_working_directory(project_folder):
  # check if your project folder exists. if not, it will be created.
  if os.path.isdir(root_dir + project_folder) == False:
    os.mkdir(root_dir + project_folder)
    print(root_dir + project_folder + ' did not exist but was created.')

  # change the OS to use your project folder as the working directory
  os.chdir(root_dir + project_folder)

  # create a test file to make sure it shows up in the right place
  !touch 'new_file_in_working_directory.txt'
  print('\nYour working directory was changed to ' + root_dir + project_folder + \
        "\n\nAn empty text file was created there. You can also run !pwd to confirm the current working directory." )

create_and_set_working_directory(project_folder)

View the code on Github or in this public Google Colab notebook.

Related Posts: