Upload ViT.py

ViT.py

@@ -0,0 +1,135 @@
+import tensorflow as tf
+from tensorflow.keras.layers import Dense,LayerNormalization,Dropout,Identity,Activation
+from tensorflow.keras import Model
+
+
+def pair(t):
+    return t if isinstance(t, tuple) else (t, t)
+
+
+class FeedForward:
+    def __init__(self, dim, hidden_dim, drop_rate = 0.):
+        self.net = tf.keras.Sequential()
+        self.net.add(LayerNormalization())
+        self.net.add(Dense(hidden_dim))
+        self.net.add(Activation('gelu'))
+        self.net.add(Dropout(drop_rate))
+        self.net.add(Dense(dim))
+        self.net.add(Dropout(drop_rate))
+
+    def __call__(self, x):
+        return self.net(x)
+
+
+class Attention:
+    def __init__(self, dim, heads = 8, dim_head = 64, drop_rate = 0.):
+        inner_dim = dim_head *  heads
+        project_out = not (heads == 1 and dim_head == dim)
+
+        self.heads = heads
+        self.scale = dim_head ** -0.5
+
+        self.norm = LayerNormalization()
+
+        self.attend = tf.nn.softmax
+        self.dropout = Dropout(drop_rate)
+
+        self.to_qkv = Dense(inner_dim * 3, use_bias = False)
+        
+        if project_out:
+            self.to_out = tf.keras.Sequential()
+            self.to_out.add(Dense(dim))
+            self.to_out.add(Dropout(drop_rate))
+        else:
+            self.to_out = Identity()
+
+    def __call__(self, x):
+        x = self.norm(x)
+
+        qkv = self.to_qkv(x)
+        q, k, v = tf.split(qkv, 3, axis=-1)
+        b = q.shape[0]
+        h = self.heads
+        n = q.shape[1]
+        d = q.shape[2] // self.heads
+        q = tf.reshape(q, (b, h, n, d))
+        k = tf.reshape(k, (b, h, n, d))
+        v = tf.reshape(v, (b, h, n, d))
+
+        dots = tf.matmul(q, tf.transpose(k, [0, 1, 3, 2])) * self.scale
+
+        attn = self.attend(dots)
+        attn = self.dropout(attn)
+
+        out = tf.matmul(attn, v)
+        out = tf.transpose(out, [0, 1, 3, 2])
+        out = tf.reshape(out, shape=[-1, n, h*d])
+        return self.to_out(out)
+
+
+class Transformer:
+    def __init__(self, dim, depth, heads, dim_head, mlp_dim, dropout = 0.):
+        self.norm = LayerNormalization()
+        self.layers = []
+        for _ in range(depth):
+            self.layers.append([Attention(dim, heads = heads, dim_head = dim_head, drop_rate = dropout),
+                                FeedForward(dim, mlp_dim, drop_rate = dropout)])
+
+    def __call__(self, x):
+        for attn, ff in self.layers:
+            x = attn(x) + x
+            x = ff(x) + x
+
+        return self.norm(x)
+
+
+class ViT(Model):
+    def __init__(self, image_size, patch_size, num_classes, dim, depth, heads, mlp_dim, pool = 'cls', channels = 3, dim_head = 64, drop_rate = 0., emb_dropout = 0.):
+        super(ViT, self).__init__()
+        image_height, image_width = pair(image_size)
+        patch_height, patch_width = pair(patch_size)
+        self.p1, self.p2 = patch_height, patch_width
+        self.dim = dim
+
+        assert image_height % patch_height == 0 and image_width % patch_width == 0, 'Image dimensions must be divisible by the patch size.'
+
+        num_patches = (image_height // patch_height) * (image_width // patch_width)
+        assert pool in {'cls', 'mean'}, 'pool type must be either cls (cls token) or mean (mean pooling)'
+
+        self.to_patch_embedding = tf.keras.Sequential()
+        self.to_patch_embedding.add(LayerNormalization())
+        self.to_patch_embedding.add(Dense(dim))
+        self.to_patch_embedding.add(LayerNormalization())
+
+        self.pos_embedding = tf.Variable(tf.random.normal((1, num_patches + 1, dim)))
+        self.cls_token = tf.Variable(tf.random.normal(((1, 1, dim))))
+        self.dropout = Dropout(emb_dropout)
+
+        self.transformer = Transformer(dim, depth, heads, dim_head, mlp_dim, drop_rate)
+
+        self.pool = pool
+        self.to_latent = Identity()
+
+        self.mlp_head = Dense(num_classes)
+
+
+    def __call__(self, data):
+        b = data.shape[0]
+        h = data.shape[1] // self.p1
+        w = data.shape[2] // self.p2
+        c = data.shape[3]
+        data = tf.reshape(data, (b, h * w, self.p1 * self.p2 * c))
+        x = self.to_patch_embedding(data)
+        b, n, _ = x.shape
+
+        cls_tokens = tf.tile(self.cls_token, multiples=[b, 1, 1])
+        x = tf.concat([cls_tokens, x], axis=1)
+        x += self.pos_embedding[:, :(n + 1)]
+        x = self.dropout(x)
+
+        x = self.transformer(x)
+
+        x = tf.reduce_mean(x, axis = 1) if self.pool == 'mean' else x[:, 0]
+
+        x = self.to_latent(x)
+        return tf.nn.softmax(self.mlp_head(x))