title image classification with swin transformers rishit daglihttps twitter comrishitdagli date created 20210908 last modified 20210908 description image classification using swin transformers a generalpurpose backbone for computer vision accelerator gpu this example implements swin transformer hierarchical vision transformer using shifted windowshttps arxiv orgabs2103 14030 by liu et al for image classification and demonstrates it on the cifar100 datasethttps www cs toronto edukrizcifar html swin transformer shifted window transformer can serve as a generalpurpose backbone for computer vision swin transformer is a hierarchical transformer whose representations are computed with shifted windows the shifted window scheme brings greater efficiency by limiting selfattention computation to nonoverlapping local windows while also allowing for crosswindow connections this architecture has the flexibility to model information at various scales and has a linear computational complexity with respect to image size this example requires tensorflow 2 5 or higher setup prepare the data we load the cifar100 dataset through tf keras datasets normalize the images and convert the integer labels to onehot encoded vectors configure the hyperparameters a key parameter to pick is the patchsize the size of the input patches in order to use each pixel as an individual input you can set patchsize to 1 1 below we take inspiration from the original paper settings for training on imagenet1k keeping most of the original settings for this example helper functions we create two helper functions to help us get a sequence of patches from the image merge patches and apply dropout window based multihead selfattention usually transformers perform global selfattention where the relationships between a token and all other tokens are computed the global computation leads to quadratic complexity with respect to the number of tokens here as the original paperhttps arxiv orgabs2103 14030 suggests we compute selfattention within local windows in a nonoverlapping manner global selfattention leads to quadratic computational complexity in the number of patches whereas windowbased selfattention leads to linear complexity and is easily scalable the complete swin transformer model finally we put together the complete swin transformer by replacing the standard multihead attention mha with shifted windows attention as suggested in the original paper we create a model comprising of a shifted windowbased mha layer followed by a 2layer mlp with gelu nonlinearity in between applying layernormalization before each msa layer and each mlp and a residual connection after each of these layers notice that we only create a simple mlp with 2 dense and 2 dropout layers often you will see models using resnet50 as the mlp which is quite standard in the literature however in this paper the s use a 2layer mlp with gelu nonlinearity in between mask array to windows model training and evaluation extract and embed patches we first create 3 layers to help us extract embed and merge patches from the images on top of which we will later use the swin transformer class we built build the model we put together the swin transformer model train on cifar100 we train the model on cifar100 here we only train the model for 40 epochs to keep the training time short in this example in practice you should train for 150 epochs to reach convergence let s visualize the training progress of the model let s display the final results of the training on cifar100 the swin transformer model we just trained has just 152k parameters and it gets us to 75 test top5 accuracy within just 40 epochs without any signs of overfitting as well as seen in above graph this means we can train this network for longer perhaps with a bit more regularization and obtain even better performance this performance can further be improved by additional techniques like cosine decay learning rate schedule other data augmentation techniques while experimenting i tried training the model for 150 epochs with a slightly higher dropout and greater embedding dimensions which pushes the performance to 72 test accuracy on cifar100 as you can see in the screenshot results of training for longerhttps i imgur com9vnqesz png the s present a top1 accuracy of 87 3 on imagenet the s also present a number of experiments to study how input sizes optimizers etc affect the final performance of this model the s further present using this model for object detection semantic segmentation and instance segmentation as well and report competitive results for these you are strongly advised to also check out the original paperhttps arxiv orgabs2103 14030 this example takes inspiration from the official pytorchhttps github commicrosoftswintransformer and tensorflowhttps github comvcampsoldiersswintransformertensorflow implementations this example implements swin transformer hierarchical vision transformer using shifted windows https arxiv org abs 2103 14030 by liu et al for image classification and demonstrates it on the cifar 100 dataset https www cs toronto edu kriz cifar html swin transformer s hifted win dow transformer can serve as a general purpose backbone for computer vision swin transformer is a hierarchical transformer whose representations are computed with _shifted windows_ the shifted window scheme brings greater efficiency by limiting self attention computation to non overlapping local windows while also allowing for cross window connections this architecture has the flexibility to model information at various scales and has a linear computational complexity with respect to image size this example requires tensorflow 2 5 or higher setup prepare the data we load the cifar 100 dataset through tf keras datasets normalize the images and convert the integer labels to one hot encoded vectors configure the hyperparameters a key parameter to pick is the patch_size the size of the input patches in order to use each pixel as an individual input you can set patch_size to 1 1 below we take inspiration from the original paper settings for training on imagenet 1k keeping most of the original settings for this example 2 by 2 sized patches dropout rate attention heads embedding dimension mlp layer size convert embedded patches to query key and values with a learnable additive value size of attention window size of shifting window initial image size helper functions we create two helper functions to help us get a sequence of patches from the image merge patches and apply dropout window based multi head self attention usually transformers perform global self attention where the relationships between a token and all other tokens are computed the global computation leads to quadratic complexity with respect to the number of tokens here as the original paper https arxiv org abs 2103 14030 suggests we compute self attention within local windows in a non overlapping manner global self attention leads to quadratic computational complexity in the number of patches whereas window based self attention leads to linear complexity and is easily scalable the complete swin transformer model finally we put together the complete swin transformer by replacing the standard multi head attention mha with shifted windows attention as suggested in the original paper we create a model comprising of a shifted window based mha layer followed by a 2 layer mlp with gelu nonlinearity in between applying layernormalization before each msa layer and each mlp and a residual connection after each of these layers notice that we only create a simple mlp with 2 dense and 2 dropout layers often you will see models using resnet 50 as the mlp which is quite standard in the literature however in this paper the s use a 2 layer mlp with gelu nonlinearity in between number of input dimensions number of embedded patches number of attention heads size of window size of window shift number of mlp nodes mask array to windows model training and evaluation extract and embed patches we first create 3 layers to help us extract embed and merge patches from the images on top of which we will later use the swin transformer class we built build the model we put together the swin transformer model train on cifar 100 we train the model on cifar 100 here we only train the model for 40 epochs to keep the training time short in this example in practice you should train for 150 epochs to reach convergence let s visualize the training progress of the model let s display the final results of the training on cifar 100 the swin transformer model we just trained has just 152k parameters and it gets us to 75 test top 5 accuracy within just 40 epochs without any signs of overfitting as well as seen in above graph this means we can train this network for longer perhaps with a bit more regularization and obtain even better performance this performance can further be improved by additional techniques like cosine decay learning rate schedule other data augmentation techniques while experimenting i tried training the model for 150 epochs with a slightly higher dropout and greater embedding dimensions which pushes the performance to 72 test accuracy on cifar 100 as you can see in the screenshot results of training for longer https i imgur com 9vnqesz png the s present a top 1 accuracy of 87 3 on imagenet the s also present a number of experiments to study how input sizes optimizers etc affect the final performance of this model the s further present using this model for object detection semantic segmentation and instance segmentation as well and report competitive results for these you are strongly advised to also check out the original paper https arxiv org abs 2103 14030 this example takes inspiration from the official pytorch https github com microsoft swin transformer and tensorflow https github com vcampsoldiers swin transformer tensorflow implementations
import matplotlib.pyplot as plt
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()
x_train, x_test = x_train / 255.0, x_test / 255.0
y_train = keras.utils.numerical_utils.to_categorical(y_train, num_classes)
y_test = keras.utils.numerical_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}")

plt.figure(figsize=(10, 10))
for i in range(25):
    plt.subplot(5, 5, i + 1)
    plt.xticks([])
    plt.yticks([])
    plt.grid(False)
    plt.imshow(x_train[i])
plt.show()



patch_size = (2, 2)  
dropout_rate = 0.03  
num_heads = 8  
embed_dim = 64  
num_mlp = 256  
qkv_bias = True  
window_size = 2  
shift_size = 1  
image_dimension = 32  

num_patch_x = input_shape[0] // patch_size[0]
num_patch_y = input_shape[1] // patch_size[1]

learning_rate = 1e-3
batch_size = 128
num_epochs = 1
validation_split = 0.1
weight_decay = 0.0001
label_smoothing = 0.1




def window_partition(x, window_size):
    _, height, width, channels = x.shape
    patch_num_y = height // window_size
    patch_num_x = width // window_size
    x = tf.reshape(
        x,
        shape=(
            -1,
            patch_num_y,
            window_size,
            patch_num_x,
            window_size,
            channels,
        ),
    )
    x = tf.transpose(x, (0, 1, 3, 2, 4, 5))
    windows = tf.reshape(x, shape=(-1, window_size, window_size, channels))
    return windows


def window_reverse(windows, window_size, height, width, channels):
    patch_num_y = height // window_size
    patch_num_x = width // window_size
    x = tf.reshape(
        windows,
        shape=(
            -1,
            patch_num_y,
            patch_num_x,
            window_size,
            window_size,
            channels,
        ),
    )
    x = tf.transpose(x, perm=(0, 1, 3, 2, 4, 5))
    x = tf.reshape(x, shape=(-1, height, width, channels))
    return x


class DropPath(layers.Layer):
    def __init__(self, drop_prob=None, **kwargs):
        super().__init__(**kwargs)
        self.drop_prob = drop_prob

    def call(self, x):
        input_shape = tf.shape(x)
        batch_size = input_shape[0]
        rank = x.shape.rank
        shape = (batch_size,) + (1,) * (rank - 1)
        random_tensor = (1 - self.drop_prob) + tf.random.uniform(
            shape, dtype=x.dtype
        )
        path_mask = tf.floor(random_tensor)
        output = tf.math.divide(x, 1 - self.drop_prob) * path_mask
        return output





class WindowAttention(layers.Layer):
    def __init__(
        self,
        dim,
        window_size,
        num_heads,
        qkv_bias=True,
        dropout_rate=0.0,
        **kwargs,
    ):
        super().__init__(**kwargs)
        self.dim = dim
        self.window_size = window_size
        self.num_heads = num_heads
        self.scale = (dim // num_heads) ** -0.5
        self.qkv = layers.Dense(dim * 3, use_bias=qkv_bias)
        self.dropout = layers.Dropout(dropout_rate)
        self.proj = layers.Dense(dim)

    def build(self, input_shape):
        num_window_elements = (2 * self.window_size[0] - 1) * (
            2 * self.window_size[1] - 1
        )
        self.relative_position_bias_table = self.add_weight(
            shape=(num_window_elements, self.num_heads),
            initializer=tf.initializers.Zeros(),
            trainable=True,
        )
        coords_h = np.arange(self.window_size[0])
        coords_w = np.arange(self.window_size[1])
        coords_matrix = np.meshgrid(coords_h, coords_w, indexing="ij")
        coords = np.stack(coords_matrix)
        coords_flatten = coords.reshape(2, -1)
        relative_coords = (
            coords_flatten[:, :, None] - coords_flatten[:, None, :]
        )
        relative_coords = relative_coords.transpose([1, 2, 0])
        relative_coords[:, :, 0] += self.window_size[0] - 1
        relative_coords[:, :, 1] += self.window_size[1] - 1
        relative_coords[:, :, 0] *= 2 * self.window_size[1] - 1
        relative_position_index = relative_coords.sum(-1)

        self.relative_position_index = tf.Variable(
            initial_value=lambda: tf.convert_to_tensor(relative_position_index),
            trainable=False,
        )

    def call(self, x, mask=None):
        _, size, channels = x.shape
        head_dim = channels // self.num_heads
        x_qkv = self.qkv(x)
        x_qkv = tf.reshape(x_qkv, shape=(-1, size, 3, self.num_heads, head_dim))
        x_qkv = tf.transpose(x_qkv, perm=(2, 0, 3, 1, 4))
        q, k, v = x_qkv[0], x_qkv[1], x_qkv[2]
        q = q * self.scale
        k = tf.transpose(k, perm=(0, 1, 3, 2))
        attn = q @ k

        num_window_elements = self.window_size[0] * self.window_size[1]
        relative_position_index_flat = tf.reshape(
            self.relative_position_index, shape=(-1,)
        )
        relative_position_bias = tf.gather(
            self.relative_position_bias_table, relative_position_index_flat
        )
        relative_position_bias = tf.reshape(
            relative_position_bias,
            shape=(num_window_elements, num_window_elements, -1),
        )
        relative_position_bias = tf.transpose(
            relative_position_bias, perm=(2, 0, 1)
        )
        attn = attn + tf.expand_dims(relative_position_bias, axis=0)

        if mask is not None:
            nW = mask.shape[0]
            mask_float = tf.cast(
                tf.expand_dims(tf.expand_dims(mask, axis=1), axis=0), tf.float32
            )
            attn = (
                tf.reshape(attn, shape=(-1, nW, self.num_heads, size, size))
                + mask_float
            )
            attn = tf.reshape(attn, shape=(-1, self.num_heads, size, size))
            attn = keras.activations.softmax(attn, axis=-1)
        else:
            attn = keras.activations.softmax(attn, axis=-1)
        attn = self.dropout(attn)

        x_qkv = attn @ v
        x_qkv = tf.transpose(x_qkv, perm=(0, 2, 1, 3))
        x_qkv = tf.reshape(x_qkv, shape=(-1, size, channels))
        x_qkv = self.proj(x_qkv)
        x_qkv = self.dropout(x_qkv)
        return x_qkv





class SwinTransformer(layers.Layer):
    def __init__(
        self,
        dim,
        num_patch,
        num_heads,
        window_size=7,
        shift_size=0,
        num_mlp=1024,
        qkv_bias=True,
        dropout_rate=0.0,
        **kwargs,
    ):
        super().__init__(**kwargs)

        self.dim = dim  
        self.num_patch = num_patch  
        self.num_heads = num_heads  
        self.window_size = window_size  
        self.shift_size = shift_size  
        self.num_mlp = num_mlp  

        self.norm1 = layers.LayerNormalization(epsilon=1e-5)
        self.attn = WindowAttention(
            dim,
            window_size=(self.window_size, self.window_size),
            num_heads=num_heads,
            qkv_bias=qkv_bias,
            dropout_rate=dropout_rate,
        )
        self.drop_path = DropPath(dropout_rate)
        self.norm2 = layers.LayerNormalization(epsilon=1e-5)

        self.mlp = keras.Sequential(
            [
                layers.Dense(num_mlp),
                layers.Activation(keras.activations.gelu),
                layers.Dropout(dropout_rate),
                layers.Dense(dim),
                layers.Dropout(dropout_rate),
            ]
        )

        if min(self.num_patch) < self.window_size:
            self.shift_size = 0
            self.window_size = min(self.num_patch)

    def build(self, input_shape):
        if self.shift_size == 0:
            self.attn_mask = None
        else:
            height, width = self.num_patch
            h_slices = (
                slice(0, -self.window_size),
                slice(-self.window_size, -self.shift_size),
                slice(-self.shift_size, None),
            )
            w_slices = (
                slice(0, -self.window_size),
                slice(-self.window_size, -self.shift_size),
                slice(-self.shift_size, None),
            )
            mask_array = np.zeros((1, height, width, 1))
            count = 0
            for h in h_slices:
                for w in w_slices:
                    mask_array[:, h, w, :] = count
                    count += 1
            mask_array = tf.convert_to_tensor(mask_array)

            
            mask_windows = window_partition(mask_array, self.window_size)
            mask_windows = tf.reshape(
                mask_windows, shape=[-1, self.window_size * self.window_size]
            )
            attn_mask = tf.expand_dims(mask_windows, axis=1) - tf.expand_dims(
                mask_windows, axis=2
            )
            attn_mask = tf.where(attn_mask != 0, -100.0, attn_mask)
            attn_mask = tf.where(attn_mask == 0, 0.0, attn_mask)
            self.attn_mask = tf.Variable(
                initial_value=attn_mask, trainable=False
            )

    def call(self, x):
        height, width = self.num_patch
        _, num_patches_before, channels = x.shape
        x_skip = x
        x = self.norm1(x)
        x = tf.reshape(x, shape=(-1, height, width, channels))
        if self.shift_size > 0:
            shifted_x = tf.roll(
                x, shift=[-self.shift_size, -self.shift_size], axis=[1, 2]
            )
        else:
            shifted_x = x

        x_windows = window_partition(shifted_x, self.window_size)
        x_windows = tf.reshape(
            x_windows, shape=(-1, self.window_size * self.window_size, channels)
        )
        attn_windows = self.attn(x_windows, mask=self.attn_mask)

        attn_windows = tf.reshape(
            attn_windows,
            shape=(-1, self.window_size, self.window_size, channels),
        )
        shifted_x = window_reverse(
            attn_windows, self.window_size, height, width, channels
        )
        if self.shift_size > 0:
            x = tf.roll(
                shifted_x, shift=[self.shift_size, self.shift_size], axis=[1, 2]
            )
        else:
            x = shifted_x

        x = tf.reshape(x, shape=(-1, height * width, channels))
        x = self.drop_path(x)
        x = x_skip + x
        x_skip = x
        x = self.norm2(x)
        x = self.mlp(x)
        x = self.drop_path(x)
        x = x_skip + x
        return x





class PatchExtract(layers.Layer):
    def __init__(self, patch_size, **kwargs):
        super().__init__(**kwargs)
        self.patch_size_x = patch_size[0]
        self.patch_size_y = patch_size[0]

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


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 = tf.range(start=0, limit=self.num_patch, delta=1)
        return self.proj(patch) + self.pos_embed(pos)


class PatchMerging(keras.layers.Layer):
    def __init__(self, num_patch, embed_dim):
        super().__init__()
        self.num_patch = num_patch
        self.embed_dim = embed_dim
        self.linear_trans = layers.Dense(2 * embed_dim, use_bias=False)

    def call(self, x):
        height, width = self.num_patch
        _, _, C = x.shape
        x = tf.reshape(x, shape=(-1, height, width, C))
        x0 = x[:, 0::2, 0::2, :]
        x1 = x[:, 1::2, 0::2, :]
        x2 = x[:, 0::2, 1::2, :]
        x3 = x[:, 1::2, 1::2, :]
        x = tf.concat((x0, x1, x2, x3), axis=-1)
        x = tf.reshape(x, shape=(-1, (height // 2) * (width // 2), 4 * C))
        return self.linear_trans(x)




input = layers.Input(input_shape)
x = layers.RandomCrop(image_dimension, image_dimension)(input)
x = layers.RandomFlip("horizontal")(x)
x = PatchExtract(patch_size)(x)
x = PatchEmbedding(num_patch_x * num_patch_y, embed_dim)(x)
x = SwinTransformer(
    dim=embed_dim,
    num_patch=(num_patch_x, num_patch_y),
    num_heads=num_heads,
    window_size=window_size,
    shift_size=0,
    num_mlp=num_mlp,
    qkv_bias=qkv_bias,
    dropout_rate=dropout_rate,
)(x)
x = SwinTransformer(
    dim=embed_dim,
    num_patch=(num_patch_x, num_patch_y),
    num_heads=num_heads,
    window_size=window_size,
    shift_size=shift_size,
    num_mlp=num_mlp,
    qkv_bias=qkv_bias,
    dropout_rate=dropout_rate,
)(x)
x = PatchMerging((num_patch_x, num_patch_y), embed_dim=embed_dim)(x)
x = layers.GlobalAveragePooling1D()(x)
output = layers.Dense(num_classes, activation="softmax")(x)



model = keras.Model(input, output)
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)}%")