rustyml 0.13.0

A high-performance machine learning & deep learning library in pure Rust, offering ML algorithms and neural network support
Documentation
//! Global max pooling layer for 2D inputs

use crate::error::Error;
use crate::neural_network::Tensor;
use crate::neural_network::layers::TrainingParameters;
use crate::neural_network::layers::layer_weight::LayerWeight;
use crate::neural_network::layers::pooling::layer_functions_global_pooling;
use crate::neural_network::layers::pooling::pooling_engine::{
    PoolKind, global_pool_backward, global_pool_forward,
};
use crate::neural_network::traits::Layer;

/// Global max pooling layer for 2D inputs
///
/// Selects the maximum value across the height and width dimensions
///
/// Input tensor shape: `[batch_size, channels, height, width]`. Output tensor shape:
/// `[batch_size, channels]`
///
/// # Examples
///
/// ```rust
/// use rustyml::neural_network::sequential::Sequential;
/// use rustyml::neural_network::layers::*;
/// use rustyml::neural_network::optimizers::*;
/// use rustyml::neural_network::losses::*;
/// use ndarray::{Array, IxDyn};
/// use approx::assert_relative_eq;
///
/// // Create a Sequential model containing multiple layers
/// let mut model = Sequential::new();
///
/// // Add a GlobalMaxPooling2D layer
/// model.add(GlobalMaxPooling2D::new());
///
/// // Create a test input tensor: [batch_size, channels, height, width]
/// let input_data = Array::from_elem(IxDyn(&[3, 4, 5, 5]), 1.0);
///
/// // Forward propagation
/// let output = model.predict(&input_data).unwrap();
///
/// // Check output shape - should be [3, 4]
/// assert_eq!(output.shape(), &[3, 4]);
///
/// // Since all input values are 1.0, all output values should also be 1.0 too
/// for b in 0..3 {
///     for c in 0..4 {
///         assert_relative_eq!(output[[b, c]], 1.0);
///     }
/// }
/// ```
///
/// # Performance
///
/// Parallel execution is used when `batch_size * channels >= 32`
#[derive(Debug)]
pub struct GlobalMaxPooling2D {
    /// Shape of the input tensor cached during the forward pass
    input_shape: Vec<usize>,
    /// Cached flat per-channel arg-max indices for backpropagation
    argmax: Option<Vec<usize>>,
}

impl GlobalMaxPooling2D {
    /// Creates a new global max pooling 2D layer
    ///
    /// # Returns
    ///
    /// - `GlobalMaxPooling2D` - New layer instance
    pub fn new() -> Self {
        GlobalMaxPooling2D {
            input_shape: Vec::new(),
            argmax: None,
        }
    }
}

impl Default for GlobalMaxPooling2D {
    fn default() -> Self {
        Self::new()
    }
}

impl Layer for GlobalMaxPooling2D {
    fn forward(&mut self, input: &Tensor) -> Result<Tensor, Error> {
        if input.ndim() != 4 {
            return Err(Error::invalid_input("input tensor is not 4D"));
        }

        // Cache input shape and arg-max positions for backpropagation
        self.input_shape = input.shape().to_vec();

        let (output, argmax) = global_pool_forward(input, PoolKind::Max);
        self.argmax = argmax;
        Ok(output)
    }

    /// Inference forward (eval mode, writes no caches). See [`Layer::predict`]
    fn predict(&self, input: &Tensor) -> Result<Tensor, Error> {
        if input.ndim() != 4 {
            return Err(Error::invalid_input("input tensor is not 4D"));
        }

        Ok(global_pool_forward(input, PoolKind::Max).0)
    }

    fn backward(&mut self, grad_output: &Tensor) -> Result<Tensor, Error> {
        if let Some(argmax) = &self.argmax {
            Ok(global_pool_backward(
                grad_output,
                &self.input_shape,
                PoolKind::Max,
                Some(argmax),
            ))
        } else {
            Err(Error::forward_pass_not_run("GlobalMaxPooling2D"))
        }
    }

    fn layer_type(&self) -> &str {
        "GlobalMaxPooling2D"
    }

    layer_functions_global_pooling!();
}