<jupyter_start><jupyter_text>Train point-estimate model<jupyter_code>def build_and_compile_point_model():
  model = tf.keras.Sequential([
      tf.keras.layers.Conv2D(filters=64, kernel_size=(5,5), padding='same',
                             activation='relu', input_shape=(28, 28, 1)),
      tf.keras.layers.MaxPool2D(strides=2),
      tf.keras.layers.Flatten(),
      tf.keras.layers.Dense(NUM_CLASSES)
  ])
  model.compile(optimizer='adam',
                loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
                metrics=['accuracy'])
  return model
print("Train point model")
model = build_and_compile_point_model()
model.fit(train_images, train_labels, epochs=10)<jupyter_output>Train point model
Epoch 1/2
1875/1875 [==============================] - 55s 29ms/step - loss: 0.7798 - accuracy: 0.8544
Epoch 2/2
1875/1875 [==============================] - 52s 28ms/step - loss: 0.3252 - accuracy: 0.8854<jupyter_text>Train Bayesian model with Reparameterization trick<jupyter_code>def define_bayesian_model():

  kl_divergence_function = (lambda q, p, _: tfp.distributions.kl_divergence(q, p) / tf.cast(NUM_TRAIN_EXAMPLES, dtype=tf.float32))

  model = tf.keras.models.Sequential([
      tfp.layers.Convolution2DReparameterization(
          64, kernel_size=5, padding='SAME',
          kernel_divergence_fn=kl_divergence_function,
          activation=tf.nn.relu),
      tf.keras.layers.MaxPooling2D(
          pool_size=[2, 2], strides=[2, 2],
          padding='SAME'),
      tf.keras.layers.Flatten(),
      tfp.layers.DenseReparameterization(
          NUM_CLASSES, kernel_divergence_fn=kl_divergence_function,
          activation=tf.nn.softmax)
  ])
  return model
def compile_bayesian_model(model):
  # Model compilation.
  optimizer = tf.keras.optimizers.Adam()
  # We use the categorical_crossentropy loss since the MNIST dataset contains
  # ten labels. The Keras API will then automatically add the
  # Kullback-Leibler divergence (contained on the individual layers of
  # the model), to the cross entropy loss, effectively
  # calcuating the (negated) Evidence Lower Bound Loss (ELBO)
  model.compile(optimizer, loss='categorical_crossentropy',
                metrics=['accuracy'], experimental_run_tf_function=False)

  model.build(input_shape=[None, 28, 28, 1])
  return model

train_labels_dense = tf.one_hot(train_labels, NUM_CLASSES)
train_labels_dense.shape
# use helper function to define the model architecture
bayesian_model = define_bayesian_model()
# use helper function to compile the model
bayesian_model = compile_bayesian_model(bayesian_model)
# initiate model training
bayesian_model.fit(train_images, train_labels_dense, epochs=10)
NUM_SAMPLES_INFERENCE = 50
softmax_predictions = tf.stack(
    [bayesian_model.predict(test_images[:NUM_SAMPLES_INFERENCE])
     for _ in range(NUM_INFERENCES)],axis=0)
# get the class predictions for the first image in the test set
image_ind = 0
# collect class predictions
class_predictions = []
for ind in range(NUM_INFERENCES):
  prediction_this_inference = np.argmax(softmax_predictions[ind][image_ind])
  class_predictions.append(prediction_this_inference)
# get class predictions in human-readable form
predicted_classes = [CLASS_NAMES[ind] for ind in class_predictions]
# define image caption
image_caption = []
for caption in range(NUM_INFERENCES):
  image_caption.append(f"Sample {caption+1}: {predicted_classes[caption]}\n")
image_caption = ' '.join(image_caption)
# visualise image and predictions
plt.figure(dpi=300)
plt.title(f"Correct class: {CLASS_NAMES[test_labels[image_ind]]}")
plt.imshow(test_images[image_ind], cmap=plt.cm.binary)
plt.xlabel(image_caption)
plt.show()
# calculate variance across model predictions
var_predictions = tf.reduce_mean(
    tf.math.reduce_variance(softmax_predictions, axis=0),
    axis=1)
# print variance
print(var_predictions[0].numpy())
# load MNIST to get OOD data
(train_images_mnist, train_labels_mnist), (test_images_mnist, test_labels_mnist) = tf.keras.datasets.mnist.load_data()
# get model predictions in MNIST data
softmax_predictions_mnist = tf.stack([bayesian_model.predict(test_images_mnist[:NUM_SAMPLES_INFERENCE])
                                      for _ in range(NUM_INFERENCES)], axis=0)
# calculate variance across model predictions in MNIST data
var_predictions_mnist = tf.reduce_mean(
    tf.math.reduce_variance(softmax_predictions_mnist, axis=0),
    axis=1)
# visualise comparison of MNIST fashion and MNIST variance
fig, ax = plt.subplots(nrows=1, ncols=2, dpi=200)
ax[0].plot(np.arange(NUM_SAMPLES_INFERENCE), var_predictions)
ax[0].set_ylim([0, 0.07])
ax[0].set_title("Fashion MNIST data")
ax[1].plot(np.arange(NUM_SAMPLES_INFERENCE), var_predictions_mnist)
ax[1].set_ylim([0, 0.07])
ax[1].set_title("MNIST data")
plt.show()<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.11855990068280571, 'fraction_numerical': 0.023380922822263603, 'mean_word_length': 5.118987341772152, 'pattern_counts': {'":': 0, '<': 8, '<?xml version=': 0, '>': 8, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '2', 'dir': 'jupyter-structured-clean-dedup', 'id': '548854', 'license': 'MIT', 'n_tokens_mistral': 1736, 'n_tokens_neox': 1614, 'n_words': 319, 'path': 'ch05/bbb/bbb_example.ipynb', 'repo_name': 'PacktPublishing/Bayesian-Deep-Learning', 'size': '4833'}