ExTorch.NN.Introspect (extorch v0.4.0)

Copy Markdown

Introspect the structure of TorchScript (JIT) models.

This module provides functions to extract model architecture information from loaded .pt models, including the computation graph, parameter shapes, submodule structure, and available methods.

Example

model = ExTorch.JIT.load("resnet18.pt")
schema = ExTorch.NN.Introspect.schema(model)
IO.inspect(schema.submodules, label: "Submodules")

Summary

Functions

graph(model)

Get the JIT computation graph IR as a string.

schema(model)

Extract a structured schema from a JIT model.

schema_expanded(model)

Extract a fully expanded schema, recursively descending into Sequential and other container modules.

to_elixir(model, module_name \\ "MyModel", opts \\ [])

Generate Elixir source code for an ExTorch.NN.Module DSL definition that mirrors the structure of a loaded JIT model.

Functions

graph(model)

@spec graph(ExTorch.JIT.Model.t()) :: String.t()

Get the JIT computation graph IR as a string.

Returns the TorchScript IR for the model's forward method, showing the operations and data flow.

schema(model)

@spec schema(ExTorch.JIT.Model.t()) :: ExTorch.NN.Introspect.Schema.t()

Extract a structured schema from a JIT model.

Returns a %ExTorch.NN.Introspect.Schema{} with parameter info, submodule info, and method names.

schema_expanded(model)

@spec schema_expanded(ExTorch.JIT.Model.t()) :: ExTorch.NN.Introspect.Schema.t()

Extract a fully expanded schema, recursively descending into Sequential and other container modules.

Returns a %Schema{} where all container modules have been expanded into their leaf layers with dotted path names (e.g., layer1.0.conv1).

to_elixir(model, module_name \\ "MyModel", opts \\ [])

@spec to_elixir(ExTorch.JIT.Model.t(), String.t(), keyword()) :: String.t()

Generate Elixir source code for an ExTorch.NN.Module DSL definition that mirrors the structure of a loaded JIT model.

Arguments

  • model - A loaded JIT model.
  • module_name - The name for the generated Elixir module (default: "MyModel").

Returns

A string containing Elixir source code.