Linear Layer and Sequential: Building Architectures

medium · neural-networks, linear-layer, sequential

Linear and Sequential: Building Network Architectures

Implement Linear, Sequential, and activation modules using the Module base class.

What you are building

1) `Linear(Module)`

Fully connected layer: y = x W^T + b

class Linear(Module):
    def __init__(self, in_features: int, out_features: int, bias: bool = True):
        # weight: (out_features, in_features)
        # bias: (out_features,) or None
        # Xavier init: k = 1/sqrt(in_features)
        pass

    def forward(self, x: List[Value]) -> List[Value]:
        # x length = in_features
        # return length = out_features
        pass

2) `Sequential(Module)`

A container that applies modules in order.

class Sequential(Module):
    def __init__(self, *modules: Module):
        # store in module_list
        # register each module (e.g., layer_0, layer_1, ...)
        pass

    def forward(self, x):
        for m in self.module_list:
            x = m(x)
        return x

3) Activation Modules

class Tanh(Module):
    def forward(self, x: List[Value]) -> List[Value]:
        pass

class ReLU(Module):
    def forward(self, x: List[Value]) -> List[Value]:
        pass

class Sigmoid(Module):
    def forward(self, x: List[Value]) -> List[Value]:
        # 1 / (1 + exp(-x))
        pass

Notes

Linear should register all weights (and bias, if used) as parameters.
Sequential should preserve order and expose parameters recursively.
Activation modules have no parameters but must still be Modules.

Example

model = Sequential(
    Linear(2, 4),
    Tanh(),
    Linear(4, 1)
)

x = [Value(1.0), Value(2.0)]
out = model(x)

from __future__ import annotations
import math
import random
from typing import List, Iterator, Tuple, Callable, Set

class Value:
    """Value class from autograd module."""

def __init__(self, data: float, _children: Tuple["Value", ...] = (), _op: str = "") -> None:
        self.data = float(data)
        self.grad = 0.0
        self._backward: Callable[[], None] = lambda: None
        self._prev: Set["Value"] = set(_children)
        self._op = _op

def __repr__(self) -> str:
        return f"Value(data={self.data:.4f})"

def __add__(self, other):
        other = other if isinstance(other, Value) else Value(other)
        out = Value(self.data + other.data, (self, other), "+")
        def _backward():
            self.grad += out.grad
            other.grad += out.grad
        out._backward = _backward
        return out

def __mul__(self, other):
        other = other if isinstance(other, Value) else Value(other)
        out = Value(self.data * other.data, (self, other), "*")
        def _backward():
            self.grad += other.data * out.grad
            other.grad += self.data * out.grad
        out._backward = _backward
        return out

def __pow__(self, n):
        out = Value(self.data ** n, (self,), f"**{n}")
        def _backward():
            self.grad += n * (self.data ** (n - 1)) * out.grad
        out._backward = _backward
        return out

def __neg__(self):
        return self * -1

def __sub__(self, other):
        return self + (-other)

def __truediv__(self, other):
        return self * (other ** -1)

def __radd__(self, other):
        return self + other

def __rmul__(self, other):
        return self * other

def tanh(self):
        t = math.tanh(self.data)
        out = Value(t, (self,), "tanh")
        def _backward():
            self.grad += (1 - t ** 2) * out.grad
        out._backward = _backward
        return out

def relu(self):
        out = Value(max(0, self.data), (self,), "relu")
        def _backward():
            self.grad += (1.0 if self.data > 0 else 0.0) * out.grad
        out._backward = _backward
        return out

def exp(self):
        out = Value(math.exp(self.data), (self,), "exp")
        def _backward():
            self.grad += out.data * out.grad
        out._backward = _backward
        return out

def backward(self):
        topo = []
        visited = set()
        def build_topo(v):
            if v not in visited:
                visited.add(v)
                for child in v._prev:
                    build_topo(child)
                topo.append(v)
        build_topo(self)
        self.grad = 1.0
        for node in reversed(topo):
            node._backward()

class Module:
    """Module base class (from previous problem)."""

def __init__(self) -> None:
        object.__setattr__(self, "_parameters", {})
        object.__setattr__(self, "_modules", {})
        object.__setattr__(self, "training", True)

def __setattr__(self, name: str, value) -> None:
        if isinstance(value, Value):
            self._parameters[name] = value
        elif isinstance(value, Module):
            self._modules[name] = value
        elif isinstance(value, (list, tuple)):
            self._register_list(name, value)
        object.__setattr__(self, name, value)

def _register_list(self, name: str, values) -> None:
        for i, v in enumerate(values):
            if isinstance(v, Value):
                self._parameters[f"{name}.{i}"] = v
            elif isinstance(v, Module):
                self._modules[f"{name}.{i}"] = v
            elif isinstance(v, (list, tuple)):
                self._register_list(f"{name}.{i}", v)

def __call__(self, *args, **kwargs):
        return self.forward(*args, **kwargs)

def forward(self, *args, **kwargs):
        raise NotImplementedError()

def parameters(self) -> List[Value]:
        params = list(self._parameters.values())
        for module in self._modules.values():
            params.extend(module.parameters())
        return params

def named_parameters(self, prefix: str = "") -> Iterator[Tuple[str, Value]]:
        for name, param in self._parameters.items():
            yield f"{prefix}{name}" if prefix else name, param
        for name, module in self._modules.items():
            subprefix = f"{prefix}{name}." if prefix else f"{name}."
            yield from module.named_parameters(subprefix)

def train(self, mode: bool = True) -> "Module":
        self.training = mode
        for module in self._modules.values():
            module.train(mode)
        return self

def eval(self) -> "Module":
        return self.train(False)

class Linear(Module):
    """TODO: Implement Linear layer."""

def __init__(self, in_features: int, out_features: int, bias: bool = True) -> None:
        """TODO: Initialize weights and bias."""
        raise NotImplementedError()

def forward(self, x: List[Value]) -> List[Value]:
        """TODO: Compute y = xW^T + b."""
        raise NotImplementedError()

class Sequential(Module):
    """TODO: Implement Sequential container."""

def __init__(self, *modules: Module) -> None:
        """TODO: Store and register modules."""
        raise NotImplementedError()

def forward(self, x):
        """TODO: Pass through each module."""
        raise NotImplementedError()

class Tanh(Module):
    """TODO: Implement Tanh activation."""

def forward(self, x: List[Value]) -> List[Value]:
        raise NotImplementedError()

class ReLU(Module):
    """TODO: Implement ReLU activation."""

def forward(self, x: List[Value]) -> List[Value]:
        raise NotImplementedError()

class Sigmoid(Module):
    """TODO: Implement Sigmoid activation."""