mirror of
https://github.com/firestar5683/StarPilot.git
synced 2026-11-07 04:53:47 +08:00
Daniel Koepping
112 lines
4.3 KiB
Python
112 lines
4.3 KiB
Python
from typing import Callable
|
|
import unittest, math
|
|
import torch
|
|
from tinygrad import Tensor
|
|
from tinygrad.dtype import dtypes
|
|
from tinygrad.uop.ops import UOp
|
|
from tinygrad.mixin.gradient import compute_gradient
|
|
|
|
class TestGradient(unittest.TestCase):
|
|
def _cmp_nan_okay(self, x, y):
|
|
if math.isnan(x) and math.isnan(y): return
|
|
self.assertAlmostEqual(x, y, places=5)
|
|
|
|
def _test_one_input_function(self, f:Callable, jf:Callable|None=None):
|
|
if jf is None: jf = f
|
|
x = UOp.variable('x', -math.inf, math.inf, dtype=dtypes.float)
|
|
gx = compute_gradient(f(x), UOp.const(1.0), set([x]))[x]
|
|
|
|
for val in [-5., -2.0, 0.0, 2.0, 5.]:
|
|
tg_out = gx.substitute({x: UOp.const(val)}).ssimplify()
|
|
tx = torch.tensor([val], dtype=torch.float, requires_grad=True)
|
|
torch_out = torch.autograd.grad(jf(tx), tx)[0].item()
|
|
self._cmp_nan_okay(tg_out, torch_out)
|
|
|
|
def _test_two_input_function(self, f:Callable, jf:Callable|None=None):
|
|
if jf is None: jf = f
|
|
x = UOp.variable('x', -math.inf, math.inf, dtype=dtypes.float)
|
|
y = UOp.variable('y', -math.inf, math.inf, dtype=dtypes.float)
|
|
grads = compute_gradient(f(x, y), UOp.const(1.0), set([x, y]))
|
|
gx, gy = grads[x], grads[y]
|
|
|
|
for valx in [-5., -2.0, 0.0, 2.0, 5.]:
|
|
for valy in [-5., -2.0, 0.0, 2.0, 5.]:
|
|
# Substitute the values into the gradient expressions
|
|
substitutions = {x: UOp.const(valx), y: UOp.const(valy)}
|
|
tg_out_x = gx.substitute(substitutions).ssimplify()
|
|
tg_out_y = gy.substitute(substitutions).ssimplify()
|
|
|
|
tx = torch.tensor([valx], dtype=torch.float, requires_grad=True)
|
|
ty = torch.tensor([valy], dtype=torch.float, requires_grad=True)
|
|
torch_grad = torch.autograd.grad(jf(tx, ty), [tx, ty])
|
|
torch_out_x, torch_out_y = [x.item() for x in torch_grad]
|
|
|
|
self._cmp_nan_okay(tg_out_x, torch_out_x)
|
|
self._cmp_nan_okay(tg_out_y, torch_out_y)
|
|
|
|
# unary ops unit
|
|
def test_recip(self): self._test_one_input_function(lambda x: 1.0/x)
|
|
def test_sin(self): self._test_one_input_function(lambda x: x.sin())
|
|
def test_sqrt(self): self._test_one_input_function(lambda x: x.sqrt())
|
|
def test_log2(self): self._test_one_input_function(lambda x: x.log2())
|
|
def test_exp2(self): self._test_one_input_function(lambda x: x.exp2())
|
|
|
|
# binary ops unit
|
|
def test_add(self): self._test_two_input_function(lambda x,y: x+y)
|
|
def test_mul(self): self._test_two_input_function(lambda x,y: x*y)
|
|
|
|
# chain rule
|
|
def test_chain(self): self._test_one_input_function(lambda x: x.sin().sqrt())
|
|
def test_chain_binop(self): self._test_two_input_function(lambda x,y: (x*y)+x*y)
|
|
def test_big_add_sin(self): self._test_two_input_function(lambda x,y: x.sin()+3.0/y)
|
|
def test_big_chain(self): self._test_two_input_function(lambda x,y: (1.0/x*y)+x*y)
|
|
def test_where(self): self._test_two_input_function(lambda x,y: (x<y).where(x,y), lambda x,y: torch.where(x<y,x,y))
|
|
|
|
class TestRealizeMeansRealize(unittest.TestCase):
|
|
def test_randn_realizes(self):
|
|
x = Tensor.randn(2, 3, 64, 64).realize()
|
|
assert x.uop is not x.uop.base
|
|
assert x.uop.is_realized
|
|
|
|
def test_uniform_realizes(self):
|
|
x = Tensor.uniform(16, 3, 3, 3).realize()
|
|
print(x.uop)
|
|
assert x.uop is not x.uop.base
|
|
assert x.uop.is_realized
|
|
|
|
def test_uniform_gradient(self):
|
|
x = Tensor.uniform(16, 3, 3, 3).realize()
|
|
y = x * 2
|
|
y.sum().gradient(x)[0].realize()
|
|
|
|
class TestTensorGradient(unittest.TestCase):
|
|
def test_cast_before_view(self):
|
|
x = Tensor([1.0, 1, 1, 1])
|
|
x_reshaped = x.reshape(2,2)
|
|
x_casted = x_reshaped.cast(dtypes.float16)
|
|
x_casted.mean().gradient(x_reshaped)
|
|
|
|
def test_non_float_tensor_raise(self):
|
|
x = Tensor([1, 2, 3])
|
|
with self.assertRaises(RuntimeError): x.sum().gradient(x)
|
|
with self.assertRaises(RuntimeError): x.float().sum().gradient(x)
|
|
|
|
def test_const_target_raise(self):
|
|
t = Tensor(2.0)
|
|
with self.assertRaises(RuntimeError): (t * 2.0).gradient(t)
|
|
|
|
def test_setitem_on_grad_used_tensor_raises(self):
|
|
x = Tensor([1.0, 2.0, 3.0, 4.0]).realize()
|
|
_ = (x * 2.0).sum()
|
|
with self.assertRaises(RuntimeError):
|
|
x[0] = 99.0
|
|
|
|
def test_bare_const_skipped_by_backward(self):
|
|
Tensor.manual_seed(0)
|
|
w = Tensor(1.0)
|
|
(Tensor.rand(()) + w).backward()
|
|
self.assertIsNone(w.grad)
|
|
|
|
if __name__ == '__main__':
|
|
unittest.main()
|