github/tinygrad
1
1
Fork 0
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-01-09 23:18:04 -05:00
Code Issues Packages Projects Releases Wiki Activity

update test_ops backward tests (#9267)

Browse Source 
instead of `(out+1).square().mean().backward()`, use forward.sum().gradient to get closer to the gradients
This commit is contained in:
chenyu
2025-02-26 12:09:24 -05:00
committed by GitHub
parent aaf0a8069f
commit 49ca90df75
1 changed files with 14 additions and 26 deletions
Download Patch File Download Diff File Expand all files Collapse all files

40
test/test_ops.py
View File

@@ -54,31 +54,19 @@ def helper_test_op(shps, torch_fxn, tinygrad_fxn=None, atol=1e-6, rtol=1e-3, gra
compare("forward pass", ret.numpy(), out.detach().cpu().numpy(), atol=atol, rtol=rtol)
torch_fbp, tinygrad_fbp = np.nan, np.nan
if not forward_only and not FORWARD_ONLY:
if not forward_only and not FORWARD_ONLY and ts and tst:
st = time.monotonic()
(out+1).square().mean().backward()
torch_grads = torch.autograd.grad(torch_fxn(*ts).sum(), ts)
torch_fbp = time.monotonic() - st
st = time.monotonic()
# NOTE: we now have to recompute the forward pass since we realized it
ret = tinygrad_fxn(*tst)
loss:Tensor = (ret+1).square().mean()
# test_ops uses new style gradient
tst_grads = loss.gradient(*tst)
if len(tst_grads): Tensor.realize(*tst_grads)
tiny_grads = tinygrad_fxn(*tst).sum().gradient(*tst)
Tensor.realize(*tiny_grads)
tinygrad_fbp = time.monotonic() - st
for i, (t, tt_grad) in enumerate(zip(ts, tst_grads)):
compare(f"backward pass tensor {i}", tt_grad.numpy(), t.grad.detach().cpu().numpy(), atol=grad_atol, rtol=grad_rtol)
"""
(ret+1).square().mean().backward()
for tt in tst: tt.grad.realize()
tinygrad_fbp = time.monotonic() - st
for i, (t, tt) in enumerate(zip(ts, tst)):
compare(f"backward pass tensor {i}", tt.grad.numpy(), t.grad.detach().numpy(), atol=grad_atol, rtol=grad_rtol)
"""
for i, (t, torch_grad) in enumerate(zip(tiny_grads, torch_grads)):
compare(f"backward pass tensor {i}", t.numpy(), torch_grad.detach().numpy(), atol=grad_atol, rtol=grad_rtol)
if not CI:
print("\ntesting %40r torch/tinygrad fp: %.2f / %.2f ms bp: %.2f / %.2f ms " % \
@@ -1339,8 +1327,8 @@ class TestOps(unittest.TestCase):
with warnings.catch_warnings():
warnings.filterwarnings("ignore", message="var\\(\\): degrees of freedom is <= 0")
helper_test_op([(1,2,3,1,5)], lambda x: x.var(axis=(0,3)))
helper_test_op([(1,2,3,1,5)], lambda x: x.var(axis=(0,3), correction=5))
# TODO: fix backward
helper_test_op([(1,2,3,1,5)], lambda x: x.var(axis=(0,3), correction=5), forward_only=True)
helper_test_op([(1,2,3,1,5)], lambda x: x.var(axis=(0,4), correction=5), forward_only=True)
helper_test_op([(1,)], lambda x: x.var(axis=(0,), correction=0))
helper_test_op([(1,2,3,1,5)], lambda x: x.var(axis=(0,3), correction=0))
@@ -1401,9 +1389,9 @@ class TestOps(unittest.TestCase):
helper_test_op([()], torch.nn.Softmax(dim=0), Tensor.softmax, atol=1e-7, grad_atol=1e-7)
helper_test_op([()], torch.nn.Softmax(dim=-1), Tensor.softmax, atol=1e-7, grad_atol=1e-7)
def test_softmax_other_axis(self):
helper_test_op([(10,10,10)], lambda x: x.softmax(0), atol=1e-7, grad_atol=1e-7)
helper_test_op([(10,10,10)], lambda x: x.softmax(1), atol=1e-7, grad_atol=1e-7)
helper_test_op([(10,10,10)], lambda x: x.softmax(2), atol=1e-7, grad_atol=1e-7)
helper_test_op([(10,10,10)], lambda x: x.softmax(0), atol=1e-7, grad_atol=2e-7)
helper_test_op([(10,10,10)], lambda x: x.softmax(1), atol=1e-7, grad_atol=2e-7)
helper_test_op([(10,10,10)], lambda x: x.softmax(2), atol=1e-7, grad_atol=2e-7)
def test_softmax_argmax(self):
helper_test_op([(45,65)], lambda x: x.softmax(0).argmax().type(torch.int32),
lambda x: x.softmax(0).argmax(), forward_only=True, atol=1e-7, grad_atol=1e-7)
@@ -1459,12 +1447,12 @@ class TestOps(unittest.TestCase):
helper_test_op([()], lambda x: torch.nn.functional.hardtanh(x, -val, val), lambda x: x.hardtanh(-val, val), grad_atol=1e-6)
def test_asinh(self):
helper_test_op([(45,65)], lambda x: x.asinh(), grad_atol=1e-6)
# NOTE: this one has larger atol
helper_test_op([(45,65)], lambda x: x.asinh(), atol=1e-2, rtol=2e-2, grad_atol=1e-6, low=-300, high=-297)
# TODO: this one has larger tol?
helper_test_op([(45,65)], lambda x: x.asinh(), atol=1e-2, rtol=2e-2, grad_rtol=2e-2, low=-300, high=-297)
helper_test_op([(45,65)], lambda x: x.asinh(), grad_atol=1e-6, low=300, high=303)
def test_acosh(self):
helper_test_op([(45,65)], lambda x: x.acosh(), grad_atol=1e-6)
helper_test_op([(45,65)], lambda x: x.acosh(), grad_atol=1e-6, low=-300, high=-297)
helper_test_op([(45,65)], lambda x: x.acosh(), grad_atol=1e-3, grad_rtol=1e-2, low=-300, high=-297)
helper_test_op([(45,65)], lambda x: x.acosh(), grad_atol=1e-6, low=300, high=303)
def test_atanh(self):
helper_test_op([(45,65)], lambda x: x.atanh(), grad_atol=1e-6)
@@ -2033,7 +2021,7 @@ class TestOps(unittest.TestCase):
helper_test_op([(bs,cin,64,64), (6,cin//groups,H,W)],
lambda x,w: torch.nn.functional.conv2d(x,w,groups=groups).relu(),
# needed to relax tolerance on NVIDIA
lambda x,w: Tensor.conv2d(x,w,groups=groups).relu(), atol=1e-4, grad_rtol=1e-5)
lambda x,w: Tensor.conv2d(x,w,groups=groups).relu(), atol=1e-4, grad_atol=1e-4, grad_rtol=1e-4)
def test_simple_grouped_conv2d(self):
bs = 1