reflectorch 1.5.1__py3-none-any.whl

reflectorch/models/encoders/conv_encoder.py

@@ -0,0 +1,219 @@
+import logging
+from pathlib import Path
+
+import torch
+from torch import nn, load
+
+from reflectorch.models.activations import activation_by_name
+from reflectorch.paths import SAVED_MODELS_DIR
+
+__all__ = [
+    "ConvEncoder",
+    "ConvDecoder",
+    "ConvAutoencoder",
+]
+
+logger = logging.getLogger(__name__)
+
+
+class ConvEncoder(nn.Module):
+    """A 1D CNN encoder / embedding network
+
+    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_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'.
+    """
+    def __init__(self,
+                 in_channels: int = 1,
+                 hidden_channels: tuple = (32, 64, 128, 256, 512),
+                 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__()
+
+        modules = []
+
+        activation = activation_by_name(activation)
+
+        for h in hidden_channels:
+            layers = [
+                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_embedding)
+
+    def forward(self, x):
+        if len(x.shape) < 3:
+            x = x.unsqueeze(1)
+        x = self.core(x)
+        x = self.avpool(x).view(x.size(0), -1)
+        x = self.fc(x)
+        return x
+
+    def load_weights(self, path: str or Path = None, strict: bool = False):
+        if not path:
+            return
+
+        if isinstance(path, str):
+            if not path.endswith('.pt'):
+                path = path + '.pt'
+            path = SAVED_MODELS_DIR / path
+
+        if not path.is_file():
+            logger.error(f'File {str(path)} is not found.')
+            return
+        try:
+            state_dict = load(path)
+            self.load_state_dict(state_dict, strict=strict)
+        except Exception as err:
+            logger.exception(err)
+
+
+class ConvDecoder(nn.Module):
+    """A 1D CNN decoder
+
+    Args:
+        hidden_dims (tuple, optional): the number of intermediate channels of each convolutional layer. Defaults to (512, 256, 128, 64, 32).
+        latent_dim (int, optional): the dimension of the input latent embedding. Defaults to 64.
+        in_size (int, optional): the initial size for upscaling. Defaults to 8.
+        use_batch_norm (bool, optional): whether to use batch normalization. Defaults to True.
+        activation (str, optional): the type of activation function. Defaults to 'relu'.
+    """
+    def __init__(self, 
+                 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',
+                 ):
+
+        super().__init__()
+
+        self.in_size = in_size
+        modules = []
+
+        self.decoder_input = nn.Linear(dim_latent, hidden_channels[0] * in_size)
+
+        activation = activation_by_name(activation)
+
+        for i in range(len(hidden_channels) - 1):
+            modules.append(
+                nn.Sequential(
+                    nn.ConvTranspose1d(
+                        hidden_channels[i],
+                        hidden_channels[i + 1],
+                        kernel_size=kernel_size, #3
+                        stride=2,
+                        padding=kernel_size // 2, #1
+                        output_padding=1,
+                    ),
+                    nn.BatchNorm1d(hidden_channels[i + 1]) if use_batch_norm else nn.Identity(),
+                    activation(),
+                )
+            )
+
+        self.decoder = nn.Sequential(*modules)
+
+        self.final_layer = nn.Sequential(
+            nn.ConvTranspose1d(hidden_channels[-1],
+                               hidden_channels[-1],
+                               kernel_size=kernel_size, #3
+                               stride=2,
+                               padding=kernel_size // 2, #1
+                               output_padding=1),
+            nn.BatchNorm1d(hidden_channels[-1]) if use_batch_norm else nn.Identity(),
+            activation(),
+            nn.Conv1d(hidden_channels[-1], out_channels=1,
+                      kernel_size=3, padding=1)
+        )
+
+    def forward(self, x):
+        batch_size = x.shape[0]
+        x = self.decoder_input(x).view(batch_size, -1, self.in_size)
+        x = self.decoder(x)
+        x = self.final_layer(x).flatten(1)
+        return x
+
+
+class ConvAutoencoder(nn.Module):
+    """A 1D convolutional denoising 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,
+                 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=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 SEBlock(nn.Module):
+    """Squeeze-and-excitation block (https://arxiv.org/abs/1709.01507) """
+    def __init__(self, in_channels, reduction=16):
+        super().__init__()
+        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):
+        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/conv_res_net.py

@@ -0,0 +1,115 @@
+ 
+
+from torch import nn
+from torch.nn import functional as F
+from torch.nn import init
+
+__all__ = [
+    'ConvResidualNet1D',
+]
+
+
+class ConvResidualBlock1D(nn.Module):
+    def __init__(
+            self,
+            channels,
+            activation=F.gelu,
+            dropout_probability=0.0,
+            use_batch_norm=False,
+            zero_initialization=True,
+            kernel_size: int = 3,
+            dilation: int = 1,
+            padding: int = 1,
+    ):
+        super().__init__()
+        self.activation = activation
+
+        self.use_batch_norm = use_batch_norm
+
+        if use_batch_norm:
+            self.batch_norm_layers = nn.ModuleList(
+                [nn.BatchNorm1d(channels, eps=1e-3) for _ in range(2)]
+            )
+        self.conv_layers = nn.ModuleList(
+            [nn.Conv1d(channels, channels, kernel_size=kernel_size, padding=padding, dilation=dilation)
+             for _ in range(2)]
+        )
+        self.dropout = nn.Dropout(p=dropout_probability)
+
+        if zero_initialization:
+            init.uniform_(self.conv_layers[-1].weight, -1e-3, 1e-3)
+            init.uniform_(self.conv_layers[-1].bias, -1e-3, 1e-3)
+
+    def forward(self, inputs):
+        temps = inputs
+        if self.use_batch_norm:
+            temps = self.batch_norm_layers[0](temps)
+        temps = self.activation(temps)
+        temps = self.conv_layers[0](temps)
+        if self.use_batch_norm:
+            temps = self.batch_norm_layers[1](temps)
+
+        temps = self.activation(temps)
+        temps = self.dropout(temps)
+        temps = self.conv_layers[1](temps)
+
+        return inputs + temps
+
+
+class ConvResidualNet1D(nn.Module):
+    def __init__(
+            self,
+            in_channels: int = 1,
+            out_channels: int = 64,
+            hidden_channels: int = 128,
+            num_blocks=5,
+            activation=F.gelu,
+            dropout_probability=0.0,
+            use_batch_norm=True,
+            kernel_size: int = 3,
+            dilation: int = 1,
+            padding: int = 1,
+            avpool: int = 8,
+
+    ):
+        super().__init__()
+
+        self.hidden_channels = hidden_channels
+
+        self.initial_layer = nn.Conv1d(
+            in_channels=in_channels,
+            out_channels=hidden_channels,
+            kernel_size=1,
+            padding=0,
+        )
+        self.blocks = nn.ModuleList(
+            [
+                ConvResidualBlock1D(
+                    channels=hidden_channels,
+                    activation=activation,
+                    dropout_probability=dropout_probability,
+                    use_batch_norm=use_batch_norm,
+                    kernel_size=kernel_size,
+                    dilation=dilation,
+                    padding=padding,
+                )
+                for _ in range(num_blocks)
+            ]
+        )
+
+        self.avpool = nn.AdaptiveAvgPool1d(avpool)
+
+        self.final_layer = nn.Linear(
+            hidden_channels * avpool, out_channels
+        )
+
+    def forward(self, x):
+        temps = self.initial_layer(x.unsqueeze(1))
+
+        for block in self.blocks:
+            temps = block(temps)
+
+        temps = self.avpool(temps).view(temps.size(0), -1)
+        outputs = self.final_layer(temps)
+
+        return outputs

reflectorch/models/encoders/fno.py

@@ -0,0 +1,134 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from reflectorch.models.activations import activation_by_name
+
+class SpectralConv1d(nn.Module):
+    def __init__(self, in_channels, out_channels, modes):
+        super().__init__()
+
+        """
+        1D Fourier layer. It does FFT, linear transform, and Inverse FFT.    
+        """
+
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.modes = modes  #Number of Fourier modes to multiply, at most floor(N/2) + 1
+
+        self.scale = (1 / (in_channels*out_channels))
+        self.weights1 = nn.Parameter(self.scale * torch.rand(in_channels, out_channels, modes, dtype=torch.cfloat))
+
+    # Complex multiplication
+    def compl_mul1d(self, input, weights):
+        # (batch, in_channel, x ), (in_channel, out_channel, x) -> (batch, out_channel, x)
+        return torch.einsum("bix,iox->box", input, weights)
+
+    def forward(self, x):
+        batchsize = x.shape[0]
+        #Compute Fourier coeffcients up to factor of e^(- something constant)
+        x_ft = torch.fft.rfft(x)
+
+        # Multiply relevant Fourier modes
+        out_ft = torch.zeros(batchsize, self.out_channels, x.size(-1)//2 + 1,  device=x.device, dtype=torch.cfloat)
+        out_ft[:, :, :self.modes] = self.compl_mul1d(x_ft[:, :, :self.modes], self.weights1)
+
+        #Return to physical space
+        x = torch.fft.irfft(out_ft, n=x.size(-1))
+        return x
+    
+
+class FnoEncoder(nn.Module):
+    """An embedding network based on the Fourier Neural Operator (FNO) architecture
+
+    .. image:: ../documentation/fig_reflectometry_embedding_networks.png
+        :width: 400px
+        :align: center
+
+    Args:
+        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): 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, 
+            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.in_channels = in_channels
+        self.dim_embedding = dim_embedding
+        
+        self.modes = modes
+        self.width_fno = width_fno
+        self.n_fno_blocks = n_fno_blocks
+        self.activation = activation_by_name(activation)()
+        self.fusion_self_attention = fusion_self_attention
+        
+
+        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(embed_dim=width_fno, hidden_dim=2*width_fno, activation=fsa_activation)
+        
+    def forward(self, x):
+        """"""
+
+        x = x.permute(0, 2, 1) #(B, D, S) -> (B, S, D)
+        x = self.fc0(x)
+        x = x.permute(0, 2, 1) #(B, S, D) -> (B, D, S) 
+
+        for i in range(self.n_fno_blocks):
+            x1 = self.spectral_convs[i](x)
+            x2 = self.w_convs[i](x)
+            
+            x = x1 + x2
+            x = self.activation(x)
+
+        if self.fusion_self_attention:
+            x = x.permute(0, 2, 1)
+            x = self.fusion(x)  
+        else:
+            x = x.mean(dim=-1)
+
+        x = self.fc_out(x)
+        
+        return x
+    
+
+class FusionSelfAttention(nn.Module):
+    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,
+                                   nn.Linear(hidden_dim, 1, bias=False))
+        
+    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)*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)