gelu with broken test
MODIFY
examples/vit.py
examples/vit.py
@@ -8,6 +8,8 @@ with tf.io.gfile.GFile(fn, "rb") as f:
g.write(dat)
"""
+import io
+from extra.utils import fetch
from tinygrad.tensor import Tensor
from models.transformer import TransformerBlock
@@ -15,10 +17,11 @@ class ViT:
def __init__(self):
self.conv_weight = Tensor.uniform(192, 3, 16, 16)
self.conv_bias = Tensor.zeros(192)
- self.cls = Tensor.ones(1, 1, 192)
+ self.cls_token = Tensor.ones(1, 1, 192)
self.tbs = [TransformerBlock(embed_dim=192, num_heads=3, ff_dim=768) for i in range(12)]
- self.pos = Tensor.ones(1, 197, 192)
- self.head = (Tensor.uniform(192, 21843), Tensor.zeros(21843))
+ self.pos_embed = Tensor.ones(1, 197, 192)
+ self.head = (Tensor.uniform(192, 1000), Tensor.zeros(1000))
+ self.norm = (Tensor.uniform(192), Tensor.zeros(192))
def forward(self, x):
print(x.shape)
@@ -27,23 +30,32 @@ class ViT:
print(x.shape)
x = x.reshape(shape=(x.shape[0], 192, -1)).transpose(order=(0,2,1))
print(x.shape)
- x = self.cls.cat(x, dim=1)
+ # TODO: expand cls_token for batch
+ x = self.cls_token.cat(x, dim=1).add(self.pos_embed)
print(x.shape)
for l in self.tbs:
x = l(x)
+ x = x.affine(self.norm)
return x[:, 0].affine(self.head)
m = ViT()
import numpy as np
-dat = np.load("cache/Ti_16-i21k-300ep-lr_0.001-aug_none-wd_0.03-do_0.0-sd_0.0.npz")
+# https://github.com/rwightman/pytorch-image-models/blob/master/timm/models/vision_transformer.py
+dat = np.load(io.BytesIO(fetch("https://storage.googleapis.com/vit_models/augreg/Ti_16-i21k-300ep-lr_0.001-aug_none-wd_0.03-do_0.0-sd_0.0--imagenet2012-steps_20k-lr_0.03-res_224.npz")))
for x in dat.keys():
print(x, dat[x].shape, dat[x].dtype)
m.conv_weight.assign(np.transpose(dat['embedding/kernel'], (3,2,1,0)))
m.conv_bias.assign(dat['embedding/bias'])
-m.cls.assign(dat['cls'])
-m.pos.assign(dat['Transformer/posembed_input/pos_embedding'])
+
+m.norm[0].assign(dat['Transformer/encoder_norm/scale'])
+m.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'])
for i in range(12):
m.tbs[i].query_dense[0].assign(dat[f'Transformer/encoderblock_{i}/MultiHeadDotProductAttention_1/query/kernel'].reshape(192, 192))
@@ -63,8 +75,36 @@ for i in range(12):
m.tbs[i].ln2[0].assign(dat[f'Transformer/encoderblock_{i}/LayerNorm_2/scale'])
m.tbs[i].ln2[1].assign(dat[f'Transformer/encoderblock_{i}/LayerNorm_2/bias'])
-test_input = Tensor.ones(1, 3, 224, 224)
-out = m.forward(test_input)
-print(out.shape)
+url = "https://upload.wikimedia.org/wikipedia/commons/4/41/Chicken.jpg"
+#url = "https://repository-images.githubusercontent.com/296744635/39ba6700-082d-11eb-98b8-cb29fb7369c0"
+
+# category labels
+import ast
+lbls = fetch("https://gist.githubusercontent.com/yrevar/942d3a0ac09ec9e5eb3a/raw/238f720ff059c1f82f368259d1ca4ffa5dd8f9f5/imagenet1000_clsidx_to_labels.txt")
+lbls = ast.literal_eval(lbls.decode('utf-8'))
+
+# junk
+from PIL import Image
+img = Image.open(io.BytesIO(fetch(url)))
+aspect_ratio = img.size[0] / img.size[1]
+img = img.resize((int(224*max(aspect_ratio,1.0)), int(224*max(1.0/aspect_ratio,1.0))))
+img = np.array(img)
+y0,x0=(np.asarray(img.shape)[:2]-224)//2
+img = img[y0:y0+224, x0:x0+224]
+img = np.moveaxis(img, [2,0,1], [0,1,2])
+img = img.astype(np.float32)[:3].reshape(1,3,224,224)
+img /= 255.0
+
+out = m.forward(Tensor(img))
+outnp = out.cpu().data.ravel()
+choice = outnp.argmax()
+print(out.shape, choice, outnp[choice])
+
+print(lbls[choice])
+
+#lookup = dict([x.split(" ") for x in open("cache/classids.txt").read().strip().split("\n")])
+#cls = open("cache/imagenet21k_wordnet_ids.txt").read().strip().split("\n")
+#print(cls[choice], lookup[cls[choice]])
+

MODIFY
models/transformer.py
models/transformer.py
@@ -51,7 +51,7 @@ class TransformerBlock:
x = inputs + attention.reshape(shape=(-1, embed_dim)).affine(self.final).dropout(0.1)
x = layernorm(x, embed_dim).affine(self.ln1)
- x = x + x.affine(self.ff1).relu().affine(self.ff2).dropout(0.1)
+ x = x + x.affine(self.ff1).gelu().affine(self.ff2).dropout(0.1)
x = layernorm(x, embed_dim).affine(self.ln2)
return x.reshape(shape=(bs, -1, embed_dim))

MODIFY
test/test_ops.py
test/test_ops.py
@@ -69,6 +69,10 @@ class TestOps(unittest.TestCase):
helper_test_op([(45,65)], lambda x: x.sigmoid(), Tensor.sigmoid)
def test_softplus(self):
helper_test_op([(45,65)], lambda x: torch.nn.functional.softplus(x), Tensor.softplus, atol=1e-6, grad_atol=1e-6)
+ def test_gelu(self):
+ pass
+ # fails?
+ #helper_test_op([(45,65)], lambda x: 0.5 * x * (1.0 + torch.tanh(x * 0.7978845608 * (1.0 + 0.044715 * x * x))), Tensor.gelu)
def test_relu6(self):
helper_test_op([(45,65)], lambda x: torch.nn.functional.relu6(x), Tensor.relu6)
def test_hardswish(self):

MODIFY
tinygrad/ops/ops_cpu.py
tinygrad/ops/ops_cpu.py
@@ -44,7 +44,7 @@ class Log(Function):
class Exp(Function):
def forward(ctx, input):
- ret = input.exp()
+ ret = np.clip(input, -88, 88).exp()
ctx.save_for_backward(ret)
return ret

MODIFY
tinygrad/tensor.py
tinygrad/tensor.py
@@ -218,8 +218,8 @@ class Tensor:
__truediv__ = div
def sigmoid(self):
- e = self.exp()
- return e.div(1 + e)
+ #e = self.exp(); return e.div(1 + e)
+ return (1.0 + (0.0-self).exp()) ** -1.0
def swish(self):
return self * self.sigmoid()
@@ -233,6 +233,10 @@ class Tensor:
def tanh(self):
return 2.0 * ((2.0 * self).sigmoid()) - 1.0
+ def gelu(x):
+ # https://github.com/huggingface/transformers/blob/master/src/transformers/activations.py
+ return 0.5 * x * (1 + (x * 0.7978845608 * (1 + 0.044715 * x * x)).tanh())
+
def leakyrelu(self, neg_slope=0.01):
return self.relu() - (-neg_slope*self).relu()

gelu with broken test