dacapo.experiments.architectures.cnnectome_unet

Classes

`CNNectomeUNet`	A U-Net architecture for 3D or 4D data. The U-Net expects 3D or 4D tensors
`CNNectomeUNetModule`	A U-Net module for 3D or 4D data. The U-Net expects 3D or 4D tensors shaped
`ConvPass`	Convolutional pass module. This module performs a series of convolutional
`Downsample`	Downsample module. This module performs downsampling of the input tensor
`Upsample`	Upsample module. This module performs upsampling of the input tensor using
`AttentionBlockModule`	Attention Block Module:

Module Contents

class dacapo.experiments.architectures.cnnectome_unet.CNNectomeUNet(architecture_config)

A U-Net architecture for 3D or 4D data. The U-Net expects 3D or 4D tensors shaped like:

``(batch=1, channels, [length,] depth, height, width)``.

This U-Net performs only “valid” convolutions, i.e., sizes of the feature maps decrease after each convolution. It will perfrom 4D convolutions as long as length is greater than 1. As soon as length is 1 due to a valid convolution, the time dimension will be dropped and tensors with (b, c, z, y, x) will be use (and returned) from there on.

fmaps_in: The number of input channels.

fmaps_out: The number of feature maps in the output layer. This is also the number of output feature maps. Stored in the channels dimension.

num_fmaps: The number of feature maps in the first layer. This is also the number of output feature maps. Stored in the channels dimension.

fmap_inc_factor: By how much to multiply the number of feature maps between layers. If layer 0 has k feature maps, layer l will have k*fmap_inc_factor**l.

downsample_factors: List of tuples (z, y, x) to use to down- and up-sample the feature maps between layers.

kernel_size_down

List of lists of kernel sizes. The number of sizes in a list determines the number of convolutional layers in the corresponding level of the build on the left side. Kernel sizes can be given as tuples or integer. If not given, each convolutional pass will consist of two 3x3x3 convolutions.

Type:: optional

kernel_size_up

List of lists of kernel sizes. The number of sizes in a list determines the number of convolutional layers in the corresponding level of the build on the right side. Within one of the lists going from left to right. Kernel sizes can be given as tuples or integer. If not given, each convolutional pass will consist of two 3x3x3 convolutions.

Type:: optional

activation: Which activation to use after a convolution. Accepts the name of any tensorflow activation function (e.g., ReLU for torch.nn.ReLU).

fov

Initial field of view in physical units

Type:: optional

voxel_size

Size of a voxel in the input data, in physical units

Type:: optional

num_heads

Number of decoders. The resulting U-Net has one single encoder path and num_heads decoder paths. This is useful in a multi-task learning context.

Type:: optional

constant_upsample

If set to true, perform a constant upsampling instead of a transposed convolution in the upsampling layers.

Type:: optional

padding

How to pad convolutions. Either ‘same’ or ‘valid’ (default).

Type:: optional

upsample_channel_contraction: When performing the ConvTranspose, whether to reduce the number of channels by the fmap_increment_factor. can be either bool or list of bools to apply independently per layer.

activation_on_upsample: Whether or not to add an activation after the upsample operation.

use_attention

Whether or not to use an attention block in the U-Net.

Methods:

forward(x):
Forward pass of the U-Net.

scale(voxel_size):
Scale the voxel size according to the upsampling factors.

input_shape:
Return the input shape of the U-Net.

num_in_channels:
Return the number of input channels.

num_out_channels:
Return the number of output channels.

eval_shape_increase:
Return the increase in shape due to the U-Net.

Note:
This class is a wrapper around the CNNectomeUNetModule class. The CNNectomeUNetModule class is the actual implementation of the U-Net architecture.

fmaps_out

fmaps_in

num_fmaps

fmap_inc_factor

downsample_factors

kernel_size_down

kernel_size_up

constant_upsample

padding

upsample_factors

use_attention

batch_norm

unet

property eval_shape_increase

The increase in shape due to the U-Net.

Returns:: The increase in shape due to the U-Net.
Raises:: AttributeError – If the increase in shape is not given.

Examples

>>> unet.eval_shape_increase
(1, 1, 128, 128, 128)

Note

The increase in shape should be given as a tuple (batch, channels, [length,] depth, height, width).

module()

Create the U-Net module.

Returns:

The U-Net module.

Raises:

AttributeError – If the number of input channels is not given.
AttributeError – If the number of output channels is not given.
AttributeError – If the number of feature maps in the first layer is not given.
AttributeError – If the factor by which the number of feature maps increases between layers is not given.
AttributeError – If the downsample factors are not given.
AttributeError – If the kernel sizes for the down pass are not given.
AttributeError – If the kernel sizes for the up pass are not given.
AttributeError – If the constant upsample flag is not given.
AttributeError – If the padding is not given.
AttributeError – If the upsample factors are not given.
AttributeError – If the activation on upsample flag is not given.
AttributeError – If the use attention flag is not given.

Examples

>>> unet.module()
CNNectomeUNetModule(
    in_channels=1,
    num_fmaps=24,
    num_fmaps_out=1,
    fmap_inc_factor=2,
    kernel_size_down=[[(3, 3, 3), (3, 3, 3)], [(3, 3, 3), (3, 3, 3)], [(3, 3, 3), (3, 3, 3)]],
    kernel_size_up=[[(3, 3, 3), (3, 3, 3)], [(3, 3, 3), (3, 3, 3)], [(3, 3, 3), (3, 3, 3)]],
    downsample_factors=[(2, 2, 2), (2, 2, 2), (2, 2, 2)],
    constant_upsample=False,
    padding='valid',
    activation_on_upsample=True,
    upsample_channel_contraction=[False, True, True],
    use_attention=False
)

Note

The U-Net module is an instance of the CNNectomeUNetModule class.

scale(voxel_size)

Scale the voxel size according to the upsampling factors.

Parameters:: voxel_size (tuple) – The size of a voxel in the input data.
Returns:: The scaled voxel size.
Raises:: ValueError – If the voxel size is not given.

Examples

>>> unet.scale((1, 1, 1))
(1, 1, 1)

Note

The voxel size should be given as a tuple (z, y, x).

property input_shape

Return the input shape of the U-Net.

Returns:: The input shape of the U-Net.
Raises:: AttributeError – If the input shape is not given.

Examples

>>> unet.input_shape
(1, 1, 128, 128, 128)

Note

The input shape should be given as a tuple (batch, channels, [length,] depth, height, width).

property num_in_channels: int

Return the number of input channels.

Returns:: The number of input channels.
Raises:: AttributeError – If the number of input channels is not given.

Examples

>>> unet.num_in_channels
1

Note

The number of input channels should be given as an integer.

property num_out_channels: int

Return the number of output channels.

Returns:: The number of output channels.
Raises:: AttributeError – If the number of output channels is not given.

Examples

>>> unet.num_out_channels
1

Note

The number of output channels should be given as an integer.

forward(x)

Forward pass of the U-Net.

Parameters:: x (Tensor) – The input tensor.
Returns:: The output tensor.
Raises:: RuntimeError – If the tensors have different dimensions.

Examples

>>> unet = CNNectomeUNet(architecture_config)
>>> x = torch.randn(1, 1, 64, 64, 64)
>>> unet(x)

Note

The input tensor should be given as a 5D tensor.

class dacapo.experiments.architectures.cnnectome_unet.CNNectomeUNetModule(in_channels, num_fmaps, fmap_inc_factor, downsample_factors, kernel_size_down=None, kernel_size_up=None, activation='ReLU', num_fmaps_out=None, num_heads=1, constant_upsample=False, padding='valid', upsample_channel_contraction=False, activation_on_upsample=False, use_attention=False, batch_norm=True)

A U-Net module for 3D or 4D data. The U-Net expects 3D or 4D tensors shaped like:

``(batch=1, channels, [length,] depth, height, width)``.

This U-Net performs only “valid” convolutions, i.e., sizes of the feature maps decrease after each convolution. It will perfrom 4D convolutions as long as length is greater than 1. As soon as length is 1 due to a valid convolution, the time dimension will be dropped and tensors with (b, c, z, y, x) will be use (and returned) from there on.

num_levels: The number of levels in the U-Net.

num_heads: The number of decoders.

in_channels: The number of input channels.

out_channels: The number of output channels.

dims: The number of dimensions.

use_attention: Whether or not to use an attention block in the U-Net.

l_conv: The left convolutional passes.

l_down: The left downsample layers.

r_up: The right up/crop/concatenate layers.

r_conv: The right convolutional passes.

kernel_size_down: The kernel sizes for the down pass.

kernel_size_up: The kernel sizes for the up pass.

fmap_inc_factor: The factor by which the number of feature maps increases between layers.

downsample_factors: The downsample factors.

constant_upsample: Whether to perform a constant upsampling instead of a transposed convolution.

padding: How to pad convolutions.

upsample_channel_contraction: Whether to reduce the number of channels by the fmap_increment_factor.

activation_on_upsample: Whether or not to add an activation after the upsample operation.

use_attention: Whether or not to use an attention block in the U-Net.

attention: The attention blocks.

rec_forward(level, f_in): Recursive forward pass of the U-Net.

forward(x): Forward pass of the U-Net.

Note

The input tensor should be given as a 5D tensor.

num_levels

num_heads

in_channels

out_channels

upsample_channel_contraction

dims

use_attention

batch_norm

kernel_size_down

kernel_size_up

crop_factors = []

factor_product = None

l_conv

l_down

r_up

r_conv

rec_forward(level, f_in)

Recursive forward pass of the U-Net.

Parameters:

level (int) – The level of the U-Net.
f_in (Tensor) – The input tensor.

Returns:

The output tensor.

Raises:

RuntimeError – If the tensors have different dimensions.

Examples

>>> unet = CNNectomeUNetModule(
...     in_channels=1,
...     num_fmaps=24,
...     num_fmaps_out=1,
...     fmap_inc_factor=2,
...     kernel_size_down=[[(3, 3, 3), (3, 3, 3)], [(3, 3, 3), (3, 3, 3)], [(3, 3, 3), (3, 3, 3)]],
...     kernel_size_up=[[(3, 3, 3), (3, 3, 3)], [(3, 3, 3), (3, 3, 3)], [(3, 3, 3), (3, 3, 3)]],
...     downsample_factors=[(2, 2, 2), (2, 2, 2), (2, 2, 2)],
...     constant_upsample=False,
...     padding='valid',
...     activation_on_upsample=True,
...     upsample_channel_contraction=[False, True, True],
...     use_attention=False
... )
>>> x = torch.randn(1, 1, 64, 64, 64)
>>> unet.rec_forward(2, x)

Note

The input tensor should be given as a 5D tensor.

forward(x)

Forward pass of the U-Net.

Parameters:: x (Tensor) – The input tensor.
Returns:: The output tensor.
Raises:: RuntimeError – If the tensors have different dimensions.

Examples

>>> unet = CNNectomeUNetModule(
...     in_channels=1,
...     num_fmaps=24,
...     num_fmaps_out=1,
...     fmap_inc_factor=2,
...     kernel_size_down=[[(3, 3, 3), (3, 3, 3)], [(3, 3, 3), (3, 3, 3)], [(3, 3, 3), (3, 3, 3)]],
...     kernel_size_up=[[(3, 3, 3), (3, 3, 3)], [(3, 3, 3), (3, 3, 3)], [(3, 3, 3), (3, 3, 3)]],
...     downsample_factors=[(2, 2, 2), (2, 2, 2), (2, 2, 2)],
...     constant_upsample=False,
...     padding='valid',
...     activation_on_upsample=True,
...     upsample_channel_contraction=[False, True, True],
...     use_attention=False
... )
>>> x = torch.randn(1, 1, 64, 64, 64)
>>> unet(x)

Note

The input tensor should be given as a 5D tensor.

class dacapo.experiments.architectures.cnnectome_unet.ConvPass(in_channels, out_channels, kernel_sizes, activation, padding='valid', batch_norm=True)

Convolutional pass module. This module performs a series of convolutional layers followed by an activation function. The module can also pad the feature maps to ensure translation equivariance. The module can perform 2D or 3D convolutions.

dims: The number of dimensions.

conv_pass: The convolutional pass module.

forward(x): Forward pass of the Conv

Note

The input tensor should be given as a 5D tensor.

layers = []

conv_pass

forward(x)

Forward pass of the ConvPass module.

Parameters:: x (Tensor) – The input tensor.
Returns:: The output tensor.
Raises:: RuntimeError – If the tensors have different dimensions.

Examples

>>> conv_pass = ConvPass(1, 1, [(3, 3, 3), (3, 3, 3)], "ReLU")
>>> x = torch.randn(1, 1, 64, 64, 64)
>>> conv_pass(x)

Note

The input tensor should be given as a 5D tensor.

class dacapo.experiments.architectures.cnnectome_unet.Downsample(downsample_factor)

Downsample module. This module performs downsampling of the input tensor using either max-pooling or average pooling. The module can also crop the feature maps to ensure translation equivariance with a stride of the downsampling factor.

dims: The number of dimensions.

downsample_factor: The downsampling factor.

down: The downsampling layer.

forward(x): Downsample the input tensor.

Note

The input tensor should be given as a 5D tensor.

dims

downsample_factor

pool

down

forward(x)

Downsample the input tensor.

Parameters:: x (Tensor) – The input tensor.
Returns:: The downsampled tensor.
Raises:: RuntimeError – If the tensors have different dimensions.

Examples

>>> downsample = Downsample((2, 2, 2))
>>> x = torch.randn(1, 1, 64, 64, 64)
>>> downsample(x)

Note

The input tensor should be given as a 5D tensor.

class dacapo.experiments.architectures.cnnectome_unet.Upsample(scale_factor, mode='transposed_conv', in_channels=None, out_channels=None, crop_factor=None, next_conv_kernel_sizes=None, activation=None)

Upsample module. This module performs upsampling of the input tensor using either transposed convolutions or nearest neighbor interpolation. The module can also crop the feature maps to ensure translation equivariance with a stride of the upsampling factor.

crop_factor: The crop factor.

next_conv_kernel_sizes: The kernel sizes for the convolutional layers.

dims: The number of dimensions.

up: The upsampling layer.

crop_to_factor(x, factor, kernel_sizes): Crop feature maps to ensure translation equivariance with stride of upsampling factor.

crop(x, shape): Center-crop x to match spatial dimensions given by shape.

forward(g_out, f_left=None): Forward pass of the Upsample module.

Note

The input tensor should be given as a 5D tensor.

crop_factor

next_conv_kernel_sizes

dims

layers = []

crop_to_factor(x, factor, kernel_sizes)

Crop feature maps to ensure translation equivariance with stride of upsampling factor. This should be done right after upsampling, before application of the convolutions with the given kernel sizes.

The crop could be done after the convolutions, but it is more efficient to do that before (feature maps will be smaller).

We need to ensure that the feature map is large enough to ensure that the translation equivariance is maintained. This is done by cropping the feature map to the largest size that is a multiple of the factor and that is large enough to ensure that the translation equivariance is maintained.

We need (spatial_shape - convolution_crop) to be a multiple of factor, i.e.: (s - c) = n*k

where s is the spatial size of the feature map, c is the crop due to the convolutions, n is the number of strides of the upsampling factor, and k is the upsampling factor.

We want to find the largest n for which s’ = n*k + c <= s

n = floor((s - c)/k)

This gives us the target shape s’

s’ = n*k + c

Parameters:

x (Tensor) – The input tensor.
factor (tuple) – The upsampling factor.
kernel_sizes (list) – The kernel sizes for the convolutional layers.

Returns:

The cropped tensor.

Raises:

RuntimeError – If the feature map is too small to ensure translation equivariance.

Examples

>>> upsample = Upsample(scale_factor=(2, 2, 2), in_channels=1, out_channels=1)
>>> x = torch.randn(1, 1, 64, 64, 64)
>>> upsample.crop_to_factor(x, (2, 2, 2), [(3, 3, 3), (3, 3, 3)])

Note

The input tensor should be given as a 5D tensor.

crop(x, shape)

Center-crop x to match spatial dimensions given by shape.

Parameters:

x (Tensor) – The input tensor.
shape (tuple) – The target shape.

Returns:

The center-cropped tensor.

Raises:

RuntimeError – If the tensors have different dimensions.

Examples

>>> upsample = Upsample(scale_factor=(2, 2, 2), in_channels=1, out_channels=1)
>>> x = torch.randn(1, 1, 64, 64, 64)
>>> upsample.crop(x, (32, 32, 32))

Note

The input tensor should be given as a 5D tensor.

forward(g_out, f_left=None)

Forward pass of the Upsample module.

Parameters:

g_out (Tensor) – The gating signal tensor.
f_left (Tensor) – The input feature tensor.

Returns:

The output feature tensor.

Raises:

RuntimeError – If the tensors have different dimensions.

Examples

>>> upsample = Upsample(scale_factor=(2, 2, 2), in_channels=1, out_channels=1)
>>> g_out = torch.randn(1, 1, 64, 64, 64)
>>> f_left = torch.randn(1, 1, 32, 32, 32)
>>> upsample(g_out, f_left)

Note

The gating signal and input feature tensors should be given as 5D tensors.

class dacapo.experiments.architectures.cnnectome_unet.AttentionBlockModule(F_g, F_l, F_int, dims, upsample_factor=None, batch_norm=True)

Attention Block Module:

The AttentionBlock uses two separate pathways to process ‘g’ and ‘x’, combines them, and applies a sigmoid activation to generate an attention map. This map is then used to scale the input features ‘x’, resulting in an output that focuses on important features as dictated by the gating signal ‘g’.

The attention block takes two inputs: ‘g’ (gating signal) and ‘x’ (input features).

[g] –> W_g – /–> psi –> * –> [output]
/

[x] –> W_x –> [+] –> relu –

Where:

W_g and W_x are 1x1 Convolution followed by Batch Normalization

[+] indicates element-wise addition

relu is the Rectified Linear Unit activation function

psi is a sequence of 1x1 Convolution, Batch Normalization, and Sigmoid activation

indicates element-wise multiplication between the output of psi and input feature ‘x’

[output] has the same dimensions as input ‘x’, selectively emphasized by attention weights

Attributes:

dims:
The number of dimensions of the input tensors.

kernel_sizes:
The kernel sizes for the convolutional layers.

upsample_factor:
The factor by which to upsample the attention map.

W_g:
The 1x1 Convolutional layer for the gating signal.

W_x:
The 1x1 Convolutional layer for the input features.

psi:
The 1x1 Convolutional layer followed by Sigmoid activation.

up:
The upsampling layer to match the dimensions of the input features.

relu:
The Rectified Linear Unit activation function.

Methods:

calculate_and_apply_padding(smaller_tensor, larger_tensor):
Calculate and apply symmetric padding to the smaller tensor to match the dimensions of the larger tensor.

forward(g, x):
Forward pass of the Attention Block.

Note:
The AttentionBlockModule is an instance of the torch.nn.Module class.

dims

kernel_sizes

batch_norm

W_g

W_x

psi

up_mode

up

relu

calculate_and_apply_padding(smaller_tensor, larger_tensor)

Calculate and apply symmetric padding to the smaller tensor to match the dimensions of the larger tensor.

Parameters:

smaller_tensor (Tensor) – The tensor to be padded.
larger_tensor (Tensor) – The tensor whose dimensions the smaller tensor needs to match.

Returns:

The padded smaller tensor with the same dimensions as the larger tensor.

Return type:

Tensor

Raises:

RuntimeError – If the tensors have different dimensions.

Examples

>>> larger_tensor = torch.randn(1, 1, 128, 128, 128)
>>> smaller_tensor = torch.randn(1, 1, 64, 64, 64)
>>> attention_block = AttentionBlockModule(F_g=1, F_l=1, F_int=1, dims=3)
>>> padded_tensor = attention_block.calculate_and_apply_padding(smaller_tensor, larger_tensor)

Note

The tensors should have the same dimensions.

forward(g, x)

Forward pass of the Attention Block.

Parameters:

g (Tensor) – The gating signal tensor.
x (Tensor) – The input feature tensor.

Returns:

The output tensor with the same dimensions as the input feature tensor.

Return type:

Tensor

Raises:

RuntimeError – If the gating signal and input feature tensors have different dimensions.

Examples

>>> g = torch.randn(1, 1, 128, 128, 128)
>>> x = torch.randn(1, 1, 128, 128, 128)
>>> attention_block = AttentionBlockModule(F_g=1, F_l=1, F_int=1, dims=3)
>>> output = attention_block(g, x)

Note

The gating signal and input feature tensors should have the same dimensions.