08 — Auditing inside a node

Generated from a notebook

This page is notebooks/08_auditing_inside_nodes.ipynb, executed, with its outputs. Run it yourself with jupyter lab notebooks/ after pip install 'somatize[viz]'. Interactive Plotly figures show here as the PNG the notebook also stores.

At graph level, gradient_audit() treats every filter as opaque: one record per node per step. With inside=, the hooks go into the node’s model and each submodule is audited under a hierarchical id (encoder/<module.path>) — the same records, flags, persistence and figures. Here it catches a textbook vanishing gradient.

import warnings

import plotly.io as pio
import torch
import torch.nn as nn

import soma

# Silencing PyTorch's benign warning about backward hooks.
# Interactivo (zoom/hover/brushing) en Jupyter y VS Code, con
# fallback PNG donde no hay JS (GitHub incluido).
pio.renderers.default = "plotly_mimetype+png"
pio.renderers["png"].scale = 2      # crisp on retina displays
pio.renderers["png"].width = 950
warnings.filterwarnings("ignore", message="Full backward hook is firing")
from soma import AuditScope, DifferentiableFilter, Graph


class DeepEncoder(DifferentiableFilter):
    _cache_version = "nb08-deep-v1"

    def __init__(self, layers=12):
        super().__init__(layers=layers)

    def build_module(self, input_shape):
        mods = []
        for _ in range(self.layers):
            lin = nn.Linear(input_shape[-1], input_shape[-1])
            with torch.no_grad():
                lin.weight.mul_(0.35)  # init contractiva: gradientes se apagan
            mods += [lin, nn.Tanh()]
        return nn.Sequential(*mods)

    def output_shape(self, input_shape):
        return input_shape


torch.manual_seed(0)
g = Graph()
enc = DeepEncoder()
g.node("encoder", enc)
x, y = torch.randn(64, 8), torch.randn(64, 8)
g.materialize(x)
g.train()
g.make_optimizer(torch.optim.Adam, lr=1e-3)

Adam (
Parameter Group 0
    amsgrad: False
    betas: (0.9, 0.999)
    capturable: False
    decoupled_weight_decay: False
    differentiable: False
    eps: 1e-08
    foreach: None
    fused: None
    lr: 0.001
    maximize: False
    weight_decay: 0
)

A materialized DifferentiableFilter draws itself too: its submodules with parameters, chained in order, with the θ count per layer.

enc

DeepEncoder (864 θ en 12 submódulos)

inside=True: no configuration at all

It selects direct children that have parameters, descending through wrappers that hold a single child. The root node keeps being audited under its plain id. The alternatives are all the same thing underneath:

g.gradient_audit(inside={"encoder": 2})            # int = depth
g.gradient_audit(inside={"encoder": ["0", "4.*"]}) # fnmatch patterns
g.gradient_audit(inside={"encoder": AuditScope(depth=3, sample_every=10)})

class DeepEncoder(DifferentiableFilter):
    _audit_scope = "auto"   # declared where the model lives

Precedence: inside[...] > the class’s _audit_scope > auto.

with g.track_run("nb08-inside", tags=["demo"]) as run:
    with g.gradient_audit(inside=True) as audit:
        for epoch in range(4):
            with g.context() as ctx:
                g.zero_grad()
                out, _ = g.forward(x)
                g.backward(ctx, nn.functional.mse_loss(out, y))
            g.step(ctx)

print(audit.report().pretty())

                  filter   steps      act|μ|       act σ      |out∂|        |θ∂|         |θ|         ∂/θ                     flags
----------------------------------------------------------------------------------------------------------------------------------
                 encoder       4   2.176e-01   2.119e-01   9.568e-02   2.778e-01   2.845e+00   9.765e-02                   HEALTHY
               encoder/0       4   2.166e-01   2.624e-01   1.931e-09   1.120e-08   7.447e-01   1.503e-08                 VANISHING
               encoder/2       4   1.993e-01   2.315e-01   1.322e-08   4.892e-08   8.452e-01   5.787e-08                 VANISHING
               encoder/4       4   1.770e-01   2.026e-01   8.374e-08   3.280e-07   8.164e-01   4.017e-07                   HEALTHY
               encoder/6       4   1.876e-01   2.247e-01   4.840e-07   1.822e-06   8.897e-01   2.048e-06                   HEALTHY
               encoder/8       4   1.320e-01   1.559e-01   2.361e-06   9.146e-06   7.380e-01   1.239e-05                   HEALTHY
              encoder/10       4   2.281e-01   2.488e-01   9.397e-06   3.364e-05   9.080e-01   3.705e-05                   HEALTHY
              encoder/12       4   2.448e-01   2.572e-01   4.302e-05   1.699e-04   8.493e-01   2.000e-04                   HEALTHY
              encoder/14       4   1.457e-01   1.886e-01   1.886e-04   7.710e-04   7.457e-01   1.034e-03                   HEALTHY
              encoder/16       4   1.837e-01   1.893e-01   6.889e-04   2.417e-03   7.794e-01   3.101e-03                   HEALTHY
              encoder/18       4   1.805e-01   2.096e-01   3.261e-03   1.047e-02   8.097e-01   1.293e-02                   HEALTHY
              encoder/20       4   2.220e-01   2.130e-01   1.783e-02   5.768e-02   7.814e-01   7.382e-02                   HEALTHY
              encoder/22       4   2.233e-01   2.180e-01   9.049e-02   2.715e-01   9.202e-01   2.951e-01                   HEALTHY

Gradient flow, layer by layer

The classic plot: submodules in real execution order on x, |∂| on a log scale on y. A vanishing gradient reads as a staircase falling towards the input — here, several orders of magnitude of it.

view = soma.RunView(run.dir)
view.plot_module_flow("encoder")

The inner architecture, annotated

The submodule tree is persisted to diagnostics/modules/encoder.json — the same Graph schema as graph.json — and rendered through the same overlay mechanism: parameters, mean |∂| and flags, per layer.

from IPython.display import SVG

SVG(view.to_svg(node="encoder"))   # flags and |∂| per layer; to_mermaid(node=...) da la fuente mermaid

Flags, rolled up

Every flagged layer emits its own HealthFlag and the parent node receives one aggregated per family, whose detail names the layers. The outer DAG marks the node; the inner views point at the exact layer.

parent = [f for f in view.health_flags() if f["node_id"] == "encoder"]
children = [f for f in view.health_flags() if f["node_id"].startswith("encoder/")]
print(f"{len(children)} flags de capa; rollup en el padre:")
[f["detail"] for f in parent]

2 flags de capa; rollup en el padre:

['in: 0, 2']

# Series temporales por capa (por defecto plot_audit solo muestra
# roots; node="encoder" would bring them all and warn about the limit of 8).
# Comparamos entrada / medio / salida:
view.plot_audit(
    filters=["encoder", "encoder/0", "encoder/10", "encoder/18"],
    metric="out_grad.norm",
)

soma report <run_id> carries all of this in the report’s Module flow section. For large models, AuditScope(sample_every=N) samples the submodules every N steps (the root always records), and max_modules caps the selection with a warning — never silently.