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:
          train_seq = MNISTSequence(
              batch_size=FLAGS.batch_size,
              fake_data_size=NUM_TRAIN_EXAMPLES,
          )
          heldout_seq = MNISTSequence(
              batch_size=FLAGS.batch_size,
              fake_data_size=NUM_HELDOUT_EXAMPLES,
          )
      else:
          train_set, heldout_set = tf_keras.datasets.mnist.load_data()
          train_seq = MNISTSequence(
              data=train_set, batch_size=FLAGS.batch_size
          )
          heldout_seq = MNISTSequence(
              data=heldout_set, batch_size=FLAGS.batch_size
          )
  
      model = create_model()
      # TODO(b/149259388): understand why Keras does not automatically build the
      # model correctly.
      model.build(input_shape=[None, 28, 28, 1])
  
      print(" ... Training convolutional neural network")
      for epoch in range(FLAGS.num_epochs):
          epoch_accuracy, epoch_loss = [], []
          for step, (batch_x, batch_y) in enumerate(train_seq):
              batch_loss, batch_accuracy = model.train_on_batch(
                  batch_x, batch_y
              )
              epoch_accuracy.append(batch_accuracy)
              epoch_loss.append(batch_loss)
  
              if step % 100 == 0:
                  print(
                      "Epoch: {}, Batch index: {}, "
                      "Loss: {:.3f}, Accuracy: {:.3f}".format(
                          epoch,
                          step,
                          tf.reduce_mean(epoch_loss),
                          tf.reduce_mean(epoch_accuracy),
                      )
                  )
  
              if (step + 1) % FLAGS.viz_steps == 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).
                  print(" ... Running monte carlo inference")
                  probs = tf.stack(
                      [
                          model.predict(heldout_seq, verbose=1)
                          for _ in range(FLAGS.num_monte_carlo)
                      ],
                      axis=0,
                  )
                  mean_probs = tf.reduce_mean(probs, axis=0)
                  heldout_log_prob = tf.reduce_mean(
                      tf.math.log(mean_probs)
                  )
                  print(
                      " ... Held-out nats: {:.3f}".format(
                          heldout_log_prob
                      )
                  )
  
                  if HAS_SEABORN:
                      names = [
                          layer.name
                          for layer in model.layers
                          if "flipout" in layer.name
                      ]
                      qm_vals = [
                          layer.kernel_posterior.mean().numpy()
                          for layer in model.layers
                          if "flipout" in layer.name
                      ]
                      qs_vals = [
                          layer.kernel_posterior.stddev().numpy()
                          for layer in model.layers
                          if "flipout" in layer.name
                      ]
                      plot_weight_posteriors(
                          names,
                          qm_vals,
                          qs_vals,
                          fname=os.path.join(
                              FLAGS.model_dir,
                              "epoch{}_step{:05d}_weights.png".format(
                                  epoch, step
                              ),
                          ),
                      )
                      plot_heldout_prediction(
                          heldout_seq.images,
                          probs.numpy(),
                          fname=os.path.join(
                              FLAGS.model_dir,
                              "epoch{}_step{}_pred.png".format(
                                  epoch, step
                              ),
                          ),
                          title="mean heldout logprob {:.2f}".format(
                              heldout_log_prob
                          ),
                      )
  
imports: import os
  import warnings
  from absl import app
  from absl import flags
  import matplotlib
  from matplotlib import figure
  from matplotlib.backends import backend_agg
  import numpy
  import tensorflow.compat.v2
  import tensorflow_probability
  from tensorflow_probability.python.internal import tf_keras
label_desc: Use image node_id 4 for calling a global function with example usage: main(argv) without return types