rustyml 0.17.0

Machine learning models, neural networks, metrics, and data utilities in Rust
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132

//! Global average pooling layer for 2D inputs
//!
//! [`GlobalAveragePooling2D`] reduces the whole height and width axes of a
//! `[batch, height, width, channels]` tensor to a single mean per channel, producing
//! `[batch, channels]`. The backward pass spreads each output gradient evenly over every
//! position of its channel, using only the cached input shape.

use crate::error::Error;
use crate::neural_network::layers::ParamCounts;
use crate::neural_network::layers::pooling::pooling_engine::{
    PoolKind, global_pool_backward, global_pool_forward,
};
use crate::neural_network::layers::pooling::{
    layer_base_functions_pooling, layer_functions_global_pooling,
};
use crate::neural_network::traits::{LayerBase, UnaryLayer};
use crate::neural_network::{Ctx, Shape, Tensor};

/// Global average pooling layer for 2D inputs
///
/// Computes the mean value across the height and width dimensions.
/// Input tensor shape: `[batch_size, height, width, channels]`. Output tensor shape:
/// `[batch_size, channels]`.
///
/// # Examples
///
/// ```rust
/// use rustyml::neural_network::Shape;
/// use rustyml::neural_network::sequential::SequentialBuilder;
/// 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 test input tensor: [batch_size, height, width, channels].
/// let input_data = Array::from_elem(IxDyn(&[3, 5, 5, 4]), 1.0);
///
/// // Build a model that holds 1 GlobalAveragePooling2D layer.
/// let model = SequentialBuilder::new()
///     .add(GlobalAveragePooling2D::new())
///     .build(&Shape::known(input_data.shape()))
///     .unwrap();
///
/// // 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.
/// for b in 0..3 {
///     for c in 0..4 {
///         assert_relative_eq!(output[[b, c]], 1.0);
///     }
/// }
/// ```
///
/// # Performance
///
/// The forward pass runs in parallel when its estimated element count (`batch * positions *
/// channels`)
/// clears the gate in [`tuning::pool`](crate::tuning::pool). The gate does not depend on any
/// fixed shape.
#[derive(Debug)]
pub struct GlobalAveragePooling2D {
    /// Shape the layer was built for, batch axis first. `None` before the build
    built: Option<Shape>,
}

impl GlobalAveragePooling2D {
    /// Creates a new global average pooling 2D layer
    ///
    /// # Returns
    ///
    /// - `GlobalAveragePooling2D` - New layer instance
    pub fn new() -> Self {
        GlobalAveragePooling2D { built: None }
    }
}

impl Default for GlobalAveragePooling2D {
    /// Creates a layer with no recorded input shape.
    ///
    /// # Default Values
    ///
    /// The layer has no build until it receives an input shape.
    ///
    /// # Returns
    ///
    /// - `GlobalAveragePooling2D` - The new layer
    fn default() -> Self {
        Self::new()
    }
}

impl LayerBase for GlobalAveragePooling2D {
    fn layer_type(&self) -> &str {
        "GlobalAveragePooling2D"
    }

    layer_base_functions_pooling!();
}

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

        let (output, _) = global_pool_forward(input, PoolKind::Average);

        if ctx.is_training() {
            ctx.push_cache("GlobalAveragePooling2D", input.shape().to_vec());
        }

        Ok(output)
    }

    fn backward(&self, grad_output: &Tensor, ctx: &mut Ctx) -> Result<Tensor, Error> {
        let input_shape: Vec<usize> = ctx.pop_cache("GlobalAveragePooling2D")?;

        Ok(global_pool_backward(
            grad_output,
            &input_shape,
            PoolKind::Average,
            None,
        ))
    }

    layer_functions_global_pooling!("GlobalAveragePooling2D", 4);
}