How to connect to private storage bucket using the Google Colab TPU

Question

I am using google colab pro and the provided TPU. I need to upload a pre-trained model into the TPU.

• TPU can load data only from a google cloud storage bucket.

Answer 1

I've been struggling with this scenario myself (although with the free version of Colab) and just got it to work. This specific use case doesn't appear to be very well-documented—it seems the official documentation mostly deals with cases involving a Compute Engine VM, rather than an auto-assigned TPU. The process that worked for me went as follows:

1. Run Google Cloud SDK authentication and set the project (these two things may be redundant—I haven't yet tried doing just one or the other)

!gcloud auth login
!gcloud config set project [Project ID of Storage Bucket]

and

from google.colab import auth
auth.authenticate_user()

2. Initialize TPU (from Tensorflow TPU docs)

resolver = tf.distribute.cluster_resolver.TPUClusterResolver(tpu='grpc://' + os.environ['COLAB_TPU_ADDR'])
tf.config.experimental_connect_to_cluster(resolver)
tf.tpu.experimental.initialize_tpu_system(resolver)
strategy = tf.distribute.experimental.TPUStrategy(resolver)

3. Try to load the model

model = tf.keras.models.load_model('gs://[Bucket name and path to saved model]')

This initially failed, but the error message included the service account of the TPU trying to access the directory, and this is the address I gave access to as described in the Cloud Storage docs. The address is in the service-[PROJECT_NUMBER]@cloud-tpu.iam.gserviceaccount.com format but the project number isn't the Project ID of the project my bucket is in, nor a value I've been able to find anywhere else.

After I gave permissions to that service account (which I was only able to find in the error message), I was able to load and save models from my private bucket.