deepextractor.models.architectures ================================== .. py:module:: deepextractor.models.architectures Module Contents --------------- .. py:data:: logger .. py:class:: SE1D(channels, r=8) Bases: :py:obj:`torch.nn.Module` Squeeze-and-Excitation channel attention for 1D features (B, C, T). .. py:attribute:: net .. py:method:: forward(x) .. py:class:: MHSA1D(channels, num_heads=4, dropout=0.0) Bases: :py:obj:`torch.nn.Module` Temporal multi-head self-attention on (B, C, T). .. py:attribute:: proj_in .. py:attribute:: attn .. py:attribute:: proj_out .. py:attribute:: norm .. py:method:: forward(x) .. py:class:: LSTMBottleneck1D(channels, hidden: Optional[int] = None, num_layers: int = 1, bidirectional: bool = True, dropout: float = 0.0) Bases: :py:obj:`torch.nn.Module` Bidirectional LSTM bottleneck for (B, C, T) features. .. py:attribute:: lstm .. py:attribute:: proj .. py:attribute:: norm .. py:method:: forward(x) .. py:class:: UNET1D_LSTM_ATT(in_channels: int = 2, out_channels: int = 4, features=(64, 128, 256, 512), use_lstm_bottleneck: bool = True, lstm_layers: int = 1, lstm_bidirectional: bool = True, lstm_dropout: float = 0.0, use_mhsa_bottleneck: bool = False, use_mhsa_encoder: bool = False, use_mhsa_decoder: bool = False, attn_heads: int = 4, use_se_on_skips: bool = True, norm: str = 'bn') Bases: :py:obj:`torch.nn.Module` Time-domain U-Net for two-detector signal/glitch separation. Input : (B, in_channels, T) — default 2 (H1 + L1 strain stacked) Output: (B, out_channels, T) — default 4 (h1_signal, h1_noise, l1_signal, l1_noise) Optional bottleneck: bidirectional LSTM, multi-head self-attention. Optional skip connections: Squeeze-and-Excitation (SE) gating. .. py:attribute:: pool .. py:attribute:: downs .. py:attribute:: ups .. py:attribute:: skip_se .. py:attribute:: use_lstm :value: True .. py:attribute:: use_mhsa_bottleneck :value: False .. py:attribute:: use_mhsa_encoder :value: False .. py:attribute:: use_mhsa_decoder :value: False .. py:attribute:: bottleneck_conv .. py:attribute:: final_conv .. py:method:: forward(x) .. py:class:: DoubleConv1D(in_channels, out_channels) Bases: :py:obj:`torch.nn.Module` Base class for all neural network modules. Your models should also subclass this class. Modules can also contain other Modules, allowing them to be nested in a tree structure. You can assign the submodules as regular attributes:: import torch.nn as nn import torch.nn.functional as F class Model(nn.Module): def __init__(self) -> None: super().__init__() self.conv1 = nn.Conv2d(1, 20, 5) self.conv2 = nn.Conv2d(20, 20, 5) def forward(self, x): x = F.relu(self.conv1(x)) return F.relu(self.conv2(x)) Submodules assigned in this way will be registered, and will also have their parameters converted when you call :meth:`to`, etc. .. note:: As per the example above, an ``__init__()`` call to the parent class must be made before assignment on the child. :ivar training: Boolean represents whether this module is in training or evaluation mode. :vartype training: bool .. py:attribute:: conv .. py:method:: forward(x) .. py:class:: UNET1D(in_channels=1, out_channels=1, features=[64, 128, 256, 512]) Bases: :py:obj:`torch.nn.Module` Base class for all neural network modules. Your models should also subclass this class. Modules can also contain other Modules, allowing them to be nested in a tree structure. You can assign the submodules as regular attributes:: import torch.nn as nn import torch.nn.functional as F class Model(nn.Module): def __init__(self) -> None: super().__init__() self.conv1 = nn.Conv2d(1, 20, 5) self.conv2 = nn.Conv2d(20, 20, 5) def forward(self, x): x = F.relu(self.conv1(x)) return F.relu(self.conv2(x)) Submodules assigned in this way will be registered, and will also have their parameters converted when you call :meth:`to`, etc. .. note:: As per the example above, an ``__init__()`` call to the parent class must be made before assignment on the child. :ivar training: Boolean represents whether this module is in training or evaluation mode. :vartype training: bool .. py:attribute:: ups .. py:attribute:: downs .. py:attribute:: pool .. py:attribute:: bottleneck .. py:attribute:: final_conv .. py:method:: forward(x) .. py:class:: DnCNN1D(depth=12, n_channels=64, image_channels=1, use_bnorm=True, kernel_size=3) Bases: :py:obj:`torch.nn.Module` Base class for all neural network modules. Your models should also subclass this class. Modules can also contain other Modules, allowing them to be nested in a tree structure. You can assign the submodules as regular attributes:: import torch.nn as nn import torch.nn.functional as F class Model(nn.Module): def __init__(self) -> None: super().__init__() self.conv1 = nn.Conv2d(1, 20, 5) self.conv2 = nn.Conv2d(20, 20, 5) def forward(self, x): x = F.relu(self.conv1(x)) return F.relu(self.conv2(x)) Submodules assigned in this way will be registered, and will also have their parameters converted when you call :meth:`to`, etc. .. note:: As per the example above, an ``__init__()`` call to the parent class must be made before assignment on the child. :ivar training: Boolean represents whether this module is in training or evaluation mode. :vartype training: bool .. py:attribute:: dncnn .. py:method:: forward(x) .. py:class:: Autoencoder1D(in_channels=1, out_channels=1, features=[64, 128, 256, 512]) Bases: :py:obj:`torch.nn.Module` Base class for all neural network modules. Your models should also subclass this class. Modules can also contain other Modules, allowing them to be nested in a tree structure. You can assign the submodules as regular attributes:: import torch.nn as nn import torch.nn.functional as F class Model(nn.Module): def __init__(self) -> None: super().__init__() self.conv1 = nn.Conv2d(1, 20, 5) self.conv2 = nn.Conv2d(20, 20, 5) def forward(self, x): x = F.relu(self.conv1(x)) return F.relu(self.conv2(x)) Submodules assigned in this way will be registered, and will also have their parameters converted when you call :meth:`to`, etc. .. note:: As per the example above, an ``__init__()`` call to the parent class must be made before assignment on the child. :ivar training: Boolean represents whether this module is in training or evaluation mode. :vartype training: bool .. py:attribute:: ups .. py:attribute:: downs .. py:attribute:: pool .. py:attribute:: bottleneck .. py:attribute:: final_conv .. py:method:: forward(x) .. py:class:: DoubleConv2D(in_channels, out_channels) Bases: :py:obj:`torch.nn.Module` Base class for all neural network modules. Your models should also subclass this class. Modules can also contain other Modules, allowing them to be nested in a tree structure. You can assign the submodules as regular attributes:: import torch.nn as nn import torch.nn.functional as F class Model(nn.Module): def __init__(self) -> None: super().__init__() self.conv1 = nn.Conv2d(1, 20, 5) self.conv2 = nn.Conv2d(20, 20, 5) def forward(self, x): x = F.relu(self.conv1(x)) return F.relu(self.conv2(x)) Submodules assigned in this way will be registered, and will also have their parameters converted when you call :meth:`to`, etc. .. note:: As per the example above, an ``__init__()`` call to the parent class must be made before assignment on the child. :ivar training: Boolean represents whether this module is in training or evaluation mode. :vartype training: bool .. py:attribute:: conv .. py:method:: forward(x) .. py:class:: UNET2D(in_channels=2, out_channels=2, features=[64, 128, 256, 512]) Bases: :py:obj:`torch.nn.Module` Base class for all neural network modules. Your models should also subclass this class. Modules can also contain other Modules, allowing them to be nested in a tree structure. You can assign the submodules as regular attributes:: import torch.nn as nn import torch.nn.functional as F class Model(nn.Module): def __init__(self) -> None: super().__init__() self.conv1 = nn.Conv2d(1, 20, 5) self.conv2 = nn.Conv2d(20, 20, 5) def forward(self, x): x = F.relu(self.conv1(x)) return F.relu(self.conv2(x)) Submodules assigned in this way will be registered, and will also have their parameters converted when you call :meth:`to`, etc. .. note:: As per the example above, an ``__init__()`` call to the parent class must be made before assignment on the child. :ivar training: Boolean represents whether this module is in training or evaluation mode. :vartype training: bool .. py:attribute:: ups .. py:attribute:: downs .. py:attribute:: pool .. py:attribute:: bottleneck .. py:attribute:: final_conv .. py:method:: forward(x) .. py:class:: Autoencoder2D(in_channels=2, out_channels=2, features=[64, 128, 256, 512]) Bases: :py:obj:`torch.nn.Module` Base class for all neural network modules. Your models should also subclass this class. Modules can also contain other Modules, allowing them to be nested in a tree structure. You can assign the submodules as regular attributes:: import torch.nn as nn import torch.nn.functional as F class Model(nn.Module): def __init__(self) -> None: super().__init__() self.conv1 = nn.Conv2d(1, 20, 5) self.conv2 = nn.Conv2d(20, 20, 5) def forward(self, x): x = F.relu(self.conv1(x)) return F.relu(self.conv2(x)) Submodules assigned in this way will be registered, and will also have their parameters converted when you call :meth:`to`, etc. .. note:: As per the example above, an ``__init__()`` call to the parent class must be made before assignment on the child. :ivar training: Boolean represents whether this module is in training or evaluation mode. :vartype training: bool .. py:attribute:: ups .. py:attribute:: downs .. py:attribute:: pool .. py:attribute:: bottleneck .. py:attribute:: final_conv .. py:method:: forward(x) .. py:class:: ModifiedAutoencoder2D(in_channels=2, out_channels=2, features=[64, 128, 256]) Bases: :py:obj:`torch.nn.Module` Base class for all neural network modules. Your models should also subclass this class. Modules can also contain other Modules, allowing them to be nested in a tree structure. You can assign the submodules as regular attributes:: import torch.nn as nn import torch.nn.functional as F class Model(nn.Module): def __init__(self) -> None: super().__init__() self.conv1 = nn.Conv2d(1, 20, 5) self.conv2 = nn.Conv2d(20, 20, 5) def forward(self, x): x = F.relu(self.conv1(x)) return F.relu(self.conv2(x)) Submodules assigned in this way will be registered, and will also have their parameters converted when you call :meth:`to`, etc. .. note:: As per the example above, an ``__init__()`` call to the parent class must be made before assignment on the child. :ivar training: Boolean represents whether this module is in training or evaluation mode. :vartype training: bool .. py:attribute:: ups .. py:attribute:: downs .. py:attribute:: pool .. py:attribute:: bottleneck .. py:attribute:: final_conv .. py:method:: forward(x)