nnschemaparser 0.1.0__py3-none-any.whl

nnschemaparser/__init__.py

@@ -0,0 +1,27 @@
+"""
+NNSchemaParser.
+
+Converts a torch nn.Module into an inspectable ModelBlockSchema graph that
+keeps the original variable names of the forward source and every sub-module,
+and executes it with a schema interpreter re-dispatching to the real torch
+callables.
+"""
+
+# Import Modules
+from nnschemaparser.core import ModelBlockSchema, TensorRepr, TensorReprWithValue
+from nnschemaparser.dump import dump_schema
+from nnschemaparser.interpreter import SchemaInterpreter, TraceListener, TraceRecorder
+from nnschemaparser.parser import NNParser
+
+__all__ = [
+    "ModelBlockSchema",
+    "TraceListener",
+    "TraceRecorder",
+    "dump_schema",
+    "NNParser",
+    "SchemaInterpreter",
+    "TensorRepr",
+    "TensorReprWithValue"
+]
+
+__version__: str = "0.1.0"

nnschemaparser/core/__init__.py

@@ -0,0 +1,16 @@
+"""
+Core data model of NNSchemaParser.
+
+The schema types describing a parsed model: the block schema tree, the tensor
+representations and the instruction types of the block bodies.
+"""
+
+# Import Modules
+from nnschemaparser.core.model_schema import ModelBlockSchema
+from nnschemaparser.core.tensor_schema import TensorRepr, TensorReprWithValue
+
+__all__ = [
+    "ModelBlockSchema",
+    "TensorRepr",
+    "TensorReprWithValue"
+]

nnschemaparser/core/container_keys.py

@@ -0,0 +1,69 @@
+"""
+Container entry key utilities shared by the lowering and the interpreter.
+
+The entries of a sub-block container are keyed by "container.index" in the
+sub-block dictionaries, and must be iterated in their declaration order.
+"""
+
+# Import Modules
+
+
+def container_entry_sort_key(entry_key: str) -> int:
+    """
+    Get the numeric sort key of a container entry key.
+
+    Args:
+        entry_key (str): The entry key (e.g. "encoders.2").
+
+    Returns:
+        int: The numeric suffix of the key, or -1 when not numeric.
+    """
+
+    # Split the key into its container prefix and its suffix
+    parts: list[str] = entry_key.rsplit(".", 1)
+
+    # A numeric suffix is sorted by its value
+    if len(parts) == 2 and parts[1].isdigit():
+        return int(parts[1])
+
+    # A non numeric suffix is sorted last
+    return -1
+
+
+def container_entry_keys(
+    sub_block_names: set[str] | dict[str, str],
+    container_name: str
+) -> list[str]:
+    """
+    Get the ordered entry keys of a container in the parsed sub-blocks.
+
+    Args:
+        sub_block_names (set[str] | dict[str, str]): The names of the parsed
+            sub-blocks, either a set of names or the sub-block dictionary.
+        container_name (str): The name of the container.
+
+    Returns:
+        list[str]: The entry keys sorted by their numeric suffix.
+    """
+
+    # The special self path holds the direct numeric entries of the module
+    if container_name == "$self":
+        return sorted(
+            [
+                sub_block_name for sub_block_name in sub_block_names
+                if sub_block_name.isdigit()
+            ],
+            key=container_entry_sort_key
+        )
+
+    # Collect the sub-block keys of the container entries
+    keys: list[str] = [
+        sub_block_name for sub_block_name in sub_block_names
+        if sub_block_name.startswith(f"{container_name}.")
+    ]
+
+    # Sort the entry keys by their numeric suffix
+    keys.sort(key=container_entry_sort_key)
+
+    # Return the ordered entry keys of the container
+    return keys

nnschemaparser/core/instructions.py

@@ -0,0 +1,399 @@
+"""
+Instruction types of the NNSchemaParser schema.
+
+Each instruction represents one step of a block body and keeps the original
+variable names of the parsed forward source.
+
+Each instruction represents one step of a model block body, and keeps the
+original variable names of the parsed Python code so the WebUI can display a
+graph that maps back to the source.
+"""
+
+# Import Modules
+from typing import Any, Callable
+
+from torch import Tensor
+
+from nnschemaparser.core.tensor_schema import TensorRepr, TensorReprWithValue
+
+
+class VarRef:
+    """
+    Reference to a variable of the current block scope.
+    """
+
+    def __init__(self, name: str) -> None:
+        """
+        Init function.
+
+        Args:
+            name (str): The name of the referenced variable.
+        """
+
+        # Store here the name of the referenced variable
+        self.name: str = name
+
+
+class BlockRef:
+    """
+    Reference to a sub-block of the current block.
+
+    The path is the attribute name of the sub-block inside its parent module
+    (e.g. "attention"). For containers (nn.ModuleList, nn.ModuleDict), the
+    index_var is the name of the scope variable holding the index or the key.
+    """
+
+    def __init__(self, path: str, index_var: str | None = None) -> None:
+        """
+        Init function.
+
+        Args:
+            path (str): The attribute path of the referenced sub-block.
+            index_var (str | None): The name of the scope variable holding the
+                container index, if the sub-block belongs to a container.
+        """
+
+        # Store here the attribute path of the referenced sub-block
+        self.path: str = path
+
+        # Store here the index variable name for container sub-blocks
+        self.index_var: str | None = index_var
+
+
+class Instruction:
+    """
+    Base class of all the instructions of a block body.
+    """
+
+    def __init__(self, out_var: str) -> None:
+        """
+        Init function.
+
+        Args:
+            out_var (str): The name of the variable the result is bound to.
+        """
+
+        # Store here the output variable name of the instruction
+        self.out_var: str = out_var
+
+        # This field is filled later by the shape resolution pass
+        self.resolved: TensorRepr | TensorReprWithValue | None = None
+
+    def resolve(self, value: Any, capture_values: bool = False) -> None:
+        """
+        Fill the resolved representation of the instruction output.
+
+        Args:
+            value (Any): The value produced by the instruction at execution.
+            capture_values (bool): If True, the tensor values are kept inside
+                the resolved representation.
+        """
+
+        # If the value is a tensor, we store a tensor representation
+        if isinstance(value, Tensor):
+
+            # If values must be captured, keep a representation with the value
+            if capture_values:
+                self.resolved = TensorReprWithValue(
+                    shape=tuple(int(dim) for dim in value.shape),
+                    dtype=str(value.dtype),
+                    value=value.detach()
+                )
+
+            # Else we only keep the shape and the dtype of the tensor
+            else:
+                self.resolved = TensorRepr(
+                    shape=tuple(int(dim) for dim in value.shape),
+                    dtype=str(value.dtype)
+                )
+
+        # Else the value is a python literal, no shape is available
+        else:
+            self.resolved = TensorRepr(shape=(), dtype=type(value).__name__)
+
+
+class OpInstruction(Instruction):
+    """
+    Call to a torch operation or function (e.g. torch.nn.functional.softmax).
+    """
+
+    def __init__(
+        self,
+        target: Callable[..., Any] | VarRef,
+        qualified_name: str,
+        args: list[VarRef | Any],
+        kwargs: dict[str, VarRef | Any],
+        out_var: str,
+        kwargs_var: str | None = None,
+        args_var: str | None = None
+    ) -> None:
+        """
+        Init function.
+
+        Args:
+            target (Callable[..., Any] | VarRef): The callable applied by the
+                instruction, or a variable reference to a callable resolved at
+                execution time (e.g. a method obtained with getattr).
+            qualified_name (str): The display name of the target
+                (e.g. "torch.nn.functional.softmax").
+            args (list[VarRef | Any]): The positional arguments, variable
+                references or literals.
+            kwargs (dict[str, VarRef | Any]): The keyword arguments, variable
+                references or literals.
+            out_var (str): The name of the variable the result is bound to.
+        kwargs_var (str | None): The name of a scope variable holding a dict
+            unpacked as extra keyword arguments (the **kwargs of the call).
+        args_var (str | None): The name of a scope variable holding a tuple
+            unpacked as extra positional arguments (the *args of the call).
+        """
+
+        # Init parent
+        super().__init__(out_var=out_var)
+
+        # Store here the callable applied by the instruction
+        self.target: Callable[..., Any] | VarRef = target
+
+        # Store here the display name of the target
+        self.qualified_name: str = qualified_name
+
+        # Store here the positional arguments of the call
+        self.args: list[VarRef | Any] = args
+
+        # Store here the keyword arguments of the call
+        self.kwargs: dict[str, VarRef | Any] = kwargs
+
+        # Store here the name of the unpacked keyword arguments variable
+        self.kwargs_var: str | None = kwargs_var
+
+        # Store here the name of the unpacked positional arguments variable
+        self.args_var: str | None = args_var
+
+
+class BlockInstruction(Instruction):
+    """
+    Call to a sub-block of the current block.
+    """
+
+    def __init__(
+        self,
+        block_ref: BlockRef,
+        args: list[VarRef | Any],
+        out_var: str,
+        kwargs: dict[str, VarRef | Any] | None = None,
+        splat_args: bool = False,
+        kwargs_var: str | None = None,
+        args_var: str | None = None
+    ) -> None:
+        """
+        Init function.
+
+        Args:
+            block_ref (BlockRef): The reference to the called sub-block.
+            args (list[VarRef | Any]): The positional arguments of the call.
+                When splat_args is True, this list holds a single variable
+                reference whose tuple value is unpacked as the arguments.
+            out_var (str): The name of the variable the result is bound to.
+            kwargs (dict[str, VarRef | Any] | None): The keyword arguments of
+                the call, bound to the sub-block inputs by name.
+            splat_args (bool): Whether the argument variable holds a tuple to
+                unpack as the positional arguments of the call.
+            kwargs_var (str | None): The name of a scope variable holding a dict
+                unpacked as extra keyword arguments (the **kwargs of the call).
+            args_var (str | None): The name of a scope variable holding a
+                tuple unpacked as extra positional arguments (the *args of
+                the call).
+        """
+
+        # Init parent
+        super().__init__(out_var=out_var)
+
+        # Store here the reference to the called sub-block
+        self.block_ref: BlockRef = block_ref
+
+        # Store here the arguments of the sub-block call
+        self.args: list[VarRef | Any] = args
+
+        # Store here the keyword arguments of the sub-block call
+        self.kwargs: dict[str, VarRef | Any] = kwargs or {}
+
+        # Store here whether the argument variable is unpacked as the arguments
+        self.splat_args: bool = splat_args
+
+        # Store here the name of the unpacked keyword arguments variable
+        self.kwargs_var: str | None = kwargs_var
+
+        # Store here the name of the unpacked positional arguments variable
+        self.args_var: str | None = args_var
+
+
+class ControlInstruction(Instruction):
+    """
+    Control flow instruction, an if statement or a for loop.
+
+    For an if statement, the condition is an OpInstruction binding a temporary
+    variable of the block scope. For a for loop, the iterable is either an
+    OpInstruction, or the sub-block container referenced by iter_container,
+    whose entries are bound to bind_var (as sub-block schemas).
+    """
+
+    def __init__(
+        self,
+        kind: str,
+        condition: OpInstruction | None,
+        body: list[Instruction],
+        else_body: list[Instruction],
+        bind_var: str | None,
+        iter_container: BlockRef | None = None,
+        index_var: str | None = None,
+        out_var: str = ""
+    ) -> None:
+        """
+        Init function.
+
+        Args:
+            kind (str): The kind of control flow, either "if" or "for".
+            condition (OpInstruction | None): The condition expression for an
+                if statement, the iterable expression for a for loop, or None
+                when iterating a container of sub-blocks.
+            body (list[Instruction]): The instructions of the main body.
+            else_body (list[Instruction]): The instructions of the else body.
+            bind_var (str | None): The name of the loop variable bound at each
+                iteration of a for loop.
+            iter_container (BlockRef | None): For a for loop over a container
+                of sub-blocks, the reference to the iterated container.
+            index_var (str | None): For a for loop with enumerate over a
+                container, the name of the variable bound to the index.
+            out_var (str): The name of the variable the result is bound to.
+        """
+
+        # Init parent
+        super().__init__(out_var=out_var)
+
+        # Only the if, for and while kinds are supported
+        if kind not in ("if", "for", "while"):
+            raise ValueError(f"Unsupported control flow kind: {kind}")
+
+        # Store here the kind of control flow
+        self.kind: str = kind
+
+        # Store here the condition or iterable expression of the control flow
+        self.condition: OpInstruction | None = condition
+
+        # Store here the instructions of the main body
+        self.body: list[Instruction] = body
+
+        # Store here the instructions of the else body
+        self.else_body: list[Instruction] = else_body
+
+        # Store here the name of the loop variable of a for loop
+        self.bind_var: str | None = bind_var
+
+        # Store here the iterated container reference of a for loop
+        self.iter_container: BlockRef | None = iter_container
+
+        # Store here the enumerate index variable of a for loop
+        self.index_var: str | None = index_var
+
+
+class TryInstruction(Instruction):
+    """
+    Try instruction of a block body, with exception handlers and a finally.
+
+    The exception types of the handlers are resolved at lowering time and
+    stored as python classes. The name of the caught exception object is not
+    bound for now.
+    """
+
+    def __init__(
+        self,
+        body: list[Instruction],
+        handlers: list[tuple[tuple[type[BaseException], ...], list[Instruction]]],
+        handler_names: list[str | None],
+        final_body: list[Instruction],
+        out_var: str = ""
+    ) -> None:
+        """
+        Init function.
+
+        Args:
+            body (list[Instruction]): The instructions of the try body.
+            handlers (list[tuple[tuple[type[BaseException], ...], list[Instruction]]]):
+                The exception types and bodies of the handlers.
+            handler_names (list[str | None]): The caught exception names of
+                the handlers, bound in the scope during their execution.
+            final_body (list[Instruction]): The instructions of the finally body.
+            out_var (str): The name of the variable the result is bound to.
+        """
+
+        # Init parent, a try instruction binds no new variable
+        super().__init__(out_var=out_var)
+
+        # Store here the instructions of the try body
+        self.body: list[Instruction] = body
+
+        # Store here the exception types and bodies of the handlers
+        self.handlers: list[
+            tuple[tuple[type[BaseException], ...], list[Instruction]]
+        ] = handlers
+
+        # Store here the caught exception names of the handlers
+        self.handler_names: list[str | None] = handler_names
+
+        # Store here the instructions of the finally body
+        self.final_body: list[Instruction] = final_body
+
+
+class ReturnInstruction(Instruction):
+    """
+    Return instruction of a block body.
+
+    A return instruction may be nested inside a control flow body, the
+    interpreter propagates its outputs up to the enclosing block.
+    """
+
+    def __init__(self, out_vars: list[str]) -> None:
+        """
+        Init function.
+
+        Args:
+            out_vars (list[str]): The names of the returned variables.
+        """
+
+        # Init parent, a return instruction binds no new variable
+        super().__init__(out_var="")
+
+        # Store here the names of the returned variables
+        self.out_vars: list[str] = out_vars
+
+
+class ContinueInstruction(Instruction):
+    """
+    Continue instruction of a loop body.
+
+    The execution of the current iteration stops and the loop goes to its
+    next iteration, like the python continue statement.
+    """
+
+    def __init__(self) -> None:
+        """
+        Init function.
+        """
+
+        # Init parent, a continue instruction binds no new variable
+        super().__init__(out_var="")
+
+
+class BreakInstruction(Instruction):
+    """
+    Break instruction of a loop body.
+
+    The execution of the loop stops, like the python break statement.
+    """
+
+    def __init__(self) -> None:
+        """
+        Init function.
+        """
+
+        # Init parent, a break instruction binds no new variable
+        super().__init__(out_var="")

nnschemaparser/core/model_schema.py

@@ -0,0 +1,94 @@
+"""
+Core schema types of NNSchemaParser.
+
+A parsed model is a tree of ModelBlockSchema blocks, each holding its inputs,
+outputs, weights, sub-blocks and the instruction list of its forward body.
+"""
+
+# Import Modules
+from typing import Any
+
+from nnschemaparser.core.instructions import Instruction
+from nnschemaparser.core.tensor_schema import TensorRepr, TensorReprWithValue
+
+
+class ModelBlockSchema:
+    """
+    Schema of one block of a parsed model.
+
+    A block mirrors one nn.Module of the source model: its forward parameters,
+    its returned values, its direct parameters and buffers, its child modules
+    as sub-blocks, and the lowered body of its forward method.
+    """
+
+    def __init__(
+        self,
+        inputs: dict[str, TensorRepr],
+        outputs: list[TensorRepr],
+        weights: dict[str, TensorReprWithValue],
+        sub_blocks: dict[str, "ModelBlockSchema"],
+        body: list[Instruction],
+        module: Any | None = None,
+        methods: dict[str, "ModelBlockSchema"] | None = None,
+        method_name: str | None = None,
+        input_defaults: dict[str, Any] | None = None,
+        skipped: dict[str, str] | None = None
+    ) -> None:
+        """
+        Init function.
+
+        Args:
+            inputs (dict[str, TensorRepr]): The list of all named input values of the model
+                                            (often one X value)
+            outputs (list[TensorRepr]): The list of all named output values of the model
+                                            (often one Y value)
+            weights: dict[str, TensorReprWithValue]: The weights of the current block.
+            sub_blocks (dict[str, ModelBlockSchema]): The list of all the sub block of the model
+                                            (all the other nn.Modules in the model class attribute)
+            body (list[Instruction]): What is happening inside the current model block
+                                            (the graph that will be shown in the webUI)
+            module (Any | None): The source module of the block, used by the
+                                            interpreter to resolve dynamic attributes.
+            methods (dict[str, ModelBlockSchema] | None): The other methods of the
+                                            source module called by the body, parsed
+                                            as block schemas (e.g. get_image_features).
+            method_name (str | None): The name of the parsed method when the block
+                                            describes a method instead of the forward.
+            input_defaults (dict[str, Any] | None): The default values of the
+                                            optional inputs, by parameter name.
+            skipped (dict[str, str] | None): The children of the module that
+                                            could not be parsed, by name with
+                                            their skip reasons. They are
+                                            reached at run time through the
+                                            source module instead.
+        """
+
+        # Store here the model input(s) representation(s)
+        self.inputs: dict[str, TensorRepr] = inputs
+
+        # Store here the model output(s) representation(s)
+        self.outputs: list[TensorRepr] = outputs
+
+        # Store here the weights of the block
+        self.weights: dict[str, TensorReprWithValue] = weights
+
+        # Store here the model sub blocks
+        self.sub_blocks: dict[str, ModelBlockSchema] = sub_blocks
+
+        # Store here the model body list of instructions / expressions
+        self.body: list[Instruction] = body
+
+        # Store here the source module of the block
+        self.module: Any | None = module
+
+        # Store here the other methods of the module called by the body
+        self.methods: dict[str, ModelBlockSchema] = methods or {}
+
+        # Store here the default values of the optional inputs
+        self.input_defaults: dict[str, Any] = input_defaults or {}
+
+        # Store here the skipped children re-dispatched at run time
+        self.skipped: dict[str, str] = skipped or {}
+
+        # Store here the name of the described method, if not the forward
+        self.method_name: str | None = method_name

nnschemaparser/core/tensor_schema.py

@@ -0,0 +1,55 @@
+"""
+Tensor representation types of NNSchemaParser.
+
+A TensorRepr describes a tensor flowing through the graph by its shape and its
+dtype, and TensorReprWithValue additionally holds the value itself.
+"""
+
+# Import Modules
+from torch import Tensor
+
+# The placeholder dtype of unresolved tensor representations
+PLACEHOLDER_DTYPE: str = ""
+
+
+class TensorRepr:
+    """
+    Representation of a tensor flowing through the graph, without its value.
+    """
+
+    def __init__(self, shape: tuple[int, ...], dtype: str) -> None:
+        """
+        Init function.
+
+        Args:
+            shape (tuple[int | str, ...]): shape of the tensor, the str type is for symbolic values
+            dtype (str): The type of the tensor's values
+        """
+
+        # Store here the shape of the tensor
+        self.shape: tuple[int, ...] = shape
+
+        # Store here the type of the tensor
+        self.dtype: str = dtype
+
+
+class TensorReprWithValue(TensorRepr):
+    """
+    Representation of a tensor keeping a snapshot of its value.
+    """
+
+    def __init__(self, shape: tuple[int, ...], dtype: str, value: Tensor) -> None:
+        """
+        Init function.
+
+        Args:
+            shape (tuple[int, ...]): Shape of the tensor.
+            dtype (str): The type of the tensor's values.
+            value (Tensor): The value of the tensor.
+        """
+
+        # Init parent
+        super().__init__(shape=shape, dtype=dtype)
+
+        # Store here the value of the tensor
+        self.value: Tensor = value