commaai
/
tinygrad
mirror of https://github.com/commaai/tinygrad.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Fix: Jacobian tests [WIP] (#1126) 

* Fix: Jacobian tests; num_jacobian either bugged or not accurate enough;

* Fix: Jacobian tests;

* Fix: Gradcheck;
This commit is contained in:
Reza Rezvan 2023-07-06 00:36:22 +02:00 committed by GitHub
parent d363d25ee2
commit d1356cac27
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 14 additions and 16 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
extra/gradcheck.py
View File

@ -26,7 +26,7 @@ def jacobian(func, input):
J[o,i] = grad
return J
def numerical_jacobian(func, input, eps = 1e-6):
def numerical_jacobian(func, input, eps = 1e-3):
output = func(input)
ji = input.numpy().reshape(-1).shape[-1]
@ -44,7 +44,7 @@ def numerical_jacobian(func, input, eps = 1e-6):
NJ[:,i] = grad_approx
return NJ
def gradcheck(func, input, eps = 1e-06, atol = 1e-5, rtol = 0.001):
def gradcheck(func, input, eps = 1e-3, atol = 1e-3, rtol = 1e-3):
NJ = numerical_jacobian(func, input, eps)
J = jacobian(func, input)
return np.allclose(J, NJ, atol=atol, rtol=rtol)
return np.allclose(J, NJ, atol = atol, rtol = rtol)

24
test/test_tensor.py
View File

@ -105,38 +105,36 @@ class TestTinygrad(unittest.TestCase):
expected = n * (1 - rate)
np.testing.assert_allclose(non_zeros, expected, rtol=2e-3)
@unittest.skip("TODO: fix")
def test_jacobian(self):
W = np.random.RandomState(1337).random((10, 5))
x = np.random.RandomState(7331).random((1, 10)) - 0.5
W = np.random.RandomState(42069).random((10, 5)).astype(np.float32)
x = np.random.RandomState(69420).random((1, 10)).astype(np.float32)
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)
tiny_W = Tensor(W)
tiny_x = Tensor(x, requires_grad=True)
tiny_W = Tensor(W, requires_grad=True)
tiny_func = lambda x: x.dot(tiny_W).relu().log_softmax()
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)
np.testing.assert_allclose(PJ, NJ, atol = 1e-3)
@unittest.skip("TODO: fix")
def test_gradcheck(self):
W = np.random.RandomState(1337).random((10, 5))
x = np.random.RandomState(7331).random((1, 10)) - 0.5
W = np.random.RandomState(1337).random((10, 5)).astype(np.float32)
x = np.random.RandomState(7331).random((1, 10)).astype(np.float32)
tiny_x = Tensor(x)
tiny_W = Tensor(W)
tiny_x = Tensor(x, requires_grad=True)
tiny_W = Tensor(W, requires_grad=True)
tiny_func = lambda x: x.dot(tiny_W).relu().log_softmax()
self.assertTrue(gradcheck(tiny_func, tiny_x))
self.assertTrue(gradcheck(tiny_func, tiny_x, eps = 1e-3))
# coarse approx. since a "big" eps and the non-linearities of the model
self.assertFalse(gradcheck(tiny_func, tiny_x, eps = 0.1))
self.assertFalse(gradcheck(tiny_func, tiny_x, eps = 1e-5))
def test_random_fns_are_deterministic_with_seed(self):
for random_fn in [Tensor.randn, Tensor.uniform, Tensor.scaled_uniform, Tensor.glorot_uniform]: