TensorFlow TPU Demos

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

TPU Demosがアップされています。

この論文も関係あるのかな。
TensorFlow Estimators: Managing Simplicity vs. Flexibility in High-Level Machine Learning Frameworks
model_fnってあるからそうなんだろうね。

Cloud TPUsで実行できるモデルは、以下の5つ。

　・alexnet
　・cifar	
　・inception_v3	
　・mnist	
　・resnet	

alexnetのメイン部のコードは、こんな感じ。

引用
def main(unused_argv):
  del unused_argv  # Unused.

  tf.logging.set_verbosity(tf.logging.INFO)

  run_config = tpu_config.RunConfig(
      master=FLAGS.master,
      model_dir=FLAGS.model_dir,
      save_checkpoints_secs=FLAGS.save_checkpoints_secs,
      session_config=tf.ConfigProto(
          allow_soft_placement=True, log_device_placement=True),
      tpu_config=tpu_config.TPUConfig(FLAGS.iterations, FLAGS.num_shards)
  )
  estimator = tpu_estimator.TPUEstimator(
      model_fn=model_fn,
      use_tpu=FLAGS.use_tpu,
      config=run_config,
      train_batch_size=FLAGS.batch_size)
  estimator.train(input_fn=input_fn, max_steps=FLAGS.train_steps)


if __name__ == "__main__":
  tf.app.run()

tpu_config.RuConfigでrun_configを、tpu_estimator.TPUEstimatorで、estimatorを作り、
estimator.trainで学習。。。

モデル作成の model_fn で

引用
  # Configuring the optimization algorithm.
  learning_rate = tf.train.exponential_decay(
      FLAGS.learning_rate, tf.train.get_global_step(), 25000, 0.97)
  if FLAGS.use_tpu:
    optimizer = tpu_optimizer.CrossShardOptimizer(
        tf.train.GradientDescentOptimizer(learning_rate=learning_rate))
  else:
    optimizer = tf.train.GradientDescentOptimizer(learning_rate=learning_rate)

  train_op = optimizer.minimize(loss, global_step=tf.train.get_global_step())

  return tf.estimator.EstimatorSpec(
      mode=mode,
      loss=loss,
      train_op=train_op)

のように、optimizerをtpu_optimizerにて生成しているところがポイント。

他のモデルも基本的には、上記のようにやっているみたい。。。