Colab Pro:使用GPU会导致会话崩溃

Question

我订阅了Google Colab，以使用其GPU并加快我的模型的训练。当我运行代码时，它可以工作一些块，然后它停止并说(你的会话因未知原因崩溃(执行在这个特定的块停止，它是TensorFlow联邦包中的一个函数。

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)

注意，当我选择运行时类型=CPU，但我想使用GPU时，这段代码工作得很好

[附日志文件截图1]

Answer 1

这是因为使用TF 2.3.4的联合硬件消耗资源。。。

[样本]：

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

[输出]：

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