Visualization API

Model explorer

visualize

visualize(model: Module, filename: str | Path = 'model_architecture.html', *, input_shape: Sequence[int | None] | None = None, open_browser: bool = True) -> ModelVisualization

Create a standalone architecture overview and full-neuron map.

Source code in MiniTorch/visualization/viewer.py

def visualize(
    model: Module,
    filename: str | Path = "model_architecture.html",
    *,
    input_shape: Sequence[int | None] | None = None,
    open_browser: bool = True,
) -> ModelVisualization:
    """Create a standalone architecture overview and full-neuron map."""
    architecture = export_architecture(model, input_shape=input_shape)
    destination = Path(filename)
    destination.parent.mkdir(parents=True, exist_ok=True)
    destination.write_text(_render_html(architecture), encoding="utf-8")
    result = ModelVisualization(destination, architecture)
    if open_browser:
        result.show()
    return result

ModelVisualization dataclass

ModelVisualization(path: Path, architecture: dict)

A generated model visualization that can be saved or opened.

show

show() -> ModelVisualization

Source code in MiniTorch/visualization/viewer.py

def show(self) -> ModelVisualization:
    webbrowser.open(self.path.resolve().as_uri())
    return self

save

save(filename: str | Path) -> ModelVisualization

Source code in MiniTorch/visualization/viewer.py

def save(self, filename: str | Path) -> ModelVisualization:
    destination = Path(filename)
    destination.parent.mkdir(parents=True, exist_ok=True)
    destination.write_text(self.path.read_text(encoding="utf-8"), encoding="utf-8")
    return ModelVisualization(destination, self.architecture)

Architecture export

export_architecture

export_architecture(model: Module, input_shape: Sequence[int | None] | None = None) -> dict[str, Any]

Export a stable, JSON-serializable neural-network description.

Source code in MiniTorch/visualization/architecture.py

def export_architecture(
    model: Module,
    input_shape: Sequence[int | None] | None = None,
) -> dict[str, Any]:
    """Export a stable, JSON-serializable neural-network description."""
    if not isinstance(model, Module):
        raise TypeError("visualize(model) expects a MiniTorch.nn.Module")

    modules = _leaf_modules(model)
    first_linear = next(
        (module for _, module in modules if isinstance(module, Linear)),
        None,
    )
    if input_shape is None:
        if first_linear is None:
            input_shape = (None,)
        else:
            input_shape = (None, first_linear.W.shape[0])

    current_shape = list(input_shape)
    input_neurons = (
        int(current_shape[-1])
        if current_shape and current_shape[-1] is not None
        else None
    )
    nodes: list[dict[str, Any]] = [
        {
            "id": "input",
            "name": "input",
            "type": "Input",
            "shape": current_shape,
            "neurons": input_neurons,
            "parameters": 0,
            "parameterTensors": [],
        }
    ]
    edges: list[dict[str, str]] = []
    previous_id = "input"

    for index, (name, module) in enumerate(modules):
        node_id = f"layer-{index}"
        parameters = module.named_parameters()
        parameter_tensors = [
            {
                "name": parameter_name.rsplit(".", 1)[-1],
                **_parameter_metadata(parameter),
            }
            for parameter_name, parameter in parameters
        ]

        node: dict[str, Any] = {
            "id": node_id,
            "name": name,
            "type": module.__class__.__name__,
            "shape": current_shape,
            "neurons": current_shape[-1] if current_shape else None,
            "parameters": int(sum(parameter.size for _, parameter in parameters)),
            "parameterTensors": parameter_tensors,
        }

        if isinstance(module, Linear):
            in_features, out_features = module.W.shape
            current_shape = [*current_shape[:-1], out_features]
            node.update(
                {
                    "inputFeatures": in_features,
                    "outputFeatures": out_features,
                    "shape": current_shape,
                    "neurons": out_features,
                    "bias": module.b is not None,
                    "dtype": str(module.W.dtype),
                    "neuronStats": _neuron_stats(module),
                }
            )

        nodes.append(node)
        edges.append(
            {
                "id": f"{previous_id}-{node_id}",
                "source": previous_id,
                "target": node_id,
            }
        )
        previous_id = node_id

    return {
        "schemaVersion": 1,
        "model": {
            "name": model.__class__.__name__,
            "totalParameters": int(sum(parameter.size for parameter in model.parameters())),
            "parameterTensors": len(model.parameters()),
            "inputShape": list(input_shape),
            "outputShape": current_shape,
        },
        "nodes": nodes,
        "edges": edges,
    }