parameters: (argv)
docstring: Run a CNN model on MNIST data to demonstrate DistributedStrategies.
docstring_summary: Run a CNN model on MNIST data to demonstrate DistributedStrategies.
docstring_tokens: Run a CNN model on MNIST data to demonstrate DistributedStrategies .
function: def main(argv):
    """Run a CNN model on MNIST data to demonstrate DistributedStrategies."""
    del argv  # Unused.
    tf.disable_v2_behavior()
  
    cluster_resolver = tf.distribute.cluster_resolver.TPUClusterResolver(
        tpu=FLAGS.tpu)
    strategy = tf.contrib.distribute.TPUStrategy(cluster_resolver)
  
    # TODO(rxsang): This doesn't work on cloud for some reason
    strategy.extended.experimental_enable_get_next_as_optional = False
  
    with strategy.scope():
      train_ds, test_ds = mnist_datasets()
      train_ds = train_ds.shuffle(NUM_TRAIN_IMAGES).batch(FLAGS.batch_size)
      test_ds = test_ds.batch(FLAGS.batch_size)
  
      model = create_model()
      optimizer = tf.keras.optimizers.SGD(FLAGS.learning_rate, FLAGS.momentum)
      training_loss = tf.keras.metrics.Mean("training_loss", dtype=tf.float32)
      training_accuracy = tf.keras.metrics.SparseCategoricalAccuracy(
          "training_accuracy", dtype=tf.float32)
      test_loss = tf.keras.metrics.Mean("test_loss", dtype=tf.float32)
      test_accuracy = tf.keras.metrics.SparseCategoricalAccuracy(
          "test_accuracy", dtype=tf.float32)
  
      def train_step(inputs):  # pylint: disable=missing-docstring
        images, labels = inputs
        with tf.GradientTape() as tape:
          logits = model(images, training=True)
          loss = compute_loss(logits, labels)
        grads = tape.gradient(loss, model.variables)
        update_vars = optimizer.apply_gradients(zip(grads, model.variables))
        update_loss = training_loss.update_state(loss)
        update_accuracy = training_accuracy.update_state(labels, logits)
  
        with tf.control_dependencies([update_vars, update_loss, update_accuracy]):
          return tf.identity(loss)
  
      def test_step(inputs):
        images, labels = inputs
        logits = model(images, training=False)
        loss = compute_loss(logits, labels)
        update_loss = test_loss.update_state(loss)
        update_accuracy = test_accuracy.update_state(labels, logits)
  
        with tf.control_dependencies([update_loss, update_accuracy]):
          return tf.identity(loss)
  
      train_iterator = strategy.experimental_distribute_dataset(
          train_ds).make_initializable_iterator()
      test_iterator = strategy.experimental_distribute_dataset(
          test_ds).make_initializable_iterator()
  
      dist_train = strategy.experimental_local_results(
          strategy.experimental_run_v2(train_step, args=(next(train_iterator),)))
      dist_test = strategy.experimental_local_results(
          strategy.experimental_run_v2(test_step, args=(next(test_iterator),)))
  
      training_loss_result = training_loss.result()
      training_accuracy_result = training_accuracy.result()
      test_loss_result = test_loss.result()
      test_accuracy_result = test_accuracy.result()
  
      tf.contrib.distribute.initialize_tpu_system(cluster_resolver)
  
      train_iterator_init = train_iterator.initialize()
      test_iterator_init = test_iterator.initialize()
  
      all_variables = (
          tf.global_variables() +
          training_loss.variables +
          training_accuracy.variables +
          test_loss.variables +
          test_accuracy.variables)
  
      config = tf.ConfigProto()
      cluster_spec = cluster_resolver.cluster_spec()
      if cluster_spec:
        config.cluster_def.CopyFrom(cluster_spec.as_cluster_def())
  
      with tf.Session(cluster_resolver.master(), config=config) as session:
        session.run([v.initializer for v in all_variables])
  
        for epoch in range(0, FLAGS.num_epochs):
          # Train
          print("Starting epoch {}".format(epoch))
          session.run(train_iterator_init)
          while True:
            try:
              session.run(dist_train)
            except tf.errors.OutOfRangeError:
              break
          print("Training loss: {:0.4f}, accuracy: {:0.2f}%".format(
              session.run(training_loss_result),
              session.run(training_accuracy_result) * 100))
          training_loss.reset_states()
          training_accuracy.reset_states()
  
          # Test
          session.run(test_iterator_init)
          while True:
            try:
              session.run(dist_test)
            except tf.errors.OutOfRangeError:
              break
          print("Test loss: {:0.4f}, accuracy: {:0.2f}%".format(
              session.run(test_loss_result),
              session.run(test_accuracy_result) * 100))
          test_loss.reset_states()
          test_accuracy.reset_states()
function_tokens: def main ( argv ) : del argv # Unused. tf . disable_v2_behavior ( ) cluster_resolver = tf . distribute . cluster_resolver . TPUClusterResolver ( tpu = FLAGS . tpu ) strategy = tf . contrib . distribute . TPUStrategy ( cluster_resolver ) # TODO(rxsang): This doesn't work on cloud for some reason strategy . extended . experimental_enable_get_next_as_optional = False with strategy . scope ( ) : train_ds , test_ds = mnist_datasets ( ) train_ds = train_ds . shuffle ( NUM_TRAIN_IMAGES ) . batch ( FLAGS . batch_size ) test_ds = test_ds . batch ( FLAGS . batch_size ) model = create_model ( ) optimizer = tf . keras . optimizers . SGD ( FLAGS . learning_rate , FLAGS . momentum ) training_loss = tf . keras . metrics . Mean ( "training_loss" , dtype = tf . float32 ) training_accuracy = tf . keras . metrics . SparseCategoricalAccuracy ( "training_accuracy" , dtype = tf . float32 ) test_loss = tf . keras . metrics . Mean ( "test_loss" , dtype = tf . float32 ) test_accuracy = tf . keras . metrics . SparseCategoricalAccuracy ( "test_accuracy" , dtype = tf . float32 ) def train_step ( inputs ) : # pylint: disable=missing-docstring images , labels = inputs with tf . GradientTape ( ) as tape : logits = model ( images , training = True ) loss = compute_loss ( logits , labels ) grads = tape . gradient ( loss , model . variables ) update_vars = optimizer . apply_gradients ( zip ( grads , model . variables ) ) update_loss = training_loss . update_state ( loss ) update_accuracy = training_accuracy . update_state ( labels , logits ) with tf . control_dependencies ( [ update_vars , update_loss , update_accuracy ] ) : return tf . identity ( loss ) def test_step ( inputs ) : images , labels = inputs logits = model ( images , training = False ) loss = compute_loss ( logits , labels ) update_loss = test_loss . update_state ( loss ) update_accuracy = test_accuracy . update_state ( labels , logits ) with tf . control_dependencies ( [ update_loss , update_accuracy ] ) : return tf . identity ( loss ) train_iterator = strategy . experimental_distribute_dataset ( train_ds ) . make_initializable_iterator ( ) test_iterator = strategy . experimental_distribute_dataset ( test_ds ) . make_initializable_iterator ( ) dist_train = strategy . experimental_local_results ( strategy . experimental_run_v2 ( train_step , args = ( next ( train_iterator ) , ) ) ) dist_test = strategy . experimental_local_results ( strategy . experimental_run_v2 ( test_step , args = ( next ( test_iterator ) , ) ) ) training_loss_result = training_loss . result ( ) training_accuracy_result = training_accuracy . result ( ) test_loss_result = test_loss . result ( ) test_accuracy_result = test_accuracy . result ( ) tf . contrib . distribute . initialize_tpu_system ( cluster_resolver ) train_iterator_init = train_iterator . initialize ( ) test_iterator_init = test_iterator . initialize ( ) all_variables = ( tf . global_variables ( ) + training_loss . variables + training_accuracy . variables + test_loss . variables + test_accuracy . variables ) config = tf . ConfigProto ( ) cluster_spec = cluster_resolver . cluster_spec ( ) if cluster_spec : config . cluster_def . CopyFrom ( cluster_spec . as_cluster_def ( ) ) with tf . Session ( cluster_resolver . master ( ) , config = config ) as session : session . run ( [ v . initializer for v in all_variables ] ) for epoch in range ( 0 , FLAGS . num_epochs ) : # Train print ( "Starting epoch {}" . format ( epoch ) ) session . run ( train_iterator_init ) while True : try : session . run ( dist_train ) except tf . errors . OutOfRangeError : break print ( "Training loss: {:0.4f}, accuracy: {:0.2f}%" . format ( session . run ( training_loss_result ) , session . run ( training_accuracy_result ) * 100 ) ) training_loss . reset_states ( ) training_accuracy . reset_states ( ) # Test session . run ( test_iterator_init ) while True : try : session . run ( dist_test ) except tf . errors . OutOfRangeError : break print ( "Test loss: {:0.4f}, accuracy: {:0.2f}%" . format ( session . run ( test_loss_result ) , session . run ( test_accuracy_result ) * 100 ) ) test_loss . reset_states ( ) test_accuracy . reset_states ( )