与GPU / CPU相比,TPU keras回归非常慢

时间:2020-09-10 18:42:34

标签: tensorflow keras google-colaboratory tpu google-cloud-tpu

我正在对数据框中的列进行回归。当我使用CPU时,每个时期约为95秒,而使用GPU时,则约为45秒,但是当使用TPU时,每个时期超过 8分钟

我基本上初始化了tpu,包装了模型定义并编译为TPU分配策略。

我*认为*问题出在我的数据集中。我已经看到了将数据放入张量的教程(出于我的gpu / cpu性能,我正在发送数据帧(下面的代码中为X_train和y_train)。我尝试同时处理数据帧和张量,两者都比cpu。我确定这是一个用户错误,我看不到我的错误。

这是我的代码:

#setup tpu 

import os
resolver = tf.distribute.cluster_resolver.TPUClusterResolver(tpu='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'))
tpu_strategy = tf.distribute.TPUStrategy(resolver)


def KerasRegression(FullDF, variableToPredict):
  
  df_train1 = FullDF[FullDF[variableToPredict].notna()].copy() #lets make train data not have na for variable we are trying to predict

  X_train = df_train1.drop(variableToPredict, axis=1)
  y_train = df_train1[variableToPredict].copy()
  x_train_shape = X_train.shape[1]

  dateset=tf.data.Dataset.from_tensor_slices((X_train, y_train)).batch(batch_size=100).prefetch(buffer_size=5000)
  
  activationLayer = 'relu'

  with tpu_strategy.scope():
    model = Sequential()
    model.add(Dense(x_train_shape, activation=activationLayer, input_dim=x_train_shape))
    model.add(Dense(x_train_shape, activation=activationLayer))
    model.add(Dense(1, activation='linear'))
    optimizer = tf.keras.optimizers.Adam()
    model.compile(loss='mse', optimizer=optimizer, metrics=['mse'],
                  experimental_steps_per_execution = 50)
  model.fit(dateset, epochs=100) 
  # model.fit(X_train, y_train, epochs=100) 
  return model

如果对我的测试数据的形状有帮助,还可以:

(590543, 209)

欢迎任何反馈!

0 个答案:

没有答案