Skip to content
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.

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

Feat: GlobalAveragePool #600

Merged
merged 2 commits into from
Apr 22, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
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
1 change: 1 addition & 0 deletions docs/framework/operators/neural-network/README.md
Original file line number Diff line number Diff line change
Expand Up @@ -39,5 +39,6 @@ Orion supports currently these `NN` types.
| [`nn.col2im`](nn.col2im.md) | Rearranges column blocks back into a multidimensional image |
| [`nn.conv_transpose`](nn.conv\_transpose.md) | Performs the convolution transpose of the input data tensor and weight tensor. |
| [`nn.conv`](nn.conv.md) | Performs the convolution of the input data tensor and weight tensor. |
| [`nn.global_average_pool`](nn.global\_average\_pool.md) | GlobalAveragePool consumes an input tensor X and applies average pooling across the values in the same channel. |
| [`nn.conv_integer`](nn.conv\_integer.md) | Performs integer convolution |

75 changes: 75 additions & 0 deletions docs/framework/operators/neural-network/nn.global_average_pool.md
Original file line number Diff line number Diff line change
@@ -0,0 +1,75 @@
# NNTrait::global_average_pool

```rust
fn global_average_pool(tensor: @Tensor<T>) -> Tensor<T>;
```

GlobalAveragePool consumes an input tensor X and applies average pooling across the values in the same channel.
This is equivalent to AveragePool with kernel size equal to the spatial dimension of input tensor.

## Args

* `tensor`(`@Tensor<T>`) - Input data tensor from the previous operator; dimensions for image case are (N x C x H x W), where N is the batch size, C is the number of channels, and H and W are the height and the width of the data. For non image case, the dimensions are in the form of (N x C x D1 x D2 ... Dn), where N is the batch size.

## Returns

* Output data tensor from pooling across the input tensor. The output tensor has the same rank as the input. The first two dimensions of output shape are the same as the input (N x C), while the other dimensions are all 1.

## Examples

```rust
use orion::operators::tensor::{FP8x23Tensor, FP8x23TensorAdd};
use core::array::{ArrayTrait, SpanTrait};
use orion::operators::tensor::{TensorTrait, Tensor};
use orion::utils::{assert_eq, assert_seq_eq};
use orion::operators::tensor::FP8x23TensorPartialEq;
use orion::numbers::{FixedTrait, FP8x23};
use orion::operators::nn::NNTrait;
use orion::operators::nn::FP8x23NN;

fn example() -> Tensor<FP8x23> {
let mut shape = ArrayTrait::<usize>::new();
shape.append(2);
shape.append(4);
shape.append(2);
shape.append(2);

let mut data = ArrayTrait::new();
data.append(FP8x23 { mag: 85392644, sign: false });
data.append(FP8x23 { mag: 61594092, sign: false });
data.append(FP8x23 { mag: 163676643, sign: true });
data.append(FP8x23 { mag: 180530738, sign: false });
data.append(FP8x23 { mag: 168048412, sign: true });
data.append(FP8x23 { mag: 5915510, sign: false });
data.append(FP8x23 { mag: 9047009, sign: true });
data.append(FP8x23 { mag: 46030420, sign: false });
data.append(FP8x23 { mag: 184797857, sign: false });
data.append(FP8x23 { mag: 129370611, sign: false });
data.append(FP8x23 { mag: 174006060, sign: true });
data.append(FP8x23 { mag: 162252480, sign: false });
data.append(FP8x23 { mag: 139240444, sign: true });
data.append(FP8x23 { mag: 168836878, sign: true });
data.append(FP8x23 { mag: 246913333, sign: true });
data.append(FP8x23 { mag: 1047194, sign: true });
data.append(FP8x23 { mag: 238599466, sign: true });
data.append(FP8x23 { mag: 216763643, sign: true });
data.append(FP8x23 { mag: 40581779, sign: true });
data.append(FP8x23 { mag: 209811161, sign: true });
data.append(FP8x23 { mag: 250078311, sign: false });
data.append(FP8x23 { mag: 31811183, sign: true });
data.append(FP8x23 { mag: 36411415, sign: true });
data.append(FP8x23 { mag: 107986324, sign: false });
data.append(FP8x23 { mag: 69727339, sign: false });
data.append(FP8x23 { mag: 223159880, sign: true });
data.append(FP8x23 { mag: 184932087, sign: true });
data.append(FP8x23 { mag: 118617436, sign: false });
data.append(FP8x23 { mag: 134825391, sign: true });
data.append(FP8x23 { mag: 217861279, sign: false });
data.append(FP8x23 { mag: 199069387, sign: false });
data.append(FP8x23 { mag: 192925915, sign: true });
let tensor1 = TensorTrait::new(shape.span(), data.span());

return NNTrait::global_average_pool(@tensor1);
}
>>> [{ mag: 40960207, sign: true } { mag: 31287372, sign: false } { mag: 75603722, sign: true } { mag: 139009462, sign: false } { mag: 176439012, sign: false } { mag: 72460509, sign: true } { mag: 54936798, sign: false } { mag: 22294840, sign: true } ]
```
110 changes: 110 additions & 0 deletions nodegen/node/global_average_pool.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,110 @@
import numpy as np
from nodegen.node import RunAll
from ..helpers import make_test, to_fp, Tensor, Dtype, FixedImpl, Trait

def global_average_pool(x: np.ndarray) -> np.ndarray:
axis = tuple(range(2, np.ndim(x)))
y = np.average(x, axis=axis)
for _ in axis:
y = np.expand_dims(y, -1)
return y # type: ignore

class Global_average_pool(RunAll):

@staticmethod
def fp8x23_2D():
x = np.random.uniform(-30, 30, (2, 4)).astype(np.float64)
y = global_average_pool(x)

x = Tensor(Dtype.FP8x23, x.shape, to_fp(
x.flatten(), FixedImpl.FP8x23))
y = Tensor(Dtype.FP8x23, y.shape, to_fp(
y.flatten(), FixedImpl.FP8x23))

name = "global_average_pool_fp8x23_2D"
make_test([x], y, "NNTrait::global_average_pool(@input_0)",
name, Trait.NN)

@staticmethod
def fp16x16_2D():
x = np.random.uniform(-30, 30, (3, 2)).astype(np.float16)
y = global_average_pool(x)

x = Tensor(Dtype.FP16x16, x.shape, to_fp(
x.flatten(), FixedImpl.FP16x16))
y = Tensor(Dtype.FP16x16, y.shape, to_fp(
y.flatten(), FixedImpl.FP16x16))

name = "global_average_pool_fp16x16_2D"
make_test([x], y, "NNTrait::global_average_pool(@input_0)",
name, Trait.NN)

@staticmethod
def fp8x23_3D():
x = np.random.uniform(-30, 30, (2, 4, 2)).astype(np.float64)
y = global_average_pool(x)

x = Tensor(Dtype.FP8x23, x.shape, to_fp(
x.flatten(), FixedImpl.FP8x23))
y = Tensor(Dtype.FP8x23, y.shape, to_fp(
y.flatten(), FixedImpl.FP8x23))

name = "global_average_pool_fp8x23_3D"
make_test([x], y, "NNTrait::global_average_pool(@input_0)",
name, Trait.NN)

@staticmethod
def fp16x16_3D():
x = np.random.uniform(-30, 30, (3, 2, 2)).astype(np.float16)
y = global_average_pool(x)

x = Tensor(Dtype.FP16x16, x.shape, to_fp(
x.flatten(), FixedImpl.FP16x16))
y = Tensor(Dtype.FP16x16, y.shape, to_fp(
y.flatten(), FixedImpl.FP16x16))

name = "global_average_pool_fp16x16_3D"
make_test([x], y, "NNTrait::global_average_pool(@input_0)",
name, Trait.NN)

@staticmethod
def fp8x23_4D():
x = np.random.uniform(-30, 30, (2, 4, 2, 2)).astype(np.float64)
y = global_average_pool(x)

x = Tensor(Dtype.FP8x23, x.shape, to_fp(
x.flatten(), FixedImpl.FP8x23))
y = Tensor(Dtype.FP8x23, y.shape, to_fp(
y.flatten(), FixedImpl.FP8x23))

name = "global_average_pool_fp8x23_4D"
make_test([x], y, "NNTrait::global_average_pool(@input_0)",
name, Trait.NN)

@staticmethod
def fp16x16_4D():
x = np.random.uniform(-30, 30, (3, 2, 2, 3)).astype(np.float16)
y = global_average_pool(x)

x = Tensor(Dtype.FP16x16, x.shape, to_fp(
x.flatten(), FixedImpl.FP16x16))
y = Tensor(Dtype.FP16x16, y.shape, to_fp(
y.flatten(), FixedImpl.FP16x16))

name = "global_average_pool_fp16x16_4D"
make_test([x], y, "NNTrait::global_average_pool(@input_0)",
name, Trait.NN)

# @staticmethod
# def fp32x32():
# x = np.random.uniform(-30, 30, (2, 4)).astype(np.float64)
# y = global_average_pool(x)

# x = Tensor(Dtype.FP32x32, x.shape, to_fp(
# x.flatten(), FixedImpl.FP32x32))
# y = Tensor(Dtype.FP32x32, y.shape, to_fp(
# y.flatten(), FixedImpl.FP32x32))

# name = "global_average_pool_fp32x32"
# make_test([x], y, "NNTrait::global_average_pool(@input_0)",
# name, Trait.NN)
79 changes: 79 additions & 0 deletions src/operators/nn/core.cairo
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,7 @@ use orion::operators::nn::{AUTO_PAD, MODE, PADDING_MODE};
/// col2im - Rearranges column blocks back into a multidimensional image
/// conv_transpose - Performs the convolution transpose of the input data tensor and weight tensor.
/// conv - Performs the convolution of the input data tensor and weight tensor.
/// global_average_pool - GlobalAveragePool consumes an input tensor X and applies average pooling across the values in the same channel.
/// conv_integer - Performs integer convolution
trait NNTrait<T> {
/// # NNTrait::relu
Expand Down Expand Up @@ -1726,4 +1727,82 @@ trait NNTrait<T> {
pads: Option<Span<usize>>,
strides: Option<Span<usize>>,
) -> Tensor<usize>;
/// # NNTrait::global_average_pool
///
/// ```rust
/// fn global_average_pool(tensor: @Tensor<T>) -> Tensor<T>;
/// ```
///
/// GlobalAveragePool consumes an input tensor X and applies average pooling across the values in the same channel.
/// This is equivalent to AveragePool with kernel size equal to the spatial dimension of input tensor.
///
/// ## Args
///
/// * `tensor`(`@Tensor<T>`) - Input data tensor from the previous operator; dimensions for image case are (N x C x H x W), where N is the batch size, C is the number of channels, and H and W are the height and the width of the data. For non image case, the dimensions are in the form of (N x C x D1 x D2 ... Dn), where N is the batch size.
///
/// ## Returns
///
/// * Output data tensor from pooling across the input tensor. The output tensor has the same rank as the input. The first two dimensions of output shape are the same as the input (N x C), while the other dimensions are all 1.
///
/// ## Examples
///
/// ```rust
/// use orion::operators::tensor::{FP8x23Tensor, FP8x23TensorAdd};
/// use core::array::{ArrayTrait, SpanTrait};
/// use orion::operators::tensor::{TensorTrait, Tensor};
/// use orion::utils::{assert_eq, assert_seq_eq};
/// use orion::operators::tensor::FP8x23TensorPartialEq;
/// use orion::numbers::{FixedTrait, FP8x23};
/// use orion::operators::nn::NNTrait;
/// use orion::operators::nn::FP8x23NN;
///
/// fn example() -> Tensor<FP8x23> {
/// let mut shape = ArrayTrait::<usize>::new();
/// shape.append(2);
/// shape.append(4);
/// shape.append(2);
/// shape.append(2);
///
/// let mut data = ArrayTrait::new();
/// data.append(FP8x23 { mag: 85392644, sign: false });
/// data.append(FP8x23 { mag: 61594092, sign: false });
/// data.append(FP8x23 { mag: 163676643, sign: true });
/// data.append(FP8x23 { mag: 180530738, sign: false });
/// data.append(FP8x23 { mag: 168048412, sign: true });
/// data.append(FP8x23 { mag: 5915510, sign: false });
/// data.append(FP8x23 { mag: 9047009, sign: true });
/// data.append(FP8x23 { mag: 46030420, sign: false });
/// data.append(FP8x23 { mag: 184797857, sign: false });
/// data.append(FP8x23 { mag: 129370611, sign: false });
/// data.append(FP8x23 { mag: 174006060, sign: true });
/// data.append(FP8x23 { mag: 162252480, sign: false });
/// data.append(FP8x23 { mag: 139240444, sign: true });
/// data.append(FP8x23 { mag: 168836878, sign: true });
/// data.append(FP8x23 { mag: 246913333, sign: true });
/// data.append(FP8x23 { mag: 1047194, sign: true });
/// data.append(FP8x23 { mag: 238599466, sign: true });
/// data.append(FP8x23 { mag: 216763643, sign: true });
/// data.append(FP8x23 { mag: 40581779, sign: true });
/// data.append(FP8x23 { mag: 209811161, sign: true });
/// data.append(FP8x23 { mag: 250078311, sign: false });
/// data.append(FP8x23 { mag: 31811183, sign: true });
/// data.append(FP8x23 { mag: 36411415, sign: true });
/// data.append(FP8x23 { mag: 107986324, sign: false });
/// data.append(FP8x23 { mag: 69727339, sign: false });
/// data.append(FP8x23 { mag: 223159880, sign: true });
/// data.append(FP8x23 { mag: 184932087, sign: true });
/// data.append(FP8x23 { mag: 118617436, sign: false });
/// data.append(FP8x23 { mag: 134825391, sign: true });
/// data.append(FP8x23 { mag: 217861279, sign: false });
/// data.append(FP8x23 { mag: 199069387, sign: false });
/// data.append(FP8x23 { mag: 192925915, sign: true });
/// let tensor1 = TensorTrait::new(shape.span(), data.span());
///
/// return NNTrait::global_average_pool(@tensor1);
/// }
/// >>> [{ mag: 40960207, sign: true } { mag: 31287372, sign: false } { mag: 75603722, sign: true } { mag: 139009462, sign: false } { mag: 176439012, sign: false } { mag: 72460509, sign: true } { mag: 54936798, sign: false } { mag: 22294840, sign: true } ]
/// ```
///
fn global_average_pool(tensor: @Tensor<T>) -> Tensor<T>;

}
1 change: 1 addition & 0 deletions src/operators/nn/functional.cairo
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@ mod conv_transpose;
mod depth_to_space;
mod space_to_depth;
mod conv;
mod global_average_pool;
mod conv_integer;
mod max_pool;
mod common_pool;
Expand Down
63 changes: 63 additions & 0 deletions src/operators/nn/functional/global_average_pool.cairo
Original file line number Diff line number Diff line change
@@ -0,0 +1,63 @@
use orion::numbers::fixed_point::core::FixedTrait;
use orion::numbers::NumberTrait;
use orion::operators::tensor::core::{Tensor, TensorTrait};
use orion::operators::tensor::helpers::{reduce_output_shape, len_from_shape, combine_indices};
use orion::operators::tensor::math::{reduce_sum::accumulate_sum, arithmetic::div_downcast};


fn global_average_pool<
T,
MAG,
impl TTensor: TensorTrait<T>,
impl TNumber: NumberTrait<T, MAG>,
impl TAdd: Add<T>,
impl TSub: Sub<T>,
impl TMul: Mul<T>,
impl TDiv: Div<T>,
impl TTensorAdd: Add<Tensor<T>>,
impl TPartialOrd: PartialOrd<T>,
impl TPartialEq: PartialEq<T>,
impl TAddEq: AddEq<T>,
impl TCopy: Copy<T>,
impl TDrop: Drop<T>,
>(
tensor: Tensor<T>
) -> Tensor<T> {
assert((tensor.shape).len() >= 2, 'Unexpected shape.');
let N = *(tensor.shape).at(0);
let C = *(tensor.shape).at(1);
let mut shape = array![N, C];
let mut i:usize = 2;
let one: T = NumberTrait::one();
let len = (tensor.shape).len();
let mut num: usize = 1;
let mut num_t: T = one;
while i < len {
shape.append(1);
let v = *(tensor.shape).at(i);
let mut v_t: T = one;
let mut j: usize = 1;
while j != v {
v_t = v_t + one;
j += 1;
};
num *= v;
num_t = num_t * v_t;
i += 1;
};
let mut arr: Array<T> = array![];
i = 0;
let tensor_len = (tensor.data).len();
while i < tensor_len {
let mut j:usize = 0;
let mut r: T = NumberTrait::zero();
while j != num {
r += *(tensor.data).at(i);
j += 1;
i += 1;
};
arr.append(r / num_t);
};

TensorTrait::<T>::new(shape.span(), arr.span())
}
4 changes: 4 additions & 0 deletions src/operators/nn/implementations/nn_fp16x16.cairo
Original file line number Diff line number Diff line change
Expand Up @@ -145,6 +145,10 @@ impl FP16x16NN of NNTrait<FP16x16> {
functional::conv::conv(X, W, B, auto_pad, dilations, group, kernel_shape, pads, strides)
}

fn global_average_pool(tensor: @Tensor<FP16x16>) -> Tensor<FP16x16> {
functional::global_average_pool::global_average_pool(*tensor)
}

fn conv_integer(
X: @Tensor<FP16x16>,
W: @Tensor<FP16x16>,
Expand Down
4 changes: 4 additions & 0 deletions src/operators/nn/implementations/nn_fp32x32.cairo
Original file line number Diff line number Diff line change
Expand Up @@ -139,6 +139,10 @@ impl FP32x32NN of NNTrait<FP32x32> {
functional::conv::conv(X, W, B, auto_pad, dilations, group, kernel_shape, pads, strides)
}

fn global_average_pool(tensor: @Tensor<FP32x32>) -> Tensor<FP32x32> {
functional::global_average_pool::global_average_pool(*tensor)
}

fn conv_integer(
X: @Tensor<FP32x32>,
W: @Tensor<FP32x32>,
Expand Down
4 changes: 4 additions & 0 deletions src/operators/nn/implementations/nn_fp64x64.cairo
Original file line number Diff line number Diff line change
Expand Up @@ -139,6 +139,10 @@ impl FP64x64NN of NNTrait<FP64x64> {
functional::conv::conv(X, W, B, auto_pad, dilations, group, kernel_shape, pads, strides)
}

fn global_average_pool(tensor: @Tensor<FP64x64>) -> Tensor<FP64x64> {
functional::global_average_pool::global_average_pool(*tensor)
}

fn conv_integer(
X: @Tensor<FP64x64>,
W: @Tensor<FP64x64>,
Expand Down
Loading
Loading