<jupyter_start><jupyter_text>Learning to tokenize in Vision Transformers**Authors:** [<NAME>](https://twitter.com/ariG23498), [<NAME>](https://twitter.com/RisingSayak) (equal contribution)**Date created:** 2021/12/10**Last modified:** 2021/12/15**Description:** Adaptively generating a smaller number of tokens for Vision Transformers.  IntroductionVision Transformers ([Dosovitskiy et al.](https://arxiv.org/abs/2010.11929)) and manyother Transformer-based architectures ([Liu et al.](https://arxiv.org/abs/2103.14030),[Yuan et al.](https://arxiv.org/abs/2101.11986), etc.) have shown strong results inimage recognition. The following provides a brief overview of the components involved in theVision Transformer architecture for image classification:* Extract small patches from input images.* Linearly project those patches.* Add positional embeddings to these linear projections.* Run these projections through a series of Transformer ([Vaswani et al.](https://arxiv.org/abs/1706.03762))blocks.* Finally, take the representation from the final Transformer block and add aclassification head.If we take 224x224 images and extract 16x16 patches, we get a total of 196 patches (alsocalled tokens) for each image. The number of patches increases as we increase theresolution, leading to higher memory footprint. Could we use a reducednumber of patches without having to compromise performance?Ryoo et al. investigate this question in[TokenLearner: Adaptive Space-Time Tokenization for Videos](https://openreview.net/forum?id=z-l1kpDXs88).They introduce a novel module called **TokenLearner** that can help reduce the numberof patches used by a Vision Transformer (ViT) in an adaptive manner. With TokenLearnerincorporated in the standard ViT architecture, they are able to reduce the amount ofcompute (measured in FLOPS) used by the model.In this example, we implement the TokenLearner module and demonstrate itsperformance with a mini ViT and the CIFAR-10 dataset. We make use of the followingreferences:* [Official TokenLearner code](https://github.com/google-research/scenic/blob/main/scenic/projects/token_learner/model.py)* [Image Classification with ViTs on keras.io](https://keras.io/examples/vision/image_classification_with_vision_transformer/)* [TokenLearner slides from NeurIPS 2021](https://nips.cc/media/neurips-2021/Slides/26578.pdf)  SetupWe need to install TensorFlow Addons to run this example. To install it, execute thefollowing:```shellpip install tensorflow-addons```  Imports<jupyter_code>import tensorflow as tf

from tensorflow import keras
from tensorflow.keras import layers
import tensorflow_addons as tfa

from datetime import datetime
import matplotlib.pyplot as plt
import numpy as np

import math<jupyter_output><empty_output><jupyter_text>HyperparametersPlease feel free to change the hyperparameters and check your results. The best way todevelop intuition about the architecture is to experiment with it.<jupyter_code># DATA
BATCH_SIZE = 256
AUTO = tf.data.AUTOTUNE
INPUT_SHAPE = (32, 32, 3)
NUM_CLASSES = 10

# OPTIMIZER
LEARNING_RATE = 1e-3
WEIGHT_DECAY = 1e-4

# TRAINING
EPOCHS = 20

# AUGMENTATION
IMAGE_SIZE = 48  # We will resize input images to this size.
PATCH_SIZE = 6  # Size of the patches to be extracted from the input images.
NUM_PATCHES = (IMAGE_SIZE // PATCH_SIZE) ** 2

# ViT ARCHITECTURE
LAYER_NORM_EPS = 1e-6
PROJECTION_DIM = 128
NUM_HEADS = 4
NUM_LAYERS = 4
MLP_UNITS = [
    PROJECTION_DIM * 2,
    PROJECTION_DIM,
]

# TOKENLEARNER
NUM_TOKENS = 4<jupyter_output><empty_output><jupyter_text>Load and prepare the CIFAR-10 dataset<jupyter_code># Load the CIFAR-10 dataset.
(x_train, y_train), (x_test, y_test) = keras.datasets.cifar10.load_data()
(x_train, y_train), (x_val, y_val) = (
    (x_train[:40000], y_train[:40000]),
    (x_train[40000:], y_train[40000:]),
)
print(f"Training samples: {len(x_train)}")
print(f"Validation samples: {len(x_val)}")
print(f"Testing samples: {len(x_test)}")

# Convert to tf.data.Dataset objects.
train_ds = tf.data.Dataset.from_tensor_slices((x_train, y_train))
train_ds = train_ds.shuffle(BATCH_SIZE * 100).batch(BATCH_SIZE).prefetch(AUTO)

val_ds = tf.data.Dataset.from_tensor_slices((x_val, y_val))
val_ds = val_ds.batch(BATCH_SIZE).prefetch(AUTO)

test_ds = tf.data.Dataset.from_tensor_slices((x_test, y_test))
test_ds = test_ds.batch(BATCH_SIZE).prefetch(AUTO)<jupyter_output><empty_output><jupyter_text>Data augmentationThe augmentation pipeline consists of:- Rescaling- Resizing- Random cropping (fixed-sized or random sized)- Random horizontal flipping<jupyter_code>data_augmentation = keras.Sequential(
    [
        layers.Rescaling(1 / 255.0),
        layers.Resizing(INPUT_SHAPE[0] + 20, INPUT_SHAPE[0] + 20),
        layers.RandomCrop(IMAGE_SIZE, IMAGE_SIZE),
        layers.RandomFlip("horizontal"),
    ],
    name="data_augmentation",
)<jupyter_output><empty_output><jupyter_text>Note that image data augmentation layers do not apply data transformations at inference time.This means that when these layers are called with `training=False` they behave differently. Refer[to the documentation](https://keras.io/api/layers/preprocessing_layers/image_augmentation/) for moredetails.  Positional embedding moduleA [Transformer](https://arxiv.org/abs/1706.03762) architecture consists of **multi-headself attention** layers and **fully-connected feed forward** networks (MLP) as the maincomponents. Both these components are _permutation invariant_: they're not aware offeature order.To overcome this problem we inject tokens with positional information. The`position_embedding` function adds this positional information to the linearly projectedtokens.<jupyter_code>def position_embedding(
    projected_patches, num_patches=NUM_PATCHES, projection_dim=PROJECTION_DIM
):
    # Build the positions.
    positions = tf.range(start=0, limit=num_patches, delta=1)

    # Encode the positions with an Embedding layer.
    encoded_positions = layers.Embedding(
        input_dim=num_patches, output_dim=projection_dim
    )(positions)

    # Add encoded positions to the projected patches.
    return projected_patches + encoded_positions<jupyter_output><empty_output><jupyter_text>MLP block for TransformerThis serves as the Fully Connected Feed Forward block for our Transformer.<jupyter_code>def mlp(x, dropout_rate, hidden_units):
    # Iterate over the hidden units and
    # add Dense => Dropout.
    for units in hidden_units:
        x = layers.Dense(units, activation=tf.nn.gelu)(x)
        x = layers.Dropout(dropout_rate)(x)
    return x<jupyter_output><empty_output><jupyter_text>TokenLearner moduleThe following figure presents a pictorial overview of the module([source](https://ai.googleblog.com/2021/12/improving-vision-transformer-efficiency.html)).![TokenLearner module GIF](https://blogger.googleusercontent.com/img/a/AVvXsEiylT3_nmd9-tzTnz3g3Vb4eTn-L5sOwtGJOad6t2we7FsjXSpbLDpuPrlInAhtE5hGCA_PfYTJtrIOKfLYLYGcYXVh1Ksfh_C1ZC-C8gw6GKtvrQesKoMrEA_LU_Gd5srl5-3iZDgJc1iyCELoXtfuIXKJ2ADDHOBaUjhU8lXTVdr2E7bCVaFgVHHkmA=w640-h208)The TokenLearner module takes as input an image-shaped tensor. It then passes it throughmultiple single-channel convolutional layers extracting different spatial attention mapsfocusing on different parts of the input. These attention maps are then element-wisemultiplied to the input and result is aggregated with pooling. This pooled output can betrated as a summary of the input and has much lesser number of patches (8, for example)than the original one (196, for example).Using multiple convolution layers helps with expressivity. Imposing a form of spatialattention helps retain relevant information from the inputs. Both of these components arecrucial to make TokenLearner work, especially when we are significantly reducing the number of patches.<jupyter_code>def token_learner(inputs, number_of_tokens=NUM_TOKENS):
    # Layer normalize the inputs.
    x = layers.LayerNormalization(epsilon=LAYER_NORM_EPS)(inputs)  # (B, H, W, C)

    # Applying Conv2D => Reshape => Permute
    # The reshape and permute is done to help with the next steps of
    # multiplication and Global Average Pooling.
    attention_maps = keras.Sequential(
        [
            # 3 layers of conv with gelu activation as suggested
            # in the paper.
            layers.Conv2D(
                filters=number_of_tokens,
                kernel_size=(3, 3),
                activation=tf.nn.gelu,
                padding="same",
                use_bias=False,
            ),
            layers.Conv2D(
                filters=number_of_tokens,
                kernel_size=(3, 3),
                activation=tf.nn.gelu,
                padding="same",
                use_bias=False,
            ),
            layers.Conv2D(
                filters=number_of_tokens,
                kernel_size=(3, 3),
                activation=tf.nn.gelu,
                padding="same",
                use_bias=False,
            ),
            # This conv layer will generate the attention maps
            layers.Conv2D(
                filters=number_of_tokens,
                kernel_size=(3, 3),
                activation="sigmoid",  # Note sigmoid for [0, 1] output
                padding="same",
                use_bias=False,
            ),
            # Reshape and Permute
            layers.Reshape((-1, number_of_tokens)),  # (B, H*W, num_of_tokens)
            layers.Permute((2, 1)),
        ]
    )(
        x
    )  # (B, num_of_tokens, H*W)

    # Reshape the input to align it with the output of the conv block.
    num_filters = inputs.shape[-1]
    inputs = layers.Reshape((1, -1, num_filters))(inputs)  # inputs == (B, 1, H*W, C)

    # Element-Wise multiplication of the attention maps and the inputs
    attended_inputs = (
        attention_maps[..., tf.newaxis] * inputs
    )  # (B, num_tokens, H*W, C)

    # Global average pooling the element wise multiplication result.
    outputs = tf.reduce_mean(attended_inputs, axis=2)  # (B, num_tokens, C)
    return outputs<jupyter_output><empty_output><jupyter_text>Transformer block<jupyter_code>def transformer(encoded_patches):
    # Layer normalization 1.
    x1 = layers.LayerNormalization(epsilon=LAYER_NORM_EPS)(encoded_patches)

    # Multi Head Self Attention layer 1.
    attention_output = layers.MultiHeadAttention(
        num_heads=NUM_HEADS, key_dim=PROJECTION_DIM, dropout=0.1
    )(x1, x1)

    # Skip connection 1.
    x2 = layers.Add()([attention_output, encoded_patches])

    # Layer normalization 2.
    x3 = layers.LayerNormalization(epsilon=LAYER_NORM_EPS)(x2)

    # MLP layer 1.
    x4 = mlp(x3, hidden_units=MLP_UNITS, dropout_rate=0.1)

    # Skip connection 2.
    encoded_patches = layers.Add()([x4, x2])
    return encoded_patches<jupyter_output><empty_output><jupyter_text>ViT model with the TokenLearner module<jupyter_code>def create_vit_classifier(use_token_learner=True, token_learner_units=NUM_TOKENS):
    inputs = layers.Input(shape=INPUT_SHAPE)  # (B, H, W, C)

    # Augment data.
    augmented = data_augmentation(inputs)

    # Create patches and project the pathces.
    projected_patches = layers.Conv2D(
        filters=PROJECTION_DIM,
        kernel_size=(PATCH_SIZE, PATCH_SIZE),
        strides=(PATCH_SIZE, PATCH_SIZE),
        padding="VALID",
    )(augmented)
    _, h, w, c = projected_patches.shape
    projected_patches = layers.Reshape((h * w, c))(
        projected_patches
    )  # (B, number_patches, projection_dim)

    # Add positional embeddings to the projected patches.
    encoded_patches = position_embedding(
        projected_patches
    )  # (B, number_patches, projection_dim)
    encoded_patches = layers.Dropout(0.1)(encoded_patches)

    # Iterate over the number of layers and stack up blocks of
    # Transformer.
    for i in range(NUM_LAYERS):
        # Add a Transformer block.
        encoded_patches = transformer(encoded_patches)

        # Add TokenLearner layer in the middle of the
        # architecture. The paper suggests that anywhere
        # between 1/2 or 3/4 will work well.
        if use_token_learner and i == NUM_LAYERS // 2:
            _, hh, c = encoded_patches.shape
            h = int(math.sqrt(hh))
            encoded_patches = layers.Reshape((h, h, c))(
                encoded_patches
            )  # (B, h, h, projection_dim)
            encoded_patches = token_learner(
                encoded_patches, token_learner_units
            )  # (B, num_tokens, c)

    # Layer normalization and Global average pooling.
    representation = layers.LayerNormalization(epsilon=LAYER_NORM_EPS)(encoded_patches)
    representation = layers.GlobalAvgPool1D()(representation)

    # Classify outputs.
    outputs = layers.Dense(NUM_CLASSES, activation="softmax")(representation)

    # Create the Keras model.
    model = keras.Model(inputs=inputs, outputs=outputs)
    return model<jupyter_output><empty_output><jupyter_text>As shown in the [TokenLearner paper](https://openreview.net/forum?id=z-l1kpDXs88), it isalmost always advantageous to include the TokenLearner module in the middle of thenetwork.  Training utility<jupyter_code>def run_experiment(model):
    # Initialize the AdamW optimizer.
    optimizer = tfa.optimizers.AdamW(
        learning_rate=LEARNING_RATE, weight_decay=WEIGHT_DECAY
    )

    # Compile the model with the optimizer, loss function
    # and the metrics.
    model.compile(
        optimizer=optimizer,
        loss="sparse_categorical_crossentropy",
        metrics=[
            keras.metrics.SparseCategoricalAccuracy(name="accuracy"),
            keras.metrics.SparseTopKCategoricalAccuracy(5, name="top-5-accuracy"),
        ],
    )

    # Define callbacks
    checkpoint_filepath = "/tmp/checkpoint"
    checkpoint_callback = keras.callbacks.ModelCheckpoint(
        checkpoint_filepath,
        monitor="val_accuracy",
        save_best_only=True,
        save_weights_only=True,
    )

    # Train the model.
    _ = model.fit(
        train_ds,
        epochs=EPOCHS,
        validation_data=val_ds,
        callbacks=[checkpoint_callback],
    )

    model.load_weights(checkpoint_filepath)
    _, accuracy, top_5_accuracy = model.evaluate(test_ds)
    print(f"Test accuracy: {round(accuracy * 100, 2)}%")
    print(f"Test top 5 accuracy: {round(top_5_accuracy * 100, 2)}%")<jupyter_output><empty_output><jupyter_text>Train and evaluate a ViT with TokenLearner<jupyter_code>vit_token_learner = create_vit_classifier()
run_experiment(vit_token_learner)<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.08925619834710743, 'fraction_numerical': 0.01893939393939394, 'mean_word_length': 3.8662868632707776, 'pattern_counts': {'":': 0, '<': 47, '<?xml version=': 0, '>': 50, 'https://': 15, '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': '11', 'dir': 'jupyter-structured-clean-dedup', 'id': '519133', 'license': 'Apache-2.0', 'n_tokens_mistral': 4800, 'n_tokens_neox': 4432, 'n_words': 1173, 'path': 'examples/vision/ipynb/token_learner.ipynb', 'repo_name': 'wariua/keras-io-ko', 'size': '14539'}