mirror of
https://github.com/tinygrad/tinygrad.git
synced 2026-06-24 02:14:17 +00:00
bcf1518309* Update all devices to be tested
ANE, CPU and OCL all now support all tests.
However tests are not currently passing on GPU and I cannot test on CPU.
Failing GPU test are not an issue caused by this update. Tests have not
been passing due to a missing "six" required installation.
OpenCL Tests have not been run since commit: 1a1c63a08b
devices have 3 types and are handle by a new DeviceTypes enum. (The goal
is to revert to Tensor.<type>, but this current setup allows for keyword
argument defaults: `device=DeviceType.CPU`)
All references to Tensor.GPU/CPU/ANE as been converted to the
corresponding `DeviceTypes` enum.
Refactor of the conversion code to allow for any device to any device
conversion.
* Add six dependency in requirements.txt
* Resolve failure to run tests
Move six into gpu required installs. Remove six from standard
installation.
* Remove repeated data conversion
* Refactor method names
Also reduce code with .to and .to_
* Dynamic device handlers
* Refactor DeviceTypes -> Device
* Add mem copy profiling back
* test_backward_pass_diamond_model passing
* Resolve Sum issue on GPU
* Revert batchnorm2d tests
* Update README with upadated API
* ANE testing with
* Last minute line gains
117 lines
3.8 KiB
Python
117 lines
3.8 KiB
Python
import numpy as np
|
|
import torch
|
|
import unittest
|
|
from tinygrad.tensor import Tensor, GPU, Device
|
|
from extra.gradcheck import numerical_jacobian, jacobian, gradcheck
|
|
from .config import ANE
|
|
|
|
|
|
x_init = np.random.randn(1,3).astype(np.float32)
|
|
U_init = np.random.randn(3,3).astype(np.float32)
|
|
V_init = np.random.randn(3,3).astype(np.float32)
|
|
W_init = np.random.randn(3,3).astype(np.float32)
|
|
m_init = np.random.randn(1,3).astype(np.float32)
|
|
|
|
class TestTinygrad(unittest.TestCase):
|
|
device = Device.CPU
|
|
|
|
def test_backward_pass(self):
|
|
def test_tinygrad():
|
|
x = Tensor(x_init, device=self.device)
|
|
W = Tensor(W_init, device=self.device)
|
|
m = Tensor(m_init, device=self.device)
|
|
out = x.dot(W).relu()
|
|
out = out.logsoftmax()
|
|
out = out.mul(m).add(m).sum()
|
|
out.backward()
|
|
return out.cpu().data, x.grad.cpu().data, W.grad.cpu().data
|
|
|
|
def test_pytorch():
|
|
x = torch.tensor(x_init, requires_grad=True)
|
|
W = torch.tensor(W_init, requires_grad=True)
|
|
m = torch.tensor(m_init)
|
|
out = x.matmul(W).relu()
|
|
out = torch.nn.functional.log_softmax(out, dim=1)
|
|
out = out.mul(m).add(m).sum()
|
|
out.backward()
|
|
return out.detach().numpy(), x.grad, W.grad
|
|
|
|
for x,y in zip(test_tinygrad(), test_pytorch()):
|
|
np.testing.assert_allclose(x, y, atol=1e-5)
|
|
|
|
def test_backward_pass_diamond_model(self):
|
|
def test_tinygrad():
|
|
u = Tensor(U_init, device=self.device)
|
|
v = Tensor(V_init, device=self.device)
|
|
w = Tensor(W_init, device=self.device)
|
|
x = u.mul(v).relu()
|
|
y = u.mul(w).relu()
|
|
out = x.add(y).mul(y).relu()
|
|
out = out.logsoftmax()
|
|
out = out.sum()
|
|
out.backward()
|
|
return out.cpu().data, u.cpu().grad.data, v.cpu().grad.data, w.cpu().grad.data
|
|
|
|
def test_pytorch():
|
|
u = torch.tensor(U_init, requires_grad=True)
|
|
v = torch.tensor(V_init, requires_grad=True)
|
|
w = torch.tensor(W_init, requires_grad=True)
|
|
x = u.mul(v).relu()
|
|
y = u.mul(w).relu()
|
|
out = x.add(y).mul(y).relu()
|
|
out = torch.nn.functional.log_softmax(out, dim=1)
|
|
out = out.sum()
|
|
out.backward()
|
|
return out.detach().numpy(), u.grad, v.grad, w.grad
|
|
|
|
for x,y in zip(test_tinygrad(), test_pytorch()):
|
|
np.testing.assert_allclose(x, y, atol=1e-5)
|
|
|
|
def test_jacobian(self):
|
|
W = np.random.RandomState(1337).random((10, 5))
|
|
x = np.random.RandomState(7331).random((1, 10)) - 0.5
|
|
|
|
torch_x = torch.tensor(x, requires_grad=True)
|
|
torch_W = torch.tensor(W, requires_grad=True)
|
|
torch_func = lambda x: torch.nn.functional.log_softmax(x.matmul(torch_W).relu(), dim=1)
|
|
PJ = torch.autograd.functional.jacobian(torch_func, torch_x).squeeze().numpy()
|
|
|
|
tiny_x = Tensor(x, device=self.device)
|
|
tiny_W = Tensor(W, device=self.device)
|
|
tiny_func = lambda x: x.dot(tiny_W).relu().logsoftmax()
|
|
J = jacobian(tiny_func, tiny_x)
|
|
NJ = numerical_jacobian(tiny_func, tiny_x)
|
|
|
|
np.testing.assert_allclose(PJ, J, atol = 1e-5)
|
|
np.testing.assert_allclose(PJ, NJ, atol = 1e-5)
|
|
|
|
def test_gradcheck(self):
|
|
W = np.random.RandomState(1337).random((10, 5))
|
|
x = np.random.RandomState(7331).random((1, 10)) - 0.5
|
|
|
|
tiny_x = Tensor(x, device=self.device)
|
|
tiny_W = Tensor(W, device=self.device)
|
|
tiny_func = lambda x: x.dot(tiny_W).relu().logsoftmax()
|
|
|
|
self.assertTrue(gradcheck(tiny_func, tiny_x))
|
|
|
|
# coarse approx. since a "big" eps and the non-linearities of the model
|
|
self.assertFalse(gradcheck(tiny_func, tiny_x, eps = 0.1))
|
|
|
|
|
|
@unittest.skipUnless(GPU, "Requires GPU")
|
|
class TestTinygradGPU(TestTinygrad):
|
|
device = Device.GPU
|
|
|
|
@unittest.skip("float64 not supported on GPU")
|
|
def test_jacobian(self): pass
|
|
|
|
@unittest.skip("float64 not supported on GPU")
|
|
def test_gradcheck(self): pass
|
|
|
|
@unittest.skipUnless(ANE, "Requires ANE")
|
|
class TestOpsANE(TestTinygrad):
|
|
device=Device.ANE
|
|
|
|
if __name__ == '__main__':
|
|
unittest.main()
|