requires grad, optim in train enet
MODIFY
examples/train_efficientnet.py
examples/train_efficientnet.py
@@ -4,12 +4,14 @@ import numpy as np
from extra.efficientnet import EfficientNet
from tinygrad.tensor import Tensor
from tinygrad.utils import get_parameters
+import tinygrad.optim as optim
if __name__ == "__main__":
Tensor.default_gpu = os.getenv("GPU") is not None
model = EfficientNet(int(os.getenv("NUM", "0")))
parameters = get_parameters(model)
print(len(parameters))
+ optimizer = optim.Adam(parameters, lr=0.001)
BS = 16
img = np.zeros((BS,3,224,224), dtype=np.float32)
@@ -29,14 +31,17 @@ if __name__ == "__main__":
y = Tensor(y)
loss = out.logsoftmax().mul(y).mean()
- # zero grad
- for p in parameters:
- p.grad = None
+ optimizer.zero_grad()
st = time.time()
loss.backward()
et = time.time()
print("backward %.2f s" % (et-st))
+ st = time.time()
+ optimizer.step()
+ et = time.time()
+ print("optimizer %.2f s" % (et-st))
+
del out, y, loss

MODIFY
tinygrad/nn.py
tinygrad/nn.py
@@ -9,7 +9,7 @@ class BatchNorm2D:
# TODO: need running_mean and running_var
self.running_mean = Tensor.zeros(sz)
self.running_var = Tensor.zeros(sz)
- self.num_batches_tracked = Tensor.zeros(1)
+ self.num_batches_tracked = Tensor.zeros(1, requires_grad=False)
def __call__(self, x):
# this work at inference?

MODIFY
tinygrad/optim.py
tinygrad/optim.py
@@ -5,7 +5,7 @@ from tinygrad.tensor import Tensor
class Optimizer:
def __init__(self, params):
- self.params = params
+ self.params = [x for x in params if x.requires_grad == True]
def zero_grad(self):
for param in self.params:
@@ -14,7 +14,7 @@ class Optimizer:
class SGD(Optimizer):
def __init__(self, params, lr=0.001):
super(SGD, self).__init__(params)
- self.lr = Tensor([lr], gpu=params[0].gpu)
+ self.lr = Tensor([lr], gpu=params[0].gpu, requires_grad=False)
def step(self):
for t in self.params:
@@ -23,14 +23,14 @@ class SGD(Optimizer):
class RMSprop(Optimizer):
def __init__(self, params, lr=0.001, decay=0.9, eps=1e-8):
super(RMSprop, self).__init__(params)
- self.lr = Tensor([lr], gpu=params[0].gpu)
- self.decay = Tensor([decay], gpu=params[0].gpu)
- self.eps = Tensor([eps], gpu=params[0].gpu)
+ self.lr = Tensor([lr], gpu=params[0].gpu, requires_grad=False)
+ self.decay = Tensor([decay], gpu=params[0].gpu, requires_grad=False)
+ self.eps = Tensor([eps], gpu=params[0].gpu, requires_grad=False)
- self.v = [Tensor(np.zeros(t.shape, dtype=np.float32), gpu=params[0].gpu) for t in self.params]
+ self.v = [Tensor(np.zeros(t.shape, dtype=np.float32), gpu=params[0].gpu, requires_grad=False) for t in self.params]
- self.one = Tensor([1], gpu=self.params[0].gpu)
- self.two = Tensor([2], gpu=self.params[0].gpu)
+ self.one = Tensor([1], gpu=self.params[0].gpu, requires_grad=False)
+ self.two = Tensor([2], gpu=self.params[0].gpu, requires_grad=False)
def step(self):
for i, t in enumerate(self.params):
@@ -40,17 +40,17 @@ class RMSprop(Optimizer):
class Adam(Optimizer):
def __init__(self, params, lr=0.001, b1=0.9, b2=0.999, eps=1e-8):
super(Adam, self).__init__(params)
- self.lr = Tensor([lr], gpu=params[0].gpu)
- self.b1 = Tensor([b1], gpu=params[0].gpu)
- self.b2 = Tensor([b2], gpu=params[0].gpu)
- self.eps = Tensor([eps], gpu=params[0].gpu)
- self.t = Tensor([0], gpu=params[0].gpu)
+ self.lr = Tensor([lr], gpu=params[0].gpu, requires_grad=False)
+ self.b1 = Tensor([b1], gpu=params[0].gpu, requires_grad=False)
+ self.b2 = Tensor([b2], gpu=params[0].gpu, requires_grad=False)
+ self.eps = Tensor([eps], gpu=params[0].gpu, requires_grad=False)
+ self.t = Tensor([0], gpu=params[0].gpu, requires_grad=False)
- self.m = [Tensor(np.zeros(t.shape, dtype=np.float32), gpu=params[0].gpu) for t in self.params]
- self.v = [Tensor(np.zeros(t.shape, dtype=np.float32), gpu=params[0].gpu) for t in self.params]
+ self.m = [Tensor(np.zeros(t.shape, dtype=np.float32), gpu=params[0].gpu, requires_grad=False) for t in self.params]
+ self.v = [Tensor(np.zeros(t.shape, dtype=np.float32), gpu=params[0].gpu, requires_grad=False) for t in self.params]
- self.one = Tensor([1], gpu=self.params[0].gpu)
- self.two = Tensor([2], gpu=self.params[0].gpu)
+ self.one = Tensor([1], gpu=self.params[0].gpu, requires_grad=False)
+ self.two = Tensor([2], gpu=self.params[0].gpu, requires_grad=False)
def step(self):
self.t = self.t + self.one

MODIFY
tinygrad/tensor.py
tinygrad/tensor.py
@@ -57,6 +57,9 @@ class GPUBuffer:
else:
self.cl = cl.Buffer(cl_ctx, cl.mem_flags.READ_WRITE, 4*np.prod(shape))
+ def __repr__(self):
+ return "<GPUBuffer with shape %r>" % (self.shape,)
+
# **** start with two base classes ****
def deepwalk(node, visited=None, nodes=None):
@@ -75,7 +78,7 @@ class Tensor:
default_gpu = False
allocated = 0
- def __init__(self, data, gpu=None):
+ def __init__(self, data, gpu=None, requires_grad=True):
if gpu is None:
gpu = Tensor.default_gpu
if isinstance(data, list):
@@ -94,6 +97,7 @@ class Tensor:
self.data = data
self.grad = None
+ self.requires_grad = requires_grad
if gpu:
self.cuda_()
@@ -122,20 +126,20 @@ class Tensor:
return self.data.dtype
@staticmethod
- def zeros(*shape, gpu=None):
- return Tensor(np.zeros(shape, dtype=np.float32), gpu)
+ def zeros(*shape, **kwargs):
+ return Tensor(np.zeros(shape, dtype=np.float32), **kwargs)
@staticmethod
- def ones(*shape, gpu=None):
- return Tensor(np.ones(shape, dtype=np.float32), gpu)
+ def ones(*shape, **kwargs):
+ return Tensor(np.ones(shape, dtype=np.float32), **kwargs)
@staticmethod
- def randn(*shape, gpu=None):
- return Tensor(np.random.randn(*shape).astype(np.float32), gpu)
+ def randn(*shape, **kwargs):
+ return Tensor(np.random.randn(*shape).astype(np.float32), **kwargs)
@staticmethod
- def eye(dim, gpu=None):
- return Tensor(np.eye(dim).astype(np.float32), gpu)
+ def eye(dim, **kwargs):
+ return Tensor(np.eye(dim).astype(np.float32), **kwargs)
def backward(self, allow_fill=True):
if self._ctx is None:
@@ -145,7 +149,7 @@ class Tensor:
# fill in the first grad with one
# this is "implicit gradient creation"
assert self.shape == (1,)
- self.grad = Tensor(np.ones(self.shape, dtype=self.dtype), gpu=self.gpu)
+ self.grad = Tensor(np.ones(self.shape, dtype=self.dtype), gpu=self.gpu, requires_grad=False)
for t0 in reversed(deepwalk(self)):
assert (t0.grad is not None)
@@ -158,8 +162,8 @@ class Tensor:
continue
assert g.shape == t.shape, \
"grad shape must match tensor shape in %r, %r != %r" % (self._ctx, g.shape, t.shape)
- t.grad = Tensor(g) if t.grad is None else (t.grad + Tensor(g))
- del t.grad._ctx # no backward pass through the add
+ gt = Tensor(g, requires_grad=False)
+ t.grad = gt if t.grad is None else (t.grad + gt)
# ***** tinygrad supports CPU and GPU *****
@@ -243,7 +247,9 @@ class Function:
for k, v in kwargs.items():
setattr(ctx, k, v)
with ProfileOp(ctx.__class__.__name__, x):
- ret = Tensor(op.forward(ctx, *[t.data for t in x], **kwargs))
+ ret = Tensor(op.forward(ctx, *[t.data for t in x], **kwargs),
+ requires_grad=any([t.requires_grad for t in x]))
+ if ret.requires_grad:
ret._ctx = ctx
return ret

requires grad, optim in train enet