clean up vit code
MODIFY
examples/vit.py
examples/vit.py
@@ -23,8 +23,8 @@ dat = np.load(io.BytesIO(fetch("https://storage.googleapis.com/vit_models/augreg
#for x in dat.keys():
# print(x, dat[x].shape, dat[x].dtype)
-m.conv_weight.assign(np.transpose(dat['embedding/kernel'], (3,2,0,1)))
-m.conv_bias.assign(dat['embedding/bias'])
+m.conv[0].assign(np.transpose(dat['embedding/kernel'], (3,2,0,1)))
+m.conv[1].assign(dat['embedding/bias'])
m.norm[0].assign(dat['Transformer/encoder_norm/scale'])
m.norm[1].assign(dat['Transformer/encoder_norm/bias'])

MODIFY
models/transformer.py
models/transformer.py
@@ -51,7 +51,6 @@ class TransformerBlock:
return x
class Transformer:
- # L = layers, H = embed_dim, A = num_heads
def __init__(self, syms, maxlen, layers, embed_dim, num_heads, ff_dim):
self.maxlen, self.syms = maxlen, syms
self.embed = Tensor.uniform(maxlen+syms, embed_dim, requires_grad=False)
@@ -75,18 +74,16 @@ class Transformer:
return x.reshape(shape=(bs, -1, x.shape[-1]))
class ViT:
- def __init__(self, embed_dim=192):
- self.conv_weight = Tensor.uniform(embed_dim, 3, 16, 16)
- self.conv_bias = Tensor.zeros(embed_dim)
+ def __init__(self, layers=12, embed_dim=192, num_heads=3):
+ self.conv = (Tensor.uniform(embed_dim, 3, 16, 16), Tensor.zeros(embed_dim))
self.cls_token = Tensor.ones(1, 1, embed_dim)
- self.tbs = [TransformerBlock(embed_dim=embed_dim, num_heads=3, ff_dim=768, prenorm=True) for i in range(12)]
self.pos_embed = Tensor.ones(1, 197, embed_dim)
- self.head = (Tensor.uniform(embed_dim, 1000), Tensor.zeros(1000))
+ self.tbs = [TransformerBlock(embed_dim=embed_dim, num_heads=num_heads, ff_dim=embed_dim*4, prenorm=True) for i in range(layers)]
self.norm = (Tensor.uniform(embed_dim), Tensor.zeros(embed_dim))
+ self.head = (Tensor.uniform(embed_dim, 1000), Tensor.zeros(1000))
def patch_embed(self, x):
- x = x.conv2d(self.conv_weight, stride=16)
- x = x.add(self.conv_bias.reshape(shape=(1,-1,1,1)))
+ x = x.conv2d(*self.conv, stride=16)
x = x.reshape(shape=(x.shape[0], x.shape[1], -1)).transpose(order=(0,2,1))
return x

clean up vit code