|
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 = self.add_weight( |
|
name='pos_embedding', |
|
shape=(1, self.num_patches + 1, self.dim), |
|
initializer=tf.keras.initializers.RandomNormal(stddev=0.02), |
|
trainable=True |
|
) |
|
self.cls_token = self.add_weight( |
|
name='cls_token', |
|
shape=(1, 1, self.dim), |
|
initializer=tf.keras.initializers.RandomNormal(stddev=0.02), |
|
trainable=True |
|
) |
|
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)) |
