make vit names match pytorch
MODIFY
models/transformer.py
models/transformer.py
@@ -8,11 +8,11 @@ class TransformerBlock:
assert self.head_size * self.num_heads == embed_dim
self.prenorm, self.act = prenorm, act
- self.query_dense = (Tensor.uniform(embed_dim, embed_dim), Tensor.zeros(embed_dim))
- self.key_dense = (Tensor.uniform(embed_dim, embed_dim), Tensor.zeros(embed_dim))
- self.value_dense = (Tensor.uniform(embed_dim, embed_dim), Tensor.zeros(embed_dim))
+ self.query = (Tensor.uniform(embed_dim, embed_dim), Tensor.zeros(embed_dim))
+ self.key = (Tensor.uniform(embed_dim, embed_dim), Tensor.zeros(embed_dim))
+ self.value = (Tensor.uniform(embed_dim, embed_dim), Tensor.zeros(embed_dim))
- self.final = (Tensor.uniform(embed_dim, embed_dim), Tensor.zeros(embed_dim))
+ self.out = (Tensor.uniform(embed_dim, embed_dim), Tensor.zeros(embed_dim))
self.ff1 = (Tensor.uniform(embed_dim, ff_dim), Tensor.zeros(ff_dim))
self.ff2 = (Tensor.uniform(ff_dim, embed_dim), Tensor.zeros(embed_dim))
@@ -23,7 +23,7 @@ class TransformerBlock:
def attn(self, x):
query, key, value = [x.linear(*y) \
.reshape(shape=(x.shape[0], -1, self.num_heads, self.head_size)) \
- for y in [self.query_dense, self.key_dense, self.value_dense]]
+ for y in [self.query, self.key, self.value]]
query = query.transpose(order=(0,2,1,3)) # (bs, num_heads, T, head_size)
key = key.transpose(order=(0,2,3,1)) # (bs, num_heads, head_size, T)
@@ -33,7 +33,7 @@ class TransformerBlock:
weights = score.softmax() # (bs, num_heads, T, T)
attention = weights.dot(value).transpose(order=(0,2,1,3)) # (bs, T, num_heads, head_size)
- return attention.reshape(shape=(x.shape[0], -1, self.num_heads * self.head_size)).linear(*self.final)
+ return attention.reshape(shape=(x.shape[0], -1, self.num_heads * self.head_size)).linear(*self.out)
def __call__(self, x):
if self.prenorm:

MODIFY
models/vit.py
models/vit.py
@@ -4,26 +4,26 @@ from models.transformer import TransformerBlock
class ViT:
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.pos_embed = Tensor.ones(1, 197, embed_dim)
+ self.embedding = (Tensor.uniform(embed_dim, 3, 16, 16), Tensor.zeros(embed_dim))
+ self.cls = Tensor.ones(1, 1, embed_dim)
+ self.pos_embedding = Tensor.ones(1, 197, embed_dim)
self.tbs = [
TransformerBlock(embed_dim=embed_dim, num_heads=num_heads, ff_dim=embed_dim*4,
prenorm=True, act=lambda x: x.gelu())
for i in range(layers)]
- self.norm = (Tensor.uniform(embed_dim), Tensor.zeros(embed_dim))
+ self.encoder_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, stride=16)
+ x = x.conv2d(*self.embedding, stride=16)
x = x.reshape(shape=(x.shape[0], x.shape[1], -1)).transpose(order=(0,2,1))
return x
def forward(self, x):
pe = self.patch_embed(x)
- x = self.cls_token.add(Tensor.zeros(pe.shape[0],1,1)).cat(pe, dim=1) + self.pos_embed
+ x = self.cls.add(Tensor.zeros(pe.shape[0],1,1)).cat(pe, dim=1) + self.pos_embedding
x = x.sequential(self.tbs)
- x = x.layernorm().linear(*self.norm)
+ x = x.layernorm().linear(*self.encoder_norm)
return x[:, 0].linear(*self.head)
def load_from_pretrained(m):
@@ -37,27 +37,27 @@ class ViT:
#for x in dat.keys():
# print(x, dat[x].shape, dat[x].dtype)
- m.conv[0].assign(np.transpose(dat['embedding/kernel'], (3,2,0,1)))
- m.conv[1].assign(dat['embedding/bias'])
+ m.embedding[0].assign(np.transpose(dat['embedding/kernel'], (3,2,0,1)))
+ m.embedding[1].assign(dat['embedding/bias'])
- m.norm[0].assign(dat['Transformer/encoder_norm/scale'])
- m.norm[1].assign(dat['Transformer/encoder_norm/bias'])
+ m.encoder_norm[0].assign(dat['Transformer/encoder_norm/scale'])
+ m.encoder_norm[1].assign(dat['Transformer/encoder_norm/bias'])
m.head[0].assign(dat['head/kernel'])
m.head[1].assign(dat['head/bias'])
- m.cls_token.assign(dat['cls'])
- m.pos_embed.assign(dat['Transformer/posembed_input/pos_embedding'])
+ m.cls.assign(dat['cls'])
+ m.pos_embedding.assign(dat['Transformer/posembed_input/pos_embedding'])
for i in range(12):
- m.tbs[i].query_dense[0].assign(dat[f'Transformer/encoderblock_{i}/MultiHeadDotProductAttention_1/query/kernel'].reshape(192, 192))
- m.tbs[i].query_dense[1].assign(dat[f'Transformer/encoderblock_{i}/MultiHeadDotProductAttention_1/query/bias'].reshape(192))
- m.tbs[i].key_dense[0].assign(dat[f'Transformer/encoderblock_{i}/MultiHeadDotProductAttention_1/key/kernel'].reshape(192, 192))
- m.tbs[i].key_dense[1].assign(dat[f'Transformer/encoderblock_{i}/MultiHeadDotProductAttention_1/key/bias'].reshape(192))
- m.tbs[i].value_dense[0].assign(dat[f'Transformer/encoderblock_{i}/MultiHeadDotProductAttention_1/value/kernel'].reshape(192, 192))
- m.tbs[i].value_dense[1].assign(dat[f'Transformer/encoderblock_{i}/MultiHeadDotProductAttention_1/value/bias'].reshape(192))
- m.tbs[i].final[0].assign(dat[f'Transformer/encoderblock_{i}/MultiHeadDotProductAttention_1/out/kernel'].reshape(192, 192))
- m.tbs[i].final[1].assign(dat[f'Transformer/encoderblock_{i}/MultiHeadDotProductAttention_1/out/bias'].reshape(192))
+ m.tbs[i].query[0].assign(dat[f'Transformer/encoderblock_{i}/MultiHeadDotProductAttention_1/query/kernel'].reshape(192, 192))
+ m.tbs[i].query[1].assign(dat[f'Transformer/encoderblock_{i}/MultiHeadDotProductAttention_1/query/bias'].reshape(192))
+ m.tbs[i].key[0].assign(dat[f'Transformer/encoderblock_{i}/MultiHeadDotProductAttention_1/key/kernel'].reshape(192, 192))
+ m.tbs[i].key[1].assign(dat[f'Transformer/encoderblock_{i}/MultiHeadDotProductAttention_1/key/bias'].reshape(192))
+ m.tbs[i].value[0].assign(dat[f'Transformer/encoderblock_{i}/MultiHeadDotProductAttention_1/value/kernel'].reshape(192, 192))
+ m.tbs[i].value[1].assign(dat[f'Transformer/encoderblock_{i}/MultiHeadDotProductAttention_1/value/bias'].reshape(192))
+ m.tbs[i].out[0].assign(dat[f'Transformer/encoderblock_{i}/MultiHeadDotProductAttention_1/out/kernel'].reshape(192, 192))
+ m.tbs[i].out[1].assign(dat[f'Transformer/encoderblock_{i}/MultiHeadDotProductAttention_1/out/bias'].reshape(192))
m.tbs[i].ff1[0].assign(dat[f'Transformer/encoderblock_{i}/MlpBlock_3/Dense_0/kernel'])
m.tbs[i].ff1[1].assign(dat[f'Transformer/encoderblock_{i}/MlpBlock_3/Dense_0/bias'])
m.tbs[i].ff2[0].assign(dat[f'Transformer/encoderblock_{i}/MlpBlock_3/Dense_1/kernel'])

make vit names match pytorch