Почему Google colab TPU медленный?

Я использую Talos для настройки гиперпараметров модели Keras. Запуск этого короткого кода в Google colab TPU выполняется очень медленно. Я думаю, это как-то связано с типом данных. Должен ли я преобразовать его в тензоры, чтобы ускорить TPU?

%tensorflow_version 2.x
import os
import tensorflow as tf
import talos as ta
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Dense
from tensorflow.keras.optimizers import Adam
from sklearn.model_selection import train_test_split

def iris_model(x_train, y_train, x_val, y_val, params):

    # Specify a distributed strategy to use TPU
    resolver = tf.distribute.cluster_resolver.TPUClusterResolver(tpu='grpc://' + os.environ['COLAB_TPU_ADDR'])
    tf.config.experimental_connect_to_host(resolver.master())
    tf.tpu.experimental.initialize_tpu_system(resolver)
    strategy = tf.distribute.experimental.TPUStrategy(resolver)

    # Use the strategy to create and compile a Keras model
    with strategy.scope():
      model = Sequential()
      model.add(Dense(32, input_shape=(4,), activation=tf.nn.relu, name="relu"))
      model.add(Dense(3, activation=tf.nn.softmax, name="softmax"))
      model.compile(optimizer=Adam(learning_rate=0.1), loss=params['losses'])

    # Convert data type to use TPU
    x_train = x_train.astype('float32')
    x_val = x_val.astype('float32')

    # Fit the Keras model on the dataset
    out = model.fit(x_train, y_train, batch_size=params['batch_size'], epochs=params['epochs'], validation_data=[x_val, y_val], verbose=0, steps_per_epoch=0)

    return out, model

# Load dataset
X, y = ta.templates.datasets.iris()

# Train and test set
x_train, x_val, y_train, y_val = train_test_split(X, y, test_size=0.30, shuffle=False)

# Create a hyperparameter distributions 
p = {'losses': ['logcosh'], 'batch_size': [128, 256, 384, 512, 1024], 'epochs': [10, 20]}

# Use Talos to scan the best hyperparameters of the Keras model
scan_object = ta.Scan(x_train, y_train, params=p, model=iris_model, experiment_name='test', x_val=x_val, y_val=y_val, fraction_limit=0.5)

keras tensorflow google-colaboratory google-cloud-tpu talos
person Sami Belkacem    schedule 12.02.2020    source источник
comment
Возможно, вам повезет больше на ai.stackexchange.com   -  person Mischa    schedule 30.03.2020

Ответы (1)


arrow_upward
3
arrow_downward

Спасибо Вам за Ваш вопрос.

К сожалению, мне не удалось запустить ваш образец кода на TensorFlow 2.2, поэтому я не знаю, какую производительность вы видели изначально. Мне удалось исправить это и запустить на TPU со следующими изменениями:

  • Заменить tf.config.experimental_connect_to_host(resolver.master()) на tf.config.experimental_connect_to_cluster(resolver)
  • Перенести инициализацию TPU за пределы iris_model().
  • Используйте tf.data.Dataset для ввода TPU.

Вот модифицированный код Colab:

# Run this to install Talos before running the rest of the code.
!pip install git+https://github.com/autonomio/[email protected]

%tensorflow_version 2.x
import os
import tensorflow as tf
import talos as ta
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Dense
from tensorflow.keras.optimizers import Adam
from sklearn.model_selection import train_test_split

print(tf.__version__) # TF 2.2.0 in my case

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

def iris_model(x_train, y_train, x_val, y_val, params):
    # Use the strategy to create and compile a Keras model
    strategy = tf.distribute.experimental.TPUStrategy(resolver)
    with strategy.scope():
      model = Sequential()
      model.add(Dense(32, input_shape=(4,), activation=tf.nn.relu, name="relu"))
      model.add(Dense(3, activation=tf.nn.softmax, name="softmax"))
      model.compile(optimizer=Adam(learning_rate=0.1), loss=params['losses'])

    train_dataset = tf.data.Dataset.from_tensor_slices((x_train, y_train)).batch(params['batch_size'])
    val_dataset = tf.data.Dataset.from_tensor_slices((x_val, y_val)).batch(params['batch_size'])

    # Fit the Keras model on the dataset
    out = model.fit(train_dataset, epochs=params['epochs'], validation_data=val_dataset)

    return out, model

# Load dataset
X, y = ta.templates.datasets.iris()

# Train and test set
x_train, x_val, y_train, y_val = train_test_split(X, y, test_size=0.30, shuffle=False)

# Create a hyperparameter distributions 
p = {'losses': ['logcosh'], 'batch_size': [128, 256, 384, 512, 1024], 'epochs': [10, 20]}

# Use Talos to scan the best hyperparameters of the Keras model
scan_object = ta.Scan(x_train, y_train, params=p, model=iris_model, experiment_name='test', x_val=x_val, y_val=y_val, fraction_limit=0.5)

Для меня последний звонок занял чуть меньше 2 минут.

Для хорошо известных наборов данных вы можете пропустить этап создания собственных tf.data.Dataset с помощью библиотеки наборов данных TensorFlow . TFDS действительно имеет в своей библиотеке набор данных iris. Полный пример использования TFDS с TPU см. В официальном руководстве по TPU TensorFlow.

person Will Cromar    schedule 08.06.2020