Skip to content

Module¤

Subclass Module, assign layers as attributes, and implement forward:

from minitorch import Module, Linear, ReLU

class MyModel(Module):
    def __init__(self):
        super().__init__()
        self.fc1 = Linear(784, 128)
        self.relu = ReLU()
        self.fc2 = Linear(128, 10)

    def forward(self, x):
        return self.fc2(self.relu(self.fc1(x)))

Module ¤

Module()

Base class for layers and models.

Subclass it and implement forward. Any Tensor or Module assigned as an attribute is discovered automatically by parameters(), state_dict(), and the train/eval switches.

Source code in minitorch/module.py 
24
25
def __init__(self):
    self._training = True

parameters ¤

parameters()

Collect every trainable Tensor in this module and its children.

Source code in minitorch/module.py 
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
def parameters(self):
    """Collect every trainable `Tensor` in this module and its children."""
    params = []
    for val in self.__dict__.values():
        if isinstance(val, Tensor) and val.requires_grad:
            params.append(val)
        elif isinstance(val, Module):
            params.extend(val.parameters())
        elif isinstance(val, (list, tuple)):
            for item in val:
                if isinstance(item, Tensor) and item.requires_grad:
                    params.append(item)
                elif isinstance(item, Module):
                    params.extend(item.parameters())
    return _dedup(params)

save ¤

save(path)

Write all parameters to a .npz file.

Source code in minitorch/module.py 
71
72
73
def save(self, path):
    """Write all parameters to a .npz file."""
    np.savez(path, **self.state_dict())

load ¤

load(path)

Load parameters from a .npz file written by save().

Source code in minitorch/module.py 
75
76
77
78
79
def load(self, path):
    """Load parameters from a .npz file written by save()."""
    with np.load(path) as data:
        self.load_state_dict({k: data[k] for k in data.files})
    return self

Sequential ¤

Sequential(*layers)

Bases: Module

Chain modules so the output of each feeds the next.

Parameter collection, train/eval, and state_dict all come from Module, which already walks the layers list.

Source code in minitorch/module.py 
127
128
129
def __init__(self, *layers):
    super().__init__()
    self.layers = list(layers)