| 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)) |