Keyboard shortcuts

Press or to navigate between chapters

Press S or / to search in the book

Press ? to show this help

Press Esc to hide this help

3.6. Pooling Layers

Pooling downsamples a feature map. It summarizes each local window with a single number. The spatial dimensions shrink, and the channel count stays the same. Pooling is the cheap counterpart to the convolutions in 3.5. Convolutional Layers. It has no weights and no matrix multiply. It only reduces a sliding window to one number.

RustyML ships pooling as a family of 12 layers. All 12 layers use one engine that works for any spatial rank. This page covers the constructors, the output shapes, and the difference between max and average pooling. It also covers how gradients route backward, and what “no learnable parameters” means for summary() and model persistence.

3.6.1. The family: 2 reductions, 3 ranks, plus global variants

Every pooling layer picks one of 2 reductions, max or average. It applies that reduction at one of 3 spatial ranks: 1D, 2D, or 3D. Each combination also has a “global” variant that collapses the whole spatial extent at once. That is 2 x 3 x 2 = 12 concrete types. All 12 types reduce to 4 functions in a private pooling_engine.

The engine derives the spatial rank at run time from input.ndim() - 2. One loop serves 1D, 2D, and 3D. The only per-layer differences are the reduction kind, and how the public pool_size/strides tuples flatten into slices.

Tensors follow the channels-last convention [batch, spatial..., channels], the same convention the convolution layers use. A 1D layer takes a 3D tensor, a 2D layer takes a 4D tensor, and a 3D layer takes a 5D tensor. Windowed pooling keeps the rank and shrinks the spatial axes. Global pooling drops the spatial axes and returns [batch, channels].

The channel axis sits innermost, so a whole position reduces at once. The window geometry (the bounds checks, and the count of real elements a Same-padded average divides by) is worked out once per output position. Every channel shares that same geometry.

LayerInput tensorWindow controlOutput
MaxPooling1D / AveragePooling1D[N, L, C]new(pool_size, input_shape) + with_stride[N, L', C]
MaxPooling2D / AveragePooling2D[N, H, W, C]new((ph, pw), input_shape) + with_strides[N, H', W', C]
MaxPooling3D / AveragePooling3D[N, D, H, W, C]new((pd, ph, pw), input_shape) + with_strides[N, D', H', W', C]
GlobalMaxPooling{1,2,3}Drank 3 / 4 / 5new()[N, C]
GlobalAveragePooling{1,2,3}Drank 3 / 4 / 5new()[N, C]

One naming point differs by rank. The 1D windowed layers expose with_stride, a single usize. The 2D and 3D layers expose with_strides, a tuple. This matches the shape of pool_size at each rank.

3.6.2. Constructors, pool size, stride, and padding

Windowed layers take the pool window and the declared input shape. For example: MaxPooling2D::new((2, 2), vec![batch, height, width, channels]). The constructor checks eagerly that the window fits inside the declared spatial dimensions. 2 builder methods override the defaults. Both are optional:

  • with_stride / with_strides sets the step between windows. If you never call it, the stride defaults to the pool size. This gives non-overlapping windows, the same default Keras uses. Pass a smaller stride for overlapping windows.
  • with_padding sets PaddingType::Valid (the default, no padding) or PaddingType::Same. This is the same PaddingType enum the convolution layers use.

Global layers take no arguments at all: GlobalMaxPooling2D::new(). There is no window, stride, or padding to configure, because the window is the whole spatial plane.

// Windowed 2D max pooling, 3x3 window, stride 2, Same padding:
let layer = MaxPooling2D::new((3, 3), vec![1, 32, 32, 16])
    .unwrap()
    .with_strides((2, 2))
    .unwrap()
    .with_padding(PaddingType::Same);

// Global 2D average pooling needs nothing:
let head = GlobalAveragePooling2D::new();

Construction validates the declared input_shape. output_shape() reports this declared shape. The forward pass itself only checks the tensor’s rank. Inside a Sequential model, the actual spatial dimensions come from the upstream layer at run time.

Keep the runtime spatial dimensions at least as large as the pool window. The engine computes output sizes with unsigned arithmetic. A plane smaller than the window causes an arithmetic underflow, instead of a clean error.

Every constructor validates its inputs and returns Result, so the failure modes are explicit. See 1.6. Error Handling for the error taxonomy:

ConditionError
input_shape has the wrong rank (e.g. 3D given to a 2D layer)Error::DimensionMismatch
any input_shape dimension is zero (including batch or channels)Error::InvalidInput
a pool dimension is zero, or exceeds the corresponding input dimensionError::InvalidParameter
a stride is zero (from with_stride/with_strides)Error::InvalidParameter

The zero-batch and zero-channel checks are deliberate. An earlier version let a [0, 1, 4, 4] shape pass the constructor. That shape only failed later, at the first forward pass. Validation now rejects it at construction time, where the stack trace points at the real problem.

3.6.3. Output-shape formulas

For Valid padding, each spatial axis shrinks by the same formula the convolution layers use. The division floors, because the trailing remainder is dropped:

out = (in - pool) / stride + 1

For Same padding, the output rounds up to ceil(in / stride). The engine pads symmetrically, with the extra cell on the trailing edge. This matches the convolution engine. Global pooling ignores padding and window size. It always produces [batch, channels].

output_shape() returns these sizes as a formatted string. Windowed layers can compute their output shape immediately, because they store input_shape at construction. Global layers return "Unknown" until a forward pass has run, because they only learn the input shape when a tensor flows through. This difference shows up directly in summary().

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

fn main() {
    // Windowed layers know their output shape at construction, from `input_shape`.
    let mp = MaxPooling2D::new((2, 2), vec![1, 6, 6, 3]).unwrap();
    println!("MaxPooling2D:     {}", mp.output_shape()); // (1, 3, 3, 3)

    let ap = AveragePooling1D::new(2, vec![1, 6, 1]).unwrap();
    println!("AveragePooling1D: {}", ap.output_shape()); // (1, 3, 1)

    // Global layers reduce every spatial axis to one value per channel, but only report a
    // concrete shape once a forward pass has cached the input shape.
    let mut gap = GlobalAveragePooling2D::new();
    println!("before forward:   {}", gap.output_shape()); // Unknown

    let x = Array::from_elem(ndarray::IxDyn(&[3, 5, 5, 4]), 1.0f32);
    let out = gap.forward(&x).unwrap();
    assert_eq!(out.shape(), &[3, 4]);
    println!("after forward:    {}", gap.output_shape()); // (3, 4)
}

Same padding is not only a shape adjustment. The padded cells are virtual. The forward and backward passes skip out-of-bounds positions instead of substituting zeros. This matters for averaging. An edge window divides by the count of real, in-bounds elements, not by the window area.

This is Keras count_include_pad=False behavior. Same-padded average pooling does not dilute edges toward zero.

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

fn main() {
    // 3x3 input, 2x2 window, stride 2. Valid padding drops the last row and column.
    // Same padding rounds the output up to ceil(3/2) = 2. The trailing windows see only
    // their in-bounds cells, because the padding is virtual.
    let x = Array::from_shape_vec((1, 3, 3, 1), (1..=9).map(|v| v as f32).collect())
        .unwrap()
        .into_dyn();

    let mut max_same = MaxPooling2D::new((2, 2), vec![1, 3, 3, 1])
        .unwrap()
        .with_strides((2, 2))
        .unwrap()
        .with_padding(PaddingType::Same);
    let m = max_same.forward(&x).unwrap();
    assert_eq!(m.shape(), &[1, 2, 2, 1]);
    // [[max(1,2,4,5), max(3,6)], [max(7,8), 9]] = [[5, 6], [8, 9]]
    assert_eq!(m.iter().copied().collect::<Vec<_>>(), vec![5.0, 6.0, 8.0, 9.0]);

    let mut avg_same = AveragePooling2D::new((2, 2), vec![1, 3, 3, 1])
        .unwrap()
        .with_strides((2, 2))
        .unwrap()
        .with_padding(PaddingType::Same);
    let a = avg_same.forward(&x).unwrap();
    // Averages divide by the count of REAL cells, not the window area (count_include_pad = False):
    // [[(1+2+4+5)/4, (3+6)/2], [(7+8)/2, 9/1]] = [[3.0, 4.5], [7.5, 9.0]]
    assert_eq!(a.iter().copied().collect::<Vec<_>>(), vec![3.0, 4.5, 7.5, 9.0]);
}

3.6.4. Max versus average: semantics and when each helps

Both reductions run over the same window. They differ only in what they keep. Max pooling records the single largest activation and discards the rest. It acts as a detector for “did this feature fire anywhere in the window”, and it tolerates translation. Average pooling keeps the mean. It preserves the overall magnitude of the region, and it smooths instead of selecting.

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

fn main() {
    // The same 4x4 plane, pooled with a non-overlapping 2x2 window (stride defaults to 2).
    let data: Vec<f32> = (0..16).map(|v| v as f32).collect();
    let x = Array::from_shape_vec((1, 4, 4, 1), data).unwrap().into_dyn();

    let mut max_pool = MaxPooling2D::new((2, 2), vec![1, 4, 4, 1]).unwrap();
    let m = max_pool.forward(&x).unwrap();
    // window maxima: [[5, 7], [13, 15]]
    assert_eq!(m.iter().copied().collect::<Vec<_>>(), vec![5.0, 7.0, 13.0, 15.0]);

    let mut avg_pool = AveragePooling2D::new((2, 2), vec![1, 4, 4, 1]).unwrap();
    let a = avg_pool.forward(&x).unwrap();
    // window means: [[2.5, 4.5], [10.5, 12.5]]
    assert_eq!(a.iter().copied().collect::<Vec<_>>(), vec![2.5, 4.5, 10.5, 12.5]);
}

Use max pooling inside the convolutional stack of a discriminative model. It keeps the strongest response through downsampling, and the exact position in the window does not matter. Use average pooling when magnitude matters more than peaks. Use the global average reduction at the end of a network. Averaging the whole plane gives a stable, smooth summary per channel.

2 engine details matter in edge cases. Max pooling breaks ties toward the first (lowest-index) maximum, using a strict > comparison. Max pooling also propagates NaN on purpose. Once a NaN enters a window, it wins and stays there. This matches PyTorch and TensorFlow, rather than dropping the NaN silently.

3.6.5. Gradient routing: backprop through pooling

Pooling has no parameters to update. It still must route the upstream gradient back to the inputs that produced its output. The 2 reductions route this gradient differently.

Max pooling is winner-takes-gradient. During the forward pass, each output records the flat index of the input element it picked. This index is the “arg-max”, and the layer caches it. During backward, each upstream gradient scatters to that one position. Every other input in the window gets zero. When windows overlap and the same input wins more than one window, the contributions add up.

Average pooling instead spreads each output gradient evenly across its window. Every in-bounds element of the window receives grad / count. Overlaps add up here too.

Global max pooling routes each channel’s gradient to its single arg-max element. Global average pooling spreads each channel’s gradient evenly across the whole plane, as grad / spatial_size.

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

fn main() {
    let data: Vec<f32> = (0..16).map(|v| v as f32).collect();
    let x = Array::from_shape_vec((1, 4, 4, 1), data).unwrap().into_dyn();
    let grad = Array::ones((1, 2, 2, 1)).into_dyn();

    // Max pooling: each window's gradient reaches only the winning cell (here the 4 maxima
    // at flat indices 5, 7, 13, 15). Every other cell receives 0.
    let mut max_pool = MaxPooling2D::new((2, 2), vec![1, 4, 4, 1]).unwrap();
    max_pool.forward(&x).unwrap();
    let gmax = max_pool.backward(&grad).unwrap();
    let expected_max = vec![
        0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 1.0,
    ];
    assert_eq!(gmax.iter().copied().collect::<Vec<_>>(), expected_max);

    // Average pooling: each window's gradient is spread evenly over its cells (1.0 / 4 = 0.25).
    let mut avg_pool = AveragePooling2D::new((2, 2), vec![1, 4, 4, 1]).unwrap();
    avg_pool.forward(&x).unwrap();
    let gavg = avg_pool.backward(&grad).unwrap();
    assert!(gavg.iter().all(|&g| (g - 0.25).abs() < 1e-6));
}

Because max pooling depends on the arg-max cache, the split between forward and predict matters. forward (training mode) records the cache. predict (inference mode, &self) writes no cache, as the Layer trait documents. So backward only works after forward. Calling backward first, or calling it after a cache-free predict, returns Error::NeuralNetwork(NnError::ForwardPassNotRun(_)).

Inside Sequential::fit, this never happens, because training always runs forward first. You trip this error only when you drive a bare layer by hand. Average pooling caches only the input shape, because it needs the geometry, not the values. Global average pooling also caches only the shape. The “run forward before backward” contract is the same across all 12 layers.

3.6.6. Global pooling as a modern head

A traditional convolutional classifier ends with Flatten, followed by a large Dense layer. This design ties the network to one exact input resolution. It also puts most of the parameters into that last matrix.

Global average pooling replaces that head with a reduction that has no parameters. Network-in-Network and ResNet made this design common. The layer collapses each channel to its mean, so each channel yields one feature. A small Dense classifier follows the pooling layer. The pooling step has no parameters. It cannot overfit. This gives the head free regularization. The classifier’s input width depends only on the channel count, not on H x W.

This size independence is real, but only for the pooling layer by itself. GlobalAveragePooling2D validates only that its input is 4D. So the same pooling layer accepts any height and width at run time. The rest of the model stack does not share this property. RustyML’s Dense fixes its input_dim at construction. The upstream convolution layers also pin their input_shape. So a saved model still expects one consistent resolution, end to end.

The payoff of the global-pooling head is the parameter-free reduction, resistant to overfitting, and the tidy [N, C] -> Dense interface. It does not give automatic variable-resolution inference.

use rustyml::neural_network::sequential::Sequential;
use rustyml::neural_network::layers::*;
use rustyml::neural_network::optimizers::*;
use rustyml::neural_network::losses::*;
use ndarray::Array;

fn main() {
    // A small classifier head: Conv -> pool -> global pool -> Dense.
    let x = Array::from_shape_fn((2, 8, 8, 1), |(b, i, j, _)| {
        (i + j) as f32 * 0.1 + b as f32
    })
    .into_dyn();
    let y = Array::ones((2, 3)).into_dyn();

    let mut model = Sequential::new();
    model
        .add(Conv2D::new(4, (3, 3), vec![2, 8, 8, 1], (1, 1), Activation::ReLU).unwrap()) // -> [2, 6, 6, 4]
        .add(MaxPooling2D::new((2, 2), vec![2, 6, 6, 4]).unwrap())                        // -> [2, 3, 3, 4]
        .add(GlobalAveragePooling2D::new())                                               // -> [2, 4]
        .add(Dense::new(4, 3, Activation::ReLU).unwrap())                                 // -> [2, 3]
        .compile(RMSprop::new(0.001, 0.9, 1e-8, 0.0).unwrap(), MeanSquaredError::new());

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

    let prediction = model.predict(&x).unwrap();
    assert_eq!(prediction.shape(), &[2, 3]);
}

The summary() call above runs before fit. This is when the global-pooling behavior from 3.6.3 appears. Its output-shape column reads Unknown, because no tensor has flowed through the layer yet. The windowed layers, in contrast, already report concrete shapes:

Model: "sequential"
┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━┓
┃ Layer (type)                    ┃ Output Shape           ┃       Param # ┃
┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━┩
│ conv2d (Conv2D)                 │ (2, 6, 6, 4)           │            40 │
│ maxpooling2d (MaxPooling2D)     │ (2, 3, 3, 4)           │             0 │
│ globalaveragepooling2d (GlobalAveragePooling2D) │ Unknown                │             0 │
│ dense (Dense)                   │ (None, 3)              │            15 │
└─────────────────────────────────┴────────────────────────┴───────────────┘
 Total params: 55 (220 B)
 Trainable params: 55 (220 B)
 Non-trainable params: 0 (0 B)

Both pooling rows show 0 in the Param # column. Neither one adds to any of the 3 totals. Call summary() again after fit, or after any forward pass, and the global-pooling row settles to (2, 4).

3.6.7. No learnable parameters: what that means for persistence

Pooling layers report TrainingParameters::NoTrainable. They expose LayerWeight::Empty. The optimizer skips them, because they yield no ParamGrad entries. This is why they never appear in the trainable or non-trainable totals above. It is also why they add nothing to a saved file’s size.

Persistence is where “no parameters” has a real consequence. save_to_path writes each layer’s type name, its reported output shape, and its weights, as metadata. For a pooling layer, the weights are LayerWeight::Empty. So the file stores nothing of substance for that layer.

load_from_path does not reconstruct the architecture. You must rebuild the model with the same layers first, then load the weights. The load walks the layers position by position. It checks that the layer count matches. It also checks that each saved type string equals the type at the same index, before it applies any weights. A pooling layer has no weights to restore, but it must still occupy its exact position, with its exact type. Otherwise the load fails with Error::Io(IoError::ModelStructureMismatch).

You cannot drop a MaxPooling2D from the rebuilt model to save space. The structural checkpoint depends on that layer being present.

use rustyml::neural_network::sequential::Sequential;
use rustyml::neural_network::layers::*;
use rustyml::neural_network::optimizers::*;
use rustyml::neural_network::losses::*;
use ndarray::Array;

fn main() {
    let x = Array::from_elem((2, 4, 4, 3), 0.5f32).into_dyn();
    let y = Array::ones((2, 2)).into_dyn();

    let mut model = Sequential::new();
    model
        .add(GlobalAveragePooling2D::new())
        .add(Dense::new(3, 2, Activation::ReLU).unwrap())
        .compile(RMSprop::new(0.001, 0.9, 1e-8, 0.0).unwrap(), MeanSquaredError::new());
    model.fit(&x, &y, 2).unwrap();
    model.save_to_path("pool_head.bin").unwrap();

    // Rebuild the same architecture, then load. The pooling layer carries no weights.
    // It must still occupy the same position, with the same type, for the structure check to pass.
    let mut restored = Sequential::new();
    restored
        .add(GlobalAveragePooling2D::new())
        .add(Dense::new(3, 2, Activation::ReLU).unwrap());
    restored.load_from_path("pool_head.bin").unwrap();

    let out = restored.predict(&x).unwrap();
    assert_eq!(out.shape(), &[2, 2]);

    std::fs::remove_file("pool_head.bin").unwrap();
}

See 3.9. Saving and Loading Weights and 7.2. Model Persistence in Depth for the full serialization story. That story covers the postcard format, why the optimizer and loss are not saved, and how weight shapes are checked on load.

3.6.8. Performance and cost

The engine avoids extra allocations. It reads the whole input as one contiguous &[f32]. Its unit of work is a single output position. For each window tap, a serial loop folds channels adjacent floats into a channels-wide accumulator. The engine then assembles the output with one from_shape_vec call, instead of writing scalars one at a time.

Cost scales with output positions, times window volume, times channels. Overlapping windows (a stride smaller than the pool size) do proportionally more work than the non-overlapping default. Global pooling is the cheapest case. It runs one linear scan per batch item.

Forward work runs in parallel with rayon, split over (batch item, block of output positions). Each block writes a disjoint output slab. Splitting on positions, rather than on the batch alone, keeps every thread busy even when batch == 1.

Backward work splits differently, over (batch item, channel slab). A slab that owns channels [j0, j1) touches only input addresses in that range, modulo channels. So the scatter needs no halo, no merge step, and no atomic operations.

Both passes go parallel only after the estimated total work clears a threshold. That estimate is batch * output_positions * channels * window_volume element operations. The threshold, POOL_PARALLEL_MIN_OPS, defaults to 12,000, and you can override it through the tuning module. The gate counts total taps, not task count. This keeps the gate honest, whether the work sits in a few wide-channel positions or in many narrow ones. Small tensors run in serial, to avoid paying the rayon task overhead for no benefit.

See 7.3. Performance Tuning and Parallelism to move that gate and measure the effect.

One warning carries over from the Layer contract. Backward is pure math, and it does not sanitize non-finite values. Max pooling also propagates NaN on purpose. So a NaN in a pooling input travels straight through, unchanged. Control large but finite gradients with optimizer-level global-norm clipping. Do not expect pooling to clamp them.