3.5. Convolutional Layers

RustyML ships 5 convolutional layers: Conv1D, Conv2D, Conv3D, DepthwiseConv2D, and SeparableConv2D. All 5 live in rustyml::neural_network::layers, and the layers glob re-exports them. They share 1 design. A layer struct holds the weights, the bias, the activation, and the forward and backward caches. The layer delegates the actual numerics for the plain convolutions to a single dimension-generic engine.

If you know Keras, you already know the layouts and weight shapes here: tensors are channels-last, and kernels carry their taps first. 1 difference catches new users early. You pass the full input shape into the constructor. This lets the layer size its weights up front, instead of inferring them lazily on the first batch.

This page covers the layouts, the constructor and padding rules, and the im2col plus GEMM engine behind the standard convolutions. It also covers the depthwise and separable factorization with its cost math, and the error types the layers return.

3.5.1. Tensor Layouts and the Layer Family

Every convolution here is channels-last. The spatial axes follow the batch axis, and the channel axis comes last. A Conv2D reads [batch, height, width, channels] and writes [batch, out_height, out_width, filters]. This is the Keras and TensorFlow NHWC convention.

A tensor built for Keras needs no permutation. If you build inputs by hand with ndarray, put the channels last. The weight tensor uses Keras’ kernel shape too: the kernel taps come first, then the input channels, then the filters.

That is not only an interface choice. With the channel axis innermost, a kernel tap at a given output position is Cin contiguous floats. This lets the engine build its im2col matrix from run copies, instead of a scalar gather. It also lets the flat weight matrix line up with the im2col matrix without any permutation.

Every shape in the table matches Keras, so a kernel exported from Keras drops in without a permutation. DepthwiseConv2D and SeparableConv2D need more explanation. DepthwiseConv2D has no filters argument at all. It applies depth_multiplier kernels to each input channel and emits C * dm channels. Input channel c’s multiplier m lands at output channel c * dm + m.

SeparableConv2D holds 2 weight tensors because it fuses 2 convolutions into 1 layer. The depthwise stage emits its channels in that same c * dm + m order. This means the pointwise weight’s rows already match the depthwise output, so the layer needs no repacking between the stages. Section 3.5.4 covers this in more detail.

The forward math is cross-correlation, the same convention Keras and PyTorch use. The engine does not flip the kernel. The bias is added last, once per filter, after the multiply-accumulate step. Weights start from Xavier/Glorot uniform bounds. Biases start at zero.

The following example sets the weights by hand and runs 1 forward pass. It shows the layout and the Valid-padding output size.

use ndarray::{Array, Array1, Array4};
use rustyml::neural_network::layers::*;
use rustyml::neural_network::traits::Layer;

fn main() {
    // Channels-last: [batch, height, width, channels].
    let mut layer = Conv2D::new(1, (2, 2), vec![1, 4, 4, 1], (1, 1), Activation::Linear).unwrap();

    // Weights are [kernel_h, kernel_w, channels, filters]. Bias is [filters].
    let weights = Array4::from_elem((2, 2, 1, 1), 1.0f32);
    let bias = Array1::zeros(1);
    layer.set_weights(weights, bias).unwrap();

    let pixels: Vec<f32> = (1..=16).map(|v| v as f32).collect();
    let x = Array::from_shape_vec((1, 4, 4, 1), pixels).unwrap().into_dyn();

    let out = layer.forward(&x).unwrap();
    // Valid padding: (4 - 2)/1 + 1 = 3 along each spatial axis -> [1, 3, 3, 1].
    assert_eq!(out.shape(), &[1, 3, 3, 1]);
    // An all-ones 2x2 kernel sums each window: the first is 1 + 2 + 5 + 6 = 14.
    assert_eq!(out[[0, 0, 0, 0]], 14.0);
    println!("output shape: {:?}", out.shape());
}

3.5.2. Constructors, Padding, and Output Shapes

The constructors take the hyperparameters as positional arguments. The kernel and stride argument shape tracks the rank. Conv1D uses a scalar kernel_size and stride. The 2D and 3D layers take tuples instead.

SeparableConv2D inserts a depth_multiplier argument before the activation. DepthwiseConv2D has no filters argument. Like Keras, it derives its output width from the input, and it takes its depth multiplier through a builder method.

Conv1D::new(filters, kernel_size: usize, input_shape: Vec<usize>, stride: usize, activation) -> Result<Conv1D, Error>
Conv2D::new(filters, kernel_size: (usize, usize), input_shape: Vec<usize>, strides: (usize, usize), activation) -> Result<Conv2D, Error>
Conv3D::new(filters, kernel_size: (usize, usize, usize), input_shape: Vec<usize>, strides: (usize, usize, usize), activation) -> Result<Conv3D, Error>
DepthwiseConv2D::new(kernel_size: (usize, usize), input_shape: Vec<usize>, strides: (usize, usize), activation) -> Result<DepthwiseConv2D, Error>
DepthwiseConv2D::with_depth_multiplier(self, depth_multiplier: usize) -> Result<DepthwiseConv2D, Error>   // builder, defaults to 1
SeparableConv2D::new(filters, kernel_size: (usize, usize), input_shape: Vec<usize>, strides: (usize, usize), depth_multiplier: usize, activation) -> Result<SeparableConv2D, Error>

activation takes impl Into<Activation>. You can pass an Activation variant (Activation::ReLU, Activation::Sigmoid, Activation::Tanh, Activation::Softmax, Activation::Linear), or a standalone activation layer such as ReLU::new(). The input_shape you supply is the full expected input, including the batch dimension. Only input_shape[1..] (the channels and spatial extents) sizes the weights. The batch value you write is informational only. See 3.2. Dense Layers and Activations for more about the activation set.

3 builder methods refine a constructed layer. Each one consumes and returns self, so you can chain them. with_padding(PaddingType) switches the padding mode. with_random_state(u64) re-runs the Xavier initialization deterministically from a seed. Call it before you assign custom weights or start training (see 7.1. Reproducibility and Random Seeds).

set_weights(...) installs weights and a bias that you control. set_weights checks every array against the layer’s expected shape. Note that DepthwiseConv2D::set_weights takes an Array1 bias. SeparableConv2D::set_weights takes 3 arrays instead: depthwise weights, pointwise weights, then bias.

Padding uses the PaddingType enum, which has 2 variants: Valid (the default) and Same. It sets the output size as follows:

• Valid applies no padding. It computes output values only where the kernel fully overlaps the input. The formula is out = (in - k) / stride + 1 (integer floor division), on each spatial axis. Every input axis must be at least the kernel size.
• Same zero-pads the borders so out = ceil(in / stride). At stride == 1 this keeps the spatial size exactly. At a larger stride, the output is the input length divided by the stride, rounded up. The layer splits the total padding evenly and puts the extra cell on the trailing edge (pad_before = pad_total / 2). This matches TensorFlow’s SAME padding.

The following numbers show the rule in practice. A [N, 8, 8, 1] input through a 3x3 kernel gives 3 results. Valid padding at stride 1 gives (8-3)/1+1 = 6, so the output is [N, 6, 6, F]. Same padding at stride 1 gives 8, so the output is [N, 8, 8, F]. Same padding at stride 2 gives ceil(8/2) = 4. A Conv1D over length 6, with kernel 3 and stride 2, under Valid padding, gives (6-3)/2+1 = 2.

Conv3D applies the same rule on depth, height, and width, each on its own. model.summary() prints these output shapes and the per-layer parameter counts. Use it to check a stack before you train it. See 3.1. Sequential Model for the full description of summary.

3.5.3. The im2col + GEMM Engine

Conv1D, Conv2D, and Conv3D do not each carry their own loop nest. A plain convolution is the same operation at every rank. Only the number of spatial axes changes. All 3 layers delegate their forward and backward numerics to 1 implementation in convolution_engine.rs.

This implementation is generic over the spatial rank R = ndim - 2. The layer wrapper keeps the public API, the weight storage, the activation, and the caches. The engine does the arithmetic.

The engine uses im2col plus GEMM, the same strategy the major frameworks use. For the forward pass, it gathers each output window into a row. This forms an [out_plane, k_plane*Cin] matrix, whose columns align with the flat weight matrix [k_plane*Cin, F]. A single matrix multiply then produces [out_plane, F], and the engine adds the bias per filter.

Under the channels-last layout, this gather is a run of Cin-wide copy_from_slice calls, not a scalar-at-a-time walk. The product lands directly on a contiguous slab of the output, with no scatter. Trading a 6-deep loop nest for a matrix multiply lets the layer use the crate’s tuned in-house GEMM, instead of a naive triple loop. See 6.2. Matrix Multiplication for more on that GEMM.

The backward pass runs 2 GEMMs per batch item. One computes the weight gradient. The other computes the input-gradient columns, which the engine then scatters back (col2im) into the input-gradient tensor.

The engine gates parallelism on estimated FLOPs, not on element counts. A 7x7x512 convolution and a 3x3x3 convolution can share the same output-element count, but their costs differ by a wide margin. The forward gate compares 2 * batch * F * out_plane * Cin*k against CONV_PARALLEL_MIN_FLOPS (default 4,000,000, tunable at runtime through rustyml::tuning).

Below this threshold, the forward pass runs serial. Above it, the forward pass parallelizes over (batch item, output-position block) tasks. This lets a single large image fill every core, even at batch == 1. Each task builds its own im2col block and runs its own GEMM into a disjoint output region.

The backward pass parallelizes over batch items. It reduces the weight and bias partials in batch order, which keeps results bit-reproducible across runs on the same machine. It also routes each item’s GEMMs through a switch. The switch keeps them parallel while the batch is too short to fill the thread pool. It flips them to serial once the batch alone saturates the pool. This way, a batch task never forks rayon again inside its own GEMM.

The parallel path and the serial path return the same numbers, so you rarely need to touch the gate. If you profile a workload that sits just under the threshold, see 7.3. Performance Tuning and Parallelism for how to move it.

3.5.4. Depthwise and Separable Convolutions

A standard convolution mixes across channels and across space in 1 step. Every output channel is a weighted sum over all input channels and all kernel taps. That coupling is where the parameters live. The weight tensor holds F * Cin * kh * kw values.

Depthwise separable convolution factors this operation into 2 cheaper stages. The first stage is a depthwise convolution that filters each input channel on its own (spatial mixing only, no cross-channel mixing). The second stage is a pointwise 1x1 convolution that recombines the channels (cross-channel mixing only, no spatial extent). This is the idea behind MobileNet and Xception. RustyML exposes both stages.

DepthwiseConv2D is the first stage alone. It carries depth_multiplier kernels of size kh x kw for each input channel, and emits C * depth_multiplier output channels. This is why it has no filters argument, exactly as in Keras. It does no channel recombination, so it cannot mix channels on its own. In practice, you almost always pair it with a 1x1 convolution downstream.

depth_multiplier defaults to 1. Set it with with_depth_multiplier, which returns a Result because 0 is rejected. SeparableConv2D takes its own depth_multiplier as a positional argument. It expands the intermediate channel count to Cin * depth_multiplier, before the pointwise stage collapses it back to filters.

The parameter counts are the point of this factorization. Consider Cin input channels, F output filters, and a kh x kw kernel:

• Standard Conv2D: F * Cin * kh * kw + F.
• DepthwiseConv2D: C * dm * kh * kw + C * dm (at the default dm = 1, this is C * kh * kw + C).
• SeparableConv2D: dm * Cin * kh * kw (depthwise) + F * Cin * dm (pointwise) + F (bias).

Ignoring the bias term, the separable-to-standard ratio is 1/F + 1/(kh*kw). The savings grow with both the filter count and the kernel area. Each parameter costs 1 multiply-accumulate per output position, so this same ratio is also the compute (FLOP) ratio.

As an example, take a 64-filter 3x3 convolution over 3 channels. The standard layer has 64*3*3*3 + 64 = 1792 parameters. The separable equivalent has 27 + 192 + 64 = 283 parameters, about 6.3 times fewer. A depthwise-only layer over those 3 channels has just 30 parameters. The following program builds all 3 layers and checks the counts.

use rustyml::neural_network::layers::*;
use rustyml::neural_network::traits::Layer;

fn count(p: TrainingParameters) -> usize {
    match p {
        TrainingParameters::Trainable(n) | TrainingParameters::NonTrainable(n) => n,
        TrainingParameters::NoTrainable => 0,
    }
}

fn main() {
    let input_shape = vec![1, 32, 32, 3]; // [batch, H, W, channels]

    // Standard 64-filter 3x3 convolution over 3 input channels.
    let standard =
        Conv2D::new(64, (3, 3), input_shape.clone(), (1, 1), Activation::ReLU).unwrap();

    // Separable equivalent: depthwise (depth_multiplier = 1) then a pointwise 1x1 to 64 filters.
    let separable =
        SeparableConv2D::new(64, (3, 3), input_shape.clone(), (1, 1), 1, Activation::ReLU).unwrap();

    // Depthwise-only emits C * depth_multiplier = 3 channels and does no cross-channel mixing.
    let depthwise = DepthwiseConv2D::new((3, 3), input_shape, (1, 1), Activation::ReLU).unwrap();

    println!("standard  Conv2D params: {}", count(standard.param_count()));
    println!("separable Conv2D params: {}", count(separable.param_count()));
    println!("depthwise Conv2D params: {}", count(depthwise.param_count()));

    assert_eq!(count(standard.param_count()), 1792); // 64*3*3*3 + 64
    assert_eq!(count(separable.param_count()), 283); // 27 + 192 + 64
    assert_eq!(count(depthwise.param_count()), 30); // 3*3*3 + 3
}

Use these layers when channel counts run high and you want to cut both parameters and compute. Wide feature maps are one example. A model meant to run on modest hardware is another. In exchange, you accept a modeling tradeoff: the factored form is strictly less expressive than a full convolution, for the same F and kernel.

Their engine differs from the standard layers’ engine. DepthwiseConv2D uses a direct loop nest, because the channels are independent and im2col buys little here. It parallelizes over (batch item, output row) tasks. Output rows are disjoint, so this split needs no merge. It gates on NAIVE_CONV_PARALLEL_MIN_FLOPS (default 1,000,000).

SeparableConv2D runs its depthwise stage through this same naive path. It then routes its pointwise 1x1 stage back through the shared im2col plus GEMM engine. A 1x1 convolution is exactly a per-position cross-channel matrix multiply.

3.5.5. Building a Small CNN

Convolutions emit rank-3, rank-4, or rank-5 tensors, but a Dense classifier head needs a rank-2 [batch, features] matrix. Flatten bridges the two. Flatten::new(input_shape: Vec<usize>) builds a parameter-free layer that reshapes [batch, ...] into [batch, product-of-the-rest]. It accepts 3D, 4D, or 5D input at forward time. You give it the shape of the tensor entering it, which is the convolution’s output shape. It then works out the flattened feature count.

The next example builds a complete stack. A Conv2D runs over synthetic single-channel 8x8 images. Flatten then feeds the result into a Dense regression head. The Sequential model trains the stack for a few epochs.

use ndarray::{Array2, Array4};
use rustyml::neural_network::layers::*;
use rustyml::neural_network::losses::*;
use rustyml::neural_network::optimizers::*;
use rustyml::neural_network::sequential::Sequential;

fn main() {
    // Synthetic "images": 6 samples, 8x8, 1 channel.
    let mut x = Array4::<f32>::zeros((6, 8, 8, 1));
    for n in 0..6 {
        for i in 0..8 {
            for j in 0..8 {
                x[[n, i, j, 0]] = ((n + i + j) as f32 * 0.1).sin();
            }
        }
    }
    let x = x.into_dyn();

    // 3 regression targets per sample.
    let y = Array2::<f32>::from_shape_fn((6, 3), |(n, k)| (n as f32) * 0.01 + (k as f32) * 0.1)
        .into_dyn();

    let mut model = Sequential::new();
    model
        // [6, 8, 8, 1] -> Conv2D(4 filters, 3x3, Valid) -> [6, 6, 6, 4]
        .add(
            Conv2D::new(4, (3, 3), vec![6, 8, 8, 1], (1, 1), Activation::ReLU)
                .unwrap()
                .with_random_state(42),
        )
        // [6, 6, 6, 4] -> Flatten -> [6, 144]
        .add(Flatten::new(vec![6, 6, 6, 4]).unwrap())
        // [6, 144] -> Dense -> [6, 3]
        .add(Dense::new(6 * 6 * 4, 3, Activation::Linear).unwrap().with_random_state(7))
        .compile(SGD::new(0.01, 0.0, false, 0.0).unwrap(), MeanSquaredError::new());

    model.summary();
    model.fit(&x, &y, 3).unwrap();

    let pred = model.predict(&x).unwrap();
    assert_eq!(pred.shape(), &[6, 3]);
    println!("prediction shape: {:?}", pred.shape());
}

The Dense input dimension is not a guess. It is the flattened feature count 6 * 6 * 4 = 144, which you can read directly from the convolution’s output shape. Get it wrong, and the Dense layer’s matrix multiply rejects the batch. summary() prints the shape at each stage and the parameter budget. This is the fastest way to catch a mis-sized head:

Model: "sequential"
┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━┓
┃ Layer (type)                    ┃ Output Shape           ┃       Param # ┃
┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━┩
│ conv2d (Conv2D)                 │ (6, 6, 6, 4)           │            40 │
│ flatten (Flatten)               │ (None, 144)            │             0 │
│ dense (Dense)                   │ (None, 3)              │           435 │
└─────────────────────────────────┴────────────────────────┴───────────────┘
 Total params: 475 (1900 B)
 Trainable params: 475 (1900 B)
 Non-trainable params: 0 (0 B)

Insert a pooling layer between the convolution and the flatten step, if you want to shrink the spatial extent before the dense head. Save the trained stack with the tools in 3.9. Saving and Loading Weights.

3.5.6. Errors and How to Avoid Them

The convolution layers validate their configuration closely. They return typed errors instead of a panic on bad input. The table below lists the common cases and their error variants.

2 of these error paths need more explanation. The engine catches a kernel larger than the input, under Valid padding, at 2 points. It checks at construction, against the declared input_shape. It checks again at forward time, against the actual tensor. The engine computes in - k in usize. Rather than let that computation underflow, it returns InvalidInput.

A channel mismatch is a recoverable error only on DepthwiseConv2D. This layer checks the runtime channel count directly, and returns DimensionMismatch when the count does not match. The standard Conv1D, Conv2D, and Conv3D layers size their weight matrix from the channel count in the declared input_shape. They do not check this count again on every forward call. You must always supply the channel count you declared. Treat the declared channel count as a contract.

The following program exercises these recoverable paths:

use ndarray::Array;
use rustyml::error::Error;
use rustyml::neural_network::layers::*;
use rustyml::neural_network::traits::Layer;

fn main() {
    // 1. Kernel larger than the declared input under Valid padding: rejected at construction.
    let too_big = Conv2D::new(1, (3, 3), vec![1, 2, 2, 1], (1, 1), Activation::Linear);
    assert!(matches!(too_big, Err(Error::InvalidInput(_))));

    // 2. A zero depth multiplier is rejected by the builder rather than panicking.
    let bad_dm = DepthwiseConv2D::new((2, 2), vec![1, 4, 4, 2], (1, 1), Activation::Linear)
        .unwrap()
        .with_depth_multiplier(0);
    assert!(matches!(bad_dm, Err(Error::InvalidParameter { .. })));

    // 3. A runtime tensor smaller than the kernel: Valid geometry returns an error, not a panic.
    let mut conv = Conv2D::new(1, (3, 3), vec![1, 5, 5, 1], (1, 1), Activation::Linear).unwrap();
    let small = Array::ones((1, 2, 5, 1)).into_dyn(); // height 2 < kernel 3
    assert!(matches!(conv.forward(&small), Err(Error::InvalidInput(_))));

    // 4. DepthwiseConv2D turns a runtime channel mismatch into a recoverable error.
    let mut dw =
        DepthwiseConv2D::new((2, 2), vec![1, 4, 4, 2], (1, 1), Activation::Linear).unwrap();
    let wrong_channels = Array::ones((1, 4, 4, 3)).into_dyn(); // 3 channels, layer expects 2
    assert!(matches!(
        dw.forward(&wrong_channels),
        Err(Error::DimensionMismatch { .. })
    ));

    println!("all error paths behaved as documented");
}

The ForwardPassNotRun error most often surprises you during training. The backward pass reads caches that forward writes. Call backward on a fresh layer, or twice in a row without a forward call between them. Either way, the layer returns this variant instead of reading stale state.

Inside Sequential, the framework handles this order for you. You meet this error only when you drive layers by hand. See 1.6. Error Handling for the full error taxonomy and the smart constructors behind these variants.