coreylowman · coreylowman · Oct 4, 2022 · Oct 4, 2022
diff --git a/src/nn/mod.rs b/src/nn/mod.rs
@@ -83,6 +83,9 @@ mod layer_norm;
 mod linear;
 mod module;
 mod npz;
+mod pool_global_avg;
+mod pool_global_max;
+mod pool_global_min;
 mod repeated;
 mod residual;
 mod split_into;
@@ -95,6 +98,9 @@ pub use layer_norm::*;
 pub use linear::*;
 pub use module::*;
 pub use npz::*;
+pub use pool_global_avg::*;
+pub use pool_global_max::*;
+pub use pool_global_min::*;
 pub use repeated::*;
 pub use residual::*;
 pub use split_into::*;

diff --git a/src/nn/pool_global_avg.rs b/src/nn/pool_global_avg.rs
@@ -0,0 +1,55 @@
+use super::{LoadFromNpz, Module, ResetParams, SaveToNpz};
+use crate::gradients::*;
+use crate::tensor::*;
+
+/// Applies average pooling over an entire image, fully reducing the height and width
+/// dimensions:
+/// - Reduces 2d (C, L) to 1d (C, )
+/// - Reduces 3d (C, H, W) to 1d (C, )
+/// - Reduces 4d (B, C, H, W) to 2d (B, C)
+///
+/// **Pytorch equivalent**: `torch.nn.AdaptiveAvgPool2d(1)` followed by a flatten.
+///
+/// Examples:
+/// ```rust
+/// # use dfdx::prelude::*;
+/// let m: AvgPoolGlobal = Default::default();
+/// let _: Tensor1D<5> = m.forward(Tensor3D::<5, 16, 8>::zeros());
+/// let _: Tensor2D<10, 5> = m.forward(Tensor4D::<10, 5, 16, 8>::zeros());
+/// ```
+#[derive(Clone, Copy, Default)]
+pub struct AvgPoolGlobal;
+
+impl ResetParams for AvgPoolGlobal {
+    fn reset_params<R: rand::Rng>(&mut self, _: &mut R) {}
+}
+impl CanUpdateWithGradients for AvgPoolGlobal {
+    fn update<G: GradientProvider>(&mut self, _: &mut G, _: &mut UnusedTensors) {}
+}
+impl SaveToNpz for AvgPoolGlobal {}
+impl LoadFromNpz for AvgPoolGlobal {}
+
+impl<const C: usize, const L: usize, T: Tape> Module<Tensor2D<C, L, T>> for AvgPoolGlobal {
+    type Output = Tensor1D<C, T>;
+    fn forward(&self, input: Tensor2D<C, L, T>) -> Self::Output {
+        input.mean()
+    }
+}
+
+impl<const C: usize, const H: usize, const W: usize, T: Tape> Module<Tensor3D<C, H, W, T>>
+    for AvgPoolGlobal
+{
+    type Output = Tensor1D<C, T>;
+    fn forward(&self, input: Tensor3D<C, H, W, T>) -> Self::Output {
+        input.mean()
+    }
+}
+
+impl<const B: usize, const C: usize, const H: usize, const W: usize, T: Tape>
+    Module<Tensor4D<B, C, H, W, T>> for AvgPoolGlobal
+{
+    type Output = Tensor2D<B, C, T>;
+    fn forward(&self, input: Tensor4D<B, C, H, W, T>) -> Self::Output {
+        input.mean()
+    }
+}
diff --git a/src/nn/pool_global_max.rs b/src/nn/pool_global_max.rs
@@ -0,0 +1,55 @@
+use super::{LoadFromNpz, Module, ResetParams, SaveToNpz};
+use crate::gradients::*;
+use crate::tensor::*;
+
+/// Applies max pooling over an entire image, fully reducing the height and width
+/// dimensions:
+/// - Reduces 2d (C, L) to 1d (C, )
+/// - Reduces 3d (C, H, W) to 1d (C, )
+/// - Reduces 4d (B, C, H, W) to 2d (B, C)
+///
+/// **Pytorch equivalent**: `torch.nn.AdaptiveMaxPool2d(1)` followed by a flatten.
+///
+/// Examples:
+/// ```rust
+/// # use dfdx::prelude::*;
+/// let m: MaxPoolGlobal = Default::default();
+/// let _: Tensor1D<5> = m.forward(Tensor3D::<5, 16, 8>::zeros());
+/// let _: Tensor2D<10, 5> = m.forward(Tensor4D::<10, 5, 16, 8>::zeros());
+/// ```
+#[derive(Clone, Copy, Default)]
+pub struct MaxPoolGlobal;
+
+impl ResetParams for MaxPoolGlobal {
+    fn reset_params<R: rand::Rng>(&mut self, _: &mut R) {}
+}
+impl CanUpdateWithGradients for MaxPoolGlobal {
+    fn update<G: GradientProvider>(&mut self, _: &mut G, _: &mut UnusedTensors) {}
+}
+impl SaveToNpz for MaxPoolGlobal {}
+impl LoadFromNpz for MaxPoolGlobal {}
+
+impl<const C: usize, const L: usize, T: Tape> Module<Tensor2D<C, L, T>> for MaxPoolGlobal {
+    type Output = Tensor1D<C, T>;
+    fn forward(&self, input: Tensor2D<C, L, T>) -> Self::Output {
+        input.max()
+    }
+}
+
+impl<const C: usize, const H: usize, const W: usize, T: Tape> Module<Tensor3D<C, H, W, T>>
+    for MaxPoolGlobal
+{
+    type Output = Tensor1D<C, T>;
+    fn forward(&self, input: Tensor3D<C, H, W, T>) -> Self::Output {
+        input.max()
+    }
+}
+
+impl<const B: usize, const C: usize, const H: usize, const W: usize, T: Tape>
+    Module<Tensor4D<B, C, H, W, T>> for MaxPoolGlobal
+{
+    type Output = Tensor2D<B, C, T>;
+    fn forward(&self, input: Tensor4D<B, C, H, W, T>) -> Self::Output {
+        input.max()
+    }
+}
diff --git a/src/nn/pool_global_min.rs b/src/nn/pool_global_min.rs
@@ -0,0 +1,55 @@
+use super::{LoadFromNpz, Module, ResetParams, SaveToNpz};
+use crate::gradients::*;
+use crate::tensor::*;
+
+/// Applies min pooling over an entire image, fully reducing the height and width
+/// dimensions:
+/// - Reduces 2d (C, L) to 1d (C, )
+/// - Reduces 3d (C, H, W) to 1d (C, )
+/// - Reduces 4d (B, C, H, W) to 2d (B, C)
+///
+/// **Pytorch equivalent**: `torch.nn.AdaptiveMinPool2d(1)` followed by a flatten.
+///
+/// Examples:
+/// ```rust
+/// # use dfdx::prelude::*;
+/// let m: MinPoolGlobal = Default::default();
+/// let _: Tensor1D<5> = m.forward(Tensor3D::<5, 16, 8>::zeros());
+/// let _: Tensor2D<10, 5> = m.forward(Tensor4D::<10, 5, 16, 8>::zeros());
+/// ```
+#[derive(Clone, Copy, Default)]
+pub struct MinPoolGlobal;
+
+impl ResetParams for MinPoolGlobal {
+    fn reset_params<R: rand::Rng>(&mut self, _: &mut R) {}
+}
+impl CanUpdateWithGradients for MinPoolGlobal {
+    fn update<G: GradientProvider>(&mut self, _: &mut G, _: &mut UnusedTensors) {}
+}
+impl SaveToNpz for MinPoolGlobal {}
+impl LoadFromNpz for MinPoolGlobal {}
+
+impl<const C: usize, const L: usize, T: Tape> Module<Tensor2D<C, L, T>> for MinPoolGlobal {
+    type Output = Tensor1D<C, T>;
+    fn forward(&self, input: Tensor2D<C, L, T>) -> Self::Output {
+        input.min()
+    }
+}
+
+impl<const C: usize, const H: usize, const W: usize, T: Tape> Module<Tensor3D<C, H, W, T>>
+    for MinPoolGlobal
+{
+    type Output = Tensor1D<C, T>;
+    fn forward(&self, input: Tensor3D<C, H, W, T>) -> Self::Output {
+        input.min()
+    }
+}
+
+impl<const B: usize, const C: usize, const H: usize, const W: usize, T: Tape>
+    Module<Tensor4D<B, C, H, W, T>> for MinPoolGlobal
+{
+    type Output = Tensor2D<B, C, T>;
+    fn forward(&self, input: Tensor4D<B, C, H, W, T>) -> Self::Output {
+        input.min()
+    }
+}