title image classification with perceiver khalid salamahttps www linkedin cominkhalidsalama24403144 date created 20210430 last modified 20210130 description implementing the perceiver model for image classification accelerator gpu introduction this example implements the perceiver general perception with iterative attentionhttps arxiv orgabs2103 03206 model by andrew jaegle et al for image classification and demonstrates it on the cifar100 dataset the perceiver model leverages an asymmetric attention mechanism to iteratively distill inputs into a tight latent bottleneck allowing it to scale to handle very large inputs in other words let s assume that your input data array e g image has m elements i e patches where m is large in a standard transformer model a selfattention operation is performed for the m elements the complexity of this operation is om2 however the perceiver model creates a latent array of size n elements where n m and performs two operations iteratively 1 crossattention transformer between the latent array and the data array the complexity of this operation is om n 2 selfattention transformer on the latent array the complexity of this operation is on2 this example requires tensorflow 2 4 or higher as well as tensorflow addonshttps www tensorflow orgaddonsoverview which can be installed using the following command python pip install u tensorflowaddons setup prepare the data configure the hyperparameters note that in order to use each pixel as an individual input in the data array set patchsize to 1 use data augmentation compute the mean and the variance of the training data for normalization implement feedforward network ffn implement patch creation as a layer implement the patch encoding layer the patchencoder layer will linearly transform a patch by projecting it into a vector of size latentdim in addition it adds a learnable position embedding to the projected vector note that the orginal perceiver paper uses the fourier feature positional encodings build the perceiver model the perceiver consists of two modules a crossattention module and a standard transformer with selfattention crossattention module the crossattention expects a latentdim projectiondim latent array and the datadim projectiondim data array as inputs to produce a latentdim projectiondim latent array as an output to apply crossattention the query vectors are generated from the latent array while the key and value vectors are generated from the encoded image note that the data array in this example is the image where the datadim is set to the numpatches recieve the latent array as an input of shape 1 latentdim projectiondim recieve the dataarray encoded image as an input of shape batchsize datadim projectiondim apply layer norm to the inputs create query tensor 1 latentdim projectiondim create key tensor batchsize datadim projectiondim create value tensor batchsize datadim projectiondim generate crossattention outputs batchsize latentdim projectiondim skip connection 1 apply layer norm apply feedforward network skip connection 2 create the keras model transformer module the transformer expects the output latent vector from the crossattention module as an input applies multihead selfattention to its latentdim elements followed by feedforward network to produce another latentdim projectiondim latent array inputshape 1 latentdim projectiondim create multiple layers of the transformer block apply layer normalization 1 create a multihead selfattention layer skip connection 1 apply layer normalization 2 apply feedforward network skip connection 2 create the keras model perceiver model the perceiver model repeats the crossattention and transformer modules numiterations timeswith shared weights and skip connectionsto allow the latent array to iteratively extract information from the input image as it is needed create latent array create patching module create patch encoder create crossattenion module create transformer module create global average pooling layer create a classification head augment data create patches encode patches prepare crossattention inputs apply the crossattention and the transformer modules iteratively apply crossattention from the latent array to the data array apply selfattention transformer to the latent array set the latent array of the next iteration apply global average pooling to generate a batchsize projectiondim repesentation tensor generate logits compile train and evaluate the mode create adam optimizer with weight decay compile the model create a learning rate scheduler callback create an early stopping callback fit the model return history to plot learning curves note that training the perceiver model with the current settings on a v100 gpus takes around 200 seconds after 40 epochs the perceiver model achieves around 53 accuracy and 81 top5 accuracy on the test data as mentioned in the ablations of the perceiverhttps arxiv orgabs2103 03206 paper you can obtain better results by increasing the latent array size increasing the projection dimensions of the latent array and data array elements increasing the number of blocks in the transformer module and increasing the number of iterations of applying the crossattention and the latent transformer modules you may also try to increase the size the input images and use different patch sizes the perceiver benefits from inceasing the model size however larger models needs bigger accelerators to fit in and train efficiently this is why in the perceiver paper they used 32 tpu cores to run the experiments introduction this example implements the perceiver general perception with iterative attention https arxiv org abs 2103 03206 model by andrew jaegle et al for image classification and demonstrates it on the cifar 100 dataset the perceiver model leverages an asymmetric attention mechanism to iteratively distill inputs into a tight latent bottleneck allowing it to scale to handle very large inputs in other words let s assume that your input data array e g image has m elements i e patches where m is large in a standard transformer model a self attention operation is performed for the m elements the complexity of this operation is o m 2 however the perceiver model creates a latent array of size n elements where n m and performs two operations iteratively 1 cross attention transformer between the latent array and the data array the complexity of this operation is o m n 2 self attention transformer on the latent array the complexity of this operation is o n 2 this example requires tensorflow 2 4 or higher as well as tensorflow addons https www tensorflow org addons overview which can be installed using the following command python pip install u tensorflow addons setup prepare the data configure the hyperparameters we ll resize input images to this size size of the patches to be extract from the input images size of the data array size of the latent array embedding size of each element in the data and latent arrays number of transformer heads size of the transformer feedforward network repetitions of the cross attention and transformer modules size of the feedforward network of the final classifier note that in order to use each pixel as an individual input in the data array set patch_size to 1 use data augmentation compute the mean and the variance of the training data for normalization implement feedforward network ffn implement patch creation as a layer implement the patch encoding layer the patchencoder layer will linearly transform a patch by projecting it into a vector of size latent_dim in addition it adds a learnable position embedding to the projected vector note that the orginal perceiver paper uses the fourier feature positional encodings build the perceiver model the perceiver consists of two modules a cross attention module and a standard transformer with self attention cross attention module the cross attention expects a latent_dim projection_dim latent array and the data_dim projection_dim data array as inputs to produce a latent_dim projection_dim latent array as an output to apply cross attention the query vectors are generated from the latent array while the key and value vectors are generated from the encoded image note that the data array in this example is the image where the data_dim is set to the num_patches recieve the latent array as an input of shape 1 latent_dim projection_dim recieve the data_array encoded image as an input of shape batch_size data_dim projection_dim apply layer norm to the inputs create query tensor 1 latent_dim projection_dim create key tensor batch_size data_dim projection_dim create value tensor batch_size data_dim projection_dim generate cross attention outputs batch_size latent_dim projection_dim skip connection 1 apply layer norm apply feedforward network skip connection 2 create the keras model transformer module the transformer expects the output latent vector from the cross attention module as an input applies multi head self attention to its latent_dim elements followed by feedforward network to produce another latent_dim projection_dim latent array input_shape 1 latent_dim projection_dim create multiple layers of the transformer block apply layer normalization 1 create a multi head self attention layer skip connection 1 apply layer normalization 2 apply feedforward network skip connection 2 create the keras model perceiver model the perceiver model repeats the cross attention and transformer modules num_iterations times with shared weights and skip connections to allow the latent array to iteratively extract information from the input image as it is needed create latent array create patching module create patch encoder create cross attenion module create transformer module create global average pooling layer create a classification head augment data create patches encode patches prepare cross attention inputs apply the cross attention and the transformer modules iteratively apply cross attention from the latent array to the data array apply self attention transformer to the latent array set the latent array of the next iteration apply global average pooling to generate a batch_size projection_dim repesentation tensor generate logits compile train and evaluate the mode create adam optimizer with weight decay compile the model create a learning rate scheduler callback create an early stopping callback fit the model return history to plot learning curves note that training the perceiver model with the current settings on a v100 gpus takes around 200 seconds after 40 epochs the perceiver model achieves around 53 accuracy and 81 top 5 accuracy on the test data as mentioned in the ablations of the perceiver https arxiv org abs 2103 03206 paper you can obtain better results by increasing the latent array size increasing the projection dimensions of the latent array and data array elements increasing the number of blocks in the transformer module and increasing the number of iterations of applying the cross attention and the latent transformer modules you may also try to increase the size the input images and use different patch sizes the perceiver benefits from inceasing the model size however larger models needs bigger accelerators to fit in and train efficiently this is why in the perceiver paper they used 32 tpu cores to run the experiments
import numpy as np
import tensorflow as tf
import keras
from keras import layers



num_classes = 100
input_shape = (32, 32, 3)

(x_train, y_train), (x_test, y_test) = keras.datasets.cifar100.load_data()

print(f"x_train shape: {x_train.shape} - y_train shape: {y_train.shape}")
print(f"x_test shape: {x_test.shape} - y_test shape: {y_test.shape}")



learning_rate = 0.001
weight_decay = 0.0001
batch_size = 64
num_epochs = 50
dropout_rate = 0.2
image_size = 64  
patch_size = 2  
num_patches = (image_size // patch_size) ** 2  
latent_dim = 256  
projection_dim = (
    256  
)
num_heads = 8  
ffn_units = [
    projection_dim,
    projection_dim,
]  
num_transformer_blocks = 4
num_iterations = (
    2  
)
classifier_units = [
    projection_dim,
    num_classes,
]  

print(f"Image size: {image_size} X {image_size} = {image_size ** 2}")
print(f"Patch size: {patch_size} X {patch_size} = {patch_size ** 2} ")
print(f"Patches per image: {num_patches}")
print(f"Elements per patch (3 channels): {(patch_size ** 2) * 3}")
print(f"Latent array shape: {latent_dim} X {projection_dim}")
print(f"Data array shape: {num_patches} X {projection_dim}")





data_augmentation = keras.Sequential(
    [
        layers.Normalization(),
        layers.Resizing(image_size, image_size),
        layers.RandomFlip("horizontal"),
        layers.RandomZoom(height_factor=0.2, width_factor=0.2),
    ],
    name="data_augmentation",
)

data_augmentation.layers[0].adapt(x_train)




def create_ffn(hidden_units, dropout_rate):
    ffn_layers = []
    for units in hidden_units[:-1]:
        ffn_layers.append(layers.Dense(units, activation=tf.nn.gelu))

    ffn_layers.append(layers.Dense(units=hidden_units[-1]))
    ffn_layers.append(layers.Dropout(dropout_rate))

    ffn = keras.Sequential(ffn_layers)
    return ffn





class Patches(layers.Layer):
    def __init__(self, patch_size):
        super().__init__()
        self.patch_size = patch_size

    def call(self, images):
        batch_size = tf.shape(images)[0]
        patches = tf.image.extract_patches(
            images=images,
            sizes=[1, self.patch_size, self.patch_size, 1],
            strides=[1, self.patch_size, self.patch_size, 1],
            rates=[1, 1, 1, 1],
            padding="VALID",
        )
        patch_dims = patches.shape[-1]
        patches = tf.reshape(patches, [batch_size, -1, patch_dims])
        return patches





class PatchEncoder(layers.Layer):
    def __init__(self, num_patches, projection_dim):
        super().__init__()
        self.num_patches = num_patches
        self.projection = layers.Dense(units=projection_dim)
        self.position_embedding = layers.Embedding(
            input_dim=num_patches, output_dim=projection_dim
        )

    def call(self, patches):
        positions = tf.range(start=0, limit=self.num_patches, delta=1)
        encoded = self.projection(patches) + self.position_embedding(positions)
        return encoded







def create_cross_attention_module(
    latent_dim, data_dim, projection_dim, ffn_units, dropout_rate
):
    inputs = {
        
        "latent_array": layers.Input(
            shape=(latent_dim, projection_dim), name="latent_array"
        ),
        
        "data_array": layers.Input(
            shape=(data_dim, projection_dim), name="data_array"
        ),
    }

    
    latent_array = layers.LayerNormalization(epsilon=1e-6)(
        inputs["latent_array"]
    )
    data_array = layers.LayerNormalization(epsilon=1e-6)(inputs["data_array"])

    
    query = layers.Dense(units=projection_dim)(latent_array)
    
    key = layers.Dense(units=projection_dim)(data_array)
    
    value = layers.Dense(units=projection_dim)(data_array)

    
    attention_output = layers.Attention(use_scale=True, dropout=0.1)(
        [query, key, value], return_attention_scores=False
    )
    
    attention_output = layers.Add()([attention_output, latent_array])

    
    attention_output = layers.LayerNormalization(epsilon=1e-6)(attention_output)
    
    ffn = create_ffn(hidden_units=ffn_units, dropout_rate=dropout_rate)
    outputs = ffn(attention_output)
    
    outputs = layers.Add()([outputs, attention_output])

    
    model = keras.Model(inputs=inputs, outputs=outputs)
    return model





def create_transformer_module(
    latent_dim,
    projection_dim,
    num_heads,
    num_transformer_blocks,
    ffn_units,
    dropout_rate,
):
    
    inputs = layers.Input(shape=(latent_dim, projection_dim))

    x0 = inputs
    
    for _ in range(num_transformer_blocks):
        
        x1 = layers.LayerNormalization(epsilon=1e-6)(x0)
        
        attention_output = layers.MultiHeadAttention(
            num_heads=num_heads, key_dim=projection_dim, dropout=0.1
        )(x1, x1)
        
        x2 = layers.Add()([attention_output, x0])
        
        x3 = layers.LayerNormalization(epsilon=1e-6)(x2)
        
        ffn = create_ffn(hidden_units=ffn_units, dropout_rate=dropout_rate)
        x3 = ffn(x3)
        
        x0 = layers.Add()([x3, x2])

    
    model = keras.Model(inputs=inputs, outputs=x0)
    return model





class Perceiver(keras.Model):
    def __init__(
        self,
        patch_size,
        data_dim,
        latent_dim,
        projection_dim,
        num_heads,
        num_transformer_blocks,
        ffn_units,
        dropout_rate,
        num_iterations,
        classifier_units,
    ):
        super().__init__()

        self.latent_dim = latent_dim
        self.data_dim = data_dim
        self.patch_size = patch_size
        self.projection_dim = projection_dim
        self.num_heads = num_heads
        self.num_transformer_blocks = num_transformer_blocks
        self.ffn_units = ffn_units
        self.dropout_rate = dropout_rate
        self.num_iterations = num_iterations
        self.classifier_units = classifier_units

    def build(self, input_shape):
        
        self.latent_array = self.add_weight(
            shape=(self.latent_dim, self.projection_dim),
            initializer="random_normal",
            trainable=True,
        )

        
        self.patcher = Patches(self.patch_size)

        
        self.patch_encoder = PatchEncoder(self.data_dim, self.projection_dim)

        
        self.cross_attention = create_cross_attention_module(
            self.latent_dim,
            self.data_dim,
            self.projection_dim,
            self.ffn_units,
            self.dropout_rate,
        )

        
        self.transformer = create_transformer_module(
            self.latent_dim,
            self.projection_dim,
            self.num_heads,
            self.num_transformer_blocks,
            self.ffn_units,
            self.dropout_rate,
        )

        
        self.global_average_pooling = layers.GlobalAveragePooling1D()

        
        self.classification_head = create_ffn(
            hidden_units=self.classifier_units, dropout_rate=self.dropout_rate
        )

        super().build(input_shape)

    def call(self, inputs):
        
        augmented = data_augmentation(inputs)
        
        patches = self.patcher(augmented)
        
        encoded_patches = self.patch_encoder(patches)
        
        cross_attention_inputs = {
            "latent_array": tf.expand_dims(self.latent_array, 0),
            "data_array": encoded_patches,
        }
        
        for _ in range(self.num_iterations):
            
            latent_array = self.cross_attention(cross_attention_inputs)
            
            latent_array = self.transformer(latent_array)
            
            cross_attention_inputs["latent_array"] = latent_array

        
        representation = self.global_average_pooling(latent_array)
        
        logits = self.classification_head(representation)
        return logits





def run_experiment(model):
    
    optimizer = keras.optimizers.Adam(
        learning_rate=learning_rate,
        weight_decay=weight_decay,
    )

    
    model.compile(
        optimizer=optimizer,
        loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
        metrics=[
            keras.metrics.SparseCategoricalAccuracy(name="acc"),
            keras.metrics.SparseTopKCategoricalAccuracy(5, name="top5-acc"),
        ],
    )

    
    reduce_lr = keras.callbacks.ReduceLROnPlateau(
        monitor="val_loss", factor=0.2, patience=3
    )

    
    early_stopping = keras.callbacks.EarlyStopping(
        monitor="val_loss", patience=15, restore_best_weights=True
    )

    
    history = model.fit(
        x=x_train,
        y=y_train,
        batch_size=batch_size,
        epochs=num_epochs,
        validation_split=0.1,
        callbacks=[early_stopping, reduce_lr],
    )

    _, accuracy, top_5_accuracy = model.evaluate(x_test, y_test)
    print(f"Test accuracy: {round(accuracy * 100, 2)}%")
    print(f"Test top 5 accuracy: {round(top_5_accuracy * 100, 2)}%")

    
    return history




perceiver_classifier = Perceiver(
    patch_size,
    num_patches,
    latent_dim,
    projection_dim,
    num_heads,
    num_transformer_blocks,
    ffn_units,
    dropout_rate,
    num_iterations,
    classifier_units,
)


history = run_experiment(perceiver_classifier)