Ниже приведен код, который я использую.Я закомментировал строку для преобразования моей модели в модель TPU.При использовании GPU для того же объема данных на эпоху уходит 7 секунд, а при использовании TPU - 90 секунд.
Inp = tf.keras.Input(name='input', shape=(input_dim,), dtype=tf.float32)
x = tf.keras.layers.Dense(900, kernel_initializer='uniform', activation='relu', input_dim=input_dim, name = 'Dense_01')(Inp)
x = tf.keras.layers.Dropout(0.3, name = 'Dropout_02')(x)
output = tf.keras.layers.Dense(stop_criteria, activation='softmax',name = 'Dense_02')(x)
model = tf.keras.Model(inputs=[Inp], outputs=[output])
opt = tf.train.AdamOptimizer(.001)
model.compile(optimizer=opt, loss='categorical_crossentropy', metrics=['acc'])
'''tpu_model = tf.contrib.tpu.keras_to_tpu_model(model,
strategy=tf.contrib.tpu.TPUDistributionStrategy(
tf.contrib.cluster_resolver.TPUClusterResolver(TPU_ADDRESS)))'''
model.fit(X_tra, y_tra, epochs=5, batch_size=batch_size, shuffle=False,
validation_split=0.1, verbose=2)
Вот ссылка на блокнот