cu.tools: summaries, layer shapes, and Netron

Model inspection tools: torchinfo summaries, layer shapes, and the Netron graph viewer.

colabutilscs4337.tools.summary(model, input_size=(1, 3, 224, 224), *, depth=3, col_names=('input_size', 'output_size', 'num_params', 'trainable'), verbose=1, **kwargs)

torchinfo.summary with computer-vision defaults; accepts fastai Learners too.

input_size is (batch, channels, height, width); pass None (with input_data=...) to summarize from real inputs instead.

Parameters:

• model (Any)

• input_size (Sequence[int] | None)

• depth (int)

• col_names (Sequence[str])

• verbose (int)

• kwargs (Any)

colabutilscs4337.tools.layer_shapes(model, layers, input_size=(1, 3, 224, 224), *, x=None, verbose=True)

Input and output shapes of one or more layers/blocks, from a real forward pass.

Parameters:

• model (Any) – An nn.Module or a fastai Learner.

• layers (Any) – A submodule (backbone.features[2][0]), a dotted name ("features.2.0"), or a list of either.

• input_size (Sequence[int]) – Shape of the dummy input, (batch, channels, height, width). Ignored if x is given.

• x (Any) – A real input batch to use instead of zeros.

• verbose (bool) – Print one line per layer.

Returns:

• LayerShapes – For a single layer; a {name: LayerShapes} dict for a list.

• The model runs once in eval mode without gradients, so BatchNorm statistics are

• not changed. If a layer runs several times in one pass, its first call is shown.

Examples

>>> backbone = torchvision.models.efficientnet_b0()
>>> s = cu.tools.layer_shapes(backbone, backbone.features[2][0], (1, 3, 160, 160))
features.2.0 (MBConv): (1, 16, 80, 80) -> (1, 24, 40, 40)
>>> s.input, s.output            # first input / output
>>> s.inputs                     # every positional input
>>> cu.tools.layer_shapes(model, ["layer1", "layer2.0"])

class colabutilscs4337.tools.LayerShapes(name, type, inputs=<factory>, kwargs=<factory>, outputs=<factory>)

Bases: object

Shapes seen by one layer during a forward pass.

inputs has one entry per positional argument (a shape tuple, a nested list for lists of tensors, or None for non-tensors); kwargs holds tensors passed by keyword. outputs has one entry per returned value.

Parameters:

• name (str)

• type (str)

• inputs (list)

• kwargs (dict)

• outputs (list)

property input

Shape of the first input (the usual case).

property output

Shape of the first output (the usual case).

colabutilscs4337.tools.netron(model_or_path, *, input_size=(1, 3, 224, 224), format='onnx', port=None, height=700)

Open Netron for a model file or an in-memory model (nn.Module / fastai Learner).

In Colab the viewer is shown inside the cell output through Colab’s port proxy. A free port is chosen unless port is given, and running it again for the same model replaces the previous viewer. Returns Netron’s (host, port) address; pass it to stop_netron() to stop.

Examples

>>> cu.tools.netron(model, input_size=(1, 1, 28, 28))   # MNIST model
>>> cu.tools.netron(learn)                              # fastai Learner, 224x224 RGB
>>> cu.tools.netron("model.onnx", height=600)           # a file you saved

Parameters:

• model_or_path (Any)

• input_size (Sequence[int])

• format (str)

• port (int | None)

• height (int)

Return type:

tuple[str, int]

colabutilscs4337.tools.stop_netron(address=None)

Stop one Netron server, or all servers started by this module.

Parameters:

address (tuple[str, int] | None)

Return type:

None

colabutilscs4337.tools.export_for_netron(model, path=None, input_size=(1, 3, 224, 224), format='onnx')

Save a model in a format Netron can draw as a graph.

format="onnx" (default) gives the cleanest graph, with the input and output named input and output. format="torchscript" writes a traced .pt file and format="export" a torch.export .pt2 archive.

Parameters:

• model (Any)

• path (str | PathLike | None)

• input_size (Sequence[int])

• format (str)

Return type:

Path

colabutilscs4337.tools.unwrap_model(model)

Return the torch.nn.Module inside a fastai Learner (or the module itself).

Parameters:

model (Any)

colabutilscs4337.tools.install(*extras, **kwargs)

Reinstall the tools (torchinfo, netron, onnx, onnxscript) if they are missing.

They come with pip install colabutilscs4337; this is only needed if you uninstalled them or installed the package with --no-deps.

Parameters:

extras (str)