title image classification with eanet external attention transformer zhiyong changhttps github comczy00000 date created 20211019 last modified 20230718 description image classification with a transformer that leverages external attention accelerator gpu converted to keras 3 muhammad anas razahttps anasrz com introduction this example implements the eanethttps arxiv orgabs2105 02358 model for image classification and demonstrates it on the cifar100 dataset eanet introduces a novel attention mechanism named external attention based on two external small learnable and shared memories which can be implemented easily by simply using two cascaded linear layers and two normalization layers it conveniently replaces selfattention as used in existing architectures external attention has linear complexity as it only implicitly considers the correlations between all samples setup prepare the data configure the hyperparameters use data augmentation compute the mean and the variance of the training data for normalization implement the patch extraction and encoding layer implement the external attention block create tensor batchsize numpatches numheads dimdimcoefficientnumheads a linear layer mk normalize attention map dobulenormalization a linear layer mv a linear layer to project original dim implement the mlp block implement the transformer block implement the eanet model the eanet model leverages external attention the computational complexity of traditional self attention is od n 2 where d is the embedding size and n is the number of patch the s find that most pixels are closely related to just a few other pixels and an nton attention matrix may be redundant so they propose as an alternative an external attention module where the computational complexity of external attention is od s n as d and s are hyperparameters the proposed algorithm is linear in the number of pixels in fact this is equivalent to a drop patch operation because a lot of information contained in a patch in an image is redundant and unimportant image augment extract patches create patch embedding create transformer block train on cifar100 let s visualize the training progress of the model let s display the final results of the test on cifar100 eanet just replaces self attention in vit with external attention the traditional vit achieved a 73 test top5 accuracy and 41 top1 accuracy after training 50 epochs but with 0 6m parameters under the same experimental environment and the same hyperparameters the eanet model we just trained has just 0 3m parameters and it gets us to 73 test top5 accuracy and 43 top1 accuracy this fully demonstrates the effectiveness of external attention we only show the training process of eanet you can train vit under the same experimental conditions and observe the test results introduction this example implements the eanet https arxiv org abs 2105 02358 model for image classification and demonstrates it on the cifar 100 dataset eanet introduces a novel attention mechanism named external attention based on two external small learnable and shared memories which can be implemented easily by simply using two cascaded linear layers and two normalization layers it conveniently replaces self attention as used in existing architectures external attention has linear complexity as it only implicitly considers the correlations between all samples setup prepare the data configure the hyperparameters size of the patches to be extracted from the input images number of patch number of hidden units number of repetitions of the transformer layer use data augmentation compute the mean and the variance of the training data for normalization implement the patch extraction and encoding layer implement the external attention block create tensor batch_size num_patches num_heads dim dim_coefficient num_heads a linear layer m_k normalize attention map dobule normalization a linear layer m_v a linear layer to project original dim implement the mlp block implement the transformer block implement the eanet model the eanet model leverages external attention the computational complexity of traditional self attention is o d n 2 where d is the embedding size and n is the number of patch the s find that most pixels are closely related to just a few other pixels and an n to n attention matrix may be redundant so they propose as an alternative an external attention module where the computational complexity of external attention is o d s n as d and s are hyper parameters the proposed algorithm is linear in the number of pixels in fact this is equivalent to a drop patch operation because a lot of information contained in a patch in an image is redundant and unimportant image augment extract patches create patch embedding create transformer block train on cifar 100 let s visualize the training progress of the model let s display the final results of the test on cifar 100 eanet just replaces self attention in vit with external attention the traditional vit achieved a 73 test top 5 accuracy and 41 top 1 accuracy after training 50 epochs but with 0 6m parameters under the same experimental environment and the same hyperparameters the eanet model we just trained has just 0 3m parameters and it gets us to 73 test top 5 accuracy and 43 top 1 accuracy this fully demonstrates the effectiveness of external attention we only show the training process of eanet you can train vit under the same experimental conditions and observe the test results
import keras
from keras import layers
from keras import ops

import matplotlib.pyplot as plt




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

(x_train, y_train), (x_test, y_test) = keras.datasets.cifar100.load_data()
y_train = keras.utils.to_categorical(y_train, num_classes)
y_test = keras.utils.to_categorical(y_test, num_classes)
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}")



weight_decay = 0.0001
learning_rate = 0.001
label_smoothing = 0.1
validation_split = 0.2
batch_size = 128
num_epochs = 50
patch_size = 2  
num_patches = (input_shape[0] // patch_size) ** 2  
embedding_dim = 64  
mlp_dim = 64
dim_coefficient = 4
num_heads = 4
attention_dropout = 0.2
projection_dropout = 0.2
num_transformer_blocks = 8  

print(f"Patch size: {patch_size} X {patch_size} = {patch_size ** 2} ")
print(f"Patches per image: {num_patches}")




data_augmentation = keras.Sequential(
    [
        layers.Normalization(),
        layers.RandomFlip("horizontal"),
        layers.RandomRotation(factor=0.1),
        layers.RandomContrast(factor=0.1),
        layers.RandomZoom(height_factor=0.2, width_factor=0.2),
    ],
    name="data_augmentation",
)

data_augmentation.layers[0].adapt(x_train)




class PatchExtract(layers.Layer):
    def __init__(self, patch_size, **kwargs):
        super().__init__(**kwargs)
        self.patch_size = patch_size

    def call(self, x):
        B, C = ops.shape(x)[0], ops.shape(x)[-1]
        x = ops.image.extract_patches(x, self.patch_size)
        x = ops.reshape(x, (B, -1, self.patch_size * self.patch_size * C))
        return x


class PatchEmbedding(layers.Layer):
    def __init__(self, num_patch, embed_dim, **kwargs):
        super().__init__(**kwargs)
        self.num_patch = num_patch
        self.proj = layers.Dense(embed_dim)
        self.pos_embed = layers.Embedding(
            input_dim=num_patch, output_dim=embed_dim
        )

    def call(self, patch):
        pos = ops.arange(start=0, stop=self.num_patch, step=1)
        return self.proj(patch) + self.pos_embed(pos)





def external_attention(
    x,
    dim,
    num_heads,
    dim_coefficient=4,
    attention_dropout=0,
    projection_dropout=0,
):
    _, num_patch, channel = x.shape
    assert dim % num_heads == 0
    num_heads = num_heads * dim_coefficient

    x = layers.Dense(dim * dim_coefficient)(x)
    
    x = ops.reshape(
        x, (-1, num_patch, num_heads, dim * dim_coefficient // num_heads)
    )
    x = ops.transpose(x, axes=[0, 2, 1, 3])
    
    attn = layers.Dense(dim // dim_coefficient)(x)
    
    attn = layers.Softmax(axis=2)(attn)
    
    attn = layers.Lambda(
        lambda attn: ops.divide(
            attn,
            ops.convert_to_tensor(1e-9) + ops.sum(attn, axis=-1, keepdims=True),
        )
    )(attn)
    attn = layers.Dropout(attention_dropout)(attn)
    
    x = layers.Dense(dim * dim_coefficient // num_heads)(attn)
    x = ops.transpose(x, axes=[0, 2, 1, 3])
    x = ops.reshape(x, [-1, num_patch, dim * dim_coefficient])
    
    x = layers.Dense(dim)(x)
    x = layers.Dropout(projection_dropout)(x)
    return x





def mlp(x, embedding_dim, mlp_dim, drop_rate=0.2):
    x = layers.Dense(mlp_dim, activation=ops.gelu)(x)
    x = layers.Dropout(drop_rate)(x)
    x = layers.Dense(embedding_dim)(x)
    x = layers.Dropout(drop_rate)(x)
    return x





def transformer_encoder(
    x,
    embedding_dim,
    mlp_dim,
    num_heads,
    dim_coefficient,
    attention_dropout,
    projection_dropout,
    attention_type="external_attention",
):
    residual_1 = x
    x = layers.LayerNormalization(epsilon=1e-5)(x)
    if attention_type == "external_attention":
        x = external_attention(
            x,
            embedding_dim,
            num_heads,
            dim_coefficient,
            attention_dropout,
            projection_dropout,
        )
    elif attention_type == "self_attention":
        x = layers.MultiHeadAttention(
            num_heads=num_heads,
            key_dim=embedding_dim,
            dropout=attention_dropout,
        )(x, x)
    x = layers.add([x, residual_1])
    residual_2 = x
    x = layers.LayerNormalization(epsilon=1e-5)(x)
    x = mlp(x, embedding_dim, mlp_dim)
    x = layers.add([x, residual_2])
    return x







def get_model(attention_type="external_attention"):
    inputs = layers.Input(shape=input_shape)
    
    x = data_augmentation(inputs)
    
    x = PatchExtract(patch_size)(x)
    
    x = PatchEmbedding(num_patches, embedding_dim)(x)
    
    for _ in range(num_transformer_blocks):
        x = transformer_encoder(
            x,
            embedding_dim,
            mlp_dim,
            num_heads,
            dim_coefficient,
            attention_dropout,
            projection_dropout,
            attention_type,
        )

    x = layers.GlobalAveragePooling1D()(x)
    outputs = layers.Dense(num_classes, activation="softmax")(x)
    model = keras.Model(inputs=inputs, outputs=outputs)
    return model





model = get_model(attention_type="external_attention")

model.compile(
    loss=keras.losses.CategoricalCrossentropy(label_smoothing=label_smoothing),
    optimizer=keras.optimizers.AdamW(
        learning_rate=learning_rate, weight_decay=weight_decay
    ),
    metrics=[
        keras.metrics.CategoricalAccuracy(name="accuracy"),
        keras.metrics.TopKCategoricalAccuracy(5, name="top-5-accuracy"),
    ],
)

history = model.fit(
    x_train,
    y_train,
    batch_size=batch_size,
    epochs=num_epochs,
    validation_split=validation_split,
)



plt.plot(history.history["loss"], label="train_loss")
plt.plot(history.history["val_loss"], label="val_loss")
plt.xlabel("Epochs")
plt.ylabel("Loss")
plt.title("Train and Validation Losses Over Epochs", fontsize=14)
plt.legend()
plt.grid()
plt.show()



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