reflectorch 1.2.0__py3-none-any.whl → 1.3.0__py3-none-any.whl

reflectorch/models/encoders/conv_encoder.py

@@ -11,7 +11,6 @@ __all__ = [
     "ConvEncoder",
     "ConvDecoder",
     "ConvAutoencoder",
-    "ConvVAE",
 ]
 
 logger = logging.getLogger(__name__)
@@ -23,7 +22,7 @@ class ConvEncoder(nn.Module):
     Args:
         in_channels (int, optional): the number of input channels. Defaults to 1.
         hidden_channels (tuple, optional): the number of intermediate channels of each convolutional layer. Defaults to (32, 64, 128, 256, 512).
-        dim_latent (int, optional): the dimension of the output latent embedding. Defaults to 64.
+        dim_embedding (int, optional): the dimension of the output latent embedding. Defaults to 64.
         dim_avpool (int, optional): the output size of the adaptive average pooling layer. Defaults to 1.
         use_batch_norm (bool, optional): whether to use batch normalization. Defaults to True.
         activation (str, optional): the type of activation function. Defaults to 'relu'.
@@ -31,9 +30,11 @@ class ConvEncoder(nn.Module):
     def __init__(self,
                  in_channels: int = 1,
                  hidden_channels: tuple = (32, 64, 128, 256, 512),
-                 dim_latent: int = 64,
+                 kernel_size: int = 3,
+                 dim_embedding: int = 64,
                  dim_avpool: int = 1,
                  use_batch_norm: bool = True,
+                 use_se: bool = False,
                  activation: str = 'relu',
                  ):
         super().__init__()
@@ -44,22 +45,24 @@ class ConvEncoder(nn.Module):
 
         for h in hidden_channels:
             layers = [
-                nn.Conv1d(in_channels, out_channels=h, kernel_size=3, stride=2, padding=1),
+                nn.Conv1d(in_channels, out_channels=h, kernel_size=kernel_size, stride=2, padding=kernel_size // 2),
                 activation(),
             ]
 
             if use_batch_norm:
                 layers.insert(1, nn.BatchNorm1d(h))
 
+            if use_se:
+                layers.insert(2, SEBlock(h))
+
             modules.append(nn.Sequential(*layers))
             in_channels = h
 
         self.core = nn.Sequential(*modules)
         self.avpool = nn.AdaptiveAvgPool1d(dim_avpool)
-        self.fc = nn.Linear(hidden_channels[-1] * dim_avpool, dim_latent)
+        self.fc = nn.Linear(hidden_channels[-1] * dim_avpool, dim_embedding)
 
     def forward(self, x):
-        """"""
         if len(x.shape) < 3:
             x = x.unsqueeze(1)
         x = self.core(x)
@@ -100,6 +103,7 @@ class ConvDecoder(nn.Module):
                  hidden_channels: tuple = (512, 256, 128, 64, 32), 
                  dim_latent: int = 64, 
                  in_size: int = 8,
+                 kernel_size: int = 3,
                  use_batch_norm: bool = True,
                  activation: str = 'relu',
                  ):
@@ -119,9 +123,9 @@ class ConvDecoder(nn.Module):
                     nn.ConvTranspose1d(
                         hidden_channels[i],
                         hidden_channels[i + 1],
-                        kernel_size=3,
+                        kernel_size=kernel_size, #3
                         stride=2,
-                        padding=1,
+                        padding=kernel_size // 2, #1
                         output_padding=1,
                     ),
                     nn.BatchNorm1d(hidden_channels[i + 1]) if use_batch_norm else nn.Identity(),
@@ -134,9 +138,9 @@ class ConvDecoder(nn.Module):
         self.final_layer = nn.Sequential(
             nn.ConvTranspose1d(hidden_channels[-1],
                                hidden_channels[-1],
-                               kernel_size=3,
+                               kernel_size=kernel_size, #3
                                stride=2,
-                               padding=1,
+                               padding=kernel_size // 2, #1
                                output_padding=1),
             nn.BatchNorm1d(hidden_channels[-1]) if use_batch_norm else nn.Identity(),
             activation(),
@@ -160,46 +164,56 @@ class ConvAutoencoder(nn.Module):
                  decoder_hidden_channels: tuple = (512, 256, 128, 64, 32),
                  dim_latent: int = 64,
                  dim_avpool: int = 1,
+                 kernel_size: int = 3,
                  use_batch_norm: bool = True,
                  activation: str = 'relu',
                  decoder_in_size: int = 8,
                  **kwargs
                  ):
         super().__init__()
-        self.encoder = ConvEncoder(in_channels, encoder_hidden_channels, dim_latent, dim_avpool, use_batch_norm, activation, **kwargs)
-        self.decoder = ConvDecoder(decoder_hidden_channels, dim_latent, decoder_in_size, use_batch_norm, activation, **kwargs)
+        self.encoder = ConvEncoder(
+            in_channels=in_channels, 
+            hidden_channels=encoder_hidden_channels,
+            kernel_size=kernel_size,
+            dim_embedding=dim_latent, 
+            dim_avpool=dim_avpool, 
+            use_batch_norm=use_batch_norm, 
+            activation=activation,
+            **kwargs)
+
+        self.decoder = ConvDecoder(
+            hidden_channels=decoder_hidden_channels, 
+            dim_latent=dim_latent, 
+            in_size=decoder_in_size, 
+            kernel_size=kernel_size,
+            use_batch_norm=use_batch_norm, 
+            activation=activation, 
+            **kwargs)
 
     def forward(self, x):
         return self.decoder(self.encoder(x))
     
-class ConvVAE(nn.Module):
-    """A 1D convolutional variational autoencoder"""
-    def __init__(self,
-                 in_channels: int = 1,
-                 encoder_hidden_channels: tuple = (32, 64, 128, 256, 512),
-                 decoder_hidden_channels: tuple = (512, 256, 128, 64, 32),
-                 dim_latent: int = 64,
-                 dim_avpool: int = 1,
-                 use_batch_norm: bool = True,
-                 activation: str = 'relu',
-                 decoder_in_size: int = 8,
-                 **kwargs
-                 ):
+class SEBlock(nn.Module):
+    """Squeeze-and-excitation block (https://arxiv.org/abs/1709.01507) """
+    def __init__(self, in_channels, reduction=16):
         super().__init__()
-        self.encoder = ConvEncoder(in_channels, encoder_hidden_channels, 2*dim_latent, dim_avpool, use_batch_norm, activation, **kwargs)
-        self.decoder = ConvDecoder(decoder_hidden_channels, dim_latent, decoder_in_size, use_batch_norm, activation, **kwargs)
+        self.fc1 = nn.Linear(in_channels, in_channels // reduction, bias=False)
+        self.fc2 = nn.Linear(in_channels // reduction, in_channels, bias=False)
+        self.relu = nn.ReLU()
+        self.sigmoid = nn.Sigmoid()
+        self.global_avg_pool = nn.AdaptiveAvgPool1d(1)
 
     def forward(self, x):
-        z_mu, z_logvar = self.encoder(x).chunk(2, dim=-1)
-        z = self.reparameterize(z_mu, z_logvar)
-
-        x_r_mu, x_r_logvar = self.decoder(z).chunk(2, dim=-1)
-        x = self.reparameterize(x_r_mu, x_r_logvar)
-
-        return x, (z_mu, z_logvar, x_r_mu, x_r_logvar)
-    
-    @staticmethod
-    def reparameterize(mu, logvar):
-        std = torch.exp(0.5 * logvar)
-        eps = torch.randn_like(std).to(std)
-        return mu + eps * std
+        batch_size, channels, _ = x.size()
+        
+        #Squeeze
+        se = self.global_avg_pool(x).view(batch_size, channels)
+        
+        #Excitation
+        se = self.fc1(se)
+        se = self.relu(se)
+        se = self.fc2(se)
+        se = self.sigmoid(se).view(batch_size, channels, 1)
+        
+        #Scale the input feature maps (channel-wise attention)
+        return x * se

reflectorch/models/encoders/fno.py

@@ -46,28 +46,30 @@ class FnoEncoder(nn.Module):
         :align: center
 
     Args:
-        ch_in (int): number of input channels
+        in_channels (int): number of input channels
         dim_embedding (int): dimension of the output embedding
         modes (int): number of Fourier modes
         width_fno (int): number of channels of the intermediate representations
         n_fno_blocks (int): number of FNO blocks
         activation (str): the activation function
-        fusion_self_attention (bool): if ``True`` a fusion layer is used after the FNO blocks to produce the final embedding
+        fusion_self_attention (bool): whether to use fusion self attention for merging the tokens (instead of mean)
+        fsa_activation (str): the activation function of the fusion self attention block
     """
     def __init__(
             self, 
-            ch_in: int = 2, 
+            in_channels: int = 2, 
             dim_embedding: int = 128, 
             modes: int = 32, 
             width_fno: int = 64, 
             n_fno_blocks: int = 6, 
             activation: str = 'gelu',
             fusion_self_attention: bool = False,
+            fsa_activation: str = 'tanh',
             ):
         super().__init__()
 
 
-        self.ch_in = ch_in
+        self.in_channels = in_channels
         self.dim_embedding = dim_embedding
         
         self.modes = modes
@@ -77,13 +79,17 @@ class FnoEncoder(nn.Module):
         self.fusion_self_attention = fusion_self_attention
         
 
-        self.fc0 = nn.Linear(ch_in, width_fno) #(r(q), q)
-        self.spectral_convs = nn.ModuleList([SpectralConv1d(in_channels=width_fno, out_channels=width_fno, modes=modes) for _ in range(n_fno_blocks)])
-        self.w_convs = nn.ModuleList([nn.Conv1d(in_channels=width_fno, out_channels=width_fno, kernel_size=1) for _ in range(n_fno_blocks)])
+        self.fc0 = nn.Linear(in_channels, width_fno) #(r(q), q)
+        self.spectral_convs = nn.ModuleList([
+            SpectralConv1d(in_channels=width_fno, out_channels=width_fno, modes=modes) for _ in range(n_fno_blocks)
+            ])
+        self.w_convs = nn.ModuleList([
+            nn.Conv1d(in_channels=width_fno, out_channels=width_fno, kernel_size=1) for _ in range(n_fno_blocks)
+            ])
         self.fc_out = nn.Linear(width_fno, dim_embedding)
 
         if fusion_self_attention:
-            self.fusion = FusionSelfAttention(width_fno, 2*width_fno)
+            self.fusion = FusionSelfAttention(embed_dim=width_fno, hidden_dim=2*width_fno, activation=fsa_activation)
         
     def forward(self, x):
         """"""
@@ -109,19 +115,20 @@ class FnoEncoder(nn.Module):
         
         return x
     
+
 class FusionSelfAttention(nn.Module):
-    def __init__(self, 
-                 embed_dim: int = 64, 
-                 hidden_dim: int = 64,
-                 activation=nn.Tanh,
-                 ):
+    def __init__(self, embed_dim: int = 64, hidden_dim: int = 64, activation: str = 'gelu'):
         super().__init__()
+        activation = activation_by_name(activation)()
         self.fuser = nn.Sequential(nn.Linear(embed_dim, hidden_dim), 
-                                   activation(),
+                                   activation,
                                    nn.Linear(hidden_dim, 1, bias=False))
         
-    def forward(self, c):  # (batch_size x seq_len x embed_dim)
+    def forward(self, 
+                c: torch.Tensor,  # (batch_size x seq_len x embed_dim)
+                mask: torch.Tensor = None, # (batch_size x seq_len)
+                ):  
         a = self.fuser(c)
-        alpha = torch.exp(a)
+        alpha = torch.exp(a)*mask.unsqueeze(-1) if mask is not None else torch.exp(a)
         alpha = alpha/alpha.sum(dim=1, keepdim=True)
         return (alpha*c).sum(dim=1)  # (batch_size x embed_dim)

reflectorch/models/networks/__init__.py

@@ -1,4 +1,5 @@
 from reflectorch.models.networks.mlp_networks import (
+    NetworkWithPriors,
     NetworkWithPriorsConvEmb,
     NetworkWithPriorsFnoEmb,
 )
@@ -7,6 +8,7 @@ from reflectorch.models.networks.residual_net import ResidualMLP
 
 __all__ = [
     "ResidualMLP",
+    "NetworkWithPriors",
     "NetworkWithPriorsConvEmb",
     "NetworkWithPriorsFnoEmb",
 ]