问题描述
我正在尝试了解使用 TPU 和 GPU 训练模型之间的区别。
这是训练模型部分:
import time
start = time.time()
tf.keras.backend.clear_session()
resolver = tf.distribute.cluster_resolver.TPUClusterResolver(tpu='grpc://' + os.environ['COLAB_TPU_ADDR'])
tf.config.experimental_connect_to_cluster(resolver)
#TPU initialization
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():
training_model = lstm_model(seq_len=100,stateful=False)
training_model.compile(
optimizer=tf.keras.optimizers.RMSprop(learning_rate=0.01),loss='sparse_categorical_crossentropy',metrics=['sparse_categorical_accuracy'])
training_model.fit(
input_fn(),steps_per_epoch=100,epochs=10
)
training_model.save_weights('/tmp/bard.h5',overwrite=True)
end = time.time()
elapsed_TPU = end - start
print(elapsed_TPU)
代码的顶部是用于 TPU 初始化。有什么办法可以改变它,让它适应在 GPU 上运行吗?
解决方法
除非您有 TPU 或多个 CPU/GPU,否则您不需要使用 tf.distribute.Strategy
。见here。您可以在没有策略的情况下将其作为标准 Tensorflow 代码运行。
import time
start = time.time()
tf.keras.backend.clear_session()
training_model = lstm_model(seq_len=100,stateful=False)
training_model.compile(
optimizer=tf.keras.optimizers.RMSprop(learning_rate=0.01),loss='sparse_categorical_crossentropy',metrics=['sparse_categorical_accuracy'])
training_model.fit(
input_fn(),steps_per_epoch=100,epochs=10
)
training_model.save_weights('/tmp/bard.h5',overwrite=True)
end = time.time()
elapsed_TPU = end - start
print(elapsed_TPU)