t.assign in optim
MODIFY
test/graph_batchnorm.py
test/graph_batchnorm.py
@@ -14,7 +14,6 @@ def model_step(lm):
optimizer.step()
Tensor.training = False
-
class TestBatchnorm(unittest.TestCase):
def test_conv(self):
class LilModel:
@@ -35,6 +34,19 @@ class TestBatchnorm(unittest.TestCase):
lm = LilModel()
model_step(lm)
+ def test_two_conv_bn(self):
+ class LilModel:
+ def __init__(self):
+ self.c = Conv2d(12, 24, 3, padding=1, bias=False)
+ self.bn = BatchNorm2D(24, track_running_stats=False)
+ self.c2 = Conv2d(24, 32, 3, padding=1, bias=False)
+ self.bn2 = BatchNorm2D(32, track_running_stats=False)
+ def forward(self, x):
+ x = self.bn(self.c(x)).relu()
+ return self.bn2(self.c2(x)).relu()
+ lm = LilModel()
+ model_step(lm)
+
def test_conv_bn(self):
class LilModel:
def __init__(self):

MODIFY
tinygrad/nn/optim.py
tinygrad/nn/optim.py
@@ -20,7 +20,7 @@ class SGD(Optimizer):
def step(self):
for t in self.params:
- t -= t.grad * self.lr
+ t.assign(t.detach() - t.grad * self.lr)
self.realize()
class RMSprop(Optimizer):
@@ -33,7 +33,7 @@ class RMSprop(Optimizer):
def step(self):
for i, t in enumerate(self.params):
self.v[i] = self.decay * self.v[i] + (1.0 - self.decay) * (t.grad * t.grad)
- t -= (t.grad * self.lr).div(self.v[i].sqrt() + self.eps)
+ t.assign(t.detach() - (t.grad * self.lr).div(self.v[i].sqrt() + self.eps))
self.realize(self.v)
class Adam(Optimizer):
@@ -50,5 +50,5 @@ class Adam(Optimizer):
for i, t in enumerate(self.params):
self.m[i] = self.b1 * self.m[i] + (1.0 - self.b1) * t.grad
self.v[i] = self.b2 * self.v[i] + (1.0 - self.b2) * (t.grad * t.grad)
- t -= a * self.m[i].div(self.v[i].sqrt() + self.eps)
+ t.assign(t.detach() - a * self.m[i].div(self.v[i].sqrt() + self.eps))
self.realize(self.m + self.v)

MODIFY
tinygrad/tensor.py
tinygrad/tensor.py
@@ -126,6 +126,7 @@ class Tensor:
if g is not None and t.requires_grad:
assert g.shape == t.shape, f"grad shape must match tensor shape in {self._ctx!r}, {g.shape!r} != {t.shape!r}"
t.grad = g if t.grad is None else (t.grad + g)
+ del t0._ctx
# ***** non first class ops (hlops) *****

t.assign in optim