Conv3d

Conv3d#

class braintrace.nn.Conv3d#

Three-dimensional convolution layer.

Applies a 3D convolution over an input signal composed of several input planes. The input should be a 5D array with the shape of [B, H, W, D, C] where B is batch size, H is height, W is width, D is depth, and C is the number of input channels (channels-last format).

This layer is commonly used for processing 3D data such as video sequences or volumetric medical imaging data.

Parameters:

  • in_size (tuple of int) – The input shape without the batch dimension. For Conv3d: (H, W, D, C) where H is height, W is width, D is depth, and C is the number of input channels. This argument is important as it is used to evaluate the output shape.

  • out_channels (int) – The number of output channels (also called filters or feature maps). These determine the depth of the output feature map.

  • kernel_size (int or tuple of int) – The shape of the convolutional kernel. Can be:

    • An integer (e.g., 3): creates a cubic kernel (3, 3, 3)

    • A tuple of three integers (e.g., (3, 5, 5)): creates a (height, width, depth) kernel

  • stride (int or tuple of int, optional) – The stride of the convolution. Controls how much the kernel moves at each step. Can be:

    • An integer: same stride for all dimensions

    • A tuple of three integers: (stride_h, stride_w, stride_d)

    Default: 1.

  • padding ({‘SAME’, ‘VALID’} or int or tuple of int or sequence of tuple, optional) – The padding strategy. Options:

    • ‘SAME’: output spatial size equals input size when stride=1

    • ‘VALID’: no padding, output size reduced by kernel size

    • int: same symmetric padding for all dimensions

    • (pad_h, pad_w[, pad_d]): symmetric padding per spatial dimension

    • [(pad_h_before, pad_h_after), (pad_w_before, pad_w_after), (pad_d_before, pad_d_after)]: explicit padding

    Default: ‘SAME’.

  • lhs_dilation (int or tuple of int, optional) – The dilation factor for the input (left-hand side). Controls spacing between input elements. A value > 1 inserts zeros between input elements, equivalent to transposed convolution. Default: 1.

  • rhs_dilation (int or tuple of int, optional) – The dilation factor for the kernel (right-hand side). Also known as atrous convolution or dilated convolution. Increases the receptive field without increasing parameters by inserting zeros between kernel elements. Particularly useful for 3D data to capture larger temporal/spatial context. Default: 1.

  • groups (int, optional) – Number of groups for grouped convolution. Must divide both in_channels and out_channels.

    • groups=1: standard convolution (all-to-all connections)

    • groups>1: grouped convolution (significantly reduces parameters and computation for 3D)

    • groups=in_channels: depthwise convolution (each input channel convolved separately)

    Default: 1.

  • w_init (Callable or ArrayLike, optional) – Weight initializer for the convolutional kernel. Can be:

    • An initializer instance (e.g., braintools.init.XavierNormal())

    • A callable that returns an array given a shape

    • A direct array matching the kernel shape

    Default: XavierNormal().

  • b_init (Callable or ArrayLike or None, optional) – Bias initializer. If None, no bias term is added to the output. Default: None.

  • w_mask (ArrayLike or Callable or None, optional) – Optional weight mask for structured sparsity or custom connectivity. The mask is element-wise multiplied with the kernel weights during the forward pass. Default: None.

  • name (str, optional) – Name identifier for this module instance. Default: None.

  • param_type (type, optional) – The parameter state class to use for managing learnable parameters. Default: ParamState.

Variables:

  • in_size (tuple of int) – The input shape (H, W, D, C) without batch dimension.

  • out_size (tuple of int) – The output shape (H_out, W_out, D_out, out_channels) without batch dimension.

  • in_channels (int) – Number of input channels.

  • out_channels (int) – Number of output channels.

  • kernel_size (tuple of int) – Size of the convolving kernel (height, width, depth).

  • weight (ParamState) – The learnable weights (and bias if specified) of the module.

Examples

>>> import braintrace as braintrace
>>> import jax.numpy as jnp
>>>
>>> # Create a 3D convolution layer for video data
>>> conv = braintrace.nn.Conv3d(in_size=(16, 64, 64, 3), out_channels=32, kernel_size=3)
>>>
>>> # Apply to input: batch_size=4, frames=16, height=64, width=64, channels=3
>>> x = jnp.ones((4, 16, 64, 64, 3))
>>> y = conv(x)
>>> print(y.shape)  # (4, 16, 64, 64, 32) with 'SAME' padding
>>>
>>> # Without batch dimension
>>> x_single = jnp.ones((16, 64, 64, 3))
>>> y_single = conv(x_single)
>>> print(y_single.shape)  # (16, 64, 64, 32)
>>>
>>> # For medical imaging with custom parameters
>>> conv = braintrace.nn.Conv3d(
...     in_size=(32, 32, 32, 1),
...     out_channels=64,
...     kernel_size=(3, 3, 3),
...     stride=2,
...     padding='VALID',
...     b_init=braintools.init.Constant(0.1)
... )

Notes

Output dimensions:

The output spatial dimensions depend on the padding mode:

  • ‘SAME’: output_size = ceil(input_size / stride)

  • ‘VALID’: output_size = ceil((input_size - kernel_size + 1) / stride)

Performance considerations:

3D convolutions are computationally expensive. Consider using:

  • Smaller kernel sizes

  • Grouped convolutions (groups > 1)

  • Separable convolutions for large-scale applications

Conv3d.__init__(in_size, out_channels, kernel_size, stride=1, padding='SAME', lhs_dilation=1, rhs_dilation=1, groups=1, w_init=XavierNormal(   scale=1.0, mode='fan_avg', in_axis=-2, out_axis=-1, distribution='truncated_normal', rng=RandomState(Array((), dtype=key<fry>) overlaying:   [2233333421 2029709265]), unit=Unit("1") ), b_init=None, w_mask=None, channel_first=False, name=None, param_type=<class 'brainstate.ParamState'>)#