xnn.cnn.models.steerable.SteerableConv3d#

class xnn.cnn.models.steerable.SteerableConv3d(fields_in, fields_out, kernel_size=5, n_gates=0, stride=1, padding=None, dilation=1, bandlimit='compromise', shell_width=0.6)[source]#

Bases: Module

SE(3)-equivariant 3D convolution between stacks of fields.

For every pair of output order j and input order l the kernel block is a learned linear combination of the B_{jl} basis kernels of steerable_kernel_basis(), with one weight per (output field, input field, basis kernel); the blocks are stacked into one (K_out, K_in, s, s, s) filter bank for conv3d (Sec. 4.4.3). The basis of a block is scaled by sqrt((2j + 1) / (B_{jl} sum_l m_l)) so that weights drawn from N(0, 1) give outputs of unit variance (the initialization of the reference implementation). There is no bias: the biases of scalar fields live in the gated nonlinearity.

Parameters:

  • fields_in (sequence of int) – Multiplicities (m_0, m_1, ...) of the input and output stacks.

  • fields_out (sequence of int) – Multiplicities (m_0, m_1, ...) of the input and output stacks.

  • kernel_size (int, optional) – Side of the cubic kernel, by default 5.

  • n_gates (int, optional) – Extra scalar output fields appended after fields_out (the gates of GatedBlock), by default 0.

  • stride (int, optional) – As in conv3d; padding=None keeps the grid size.

  • padding (int, optional) – As in conv3d; padding=None keeps the grid size.

  • dilation (int, optional) – As in conv3d; padding=None keeps the grid size.

  • bandlimit (str or sequence of int, optional) – Shell bandlimits, see shell_bandlimits().

  • shell_width (float, optional) – Width of the radial Gaussian shells in voxels, by default 0.6.

Variables:

weight (torch.nn.ParameterDict) – "{i}_{j}" -> (m_out, m_in, B) weights of output block i and input block j.

kernel()[source]#

The filter bank spanned by the current weights, (K_out, K_in, s, s, s).

Return type:

Tensor

forward(x)[source]#

Convolve (B, in_channels, X, Y, Z) fields into (B, out_channels, X', Y', Z').

Parameters:

x (Tensor)

Return type:

Tensor

extra_repr()[source]#

Set the extra representation of the module.

To print customized extra information, you should re-implement this method in your own modules. Both single-line and multi-line strings are acceptable.

Return type:

str