more optim cleanup
MODIFY
tinygrad/optim.py
tinygrad/optim.py
@@ -23,26 +23,26 @@ 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, self.decay, self.eps = Tensor([lr], gpu=self.params[0].gpu, requires_grad=False), decay, eps
+ self.lr, self.decay, self.eps = lr, decay, eps
self.v = [Tensor(np.zeros(t.shape, dtype=np.float32), gpu=params[0].gpu, requires_grad=False) for t in self.params]
def step(self):
for i, t in enumerate(self.params):
self.v[i] = self.decay * self.v[i] + (1.0 - self.decay) * t.grad.pow(2.0)
- t -= self.lr.div(self.v[i].sqrt() + self.eps) * t.grad
+ t -= (t.grad * self.lr).div(self.v[i].sqrt() + self.eps)
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, self.b1, self.b2, self.eps, self.t = lr, b1, b2, eps, Tensor([0.0], gpu=self.params[0].gpu, requires_grad=False)
+ self.lr, self.b1, self.b2, self.eps, self.t = lr, b1, b2, eps, 0
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]
def step(self):
- self.t = self.t + 1.0
- a = self.lr * ((1.0 - self.b2**self.t)**0.5).div(1.0 - self.b1**self.t)
+ self.t = self.t + 1
+ a = self.lr * ((1.0 - self.b2**self.t)**0.5) / (1.0 - self.b1**self.t)
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.pow(2.0)

more optim cleanup