You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
88 lines
2.5 KiB
88 lines
2.5 KiB
1 month ago
|
#!/usr/bin/env python
|
||
|
|
import time
|
||
|
|
import unittest
|
||
|
|
import torch
|
||
|
|
from tinygrad import Tensor, Device
|
||
|
|
from tinygrad.helpers import Profiling, CI
|
||
|
|
|
||
|
|
@unittest.skipIf(CI and Device.DEFAULT in {"CUDA", "NV"}, "slow")
|
||
|
|
class TestConvSpeed(unittest.TestCase):
|
||
|
|
|
||
|
|
def test_mnist(self):
|
||
|
|
# https://keras.io/examples/vision/mnist_convnet/
|
||
|
|
conv = 3
|
||
|
|
inter_chan, out_chan = 32, 64
|
||
|
|
|
||
|
|
# ****** torch baseline *******
|
||
|
|
|
||
|
|
torch.backends.mkldnn.enabled = False
|
||
|
|
|
||
|
|
conv = 3
|
||
|
|
inter_chan, out_chan = 32, 64
|
||
|
|
c1 = torch.randn(inter_chan,1,conv,conv, requires_grad=True)
|
||
|
|
c2 = torch.randn(out_chan,inter_chan,conv,conv, requires_grad=True)
|
||
|
|
l1 = torch.randn(out_chan*5*5, 10, requires_grad=True)
|
||
|
|
|
||
|
|
c2d = torch.nn.functional.conv2d
|
||
|
|
mp = torch.nn.MaxPool2d((2,2))
|
||
|
|
lsm = torch.nn.LogSoftmax(dim=1)
|
||
|
|
|
||
|
|
cnt = 5
|
||
|
|
fpt, bpt = 0.0, 0.0
|
||
|
|
for i in range(cnt):
|
||
|
|
et0 = time.time()
|
||
|
|
x = torch.randn(128, 1, 28, 28, requires_grad=True)
|
||
|
|
x = mp(c2d(x,c1).relu())
|
||
|
|
x = mp(c2d(x,c2).relu())
|
||
|
|
x = x.reshape(x.shape[0], -1)
|
||
|
|
out = lsm(x.matmul(l1))
|
||
|
|
out = out.mean()
|
||
|
|
et1 = time.time()
|
||
|
|
out.backward()
|
||
|
|
et2 = time.time()
|
||
|
|
fpt += (et1-et0)
|
||
|
|
bpt += (et2-et1)
|
||
|
|
|
||
|
|
fpt_baseline = (fpt*1000/cnt)
|
||
|
|
bpt_baseline = (bpt*1000/cnt)
|
||
|
|
print("torch forward pass: %.3f ms" % fpt_baseline)
|
||
|
|
print("torch backward pass: %.3f ms" % bpt_baseline)
|
||
|
|
|
||
|
|
# ****** tinygrad compare *******
|
||
|
|
|
||
|
|
c1 = Tensor(c1.detach().numpy(), requires_grad=True)
|
||
|
|
c2 = Tensor(c2.detach().numpy(), requires_grad=True)
|
||
|
|
l1 = Tensor(l1.detach().numpy(), requires_grad=True)
|
||
|
|
|
||
|
|
cnt = 5
|
||
|
|
fpt, bpt = 0.0, 0.0
|
||
|
|
for i in range(1+cnt):
|
||
|
|
et0 = time.time()
|
||
|
|
x = Tensor.randn(128, 1, 28, 28)
|
||
|
|
x = x.conv2d(c1).relu().avg_pool2d()
|
||
|
|
x = x.conv2d(c2).relu().max_pool2d()
|
||
|
|
x = x.reshape(shape=(x.shape[0], -1))
|
||
|
|
out = x.dot(l1).log_softmax()
|
||
|
|
out = out.mean()
|
||
|
|
out.backward() # NOTE: we have to now compute this here, but it doesn't realize
|
||
|
|
out.realize()
|
||
|
|
et1 = time.time()
|
||
|
|
[x.grad.realize() for x in [c1, c2, l1]]
|
||
|
|
et2 = time.time()
|
||
|
|
if i == 0:
|
||
|
|
pr = Profiling(sort='time', frac=0.2)
|
||
|
|
pr.__enter__()
|
||
|
|
else:
|
||
|
|
fpt += (et1-et0)
|
||
|
|
bpt += (et2-et1)
|
||
|
|
|
||
|
|
pr.__exit__()
|
||
|
|
fpt = (fpt*1000/cnt)
|
||
|
|
bpt = (bpt*1000/cnt)
|
||
|
|
print("forward pass: %.3f ms, %.2fx off baseline %.3f ms" % (fpt, fpt/fpt_baseline, fpt_baseline))
|
||
|
|
print("backward pass: %.3f ms, %.2fx off baseline %.3f ms" % (bpt, bpt/bpt_baseline, bpt_baseline))
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == '__main__':
|
||
|
|
unittest.main()
|