Skip to content

English · Español

01 — Parameter y Module: la mecánica de registro

🇪🇸 El truco de toda la fase cabe en treinta líneas. Parameter es un marcador (Tensor con requires_grad=True y nada más); Module usa __setattr__ para detectar cuándo le asignas un Parameter o un sub-Module y los guarda en diccionarios internos. El método parameters() recorre esos diccionarios recursivamente. Una vez tienes este esqueleto, Linear, Sequential, Adam y cualquier red profunda son extensiones triviales.

La clase Parameter de cinco líneas

class Parameter(Tensor):
    """A Tensor that is an owned, learnable weight of a Module."""
    def __init__(self, data, requires_grad: bool = True) -> None:
        super().__init__(data, requires_grad=requires_grad)

¿Por qué molestarse con una subclase en vez de usar simplemente Tensor(..., requires_grad=True)?

  • isinstance los distingue. Module.__setattr__ usa isinstance(value, Parameter) para decidir qué registrar. Si usáramos Tensors planos, tendríamos que comprobar requires_grad=True — pero un intermedio de forward (out = self.fc1(x)) tiene requires_grad=True también porque hereda de un parámetro padre. Necesitamos una manera de decir "este tensor concreto es un peso aprendible raíz". Subclase es el discriminador más simple.
  • Claridad de API. Cuando un usuario escribe self.W = Parameter(rng.standard_normal(shape)), queda inmediatamente claro que esto es un peso, no una activación.

PyTorch usa el mismo patrón: torch.nn.Parameter es una subclase delgada de torch.Tensor.

La clase base Module

class Module:
    """Base class for all neural network modules."""

    def __init__(self) -> None:
        # Initialize the registration dicts via object.__setattr__ to avoid
        # infinite recursion in our own __setattr__.
        object.__setattr__(self, "_parameters", {})
        object.__setattr__(self, "_modules", {})
        object.__setattr__(self, "training", True)

    def __setattr__(self, name: str, value: Any) -> None:
        # If the new value is a Parameter, register it.
        if isinstance(value, Parameter):
            # If we're overwriting an existing Parameter or submodule, drop the old registration.
            self._parameters.pop(name, None)
            self._modules.pop(name, None)
            self._parameters[name] = value
        elif isinstance(value, Module):
            self._parameters.pop(name, None)
            self._modules.pop(name, None)
            self._modules[name] = value
        else:
            # If we're overwriting something previously registered, drop the registration.
            self._parameters.pop(name, None)
            self._modules.pop(name, None)
        # In all cases, still set the attribute on the instance.
        object.__setattr__(self, name, value)

    def parameters(self) -> Iterator[Parameter]:
        """Yield all Parameters in this module and its submodules, depth-first."""
        yield from self._parameters.values()
        for submodule in self._modules.values():
            yield from submodule.parameters()

    def zero_grad(self) -> None:
        """Set the gradient of every parameter to None (lazy convention)."""
        for p in self.parameters():
            p.grad = None

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

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

Veinticinco líneas. Repasa qué te aporta cada parte.

¿Por qué object.__setattr__ en __init__?

Si __init__ hiciera self._parameters = {}, eso llamaría a nuestro propio __setattr__. Pero en ese momento self._parameters aún no existe, así que self._parameters.pop(name, None) (dentro de __setattr__) reventaría. El huevo y la gallina. Sortea nuestra propia sobrescritura llamando directamente a object.__setattr__.

Este es el único trozo de magia de Python en todo el módulo. Una vez pasado __init__, la sintaxis ordinaria de atributos funciona.

¿Por qué hacer pop de _parameters y _modules antes de asignar?

Supón que un usuario hace:

self.layer = Linear(2, 3)  # registra en _modules
self.layer = Parameter(rng.randn(2, 3))  # ahora lo queremos en _parameters, no en _modules

Si no eliminamos el registro antiguo, ambos dicts contienen "layer", y parameters() recorre el submódulo (ahora obsoleto). Siempre haz pop de ambos para mantener los dicts consistentes.

¿Por qué parameters() recurre vía submodule.parameters()?

Para manejar anidamientos arbitrariamente profundos:

class Block(Module):
    def __init__(self):
        super().__init__()
        self.fc = Linear(10, 10)

class Net(Module):
    def __init__(self):
        super().__init__()
        self.block1 = Block()
        self.block2 = Block()

Net().parameters() debe producir 4 tensores (2 pesos, 2 biases). La recursión es Net → block1.parameters() → fc.parameters() → [W, b], e igual para block2.

¿Por qué zero_grad() es p.grad = None?

Convención perezosa de la Fase 8. Poner a None hace que el siguiente _backward aloje fresco. La alternativa (np.zeros_like) es un modo de memoria ansiosa útil en algunos bucles de entrenamiento de larga duración; cambiaremos a él en la Fase 18 si el profiling muestra que el overhead de asignación importa.

¿Por qué forward() lanza NotImplementedError?

Fuerza a las subclases a definir un forward. El patrón es:

class TenseMLP(Module):
    def __init__(self):
        super().__init__()
        self.fc1 = Linear(23, 16)
        self.fc2 = Linear(16, 5)

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

model(x) llama a model.__call__(x) que llama a model.forward(x). PyTorch hace esto para poder interceptar __call__ para hooks (Fase 18). El Module de la Fase 9 no necesita hooks, pero mantenemos la indirección por paridad de API.

state_dict y load_state_dict — el contrato de serialización

La spec de la Fase 9 no obliga al checkpointing completo, pero la superficie de la API debería anticiparlo. Esbozo:

class Module:
    def state_dict(self, prefix: str = "") -> dict[str, np.ndarray]:
        """Return a flat dict of named parameter arrays."""
        out = {}
        for name, p in self._parameters.items():
            out[prefix + name] = p.data
        for name, m in self._modules.items():
            out.update(m.state_dict(prefix + name + "."))
        return out

    def load_state_dict(self, state: dict[str, np.ndarray], prefix: str = "") -> None:
        """Load arrays into our parameters (by name)."""
        for name, p in self._parameters.items():
            if prefix + name in state:
                p.data[...] = state[prefix + name]  # in-place copy into the Parameter's array
        for name, m in self._modules.items():
            m.load_state_dict(state, prefix + name + ".")

Decisiones clave:

  • Claves planas separadas por puntos ("fc1.weight", no dicts anidados). Coincide con PyTorch.
  • Copia de datos in-place en load_state_dict. Preserva la identidad del objeto Parameter — la referencia del optimizador a ese Parameter sigue funcionando después de cargar.
  • Formato de serialización. Pickle es el antigoal de la spec (§5.1 de LYNX_CORTEX.md). Usa safetensors desde el primer día para la persistencia en disco real. El método state_dict() devuelve un dict plano; safetensors.save_file(model.state_dict(), "path.safetensors") hace la E/S. Cargar: model.load_state_dict(safetensors.load_file("path.safetensors")).

El __repr__ para depuración

def __repr__(self) -> str:
    lines = [self.__class__.__name__ + "("]
    for name, m in self._modules.items():
        sub_repr = repr(m).replace("\n", "\n  ")
        lines.append(f"  ({name}): {sub_repr}")
    for name, p in self._parameters.items():
        lines.append(f"  ({name}): Parameter(shape={p.shape})")
    lines.append(")")
    return "\n".join(lines)

Llamar a print(model) debería producir:

TenseMLP(
  (fc1): Linear(
    (weight): Parameter(shape=(16, 23))
    (bias): Parameter(shape=(16,))
  )
  (fc2): Linear(
    (weight): Parameter(shape=(5, 16))
    (bias): Parameter(shape=(5,))
  )
)

Útil para verificar formas de parámetros antes del entrenamiento. El print(model) de PyTorch se ve esencialmente idéntico.

Casos límite a probar

  1. Un Module sin parámetros. Module().parameters() no produce nada. No lanza.
  2. Un Module cuyo forward nunca llama a ningún submódulo (degenerado). parameters() sigue produciendo los registrados.
  3. Un Module que registra el mismo Parameter dos veces bajo nombres distintos (parámetros compartidos — embeddings atados). 
    shared = Parameter(...)
self.in_embedding = shared
self.out_embedding = shared
    parameters() produce el mismo objeto dos veces. Ese es también el comportamiento de PyTorch — pero es un footgun: el optimizador aplica la actualización dos veces. Lo documentamos; la Fase 17 puede añadir un paso de deduplicación.
  4. Una lista de Parameters. self.weights = [Parameter(...), Parameter(...)] no registra ningún parámetro (la lista no es un Parameter ni un Module). Para arreglarlo, PyTorch tiene ParameterList. La Fase 9 documenta el gotcha; difiere ParameterList a la Fase 14 (si es necesario para atención multi-head).
  5. Reasignar a None. self.fc1 = None no lanza, pero des-registra el módulo previo (el pop ocurre). Comportamiento sutil — pruébalo.

El invariante hashable-por-identidad (de la Fase 8)

Parameter hereda __hash__ de Tensor (identidad por defecto). El optimizador usa id(p) implícitamente vía la semántica de set/dict de Python. No sobreescribas Parameter.__eq__ para comparación elemento a elemento — rompería la recolección de parameters() si alguna vez los almacenas en un set.

Escollos (morderán en el lab)

  1. Olvidar super().__init__(). Primera línea del __init__ de cada subclase de Module. Sin ello, _parameters no existe, y la primera self.W = Parameter(...) revienta.
  2. Asignar un Tensor plano pensando que será un parámetro. No lo será — __setattr__ comprueba isinstance(value, Parameter). Test: model = MyModel(); list(model.parameters()) debe ser exactamente la cuenta esperada.
  3. Capturar el iterador en vez de la lista. params = model.parameters() es un generador; una vez agotado, está vacío. El optimizador acepta una lista. Conversión: optim = SGD(list(model.parameters()), lr=0.01).
  4. Recrear una capa en forward(). def forward(self, x): self.fc = Linear(...); return self.fc(x) crea un Linear fresco (con pesos aleatorios) en cada forward pass. El Parameter se registra, pero el nuevo cada vez. El entrenamiento nunca funciona. Bug difícil porque la pérdida sí cambia entre forward passes — solo a valores aleatorios.
  5. Llamar a model.forward(x) en vez de model(x). Funciona en la Fase 9 (sin hooks). Saltará silenciosamente los hooks en la Fase 18.

Ancla temática (§A13)

El capstone de la Fase 9 TenseMLP(Module) tiene exactamente dos submódulos (fc1, fc2) y cuatro parámetros en total. model.parameters() debe producirlos en este orden: fc1.weight, fc1.bias, fc2.weight, fc2.bias. Verifica con un test antes del entrenamiento. Si el orden es incorrecto, el estado del optimizador (m, v para Adam) se asocia al parámetro equivocado — Adam aún "entrena" pero las estimaciones de momento están desalineadas. Bug silencioso difícil de diagnosticar.

Lo que esta página NO cubre

  • Modos Module.train() / eval(). Stub como no-op en el Lab 00; activados en la Fase 10.
  • Buffers (estado persistente no aprendible). Fase 11.
  • Hooks (register_forward_hook). Fase 18.

Recapitulación de un párrafo

Parameter es una subclase marcadora de Tensor. Module.__init__ inicializa dos dicts (_parameters y _modules) vía object.__setattr__ para evitar recursión. Module.__setattr__ inspecciona cada valor asignado: los Parameters aterrizan en _parameters, los Modules en _modules, cualquier otra cosa es plana. parameters() recorre los dos dicts recursivamente. zero_grad, state_dict, load_state_dict y __repr__ usan todos el mismo recorrido. ~25 líneas de código; el resto del framework (Linear, Sequential, optimizadores) se asienta limpiamente encima.

Siguiente: 02-linear-and-sequential.md