`elasticai.creator.torch2ir`#

Transform PyTorch models to IR.

The translation process is highly customizable. (Torch2Ir)[elasticai.creator.torch2ir.Torch2Ir] is responsible for the translation process and will call module handlers to extract attributes from modules. These attributes are then attached to the corresponding nodes in the IR.

Each module handler is a function that takes a PyTorch module and returns a dictionary with the extracted attributes.

The Torch2Ir class features some factory methods as class methods, e.g., (Torch2Ir.get_default_converter)[elasticai.creator.torch2ir.Torch2Ir.get_default_converter]. Those will create a new Torch2Ir instance and register some preconfigured module handlers. However, you are free to extend or alter the behaviour of the translation process by registering your own module handlers.

Package Contents#

Classes#

`Torch2Ir`
`Implementation`
`Node`

Functions#

`get_default_converter`
`new_node`
`input_node`
`output_node`
`get_default_converter`

Data#

handlers

API#

class elasticai.creator.torch2ir.Torch2Ir(tracer: torch.fx.Tracer = _DefaultTracer())[source]#

Initialization

register(module_type: str, handler: collections.abc.Callable[[torch.nn.Module], dict]) → collections.abc.Callable[[torch.nn.Module], dict][source]#: The handlers are used to extract the attributes of the module

register_handlers(handlers: collections.abc.Iterable[collections.abc.Callable[[torch.nn.Module], dict]]) → elasticai.creator.torch2ir.torch2ir.Torch2Ir[source]#

convert(model: torch.nn.Module) → collections.abc.Iterator[elasticai.creator.torch2ir.core.Implementation][source]#

__call__(model: torch.nn.Module) → collections.abc.Iterator[elasticai.creator.torch2ir.core.Implementation][source]#

elasticai.creator.torch2ir.get_default_converter() → elasticai.creator.torch2ir.torch2ir.Torch2Ir[source]#

class elasticai.creator.torch2ir.Implementation(*, graph: elasticai.creator.ir.GraphProtocol, name: str | None = None, type: str | None = None, data: dict[str, elasticai.creator.ir.Attribute] | None = None)[source]#

Bases: elasticai.creator.ir.Implementation[elasticai.creator.torch2ir.core.Node, elasticai.creator.ir.Edge]

name: str#: None

type: str#: None

class elasticai.creator.torch2ir.Node(name: str, data: dict[str, elasticai.creator.ir.base.Attribute])[source]#

Bases: elasticai.creator.ir.Node

implementation: str#: None

elasticai.creator.torch2ir.new_node(name: str, type: str, implementation: str, attributes: dict[str, Any] | None = None) → elasticai.creator.torch2ir.core.Node[source]#

elasticai.creator.torch2ir.input_node(attributes: dict[str, Any] | None = None) → elasticai.creator.torch2ir.core.Node[source]#

elasticai.creator.torch2ir.output_node(attributes: dict[str, Any] | None = None) → elasticai.creator.torch2ir.core.Node[source]#

elasticai.creator.torch2ir.get_default_converter() → elasticai.creator.torch2ir.torch2ir.Torch2Ir[source]#

elasticai.creator.torch2ir.handlers#: []

elasticai.creator.torch2ir#