<jupyter_start><jupyter_text>Deep GamblersEnd-to-end check of DeepGamblers https://arxiv.org/abs/1907.00208.  Check only MNIST experiment which is not provided on [the original implementation](https://github.com/Z-T-WANG/NIPS2019DeepGamblers).<jupyter_code>from __future__ import absolute_import, division, print_function, unicode_literals

try:
    # %tensorflow_version only exists in Colab.
    %tensorflow_version 2.x
except Exception:
    pass

import tensorflow as tf
tf.__version__
from tensorflow.keras import layers
import numpy as np<jupyter_output><empty_output><jupyter_text>MNIST data<jupyter_code>(train_images, train_labels), (test_images, test_labels) = tf.keras.datasets.mnist.load_data()
train_images = train_images.reshape(train_images.shape[0], 28, 28, 1).astype('float32')
test_images = test_images.reshape(test_images.shape[0], 28, 28, 1).astype('float32')

# Normalizing the images to the range of [0., 1.]
train_images /= 255.
test_images /= 255.
train_images.shape
TRAIN_BUF = 60000
TEST_BUF = 10000
BATCH_SIZE = 100
train_dataset = tf.data.Dataset.zip((
    tf.data.Dataset.from_tensor_slices(train_images),
    tf.data.Dataset.from_tensor_slices(tf.one_hot(train_labels, depth=10))
))

test_dataset = tf.data.Dataset.zip((
    tf.data.Dataset.from_tensor_slices(test_images),
    tf.data.Dataset.from_tensor_slices(tf.one_hot(test_labels, depth=10))
))
import matplotlib.pyplot as plt

plt.imshow(train_images[0].reshape((28, 28)), cmap="gray")<jupyter_output><empty_output><jupyter_text>ModelThe original paper uses `a simple 4-layer CNN`, here I use simple 2 conv + 1 dense network.  I found this simpler model shows better results.<jupyter_code>class CNN(tf.keras.Model):
    """Simple 2 cov layers model"""

    def __init__(self, o=5.3, output_dim=11, name='cnn', **kwargs) -> None:
        super(CNN, self).__init__(name=name, **kwargs)
        self.o = o
        self.output_dim = output_dim
        self.input_layer = layers.InputLayer(input_shape=(28, 28, 1))
        self.conv2d_1 = layers.Conv2D(filters=32, kernel_size=3, strides=(1, 1), activation='relu')
        self.conv2d_2 = layers.Conv2D(filters=32, kernel_size=3, strides=(1, 1), activation='relu')
        self.flatten = layers.Flatten()
        self.dense = layers.Dense(self.output_dim)
        self.prob = layers.Softmax()

    def call(self, inputs):
        x = self.input_layer(inputs)
        x = self.conv2d_1(x)
        x = self.conv2d_2(x)
        x = self.flatten(x)
        x = self.dense(x)
        prob = self.prob(x)
        return prob
    
    def weighted_prob(self, inputs):
        prob = self.call(inputs)
        class_pred, abstention = tf.split(prob, [self.output_dim - 1, 1], 1)
        abstention /= self.o
        return tf.concat([class_pred, abstention], 1)
def compute_loss(model, x, y):
    weighted_prob = model.weighted_prob(x)
    label_shape = y.shape
    extended_label = tf.concat([y, tf.constant(1.0, shape=[label_shape[0], 1])], 1)
    log_arg = tf.reduce_sum(extended_label * weighted_prob, 1)
    cross_ent = -tf.reduce_sum(tf.math.log(log_arg))
    return cross_ent
def train(model, optimizer, test_loss, train_dataset, test_dataset):
    for step, (x_batch_train, y_batch_train) in enumerate(train_dataset):
        with tf.GradientTape() as tape:
            loss = compute_loss(model, x_batch_train, y_batch_train)

        grads = tape.gradient(loss, model.trainable_weights)
        optimizer.apply_gradients(zip(grads, model.trainable_weights))

    for (x_batch_test, y_batch_test) in test_dataset:
        test_loss(compute_loss(model, x_batch_test, y_batch_test))<jupyter_output><empty_output><jupyter_text>TrainingSince the test accuracy itslef is not a matter of this notebook,print the test accuracy during the training just for checking.<jupyter_code>import time
train_dataset = tf.data.Dataset.zip((
    tf.data.Dataset.from_tensor_slices(train_images),
    tf.data.Dataset.from_tensor_slices(tf.one_hot(train_labels, depth=10))
)).shuffle(TRAIN_BUF).batch(BATCH_SIZE)

test_dataset = tf.data.Dataset.zip((
    tf.data.Dataset.from_tensor_slices(test_images),
    tf.data.Dataset.from_tensor_slices(tf.one_hot(test_labels, depth=10))
)).batch(BATCH_SIZE)
start = time.time()
cnn = CNN()
optimizer = tf.keras.optimizers.Adam(learning_rate=1e-4)

for epoch in range(3):
    print(f"Start of epoch {epoch + 1}")
    test_loss = tf.keras.metrics.Mean()
    train(cnn, optimizer, test_loss, train_dataset, test_dataset)
    print(f"  test loss: {test_loss.result()}")

end = time.time()
print(f"TRAINING TIME: {end - start} [sec]")<jupyter_output>Start of epoch 1
  test loss: 23.60125732421875
Start of epoch 2
  test loss: 16.859041213989258
Start of epoch 3
  test loss: 11.554091453552246
TRAINING TIME: 374.09818053245544 [sec]<jupyter_text>Check the trained model<jupyter_code>%%time

predictions = np.array([], dtype=np.float32).reshape(0, 11)

for (x_batch_test, y_batch_test) in test_dataset:
    pred = cnn(x_batch_test)
    predictions = np.vstack([predictions, pred])
predictions.shape
abstention_scores = []

for idx, elem in enumerate(predictions):
    abstention_scores.append((elem[-1], idx))

abstention_scores = sorted(abstention_scores, reverse=True)
def draw_digit(data, title, row, col, n):
    plt.subplot(row, col, n)
    plt.imshow(data.reshape(28, 28), cmap="gray")
    plt.title(title)
    plt.axis("off")

show_size = 10
plt.figure(figsize=(15,20))
total = 0
for i, (score, idx) in enumerate(abstention_scores[:10]):
    draw_digit(test_images[idx], str(score), 1, show_size, total+1)
    total += 1
plt.show()<jupyter_output><empty_output><jupyter_text>Top 10 high abstention score images.  It seems the abstention mechanism works well.<jupyter_code>def acc_with_thld(thld):
    result = []
    for pred, ans in zip(predictions, test_labels):
        if pred[-1] > thld:
            continue
        pred_label = np.argmax(pred)
        result.append(pred_label == ans)
    coverage = len(result) / len(test_labels)
    acc = sum(result) / len(result)
    return coverage, acc
for thld in np.arange(1.0, 0.0, -0.1):
    coverage, acc = acc_with_thld(thld)
    print(f"Threshold: {thld:.1f}, Coverage : {coverage:.4f}, Accuracy : {acc:.4f}")<jupyter_output>Threshold: 1.0, Coverage : 1.0000, Accuracy : 0.9654
Threshold: 0.9, Coverage : 1.0000, Accuracy : 0.9654
Threshold: 0.8, Coverage : 0.9998, Accuracy : 0.9656
Threshold: 0.7, Coverage : 0.9996, Accuracy : 0.9658
Threshold: 0.6, Coverage : 0.9995, Accuracy : 0.9659
Threshold: 0.5, Coverage : 0.9992, Accuracy : 0.9662
Threshold: 0.4, Coverage : 0.9991, Accuracy : 0.9663
Threshold: 0.3, Coverage : 0.9986, Accuracy : 0.9668
Threshold: 0.2, Coverage : 0.9969, Accuracy : 0.9677
Threshold: 0.1, Coverage : 0.9929, Accuracy : 0.9693<jupyter_text>Rotated "9" experiment<jupyter_code>from tensorflow.keras.preprocessing.image import apply_affine_transform
nines_idx = []

for idx, num in enumerate(test_labels):
    if num == 9:
        nines_idx.append(idx)
nines_idx[:5]
plt.imshow(test_images[16].reshape(28, 28), cmap="gray")
def rotation(img, degree):
    rotated_img = apply_affine_transform(img, channel_axis=2, theta=degree, fill_mode="nearest", cval=0.)
    return rotated_img
def draw_digit(data, score, label, row, col, n):
    plt.subplot(row, col, n)
    plt.imshow(data.reshape(28, 28), cmap="gray")
    plt.title(score)
    plt.xlabel("predict: " + label)
    # plt.axis("off")

show_size = 10
plt.figure(figsize=(15,20))

total = 0
for i, degree in enumerate(range(0, 200, 20)):
    rotated_img = rotation(test_images[16], degree)
    pred = cnn(tf.expand_dims(rotated_img, 0)).numpy().squeeze()
    label = np.argmax(pred)
    confidence = np.max(pred)
    abstention_score = pred[-1].round(5)
    draw_digit(rotated_img, str(abstention_score), str(label), 1, show_size, total+1)

    total += 1
plt.show()<jupyter_output><empty_output><jupyter_text>NOTE: Class `10` is the abstention class.<jupyter_code><jupyter_output><empty_output><jupyter_text>Trial and Errors<jupyter_code>a = [1, 2, 3, 4, 5]
a[:-1]
a[-1]
tf.one_hot(train_labels, depth=10)[0:3]
test_model = CNN(1.2)
check_data = tf.constant(train_images[0])
check_data.shape
tf.shape(tf.expand_dims(check_data, 0))  # Add batch dim.
prob = test_model(tf.expand_dims(check_data, 0))
prob.shape
abstention = test_model.weighted_prob(tf.expand_dims(check_data, 0))
a = tf.constant([[0,1,2,3,4,5,6,7,8,9,10]], dtype=tf.float32)
a.shape
tf.slice(a, [:, 0], [:, 3])
b, c = tf.split(a, [10, 1], 1)
b, c
c / 2.0
tf.concat([b, c], 1)
test_model.weighted_prob(tf.expand_dims(check_data, 0))
tf.one_hot(train_labels, depth=10)[0:3]
tf.data.Dataset.from_tensors(tf.one_hot(train_labels, depth=10))
train_dataset
shape = label.shape
shape[0]
class fourLayerCNN(tf.keras.Model):
    """Simple 4-layer CNN"""

    def __init__(self, o=4.4, output_dim=11, name='cnn', **kwargs) -> None:
        super(fourLayerCNN, self).__init__(name=name, **kwargs)
        self.o = o
        self.output_dim = output_dim
        self.input_layer = layers.InputLayer(input_shape=(28, 28, 1))
        self.conv2d_1 = layers.Conv2D(filters=32, kernel_size=3, strides=(1, 1), activation='relu')
        self.conv2d_2 = layers.Conv2D(filters=64, kernel_size=3, strides=(1, 1), activation='relu')
        self.dropout = layers.Dropout(rate=0.2)
        self.flatten = layers.Flatten()
        self.dense = layers.Dense(self.output_dim)
        self.prob = layers.Softmax()

    def call(self, inputs):
        x = self.input_layer(inputs)
        x = self.conv2d_1(x)
        x = self.conv2d_2(x)
        x = self.conv2d_3(x)
        x = self.conv2d_4(x)
        x = self.flatten(x)
        x = self.dense(x)
        prob = self.prob(x)
        return prob
def compute_loss(model, x, y):
    weighted_prob = model(x)
    cross_ent = tf.nn.sigmoid_cross_entropy_with_logits(logits=weighted_prob, labels=y)
    return cross_ent
abstention_scores[:3]
plt.imshow(test_images[9982].reshape(28, 28), cmap="gray")
test = rotation(test_images[16], 180)
plt.imshow(test.reshape(28, 28), cmap="gray")<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.1378011301675424, 'fraction_numerical': 0.047090314265886786, 'mean_word_length': 4.35740839086564, 'pattern_counts': {'":': 0, '<': 35, '<?xml version=': 0, '>': 38, 'https://': 2, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '9', 'dir': 'jupyter-structured-clean-dedup', 'id': '392353', 'license': 'MIT', 'n_tokens_mistral': 4077, 'n_tokens_neox': 3748, 'n_words': 748, 'path': 'colab/experiment_MNIST.ipynb', 'repo_name': 'yoheikikuta/DeepGamblers-tensorflow', 'size': '10087'}