xnn.cnn.models.steerable.GatedBlock#

class xnn.cnn.models.steerable.GatedBlock(fields_in, fields_out, kernel_size=5, stride=1, padding=None, activation='relu', gate_activation='sigmoid', normalization=None, smooth_stride=True, bias=True, bandlimit='compromise', shell_width=0.6)[source]#

Bases: Module

Steerable convolution with the gated nonlinearity (Sec. 4.3, Fig. 5).

The convolution produces the output stack plus one extra scalar field per non-scalar output field. Scalar outputs pass through activation (with a bias), the extra scalars through gate_activation (with a bias) and multiply, as gates, the non-scalar fields they belong to. With a stride the result is low-pass filtered and subsampled (Sec. 4.4.2; smooth_stride=False strides the convolution instead). normalization="batch" normalizes the input stack with SteerableBatchNorm, the “batch normalization merged with the convolution” of the supplement.

Parameters:

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

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

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

  • stride (int, optional) – Downsampling factor, by default 1.

  • padding (int or None, optional) – Zero padding; None keeps the grid size.

  • activation (str or None, optional) – Activation of the scalar fields, by default "relu" (the paper); None leaves them linear.

  • gate_activation (str or None, optional) – Activation of the gates, by default "sigmoid"; None leaves the non-scalar fields ungated (a linear block).

  • normalization (str or None, optional) – "batch" or None.

  • smooth_stride (bool, optional) – Low-pass filter before subsampling, by default True.

  • bias (bool, optional) – Biases on the scalar fields and gates, by default True.

  • bandlimit – See SteerableConv3d.

  • shell_width (float) – See SteerableConv3d.

Variables:

conv (SteerableConv3d) – The convolution (its n_gates extra scalar outputs are the gates).

property out_channels: int#

Channels of the output stack (gates excluded).

Type:

int

forward(x)[source]#

Map (B, field_dim(fields_in), X, Y, Z) to (B, field_dim(fields_out), X', Y', Z').

Parameters:

x (Tensor)

Return type:

Tensor