ModuleTreeWidget API#

ModuleTreeWidget turns a PyTorch nn.Module into an expandable tree: every submodule with its parameter count and per-tensor shapes, plus trainable, frozen and buffer badges and a density indicator, so you can see where the weights in a model actually sit. If you work in marimo you do not need it — returning an nn.Module from a cell renders the same view, since this widget has graduated to marimo core — but it still works in plain Jupyter and other anywidget hosts.

PyTorch is too heavy to run in the browser, so this page has no in-browser demo — run it on molab instead. See also: NestedTable for the same collapsing-tree layout over any hierarchy, Treemap for sizing a hierarchy by value, and LiveEdit for tracing what a Python function does line by line.

import marimo as mo
import torch.nn as nn
from wigglystuff import ModuleTreeWidget

model = nn.Sequential(
    nn.Linear(784, 256),
    nn.ReLU(),
    nn.Linear(256, 10),
)
mo.ui.anywidget(ModuleTreeWidget(model, initial_expand_depth=2))

169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188

def __init__(
    self,
    module: Any = None,
    *,
    initial_expand_depth: int = 1,
):
    """Create a ModuleTreeWidget.

    Args:
        module: A PyTorch ``nn.Module`` to visualise.
        initial_expand_depth: Number of tree levels to expand initially.
    """
    warn_if_in_marimo(
        "ModuleTreeWidget",
        "marimo's built-in PyTorch formatter will render an "
        "<code>nn.Module</code> returned from a cell.",
    )
    super().__init__(initial_expand_depth=initial_expand_depth)
    if module is not None:
        self.tree = _extract_tree(module)

total_param_count

total_size_bytes

total_trainable_count

Synced traitlets#