added conv profile test
ADD
test/test_conv_speed.py
+#!/usr/bin/env python
+import cProfile
+import unittest
+from tinygrad.tensor import Tensor
+
+def profile_conv(bs, chans, conv, cnt=100):
+ img = Tensor.zeros(bs, 1, 28, 28)
+ conv = Tensor.randn(chans, 1, conv, conv)
+ for i in range(cnt):
+ out = img.conv2d(conv)
+
+class TestConvSpeed(unittest.TestCase):
+ def test_forward_3x3(self):
+ pr = cProfile.Profile()
+ pr.enable()
+ profile_conv(128, 16, 3)
+ pr.disable()
+ pr.print_stats(sort='time')
+
+if __name__ == '__main__':
+ unittest.main()
+

MODIFY
tinygrad/tensor.py
tinygrad/tensor.py
# inspired by https://github.com/karpathy/micrograd/blob/master/micrograd/engine.py
from functools import partialmethod
import numpy as np
+from tinygrad.utils import im2col, col2im
# **** start with two base classes ****
@@ -26,6 +27,14 @@ class Tensor:
def shape(self):
return self.data.shape
+ @staticmethod
+ def zeros(*shape):
+ return Tensor(np.zeros(shape, dtype=np.float32))
+
+ @staticmethod
+ def randn(*shape):
+ return Tensor(np.random.randn(*shape).astype(np.float32))
+
def backward(self, allow_fill=True):
#print("running backward on", self)
if self._ctx is None:
@@ -216,13 +225,9 @@ class FastConv2D(Function):
bs,oy,ox = x.shape[0], x.shape[2]-(H-1), x.shape[3]-(W-1)
# im2col
- tx = np.empty((oy, ox, bs, cin*W*H), dtype=x.dtype)
- for Y in range(oy):
- for X in range(ox):
- tx[Y, X] = x[:, :, Y:Y+H, X:X+W].reshape(bs, -1)
- tx = tx.reshape(-1, cin*W*H)
+ tx = im2col(x, H, W)
- # save the im2col output
+ # save the im2col output (OMG it's bigger!)
ctx.save_for_backward(tx, w)
# now the conv is a GEMM
@@ -245,13 +250,11 @@ class FastConv2D(Function):
dw = gg.T.dot(tx).reshape(w.shape)
# dx is harder
- dxi = gg.dot(tw).reshape(oy, ox, bs, cin, H, W)
+ dxi = gg.dot(tw)
+
+ # if we im2col on the forward, we col2im on the backward
+ dx = col2im(dxi, H, W, oy+(H-1), ox+(W-1))
- # unim2col (is there a faster way to do this?)
- dx = np.zeros((bs, cin, oy+(H-1), ox+(W-1)), dtype=dxi.dtype)
- for Y in range(oy):
- for X in range(ox):
- dx[:, :, Y:Y+H, X:X+W] += dxi[Y, X]
return dx, dw
register('conv2d', FastConv2D)

MODIFY
tinygrad/utils.py
tinygrad/utils.py
@@ -26,3 +26,28 @@ def fetch_mnist():
X_test = fetch("http://yann.lecun.com/exdb/mnist/t10k-images-idx3-ubyte.gz")[0x10:].reshape((-1, 28, 28))
Y_test = fetch("http://yann.lecun.com/exdb/mnist/t10k-labels-idx1-ubyte.gz")[8:]
return X_train, Y_train, X_test, Y_test
+
+
+# these are matlab functions used to speed up convs
+# write them fast and the convs will be fast?
+
+def im2col(x, H, W):
+ bs,cin,oy,ox = x.shape[0], x.shape[1], x.shape[2]-(H-1), x.shape[3]-(W-1)
+
+ tx = np.empty((oy, ox, bs, cin*W*H), dtype=x.dtype)
+ for Y in range(oy):
+ for X in range(ox):
+ tx[Y, X] = x[:, :, Y:Y+H, X:X+W].reshape(bs, -1)
+ return tx.reshape(-1, cin*W*H)
+
+def col2im(tx, H, W, OY, OX):
+ oy, ox = OY-(H-1), OX-(W-1)
+ bs = tx.shape[0] // (oy * ox)
+ cin = tx.shape[1] // (H * W)
+ tx = tx.reshape(oy, ox, bs, cin, H, W)
+
+ x = np.zeros((bs, cin, OY, OX), dtype=tx.dtype)
+ for Y in range(oy):
+ for X in range(ox):
+ x[:, :, Y:Y+H, X:X+W] += tx[Y, X]
+ return x

added conv profile test