code: def main(argv):
      del argv  # unused
      if tf.io.gfile.exists(FLAGS.model_dir):
          tf.compat.v1.logging.warning(
              "Warning: deleting old log directory at {}".format(
                  FLAGS.model_dir
              )
          )
          tf.io.gfile.rmtree(FLAGS.model_dir)
      tf.io.gfile.makedirs(FLAGS.model_dir)
  
      if FLAGS.fake_data:
          (x_train, y_train), (x_test, y_test) = build_fake_data()
      else:
          (x_train, y_train), (
              x_test,
              y_test,
          ) = tf_keras.datasets.cifar10.load_data()
  
      (
          images,
          labels,
          handle,
          training_iterator,
          heldout_iterator,
      ) = build_input_pipeline(
          x_train, x_test, y_train, y_test, FLAGS.batch_size, 500
      )
  
      if FLAGS.architecture == "resnet":
          model_fn = bayesian_resnet
      else:
          model_fn = bayesian_vgg
  
      model = model_fn(
          IMAGE_SHAPE,
          num_classes=10,
          kernel_posterior_scale_mean=FLAGS.kernel_posterior_scale_mean,
          kernel_posterior_scale_constraint=FLAGS.kernel_posterior_scale_constraint,
      )
      logits = model(images)
      labels_distribution = tfd.Categorical(logits=logits)
  
      # Perform KL annealing. The optimal number of annealing steps
      # depends on the dataset and architecture.
      t = tf.compat.v2.Variable(0.0)
      kl_regularizer = t / (
          FLAGS.kl_annealing * len(x_train) / FLAGS.batch_size
      )
  
      # Compute the -ELBO as the loss. The kl term is annealed from 0 to 1 over
      # the epochs specified by the kl_annealing flag.
      log_likelihood = labels_distribution.log_prob(labels)
      neg_log_likelihood = -tf.reduce_mean(input_tensor=log_likelihood)
      kl = (
          sum(model.losses)
          / len(x_train)
          * tf.minimum(1.0, kl_regularizer)
      )
      loss = neg_log_likelihood + kl
  
      # Build metrics for evaluation. Predictions are formed from a single forward
      # pass of the probabilistic layers. They are cheap but noisy
      # predictions.
      predictions = tf.argmax(input=logits, axis=1)
      with tf.compat.v1.name_scope("train"):
          (
              train_accuracy,
              train_accuracy_update_op,
          ) = tf.compat.v1.metrics.accuracy(
              labels=labels, predictions=predictions
          )
          opt = tf.compat.v1.train.AdamOptimizer(FLAGS.learning_rate)
          train_op = opt.minimize(loss)
          update_step_op = tf.compat.v1.assign(t, t + 1)
  
      with tf.compat.v1.name_scope("valid"):
          (
              valid_accuracy,
              valid_accuracy_update_op,
          ) = tf.compat.v1.metrics.accuracy(
              labels=labels, predictions=predictions
          )
  
      init_op = tf.group(
          tf.compat.v1.global_variables_initializer(),
          tf.compat.v1.local_variables_initializer(),
      )
  
      stream_vars_valid = [
          v
          for v in tf.compat.v1.local_variables()
          if "valid/" in v.name
      ]
      reset_valid_op = tf.compat.v1.variables_initializer(
          stream_vars_valid
      )
  
      with tf.compat.v1.Session() as sess:
          sess.run(init_op)
  
          # Run the training loop
          train_handle = sess.run(training_iterator.string_handle())
          heldout_handle = sess.run(heldout_iterator.string_handle())
          training_steps = int(
              round(FLAGS.epochs * (len(x_train) / FLAGS.batch_size))
          )
          for step in range(training_steps):
              _ = sess.run(
                  [train_op, train_accuracy_update_op, update_step_op],
                  feed_dict={handle: train_handle},
              )
  
              # Manually print the frequency
              if step % 100 == 0:
                  loss_value, accuracy_value, kl_value = sess.run(
                      [loss, train_accuracy, kl],
                      feed_dict={handle: train_handle},
                  )
                  print(
                      "Step: {:>3d} Loss: {:.3f} Accuracy: {:.3f} KL: {:.3f}".format(
                          step, loss_value, accuracy_value, kl_value
                      )
                  )
  
              if (step + 1) % FLAGS.eval_freq == 0:
                  # Compute log prob of heldout set by averaging draws from the model:
                  # p(heldout | train) = int_model p(heldout|model) p(model|train)
                  #                   ~= 1/n * sum_{i=1}^n p(heldout | model_i)
                  # where model_i is a draw from the posterior
                  # p(model|train).
                  probs = np.asarray(
                      [
                          sess.run(
                              (labels_distribution.probs),
                              feed_dict={handle: heldout_handle},
                          )
                          for _ in range(FLAGS.num_monte_carlo)
                      ]
                  )
                  mean_probs = np.mean(probs, axis=0)
  
                  _, label_vals = sess.run(
                      (images, labels),
                      feed_dict={handle: heldout_handle},
                  )
                  heldout_lp = np.mean(
                      np.log(
                          mean_probs[
                              np.arange(mean_probs.shape[0]),
                              label_vals.flatten(),
                          ]
                      )
                  )
                  print(" ... Held-out nats: {:.3f}".format(heldout_lp))
  
                  # Calculate validation accuracy
                  for _ in range(20):
                      sess.run(
                          valid_accuracy_update_op,
                          feed_dict={handle: heldout_handle},
                      )
                  valid_value = sess.run(
                      valid_accuracy, feed_dict={handle: heldout_handle}
                  )
  
                  print(
                      " ... Validation Accuracy: {:.3f}".format(
                          valid_value
                      )
                  )
  
                  sess.run(reset_valid_op)
  
file_path: probability/tensorflow_probability/examples/cifar10_bnn.py
imports: import os
  import warnings
  from absl import flags
  import matplotlib
  import numpy
  import tensorflow.compat.v1
  import tensorflow_probability
  from tensorflow_probability.examples.models.bayesian_resnet import bayesian_resnet
  from tensorflow_probability.examples.models.bayesian_vgg import bayesian_vgg
  from tensorflow_probability.python.internal import tf_keras
label_desc: Use image node_id 3 for calling a global function with example usage: main(argv) without return types