rustyml 0.14.0

A high-performance machine learning & deep learning library in pure Rust, offering ML algorithms and neural network support
Documentation
# 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.

| Layer | Input tensor | Weight tensor | Bias | Output tensor |
| --- | --- | --- | --- | --- |
| `Conv1D` | `[N, L, Cin]` | `[k, Cin, F]` | `[F]` | `[N, L', F]` |
| `Conv2D` | `[N, H, W, Cin]` | `[kh, kw, Cin, F]` | `[F]` | `[N, H', W', F]` |
| `Conv3D` | `[N, D, H, W, Cin]` | `[kd, kh, kw, Cin, F]` | `[F]` | `[N, D', H', W', F]` |
| `DepthwiseConv2D` | `[N, H, W, C]` | `[kh, kw, C, dm]` | `[C*dm]` | `[N, H', W', C*dm]` |
| `SeparableConv2D` | `[N, H, W, Cin]` | depthwise `[kh, kw, Cin, dm]`, pointwise `[1, 1, Cin*dm, F]` | `[F]` | `[N, H', W', F]` |

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.

```rust
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.

```rust,ignore
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](./3.2._Dense_Layers_and_Activations.md) 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](../Chapter-07/7.1._Reproducibility_and_Random_Seeds.md)).

`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](./3.1._The_Sequential_Model.md) 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](../Chapter-06/6.2._Matrix_Multiplication.md) 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](../Chapter-07/7.3._Performance_Tuning_and_Parallelism.md) 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.

```rust
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](./3.1._The_Sequential_Model.md) trains the stack for a few epochs.

```rust
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:

```text
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](./3.6._Pooling_Layers.md) 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.9._Saving_and_Loading_Weights.md).

## 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.

| Condition | Error |
| --- | --- |
| `filters` is `0`, for `Conv1D`, `Conv2D`, `Conv3D`, or `SeparableConv2D` (`DepthwiseConv2D` has no `filters` argument) | `Error::InvalidParameter` |
| A kernel dim or a stride is `0` | `Error::InvalidParameter` |
| `depth_multiplier == 0` (`SeparableConv2D::new`, `DepthwiseConv2D::with_depth_multiplier`) | `Error::InvalidParameter` |
| `input_shape` wrong rank, zero channels, or smaller than the kernel (constructor) | `Error::InvalidInput` |
| `forward` handed a tensor of the wrong rank | `Error::InvalidInput` |
| Valid convolution whose runtime spatial dim is below the kernel | `Error::InvalidInput` |
| `DepthwiseConv2D` runtime channel count differs from the declared channel count | `Error::DimensionMismatch` |
| `set_weights` shape does not match | `Error::NeuralNetwork(NnError::WeightShape)` |
| `backward` called before `forward` | `Error::NeuralNetwork(NnError::ForwardPassNotRun)` |

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:

```rust
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](../Chapter-01/1.6._Error_Handling.md) for the full error taxonomy and the smart constructors behind these variants.