使用 TPU 运行时在 Google Colab 上训练 Keras 模型时出错

xEt*_*eal 4 machine-learning keras tensorflow google-colaboratory google-cloud-tpu

我正在尝试使用 Google Colab 中的 TPU 创建和训练我的 CNN 模型。我打算用它来对狗和猫进行分类。该模型使用 GPU/CPU 运行时运行,但在 TPU 运行时运行时遇到问题。这是创建模型的代码。

我使用 flow_from_directory() 函数输入我的数据集,这是它的代码

train_datagen = ImageDataGenerator(rescale=1./255)
train_generator = train_datagen.flow_from_directory(
    MAIN_DIR,
    target_size = (128,128),
    batch_size = 50,
    class_mode = 'binary'
)
Run Code Online (Sandbox Code Playgroud)
def create_model():

  model=Sequential()
  model.add(Conv2D(32,(3,3),activation='relu',input_shape=(128,128,3)))
  model.add(BatchNormalization())
  model.add(MaxPooling2D(pool_size=(2,2)))
  model.add(Dropout(0.25))
  model.add(Conv2D(64,(3,3),activation='relu'))
  model.add(BatchNormalization())
  model.add(MaxPooling2D(pool_size=(2,2)))
  model.add(Dropout(0.25))
  model.add(Conv2D(128,(3,3),activation='relu'))
  model.add(BatchNormalization())
  model.add(MaxPooling2D(pool_size=(2,2)))
  model.add(Dropout(0.25))
  model.add(Flatten())
  model.add(Dense(512,activation='relu'))
  model.add(BatchNormalization())
  model.add(Dropout(0.5))
  model.add(Dense(2,activation='softmax'))
  
  return model
Run Code Online (Sandbox Code Playgroud)

这是用于在 google Colab 上启动 TPU 的代码

tf.keras.backend.clear_session()

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

# This is the TPU initialization code that has to be at the beginning.
tf.tpu.experimental.initialize_tpu_system(resolver)
print("All devices: ", tf.config.list_logical_devices('TPU'))

strategy = tf.distribute.experimental.TPUStrategy(resolver)

with strategy.scope():
  model = create_model()
  model.compile(
      optimizer=tf.keras.optimizers.Adam(learning_rate=1e-3, ),
      loss='sparse_categorical_crossentropy',
      metrics=['sparse_categorical_accuracy'])


model.fit(
    train_generator, 
    epochs = 5,
    
)
Run Code Online (Sandbox Code Playgroud)

但是当我运行这段代码时,我遇到了这个错误:

UnavailableError                          Traceback (most recent call last)
<ipython-input-15-1970b3405ba3> in <module>()
     20 model.fit(
     21     train_generator,
---> 22     epochs = 5,
     23 
     24 )

14 frames
/usr/local/lib/python3.6/dist-packages/six.py in raise_from(value, from_value)

UnavailableError: 5 root error(s) found.
  (0) Unavailable: {{function_node __inference_train_function_42823}} failed to connect to all addresses
Additional GRPC error information from remote target /job:localhost/replica:0/task:0/device:CPU:0:
:{"created":"@1598016644.748265484","description":"Failed to pick subchannel","file":"third_party/grpc/src/core/ext/filters/client_channel/client_channel.cc","file_line":3948,"referenced_errors":[{"created":"@1598016644.748262999","description":"failed to connect to all addresses","file":"third_party/grpc/src/core/ext/filters/client_channel/lb_policy/pick_first/pick_first.cc","file_line":394,"grpc_status":14}]}
     [[{{node MultiDeviceIteratorGetNextFromShard}}]]
     [[RemoteCall]]
     [[IteratorGetNextAsOptional]]
     [[cond_11/switch_pred/_107/_78]]
  (1) Unavailable: {{function_node __inference_train_function_42823}} failed to connect to all addresses
Additional GRPC error information from remote target /job:localhost/replica:0/task:0/device:CPU:0:
:{"created":"@1598016644.748265484","description":"Failed to pick subchannel","file":"third_party/grpc/src/core/ext/filters/client_channel/client_channel.cc","file_line":3948,"referenced_errors":[{"created":"@1598016644.748262999","description":"failed to connect to all addresses","file":"third_party/grpc/src/core/ext/filters/client_channel/lb_policy/pick_first/pick_first.cc","file_line":394,"grpc_status":14}]}
     [[{{node MultiDeviceIteratorGetNextFromShard}}]]
     [[RemoteCall]]
     [[IteratorGetNextAsOptional]]
     [[cond_12/switch_pred/_118/_82]]
  (2) Unavailable: {{function_node __inference_train_function_42823}} failed to connect to all addresses
Additional GRPC error information from remote target /job:localhost/replica:0/task:0/device:CPU:0:
:{"created":"@1598016644.748265484","description":"Failed to pick subchannel","file":"third_party/grpc/src/core/ext/filters/client_channel/client_channel.cc","file_line":3948,"referenced_errors":[{"created":"@1598016644.748262999","description":"failed to connect to all addresses","file":"third_party/grpc/src/core/ext/filters/client_channel/lb_policy/pick_first/pick_first.cc","file_line":394,"grpc_status":14}]}
     [[{{node MultiDeviceIteratorGetNextFromShard}}]]
     [[RemoteCall]]
     [[IteratorGetNextAsOptional]]
     [[TPUReplicate/_compile/_7955920754087029306/_4/_266]]
  (3) Unavailable: {{function_node __inference_train_function_42823}} failed to connect to all addresses
Additional GRPC error information from remote target /job:localhost/replica:0/task:0/device:CPU:0:
:{"created":"@1598016644.748265484","description":"Failed to pick subchannel","file":"third_party/grpc/src/core/ext/filters/client_channel/client_channel.cc","file_line":3948,"referenced_errors":[{"created":"@1598016644.748262999","description":"failed to connect to all addresses","file":"third_party/grpc/src/core/ext/filters/client_channel/lb_policy/pick_first/pick_first.cc","file_line":394,"grpc_status":14}]}
     [[{{node MultiDeviceIteratorGetNextFromShard}}]]
     [[RemoteCall]]
     [[IteratorGetNextAsOptional]]
     [[Shape_7/_104]]
  (4) Unavailable: {{functi ... [truncated]
Run Code Online (Sandbox Code Playgroud)

我真的不知道,我该如何解决这个问题。我也不知道这些错误意味着什么。

Bjö*_*ist 8

您遇到了 TPU 的一个已知问题 - 它们不支持 PyFunction。详细信息在这里:#38762#34346#39099

抱歉这个问题。Dataset.from_generator 预计无法与 TPU 一起使用,因为它使用下面的 py_function ,这与 Cloud TPU 2VM 设置不兼容。如果您想读取大型数据集,可以尝试将其具体化在磁盘上并使用 TFRecordDataest 代替。

由于ImageDataGenerator也在底层使用 PyFunction,因此它与 TPU 不兼容。相反,您必须使用 tf.data API 来加载图像。本教程解释了如何操作。