Keras TPU Integration

@Vengineerの戯言 : Twitter
SystemVerilogの世界へようこそ、すべては、SystemC v0.9公開から始まった

なんか、いつものように、TensorFlowのソースコード眺めていたら、見つけました。

Keras TPU Integration

そして、これに関してのあたしのツイート、結構人気あった。

MNIST と ResNet-50 のサンプルコードがある。。。

どうやら、Google Cloud Next'18では、とあるセッションでこっそり発表されたんだって。

MNISTのサンプルコードを見てみると

とりあえず、kera でモデルを作って、、、

引用
  model = tf.keras.models.Sequential()
  model.add(
      tf.keras.layers.Conv2D(
          32, kernel_size=(3, 3), activation='relu', input_shape=input_shape))
  model.add(tf.keras.layers.Conv2D(64, (3, 3), activation='relu'))
  model.add(tf.keras.layers.MaxPooling2D(pool_size=(2, 2)))
  model.add(tf.keras.layers.Dropout(0.25))
  model.add(tf.keras.layers.Flatten())
  model.add(tf.keras.layers.Dense(128, activation='relu'))
  model.add(tf.keras.layers.Dropout(0.5))
  model.add(tf.keras.layers.Dense(num_classes, activation='softmax'))

そのあと、TPU を使う場合は、

引用
  if use_tpu:
    strategy = keras_support.TPUDistributionStrategy(num_cores_per_host=8)
    model = keras_support.tpu_model(model,
                                    strategy=strategy,
                                    tpu_name_or_address=flags.FLAGS.tpu)

感じに、keras_support.TPUDistributionStrategyなるAPIでホスト1つに対して、TPUを何個使うかを指定。
このは、keras_support.TPUDistributionStrategyは、

def TPUDistributionStrategy(*args, **kw):  # pylint: disable=invalid-name
  from tensorflow.contrib.distribute.python import tpu_strategy  # pylint: disable=g-import-not-at-top
  return tpu_strategy.TPUStrategy(*args, **kw)

感じになっていて、tpu_strategy.TPUStrategyを呼んでいるだけ。。。
で、TPUStrategyは、

引用
class TPUStrategy(one_device_strategy.OneDeviceStrategy):
  """Experimental TPU distribution strategy implementation."""

でクラスになっている。

サンプルコードに戻って。。。

引用
    model = keras_support.tpu_model(model,
                                    strategy=strategy,
                                    tpu_name_or_address=flags.FLAGS.tpu)

モデル(model)をTPUのモデルに変換？
tpu_modelを見てみると。。。

引用
def tpu_model(model, tpu_name_or_address=None, strategy=None):
  """Copy `model` along with weights to the TPU.  Returns a TPU model.

とあって、モデルをウェイトと共に、TPUにコピーして、TPUモデルが戻ってくると。。。

使い方は、

引用
  Args:
    model: A `KerasTPUModel`.
    tpu_name_or_address: A string that is either the name of the Cloud TPU,
      the grpc address of the Cloud TPU, or (Googlers only) the BNS name of the
      Cloud TPU. If tpu_name_or_address is None, the TPUClusterResolver will
      examine the environment to determine a potential Cloud TPU to use.
    strategy: `TPUDistributionStrategy`.  The strategy to use for replicating
              model across multiple TPU cores.
  Returns:
    A new `KerasTPUModel` instance.

とあります。

この後のコードは、compile => fit => evaluate と。。。

引用
  model.compile(
      loss=tf.keras.losses.categorical_crossentropy,
      optimizer=tf.train.GradientDescentOptimizer(learning_rate=0.05),
      metrics=['accuracy'])

  model.fit(
      x_train,
      y_train,
      batch_size=batch_size,
      epochs=epochs,
      verbose=1,
      validation_data=(x_test, y_test))
  score = model.evaluate(x_test, y_test, verbose=0)

。