Не работает базовый оптимизатор перекрестных осколков ТПУ - PullRequest
Новая
       24

Не работает базовый оптимизатор перекрестных осколков ТПУ

1 голос
/ 14 марта 2019

В общем, есть несколько хороших примеров, в которых используются TF-оптимизаторы для решения общих проблем (без глубокого обучения). Дано:

https://databricks.com/tensorflow/training-and-convergence https://colab.research.google.com/notebooks/tpu.ipynb#scrollTo=a_rjVo-RAoYd

Мы хотим иметь возможность объединить два вышеупомянутых подхода и использовать оптимизацию на основе TPU для решения задач больших размеров.

Для этого у меня есть простой код colab, который делает это, объединяя два примера выше: 

import tensorflow as tf
import numpy as np
from tensorflow.contrib.tpu.python.tpu import tpu_function
import os
import pprint
import tensorflow as tf

if 'COLAB_TPU_ADDR' not in os.environ:
  print('ERROR: Not connected to a TPU runtime; please see the first cell in this notebook for instructions!')
else:
  tpu_address = 'grpc://' + os.environ['COLAB_TPU_ADDR']
  print ('TPU address is', tpu_address)

  with tf.Session(tpu_address) as session:
    devices = session.list_devices()

  print('TPU devices:')
  pprint.pprint(devices)

# Add this somewhere at the top
tpu_function.get_tpu_context().set_number_of_shards(8)

# x and y are placeholders for our training data
x = tf.placeholder("float")
y = tf.placeholder("float")
# w is the variable storing our values. It is initialised with starting "guesses"
# w[0] is the "a" in our equation, w[1] is the "b"
w = tf.Variable([1.0, 2.0,3.0, 4.0], name="w")
# Our model of y = a*x + b
y_model = tf.multiply(x, w[0]) + w[1] + w[2] +3

# Our error is defined as the square of the differences
error = tf.square(y - y_model)
# The Gradient Descent Optimizer does the heavy lifting
train_op = tf.train.AdamOptimizer(0.01)
optimizer = tf.contrib.tpu.CrossShardOptimizer(train_op).minimize(error) # TPU change 1





# Normal TensorFlow - initialize values, create a session and run the model
model = tf.global_variables_initializer()

with tf.Session(tpu_address) as session:
    session.run(tf.contrib.tpu.initialize_system())
    print('init')
    session.run(model)
    for i in range(10000):
        print(i)
        x_value = np.random.rand()
        y_value = x_value * 2 + 6 + 5 + 3
        session.run(optimizer, feed_dict={x: x_value, y: y_value})

    w_value = session.run(w)
    print("Predicted model: {a:.3f}x + {b:.3f}+{c:.3f}x + {d:.3f}".format(a=w_value[0], b=w_value[1], c=w_value[2], d=w_value[3]))
    session.run(tpu.shutdown_system())

Когда я запускаю его (в colab) как есть, он просто запускает первую печать цикла: 

init
0

, а затем ничего не делает, и колаб просто продолжает охватывать.

Если я не использую 

optimizer = tf.contrib.tpu.CrossShardOptimizer(train_op).minimize(error)

И другие функции TPU, тогда он отлично работает, оценивая переменную w.

Вопросы:

  1. Почему это не работает и как мы можем заставить репликатор кросс-шард оптимизировать эту простую функцию?
  2. Как мне сформировать переменную w, чтобы использовать параллельные партии / осколки на TPU?
  3. Как мы можем сделать это еще более эффективным путем использования эквивалентной операции набора данных prefetch или использования очередей подачи?

Цель состоит в том, чтобы использовать lower level API-интерфейсы TPU без TPUEstimator, например, чтобы помочь в решении нестандартных задач за счет использования мощности TPU с использованием только тензоров, очередей и сегментов.

...