qualia_codegen_plugin_snn.graph package

Subpackages

Submodules

Module contents

class qualia_codegen_plugin_snn.graph.TorchModelGraph(model: Module)[source]

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[source]