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;
#[derive(Debug)]
pub struct GlobalMaxPooling2D {
input_shape: Vec<usize>,
argmax: Option<Vec<usize>>,
}
impl GlobalMaxPooling2D {
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"));
}
self.input_shape = input.shape().to_vec();
let (output, argmax) = global_pool_forward(input, PoolKind::Max);
self.argmax = argmax;
Ok(output)
}
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!();
}