mirror of
https://github.com/firestar5683/StarPilot.git
synced 2026-11-08 13:33:42 +08:00
firestar5683
22a8575199
Keep official openpilot ancestry and dependency licenses. Store resources directly in Git, remove LFS setup requirements, and leave the downloadable big model outside the source tree.
84 lines
3.4 KiB
Python
84 lines
3.4 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()
|
|
|
|
if __name__ == '__main__':
|
|
unittest.main()
|