Skip to content

Adding min/max/avg pool #214

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 9 commits into from
Oct 7, 2022
Merged

Adding min/max/avg pool #214

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
9 commits
Select commit Hold shift + click to select a range
a37cebe
#157 Adding min/max/avg pool
coreylowman Oct 4, 2022
83a8b58
Merge branch 'main' into pooling
coreylowman Oct 6, 2022
578a6e5
Renaming to pool2d
coreylowman Oct 6, 2022
8dc0e51
tests passing
coreylowman Oct 6, 2022
adeead8
Testing equal gradients
coreylowman Oct 6, 2022
15e50bc
Merge branch 'main' into pooling
coreylowman Oct 7, 2022
a1d0c1e
Adding batched pool test
coreylowman Oct 7, 2022
8b97f68
Adding docstrings
coreylowman Oct 7, 2022
1eaf56c
Adding nn tests for pooling
coreylowman Oct 7, 2022
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 4 additions & 0 deletions src/devices/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,10 @@ pub use select::*;
mod conv;
#[cfg(feature = "nightly")]
pub use conv::*;
#[cfg(feature = "nightly")]
mod pool2d;
#[cfg(feature = "nightly")]
pub use pool2d::*;

use std::ops::*;

Expand Down
222 changes: 222 additions & 0 deletions src/devices/pool2d.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,222 @@
use super::Cpu;

pub struct PoolMax;
pub struct PoolMin;
pub struct PoolAvg;

/// **Requires nightly** 2d convolution with stride and padding specified at trait level.
///
/// This allows the rest of the parameters to be inferred by inputs.
pub trait DevicePool2D<const K: usize, const S: usize, const P: usize, Kind> {
/// Forward operation that modifies the `out` image.
fn pool_forward<const C: usize, const H: usize, const W: usize>(
inp: &[[[f32; W]; H]; C],
out: &mut [[[f32; (W + 2 * P - K) / S + 1]; (H + 2 * P - K) / S + 1]; C],
);

/// Backward operation that modifies the gradients of img, weight, and bias.
fn pool_backward<const C: usize, const H: usize, const W: usize>(
inp: &[[[f32; W]; H]; C],
out_g: &[[[f32; (W + 2 * P - K) / S + 1]; (H + 2 * P - K) / S + 1]; C],
inp_g: &mut [[[f32; W]; H]; C],
);
}

impl<const K: usize, const S: usize, const P: usize> DevicePool2D<K, S, P, PoolMax> for Cpu {
fn pool_forward<const C: usize, const H: usize, const W: usize>(
inp: &[[[f32; W]; H]; C],
out: &mut [[[f32; (W + 2 * P - K) / S + 1]; (H + 2 * P - K) / S + 1]; C],
) {
let out_height = (H + 2 * P - K) / S + 1;
let out_width = (W + 2 * P - K) / S + 1;
for c in 0..C {
for oh in 0..out_height {
for ow in 0..out_width {
let o = &mut out[c][oh][ow];
let mut tmp = f32::NEG_INFINITY;
for k1 in 0..K {
let y = (oh * S + k1).checked_sub(P);
for k2 in 0..K {
let x = (ow * S + k2).checked_sub(P);
if let Some((y, x)) = y.zip(x) {
if y < H && x < W {
tmp = tmp.max(inp[c][y][x]);
}
}
}
}
*o = tmp;
}
}
}
}

fn pool_backward<const C: usize, const H: usize, const W: usize>(
inp: &[[[f32; W]; H]; C],
out_g: &[[[f32; (W + 2 * P - K) / S + 1]; (H + 2 * P - K) / S + 1]; C],
inp_g: &mut [[[f32; W]; H]; C],
) {
let out_height = (H + 2 * P - K) / S + 1;
let out_width = (W + 2 * P - K) / S + 1;
for c in 0..C {
for oh in 0..out_height {
for ow in 0..out_width {
let o_g = &out_g[c][oh][ow];
let mut tmp = f32::NEG_INFINITY;
for k1 in 0..K {
let y = (oh * S + k1).checked_sub(P);
for k2 in 0..K {
let x = (ow * S + k2).checked_sub(P);
if let Some((y, x)) = y.zip(x) {
if y < H && x < W {
tmp = tmp.max(inp[c][y][x]);
}
}
}
}

for k1 in 0..K {
let y = (oh * S + k1).checked_sub(P);
for k2 in 0..K {
let x = (ow * S + k2).checked_sub(P);
if let Some((y, x)) = y.zip(x) {
if y < H && x < W && inp[c][y][x] == tmp {
inp_g[c][y][x] += o_g;
}
}
}
}
}
}
}
}
}

impl<const K: usize, const S: usize, const P: usize> DevicePool2D<K, S, P, PoolMin> for Cpu {
fn pool_forward<const C: usize, const H: usize, const W: usize>(
inp: &[[[f32; W]; H]; C],
out: &mut [[[f32; (W + 2 * P - K) / S + 1]; (H + 2 * P - K) / S + 1]; C],
) {
let out_height = (H + 2 * P - K) / S + 1;
let out_width = (W + 2 * P - K) / S + 1;
for c in 0..C {
for oh in 0..out_height {
for ow in 0..out_width {
let o = &mut out[c][oh][ow];
let mut tmp = f32::INFINITY;
for k1 in 0..K {
let y = (oh * S + k1).checked_sub(P);
for k2 in 0..K {
let x = (ow * S + k2).checked_sub(P);
if let Some((y, x)) = y.zip(x) {
if y < H && x < W {
tmp = tmp.min(inp[c][y][x]);
}
}
}
}
*o = tmp;
}
}
}
}

fn pool_backward<const C: usize, const H: usize, const W: usize>(
inp: &[[[f32; W]; H]; C],
out_g: &[[[f32; (W + 2 * P - K) / S + 1]; (H + 2 * P - K) / S + 1]; C],
inp_g: &mut [[[f32; W]; H]; C],
) {
let out_height = (H + 2 * P - K) / S + 1;
let out_width = (W + 2 * P - K) / S + 1;
for c in 0..C {
for oh in 0..out_height {
for ow in 0..out_width {
let o_g = &out_g[c][oh][ow];
let mut tmp = f32::INFINITY;
for k1 in 0..K {
let y = (oh * S + k1).checked_sub(P);
for k2 in 0..K {
let x = (ow * S + k2).checked_sub(P);
if let Some((y, x)) = y.zip(x) {
if y < H && x < W {
tmp = tmp.min(inp[c][y][x]);
}
}
}
}

for k1 in 0..K {
let y = (oh * S + k1).checked_sub(P);
for k2 in 0..K {
let x = (ow * S + k2).checked_sub(P);
if let Some((y, x)) = y.zip(x) {
if y < H && x < W && inp[c][y][x] == tmp {
inp_g[c][y][x] += o_g;
}
}
}
}
}
}
}
}
}

impl<const K: usize, const S: usize, const P: usize> DevicePool2D<K, S, P, PoolAvg> for Cpu {
fn pool_forward<const C: usize, const H: usize, const W: usize>(
inp: &[[[f32; W]; H]; C],
out: &mut [[[f32; (W + 2 * P - K) / S + 1]; (H + 2 * P - K) / S + 1]; C],
) {
let out_height = (H + 2 * P - K) / S + 1;
let out_width = (W + 2 * P - K) / S + 1;
let inv_k2 = 1.0 / (K * K) as f32;
for c in 0..C {
for oh in 0..out_height {
for ow in 0..out_width {
let o = &mut out[c][oh][ow];
let mut tmp = 0.0;
for k1 in 0..K {
let y = (oh * S + k1).checked_sub(P);
for k2 in 0..K {
let x = (ow * S + k2).checked_sub(P);
if let Some((y, x)) = y.zip(x) {
if y < H && x < W {
tmp += inp[c][y][x];
}
}
}
}
*o = tmp * inv_k2;
}
}
}
}

fn pool_backward<const C: usize, const H: usize, const W: usize>(
_inp: &[[[f32; W]; H]; C],
out_g: &[[[f32; (W + 2 * P - K) / S + 1]; (H + 2 * P - K) / S + 1]; C],
inp_g: &mut [[[f32; W]; H]; C],
) {
let out_height = (H + 2 * P - K) / S + 1;
let out_width = (W + 2 * P - K) / S + 1;
let inv_k2 = 1.0 / (K * K) as f32;
for c in 0..C {
for oh in 0..out_height {
for ow in 0..out_width {
let g = out_g[c][oh][ow] * inv_k2;
for k1 in 0..K {
let y = (oh * S + k1).wrapping_sub(P);
if y < H {
for k2 in 0..K {
let x = (ow * S + k2).wrapping_sub(P);
if x < W {
inp_g[c][y][x] += g;
}
}
}
}
}
}
}
}
}
15 changes: 15 additions & 0 deletions src/nn/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -122,6 +122,21 @@ mod conv;
#[cfg(feature = "nightly")]
pub use conv::*;

#[cfg(feature = "nightly")]
mod pool2d_max;
#[cfg(feature = "nightly")]
pub use pool2d_max::*;

#[cfg(feature = "nightly")]
mod pool2d_min;
#[cfg(feature = "nightly")]
pub use pool2d_min::*;

#[cfg(feature = "nightly")]
mod pool2d_avg;
#[cfg(feature = "nightly")]
pub use pool2d_avg::*;

#[cfg(test)]
mod tests {
use crate::gradients::{GradientProvider, Gradients};
Expand Down
108 changes: 108 additions & 0 deletions src/nn/pool2d_avg.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,108 @@
use super::{LoadFromNpz, SaveToNpz};
use super::{Module, ResetParams};
use crate::gradients::*;
use crate::tensor::*;
use rand::Rng;

/// Average pool with 2d kernel that operates on images (3d) and batches of images (4d).
/// Each patch reduces to the average of the values in the patch.
///
/// Generics:
/// - `KERNEL_SIZE`: The size of the kernel applied to both width and height of the images.
/// - `STRIDE`: How far to move the kernel each step. Defaults to `1`
/// - `PADDING`: How much zero padding to add around the images. Defaults to `0`.
#[derive(Debug, Default, Clone)]
pub struct AvgPool2D<const KERNEL_SIZE: usize, const STRIDE: usize = 1, const PADDING: usize = 0>;

impl<const K: usize, const S: usize, const P: usize> CanUpdateWithGradients for AvgPool2D<K, S, P> {
fn update<G: GradientProvider>(&mut self, _: &mut G, _: &mut UnusedTensors) {}
}

impl<const K: usize, const S: usize, const P: usize> ResetParams for AvgPool2D<K, S, P> {
fn reset_params<R: Rng>(&mut self, _: &mut R) {}
}

impl<const K: usize, const S: usize, const P: usize> SaveToNpz for AvgPool2D<K, S, P> {}
impl<const K: usize, const S: usize, const P: usize> LoadFromNpz for AvgPool2D<K, S, P> {}

impl<
const K: usize,
const S: usize,
const P: usize,
const C: usize,
const H: usize,
const W: usize,
T: Tape,
> Module<Tensor3D<C, H, W, T>> for AvgPool2D<K, S, P>
where
[(); (W + 2 * P - K) / S + 1]:,
[(); (H + 2 * P - K) / S + 1]:,
{
type Output = Tensor3D<C, { (H + 2 * P - K) / S + 1 }, { (W + 2 * P - K) / S + 1 }, T>;

fn forward(&self, x: Tensor3D<C, H, W, T>) -> Self::Output {
x.avg2d::<K, S, P>()
}
}

impl<
const K: usize,
const S: usize,
const P: usize,
const B: usize,
const C: usize,
const H: usize,
const W: usize,
T: Tape,
> Module<Tensor4D<B, C, H, W, T>> for AvgPool2D<K, S, P>
where
[(); (W + 2 * P - K) / S + 1]:,
[(); (H + 2 * P - K) / S + 1]:,
{
type Output = Tensor4D<B, C, { (H + 2 * P - K) / S + 1 }, { (W + 2 * P - K) / S + 1 }, T>;

fn forward(&self, x: Tensor4D<B, C, H, W, T>) -> Self::Output {
x.avg2d::<K, S, P>()
}
}

#[cfg(test)]
mod tests {
use super::*;

#[test]
fn test_forward_3d_sizes() {
type Img = Tensor3D<3, 10, 10>;
let _: Tensor3D<3, 8, 8> = AvgPool2D::<3>::default().forward(Img::zeros());
let _: Tensor3D<3, 9, 9> = AvgPool2D::<2>::default().forward(Img::zeros());
let _: Tensor3D<3, 7, 7> = AvgPool2D::<4>::default().forward(Img::zeros());
let _: Tensor3D<3, 4, 4> = AvgPool2D::<3, 2>::default().forward(Img::zeros());
let _: Tensor3D<3, 3, 3> = AvgPool2D::<3, 3>::default().forward(Img::zeros());
let _: Tensor3D<3, 10, 10> = AvgPool2D::<3, 1, 1>::default().forward(Img::zeros());
let _: Tensor3D<3, 12, 12> = AvgPool2D::<3, 1, 2>::default().forward(Img::zeros());
let _: Tensor3D<3, 6, 6> = AvgPool2D::<3, 2, 2>::default().forward(Img::zeros());
}

#[test]
fn test_forward_4d_sizes() {
type Img = Tensor4D<5, 3, 10, 10>;
let _: Tensor4D<5, 3, 7, 7> = AvgPool2D::<4>::default().forward(Img::zeros());
let _: Tensor4D<5, 3, 8, 8> = AvgPool2D::<3>::default().forward(Img::zeros());
let _: Tensor4D<5, 3, 9, 9> = AvgPool2D::<2>::default().forward(Img::zeros());
let _: Tensor4D<5, 3, 4, 4> = AvgPool2D::<3, 2>::default().forward(Img::zeros());
let _: Tensor4D<5, 3, 3, 3> = AvgPool2D::<3, 3>::default().forward(Img::zeros());
let _: Tensor4D<5, 3, 10, 10> = AvgPool2D::<3, 1, 1>::default().forward(Img::zeros());
let _: Tensor4D<5, 3, 12, 12> = AvgPool2D::<3, 1, 2>::default().forward(Img::zeros());
let _: Tensor4D<5, 3, 6, 6> = AvgPool2D::<3, 2, 2>::default().forward(Img::zeros());
}

#[test]
fn test_tuple_pool_sizes() {
type A = AvgPool2D<3>;
type B = AvgPool2D<1, 1, 1>;
type Img = Tensor3D<1, 10, 10>;

let _: Tensor3D<1, 6, 6> = <(A, A)>::default().forward(Tensor3D::<1, 10, 10>::zeros());
let _: Tensor3D<1, 8, 8> = <(A, A, B)>::default().forward(Img::zeros());
}
}
Loading