model touchups
MODIFY
examples/efficientnet.py
examples/efficientnet.py
@@ -49,7 +49,7 @@ def infer(model, img):
if __name__ == "__main__":
# instantiate my net
model = EfficientNet(int(os.getenv("NUM", "0")))
- model.load_weights_from_torch()
+ model.load_from_pretrained()
# category labels
import ast

MODIFY
examples/train_efficientnet.py
examples/train_efficientnet.py
@@ -33,7 +33,7 @@ if __name__ == "__main__":
model = TinyConvNet(classes)
elif TRANSFER:
model = EfficientNet(int(os.getenv("NUM", "0")), classes, has_se=True)
- model.load_weights_from_torch()
+ model.load_from_pretrained()
else:
model = EfficientNet(int(os.getenv("NUM", "0")), classes, has_se=False)

MODIFY
models/efficientnet.py
models/efficientnet.py
@@ -122,7 +122,7 @@ class EfficientNet:
x = x.reshape(shape=(-1, x.shape[1]))
return x.linear(self._fc, self._fc_bias)
- def load_weights_from_torch(self):
+ def load_from_pretrained(self):
model_urls = {
0: "https://github.com/lukemelas/EfficientNet-PyTorch/releases/download/1.0/efficientnet-b0-355c32eb.pth",
1: "https://github.com/lukemelas/EfficientNet-PyTorch/releases/download/1.0/efficientnet-b1-f1951068.pth",
@@ -157,7 +157,7 @@ class EfficientNet:
vnp = vnp if vnp.shape != () else np.array([vnp])
if mv.shape == vnp.shape:
- mv.assign(Tensor(vnp))
+ mv.assign(vnp)
else:
print("MISMATCH SHAPE IN %s, %r %r" % (k, mv.shape, vnp.shape))

MODIFY
models/transformer.py
models/transformer.py
@@ -2,14 +2,12 @@ import numpy as np
from tinygrad.tensor import Tensor
class TransformerBlock:
- def __init__(self, embed_dim, num_heads, ff_dim, prenorm=False):
- # Multi-Head Attention
+ def __init__(self, embed_dim, num_heads, ff_dim, prenorm=False, act=lambda x: x.relu()):
self.num_heads = num_heads
self.head_size = embed_dim // num_heads
assert self.head_size * self.num_heads == embed_dim
- self.prenorm = prenorm
+ self.prenorm, self.act = prenorm, act
- # added bias
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))
@@ -23,8 +21,6 @@ class TransformerBlock:
self.ln2 = (Tensor.ones(embed_dim), Tensor.zeros(embed_dim))
def attn(self, x):
- embed_dim = self.num_heads * self.head_size
-
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]]
@@ -37,16 +33,16 @@ 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, embed_dim)).linear(*self.final)
+ return attention.reshape(shape=(x.shape[0], -1, self.num_heads * self.head_size)).linear(*self.final)
def __call__(self, x):
if self.prenorm:
x = x + self.attn(x.layernorm().linear(*self.ln1)).dropout(0.1)
- x = x + x.layernorm().linear(*self.ln2).linear(*self.ff1).gelu().linear(*self.ff2).dropout(0.1)
+ x = x + self.act(x.layernorm().linear(*self.ln2).linear(*self.ff1)).linear(*self.ff2).dropout(0.1)
else:
x = x + self.attn(x).dropout(0.1)
x = x.layernorm().linear(*self.ln1)
- x = x + x.linear(*self.ff1).relu().linear(*self.ff2).dropout(0.1)
+ x = x + self.act(x.linear(*self.ff1)).linear(*self.ff2).dropout(0.1)
x = x.layernorm().linear(*self.ln2)
return x
@@ -78,7 +74,10 @@ class ViT:
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.tbs = [TransformerBlock(embed_dim=embed_dim, num_heads=num_heads, ff_dim=embed_dim*4, prenorm=True) for i in range(layers)]
+ 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.head = (Tensor.uniform(embed_dim, 1000), Tensor.zeros(1000))

MODIFY
test/test_efficientnet.py
test/test_efficientnet.py
@@ -42,7 +42,7 @@ class TestEfficientNet(unittest.TestCase):
def test_chicken(self):
chicken_img = Image.open(pathlib.Path(__file__).parent / 'efficientnet/Chicken.jpg')
model = EfficientNet(number=0)
- model.load_weights_from_torch()
+ model.load_from_pretrained()
label = _infer(model, chicken_img)
self.assertEqual(label, "hen", f"Expected hen but got {label} for number=0")

model touchups