Model Summaries

Structured model-summary rendering for checked sequential models.

structure TorchLean.nn.LayerSummary : Type

Structured per-layer summary derived from a checked sequential model.

• index : ℕ

  Zero-based position in the sequential layer list.

• kind : String

  User-facing layer kind string.

• input : Shape

  Checked input shape for this layer.

• output : Shape

  Checked output shape for this layer.

• stateShapes : Array Shape

  Trainable parameter and persistent-buffer shapes owned by this layer.

• parameterCount : ℕ

  Number of trainable scalar parameters in this layer.

• stateElementCount : ℕ

  Number of scalar elements in this layer's complete state.

def TorchLean.nn.LayerSummary.render (summary : LayerSummary) : String

One-line rendering of one layer summary.

@[instance_reducible]

instance TorchLean.nn.LayerSummary.instToString : ToString LayerSummary

structure TorchLean.nn.ModelSummary : Type

Structured whole-model summary derived from a checked sequential model.

• input : Shape

  Checked input shape for the full model.

• output : Shape

  Checked output shape for the full model.

• layers : Array LayerSummary

  Per-layer summaries in order.

• layerCount : ℕ

  Total number of layers in the sequential model.

• totalParameterCount : ℕ

  Total scalar parameter count across all layers.

• totalStateElementCount : ℕ

  Total scalar element count across trainable parameters and persistent buffers.

def TorchLean.nn.ModelSummary.header (summary : ModelSummary) : String

Header line for the model summary.

def TorchLean.nn.ModelSummary.render (summary : ModelSummary) : String

Multi-line rendering of the structured model summary.

@[instance_reducible]

instance TorchLean.nn.ModelSummary.instToString : ToString ModelSummary

def TorchLean.nn.summary {σ τ : Shape} (model : Sequential σ τ) : Except String ModelSummary

Validate a sequential model and construct its structured summary.

def TorchLean.nn.printSummary {σ τ : Shape} (model : Sequential σ τ) : IO Unit

Print the structured summary of a checked sequential model.

Example:

def model : nn.Sequential [2] [1] :=
  nn.build 0 nn.Sequential![nn.linear 2 8, nn.relu, nn.linear 8 1]

-- Prints one row per layer with its kind, shapes, and parameter count, then the totals. The
-- counterpart of `print(model)` plus `torchinfo.summary`.
def main : IO Unit := nn.printSummary model