diff --git a/README.md b/README.md index 091de44..c412882 100644 --- a/README.md +++ b/README.md @@ -30,6 +30,11 @@ LeanPass is a lightweight, transparent NumPy-based autodiff library for small ne - Linear layers and multilayer perceptrons - Loss utilities: `mse_loss`, `cross_entropy_loss`, and `binary_cross_entropy_loss` - SGD and Adam optimizers +- SGD now supports momentum and L2 weight decay: `optim.SGD(params, lr=..., momentum=0.9, weight_decay=1e-4)` +- Adam supports an optional L2 weight decay: `optim.Adam(params, lr=..., weight_decay=1e-4)` +- Dropout layer: `nn.Dropout(p=0.5)` for inverted dropout during training +- Tensor utilities: `Tensor.clip(a_min, a_max)` / `Tensor.clamp(min, max)` for value clamping +- Basic indexing: `Tensor.__getitem__` supports slicing/indexing that participates in autodiff - A simple demo script that trains a toy model ## Installation diff --git a/leanpass/README.md b/leanpass/README.md index 972efef..fb953a4 100644 --- a/leanpass/README.md +++ b/leanpass/README.md @@ -37,6 +37,8 @@ Supported operations: - `Tensor.gelu()` - `Tensor.softmax()` - `Tensor.backward()` + - `Tensor.clip(a_min, a_max)` / `Tensor.clamp(min, max)` + - Basic indexing via `Tensor[...]` which participates in autodiff ### Neural network layers @@ -51,6 +53,7 @@ Available components: - `nn.Linear(in_features, out_features)` creates a linear layer with weights and bias. - `nn.MLP(layer_sizes)` creates a multilayer perceptron with ReLU activations between layers. + - `nn.Dropout(p=0.5)` creates an inverted-dropout layer for training-time regularization. - `nn.mse_loss(predictions, targets)` computes mean squared error. - `nn.cross_entropy_loss(predictions, targets)` computes categorical cross-entropy for multi-class targets. - `nn.binary_cross_entropy_loss(predictions, targets)` computes binary cross-entropy for binary classification. @@ -65,10 +68,10 @@ optimizer = optim.SGD(model.parameters(), lr=0.01) optimizer = optim.Adam(model.parameters(), lr=0.001) ``` -Available optimizers: +Available optimizers (options): -- `optim.SGD(parameters, lr=...)` performs simple gradient descent. -- `optim.Adam(parameters, lr=...)` performs Adam optimization with bias correction. +- `optim.SGD(parameters, lr=..., momentum=0.0, weight_decay=0.0)` supports momentum and L2 weight decay. +- `optim.Adam(parameters, lr=..., weight_decay=0.0)` supports an optional L2 weight decay term. ### Example diff --git a/leanpass/nn.py b/leanpass/nn.py index f46ed77..415b814 100644 --- a/leanpass/nn.py +++ b/leanpass/nn.py @@ -38,6 +38,7 @@ def forward(self, x: Tensor) -> Tensor: def __call__(self, x: Tensor) -> Tensor: return self.forward(x) + class MLP(Module): @@ -67,6 +68,31 @@ def parameters(self): return params +class Dropout(Module): + """Simple dropout layer. + + Usage: `Dropout(p=0.5)(x, training=True)` or call with `training=False` to + disable dropout at evaluation time. + """ + + def __init__(self, p: float = 0.5): + if not 0 <= p < 1: + raise ValueError("p must be in the interval [0, 1)") + self.p = float(p) + + def forward(self, x: Tensor, training: bool = True) -> Tensor: + if not training or self.p == 0.0: + return x + # create a binary mask and scale to preserve expectation (inverted dropout) + mask = (np.random.rand(*x.data.shape) > self.p).astype(np.float64) / (1.0 - self.p) + return x * Tensor(mask) + + def __call__(self, x: Tensor, training: bool = True) -> Tensor: + return self.forward(x, training=training) + + + + def mse_loss(prediction: Tensor, target: Tensor) -> Tensor: """Mean squared error loss used for regression training.""" return ((prediction - target) ** 2).mean() diff --git a/leanpass/optim.py b/leanpass/optim.py index ec72bf9..f97362c 100644 --- a/leanpass/optim.py +++ b/leanpass/optim.py @@ -2,18 +2,40 @@ class SGD: - """Simple stochastic gradient descent optimizer.""" + """Stochastic gradient descent with optional momentum and L2 weight decay. - def __init__(self, params, lr=1e-2): + Args: + params: iterable of `Tensor` parameters + lr: learning rate + momentum: momentum factor (0.0 means no momentum) + weight_decay: L2 penalty factor applied to parameters (aka weight decay) + """ + + def __init__(self, params, lr=1e-2, momentum=0.0, weight_decay=0.0): self.params = list(params) self.lr = lr + self.momentum = float(momentum) + self.weight_decay = float(weight_decay) + # velocity buffers for momentum (kept even if momentum==0 for simplicity) + self.v = [np.zeros_like(p.data) for p in self.params] def step(self): - """Apply a plain gradient descent update to each parameter.""" - for param in self.params: + """Perform a parameter update step.""" + for i, param in enumerate(self.params): if param.grad is None: continue - param.data = param.data - self.lr * param.grad + # apply L2 weight decay directly to the gradient (common choice) + g = param.grad + if self.weight_decay: + g = g + self.weight_decay * param.data + + if self.momentum: + self.v[i] = self.momentum * self.v[i] + g + update = self.v[i] + else: + update = g + + param.data = param.data - self.lr * update def zero_grad(self): """Zero out gradients so the next backward pass starts clean.""" @@ -22,13 +44,14 @@ def zero_grad(self): class Adam: - """Adam optimizer with bias-corrected moment estimates.""" + """Adam optimizer with bias-corrected moment estimates and optional L2 weight decay.""" - def __init__(self, params, lr=1e-3, betas=(0.9, 0.999), eps=1e-8): + def __init__(self, params, lr=1e-3, betas=(0.9, 0.999), eps=1e-8, weight_decay=0.0): self.params = list(params) self.lr = lr self.b1, self.b2 = betas self.eps = eps + self.weight_decay = float(weight_decay) self.m = [np.zeros_like(p.data) for p in self.params] self.v = [np.zeros_like(p.data) for p in self.params] self.t = 0 @@ -39,7 +62,11 @@ def step(self): for i, param in enumerate(self.params): if param.grad is None: continue + # apply L2 weight decay to the gradient g = param.grad + if self.weight_decay: + g = g + self.weight_decay * param.data + self.m[i] = self.b1 * self.m[i] + (1 - self.b1) * g self.v[i] = self.b2 * self.v[i] + (1 - self.b2) * (g ** 2) diff --git a/leanpass/tensor.py b/leanpass/tensor.py index 6da99b0..6ea1cfc 100644 --- a/leanpass/tensor.py +++ b/leanpass/tensor.py @@ -238,6 +238,31 @@ def _backward(): out._backward = _backward return out + def clip(self, a_min=None, a_max=None): + """Clamp tensor values to the interval [a_min, a_max]. + + Both bounds are optional. The backward pass only propagates gradient + for elements that were not clipped (standard straight-through behavior). + """ + out_data = np.clip(self.data, a_min, a_max) + out = self._create_child(out_data, "clip", (self,), meta={"min": a_min, "max": a_max}) + + def _backward(): + if self.requires_grad: + mask = np.ones_like(self.data, dtype=np.float64) + if a_min is not None: + mask = mask * (self.data > a_min) + if a_max is not None: + mask = mask * (self.data < a_max) + self.grad += _sum_to_shape(out.grad * mask, self.data.shape) + + out._backward = _backward + return out + + # alias common name + def clamp(self, min=None, max=None): + return self.clip(min, max) + def sum(self, axis=None, keepdims=False): out_data = self.data.sum(axis=axis, keepdims=keepdims) out = Tensor(out_data, requires_grad=self.requires_grad, name="sum") @@ -257,6 +282,26 @@ def _backward(): out._backward = _backward return out + def __getitem__(self, idx): + """Basic indexing / slicing returning a new Tensor view (not a view in-place). + + The returned tensor participates in autodiff; gradients are placed back + into the source tensor at the same indices during the backward pass. + """ + out_data = self.data[idx] + out = self._create_child(out_data, "getitem", (self,), meta={"index": idx}) + + def _backward(): + if self.requires_grad: + if out.grad is None: + return + grad_buf = np.zeros_like(self.data) + grad_buf[idx] = out.grad + self.grad += _sum_to_shape(grad_buf, self.data.shape) + + out._backward = _backward + return out + def mean(self, axis=None, keepdims=False): out_data = self.data.mean(axis=axis, keepdims=keepdims) out = Tensor(out_data, requires_grad=self.requires_grad, name="mean") diff --git a/scripts/debug_inspect.py b/scripts/debug_inspect.py new file mode 100644 index 0000000..ed24d25 --- /dev/null +++ b/scripts/debug_inspect.py @@ -0,0 +1,36 @@ +import sys +sys.path.insert(0, '.') +from leanpass import Tensor, nn, optim +import numpy as np + +model = nn.MLP([2,4,1]) +opt = optim.SGD(model.parameters(), lr=0.1) + +print('model __dict__ keys and types:') +for k, v in model.__dict__.items(): + print(k, type(v)) + +params = model.parameters() +print('model.parameters() ->', params) +print('len=', len(params)) + +x = Tensor([[1.0, 1.0]], requires_grad=False) +y_true = Tensor([[2.0]], requires_grad=False) + +pred = model(x) +loss = nn.mse_loss(pred, y_true) + +print('param data before:') +for p in model.parameters(): + print(repr(p), p.data) + +model.zero_grad() +loss.backward() +print('\nparam grads after backward:') +for p in model.parameters(): + print(repr(p), p.grad) + +opt.step() +print('\nparam data after step:') +for p in model.parameters(): + print(repr(p), p.data)