Colab Pro: using GPU crashes the session

Question

I have a Google Colab subscription to use its GPU and to fasten the training of my model. When I ran the code it works for some blocks and then it stops and says (Your session crashed for an unknown reason) The execution stopped at this particular block, which is a function in TensorFlow federated package.

source = tff.simulation.datasets.FilePerUserClientData(
  dataset_paths, create_tf_dataset_for_client_fn) 
source._client_ids = [tf.cast(c, tf.string) for c in source.client_ids] 
source = source.preprocess(add_parsing)
train, test = source.train_test_client_split(source, 2)

Notice that, this code is working perfectly when I select the runtime type = CPU but I want to use the GPU

![Screenshot of the log file is attached 1]

Answer 1

that is because of the federation using TF 2.3.4 Hardware consuming resources ...

[ Sample ]:

import tensorflow as tf
import tensorflow_federated as tff
import tensorflow_datasets as tfds

client_data, _ = tff.simulation.datasets.emnist.load_data()
sampled_client_ids = client_data.client_ids[:5]

dataset_paths = "F:\temp\Python\client_data\client_ids.txt"

for client_id in sampled_client_ids[:5]:
    client_local_dataset = tfds.as_numpy(
        client_data.create_tf_dataset_for_client(client_id))

print( '----------------------------------------------------------------')
print( client_data )

[ Output ]:

# ['f0000_14', 'f0001_41', 'f0005_26', 'f0006_12', 'f0008_45']
# <tensorflow_federated.python.simulation.hdf5_client_data.HDF5ClientData object at 0x0000021733D03DC0>