qualia_codegen_plugin_snn.graph.TorchModelGraph module

class qualia_codegen_plugin_snn.graph.TorchModelGraph.SequentialForward(*args, **kwargs)

Bases: Protocol

Type Sequential.forward with torch.Tensor as the original Sequential.forward is untyped and makes mypy unhappy.

forward(x: Tensor) → Tensor

class qualia_codegen_plugin_snn.graph.TorchModelGraph.TorchModelGraph(model: Module)

Bases: TorchModelGraph

MODULE_MAPPING: ClassVar[dict[type[Module], Callable[[Module, TBaseLayer], tuple[type[TBaseLayer], list[Any]]]]] = {<class 'spikingjelly.activation_based.neuron.IFNode'>: <function TorchModelGraph.<lambda>>, <class 'spikingjelly.activation_based.neuron.LIFNode'>: <function TorchModelGraph.<lambda>>, <class 'spikingjelly.activation_based.neuron.ParametricLIFNode'>: <function TorchModelGraph.<lambda>>, <class 'torch.nn.modules.activation.ReLU'>: <function TorchModelGraph.<lambda>>, <class 'torch.nn.modules.activation.ReLU6'>: <function TorchModelGraph.<lambda>>, <class 'torch.nn.modules.batchnorm.BatchNorm1d'>: <function TorchModelGraph.<lambda>>, <class 'torch.nn.modules.batchnorm.BatchNorm2d'>: <function TorchModelGraph.<lambda>>, <class 'torch.nn.modules.conv.Conv1d'>: <function TorchModelGraph.<lambda>>, <class 'torch.nn.modules.conv.Conv2d'>: <function TorchModelGraph.<lambda>>, <class 'torch.nn.modules.dropout.Dropout'>: <function TorchModelGraph.<lambda>>, <class 'torch.nn.modules.flatten.Flatten'>: <function TorchModelGraph.<lambda>>, <class 'torch.nn.modules.linear.Identity'>: <function TorchModelGraph.<lambda>>, <class 'torch.nn.modules.linear.Linear'>: <function TorchModelGraph.<lambda>>, <class 'torch.nn.modules.pooling.AdaptiveAvgPool1d'>: <function TorchModelGraph.<lambda>>, <class 'torch.nn.modules.pooling.AvgPool1d'>: <function TorchModelGraph.<lambda>>, <class 'torch.nn.modules.pooling.AvgPool2d'>: <function TorchModelGraph.<lambda>>, <class 'torch.nn.modules.pooling.MaxPool1d'>: <function TorchModelGraph.<lambda>>, <class 'torch.nn.modules.pooling.MaxPool2d'>: <function TorchModelGraph.<lambda>>}

convert(custom_layers: dict[type[Module], Callable[[Module, TBaseLayer], tuple[type[TBaseLayer], list[Any]]]] | None = None) → ModelGraph | None