Neural Nets

Fully convolutional neural networks

class atomai.nets.Unet(nb_classes=1, nb_filters=16, dropout=False, batch_norm=True, upsampling_mode='bilinear', with_dilation=False, **kwargs)[source]

Bases: torch.nn.Module

Builds a fully convolutional Unet-like neural network model

Parameters
  • nb_classes (int) – Number of classes in the ground truth

  • nb_filters (int) – Number of filters in 1st convolutional block (gets multiplied by 2 in each next block)

  • dropout (bool) – Use dropouts to the 3 inner layers (Default: False)

  • batch_norm (bool) – Use batch normalization after each convolutional layer (Default: True)

  • upsampling_mode (str) – Select between “bilinear” or “nearest” upsampling method. Bilinear is usually more accurate,but adds additional (small) randomness. For full reproducibility, consider using ‘nearest’ (this assumes that all other sources of randomness are fixed)

  • with_dilation (bool) – Use dilated convolutions instead of regular ones in the bottleneck layers (Default: False)

  • **layers (list) – List with a number of layers in each block. The first 4 elements in the list are used to determine the number of layers in each block of the encoder (incluidng bottleneck layers), and the number of layers in the decoder is chosen accordingly (to maintain symmetry between encoder and decoder)

class atomai.nets.dilnet(nb_classes=1, nb_filters=25, dropout=False, batch_norm=True, upsampling_mode='bilinear', **kwargs)[source]

Bases: torch.nn.Module

Builds a fully convolutional neural network model by utilizing a combination of regular and dilated convolutions

Parameters
  • nb_classes (int) – Number of classes in the ground truth

  • nb_filters (int) – Number of filters in first and last convolutional blocks (gets multiplied by 2 for the bottleneck layer)

  • dropout (bool) – Add dropouts to the bottleneck layers (Default: False)

  • batch_norm (bool) – Add batch normalization for each convolutional layer (Default: True)

  • upsampling_mode (str) – Select between “bilinear” or “nearest” upsampling method. Bilinear is usually more accurate,but adds additional (small) randomness. For full reproducibility, consider using ‘nearest’ (this assumes that all other sources of randomness are fixed)

  • **layers (list) – List with a number of layers for each block (Default: [3, 3, 3, 3])

class atomai.nets.ResHedNet(nb_classes=1, nb_filters=64, upsampling_mode='bilinear', **kwargs)[source]

Bases: torch.nn.Module

Holistically nested edge detector with residual connections in each block

Parameters
  • nb_classes (int) – Number of classes in the ground truth

  • nb_filters (int) – Number of filters in 1st residual block (gets multiplied by 2 in each next block)

  • upsampling_mode (str) – Select between “bilinear” or “nearest” upsampling method. Bilinear is usually more accurate,but adds additional (small) randomness. For full reproducibility, consider using ‘nearest’ (this assumes that all other sources of randomness are fixed)

  • **layers (list) – 3-element list with a number of residual blocks in each segment (Default: [3, 4, 5])

class atomai.nets.SegResNet(nb_classes=1, nb_filters=32, batch_norm=True, upsampling_mode='bilinear', **kwargs)[source]

Bases: torch.nn.Module

Builds a fully convolutional neural network based on SegNet architecture with residual blocks for semantic segmentation

Parameters
  • nb_classes (int) – Number of classes in the ground truth

  • nb_filters (int) – Number of filters in 1st residual block (gets multiplied by 2 in each next block)

  • batch_norm (bool) – Use batch normalization after each convolutional layer (Default: True)

  • upsampling_mode (str) – Select between “bilinear” or “nearest” upsampling method. Bilinear is usually more accurate,but adds additional (small) randomness. For full reproducibility, consider using ‘nearest’ (this assumes that all other sources of randomness are fixed)

  • **layers (list) – 3-element list with a number of residual blocks in each residual segment (Default: [2, 2])

Neural Networks for VAE

class atomai.nets.fcEncoderNet(in_dim, latent_dim=2, num_layers=2, hidden_dim=32, **kwargs)[source]

Bases: torch.nn.Module

Encoder/inference network (for variational autoencoder)

Parameters
  • in_dim (Tuple[int]) – Input dimensions. For images, it is (height, width) or (height, width, channels). For spectra, it is (length,)

  • latent_dim (int) – number of latent dimensions (the first 3 latent dimensions are angle & translations by default)

  • num_layers (int) – number of NN layers

  • hidden_dim (int) – number of neurons in each fully connnected layer (for mlp=True) or number of filters in each convolutional layer (for mlp=False)

  • **softplus_out (bool) – Optionally applies a softplus activation to the output associated with standard deviation of the encoded distribution

class atomai.nets.convEncoderNet(in_dim, latent_dim=2, num_layers=2, hidden_dim=32, **kwargs)[source]

Bases: torch.nn.Module

Convolutional encoder/inference network (for variational autoencoder)

Parameters
  • in_dim (Tuple[int]) – Input dimensions. For images, it is (height, width) or (height, width, channels). For spectra, it is (length,)

  • latent_dim (int) – number of latent dimensions (the first 3 latent dimensions are angle & translations by default)

  • num_layers (int) – number of NN layers

  • hidden_dim (int) – number of neurons in each fully connnected layer (for mlp=True) or number of filters in each convolutional layer (for mlp=False)

  • **softplus_out (bool) – Optionally applies a softplus activation to the output associated with standard deviation of the encoded distribution

class atomai.nets.fcDecoderNet(out_dim, latent_dim, num_layers=2, hidden_dim=32)[source]

Bases: torch.nn.Module

Decoder network (for variational autoencoder)

Parameters
  • out_dim (Tuple[int]) – Output dimensions. For images, it is (height, width) or (height, width, channels). For spectra, it is (length,)

  • latent_dim (int) – number of latent dimensions associated with images content

  • num_layers (int) – number of fully connected layers

  • hidden_dim (int) – number of neurons in each fully connected layer

class atomai.nets.convDecoderNet(out_dim, latent_dim, num_layers=2, hidden_dim=32, **kwargs)[source]

Bases: torch.nn.Module

Convolutional decoder network (for variational autoencoder)

Parameters
  • out_dim (Tuple[int]) – Output dimensions. For images, it is (height, width) or (height, width, channels). For spectra, it is (length,)

  • latent_dim (int) – number of latent dimensions associated with images content

  • num_layers (int) – number of fully connected layers

  • hidden_dim (int) – number of neurons in each fully connected layer

class atomai.nets.rDecoderNet(out_dim, latent_dim, num_layers, hidden_dim, skip=False)[source]

Bases: torch.nn.Module

Spatial decoder network with (optional) skip connections

Parameters
  • out_dim (Tuple[int]) – output dimensions: (height, width) or (height, width, channels)

  • latent_dim (int) – number of latent dimensions associated with images content

  • num_layers (int) – number of fully connected layers

  • hidden_dim (int) – number of neurons in each fully connected layer

  • skip (bool) – Use skip connections to propagate latent variables through decoder network (Default: False)

Neural Networks for ImSpec

class atomai.nets.SignalEncoder(signal_dim, z_dim, nb_layers, nb_filters, **kwargs)[source]

Bases: torch.nn.Module

Encodes 1D/2D signal into a latent vector

Parameters
  • signal_dim (Tuple[int]) – Size of input signal. For images, it is (height, width). For spectra, it is (length,)

  • z_dim (int) – Number of fully-connected neurons in a “bottleneck layer” (latent dimensions)

  • nb_layers (int) – Number of convolutional layers

  • nb_filters (int) – Number of convolutional filters (aka “kernels”) in each layer

  • **batch_norm (bool) – Apply batch normalization after each convolutional layer (Default: True)

  • **downsampling (int) – Downsamples input data by this factor before passing to convolutional layers (Default: no downsampling)

class atomai.nets.SignalDecoder(signal_dim, z_dim, nb_layers, nb_filters, **kwargs)[source]

Bases: torch.nn.Module

Decodes a latent vector into 1D/2D signal

Parameters
  • signal_dim (Tuple[int]) – Size of input signal. For images, it is (height, width). For spectra, it is (length,)

  • z_dim (int) – Number of fully-connected neurons in a “bottleneck layer” (latent dimensions)

  • nb_layers (int) – Number of convolutional layers

  • nb_filters (int) – Number of convolutional filters (aka “kernels”) in each layer

  • **batch_norm (bool) – Apply batch normalization after each convolutional layer (Default: True)

  • **upsampling (bool) – Performs upsampling+convolution operation twice on the reshaped latent vector (starting from image/spectra dims 4x smaller than the target dims) before passing to the decoder

class atomai.nets.SignalED(feature_dim, target_dim, latent_dim, nblayers_encoder=2, nblayers_decoder=2, nbfilters_encoder=64, nbfilters_decoder=2, batch_norm=True, encoder_downsampling=0, decoder_upsampling=False)[source]

Bases: torch.nn.Module

Transforms image into spectra (im2spec) and vice versa (spec2im)

Parameters
  • feature_dim (Tuple[int]) – Input data dimensions. (height, width) for images or (length,) for spectra

  • target_dim (Tuple[int]) – Output dimensions. (length,) for spectra or (height, width) for images

  • latent_dim (int) – Dimensionality of the latent space (number of neurons in a fully connected “bottleneck” layer)

  • nblayers_encoder (int) – Number of convolutional layers in the encoder

  • nblayers_decoder (int) – Number of convolutional layers in the decoder

  • nbfilters_encoder (int) – number of convolutional filters in each layer of the encoder

  • nbfilters_decoder (int) – Number of convolutional filters in each layer of the decoder

  • batch_norm (bool) – Apply batch normalization after each convolutional layer (Default: True)

  • encoder_downsampling (int) – Downsamples input data by this factor before passing to convolutional layers (Default: no downsampling)

  • decoder_upsampling (bool) – Performs upsampling+convolution operation twice on the reshaped latent vector (starting from image/spectra dims 4x smaller than the target dims) before passing to the decoder

encode(features)[source]

Embeddes the input image into a latent vector

Return type

Tensor

decode(latent)[source]

Generates signal from the embedded features

Return type

Tensor

Building blocks

class atomai.nets.ConvBlock(ndim, nb_layers, input_channels, output_channels, kernel_size=3, stride=1, padding=1, batch_norm=False, lrelu_a=0.01, dropout_=0)[source]

Bases: torch.nn.Module

Creates block of layers each consisting of convolution operation, leaky relu and (optionally) dropout and batch normalization

Parameters
  • ndim (int) – Data dimensionality (1D or 2D)

  • nb_layers (int) – Number of layers in the block

  • input_channels (int) – Number of input channels for the block

  • output_channels (int) – Number of the output channels for the block

  • kernel_size (Union[Tuple[int], int]) – Size of convolutional filter (in pixels)

  • stride (Union[Tuple[int], int]) – Stride of convolutional filter

  • padding (Union[Tuple[int], int]) – Value for edge padding

  • batch_norm (bool) – Add batch normalization to each layer in the block

  • lrelu_a (float) – Value of alpha parameter in leaky ReLU activation for each layer in the block

  • dropout – Dropout value for each layer in the block

class atomai.nets.DilatedBlock(ndim, input_channels, output_channels, dilation_values, padding_values, kernel_size=3, stride=1, lrelu_a=0.01, batch_norm=False, dropout_=0)[source]

Bases: torch.nn.Module

Creates a “cascade” with dilated convolutional layers (aka atrous convolutions)

Parameters
  • ndim (int) – Data dimensionality (1D or 2D)

  • input_channels (int) – Number of input channels for the block

  • output_channels (int) – Number of the output channels for the block

  • dilation_values (List[int]) – List of dilation rates for each convolution layer in the block (for example, dilation_values = [2, 4, 6] means that the dilated block will 3 layers with dilation values of 2, 4, and 6).

  • padding_values (List[int]) – Edge padding for each dilated layer. The number of elements in this list should be equal to that in the dilated_values list and typically they can have the same values.

  • kernel_size (Union[Tuple[int], int]) – Size of convolutional filter (in pixels)

  • stride (Union[Tuple[int], int]) – Stride of convolutional filter

  • batch_norm (bool) – Add batch normalization to each layer in the block

  • lrelu_a (float) – Value of alpha parameter in leaky ReLU activation for each layer in the block

  • dropout – Dropout value for each layer in the block

class atomai.nets.UpsampleBlock(ndim, input_channels, output_channels, scale_factor=2, mode='bilinear')[source]

Bases: torch.nn.Module

Defines upsampling block performed using bilinear or nearest-neigbor interpolation followed by 1-by-1 convolution (the latter can be used to reduce a number of feature channels)

Parameters
  • ndim (int) – Data dimensionality (1D or 2D)

  • input_channels (int) – Number of input channels for the block

  • output_channels (int) – Number of the output channels for the block

  • scale_factor (int) – Scale factor for upsampling

  • mode (str) – Upsampling mode. Select between “bilinear” and “nearest”

class atomai.nets.ResBlock(ndim, input_channels, output_channels, kernel_size=3, stride=1, padding=1, batch_norm=True, lrelu_a=0.01)[source]

Bases: torch.nn.Module

Builds a residual block

Parameters
  • ndim (int) – Data dimensionality (1D or 2D)

  • nb_layers – Number of layers in the block

  • input_channels (int) – Number of input channels for the block

  • output_channels (int) – Number of the output channels for the block

  • kernel_size (Union[Tuple[int], int]) – Size of convolutional filter (in pixels)

  • stride (Union[Tuple[int], int]) – Stride of convolutional filter

  • padding (Union[Tuple[int], int]) – Value for edge padding

  • batch_norm (bool) – Add batch normalization to each layer in the block

  • lrelu_a (float) – Value of alpha parameter in leaky ReLU activation for each layer in the block

class atomai.nets.ResModule(ndim, res_depth, input_channels, output_channels, batch_norm=True, lrelu_a=0.01)[source]

Bases: torch.nn.Module

Stitches multiple convolutional blocks with residual connections together

Parameters
  • ndim (int) – Data dimensionality (1D or 2D)

  • res_depth (int) – Number of residual blocks in a residual module

  • input_channels (int) – Number of filters in the input layer

  • output_channels (int) – Number of channels in the output layer

  • batch_norm (bool) – Batch normalization for non-unity layers in the block

  • lrelu_a (float) – value of negative slope for LeakyReLU activation