Merge pull request #595 from chachaleo/feat/ConvInteger

feat: conv integer
gizatechxyz · Apr 22, 2024 · bac0053 · bac0053
2 parents 2a2c8e5 + c67cb4a
commit bac0053
Show file tree

Hide file tree

Showing 32 changed files with 876 additions and 6 deletions.
diff --git a/docs/SUMMARY.md b/docs/SUMMARY.md
@@ -177,6 +177,7 @@
     * [nn.col2im](framework/operators/neural-network/nn.col2im.md)
     * [nn.conv\_transpose](framework/operators/neural-network/nn.conv\_transpose.md)
     * [nn.conv](framework/operators/neural-network/nn.conv.md)
+    * [nn.conv_integer](framework/operators/neural-network/nn.conv\_integer.md)
     * [nn.depth_to_space](framework/operators/neural-network/nn.depth_to_space.md)
     * [nn.space_to_depth](framework/operators/neural-network/nn.space_to_depth.md)
     * [nn.max\_pool](framework/operators/neural-network/nn.max\_pool.md)

diff --git a/docs/framework/compatibility.md b/docs/framework/compatibility.md
@@ -47,6 +47,7 @@ You can see below the list of current supported ONNX Operators:
 |            [Col2im](operators/neural-network/nn.col2im\_sigmoid.md)           | :white\_check\_mark: |
 |       [ConvTranspose](operators/neural-network/nn.conv\_transpose_.md)        | :white\_check\_mark: |
 |                  [Conv](operators/neural-network/nn.conv.md)                  | :white\_check\_mark: |
+|         [ConvInteger](operators/neural-network/nn.conv\_integer_.md)          | :white\_check\_mark: |
 |              [MaxPool](operators/neural-network/nn.max\_pool.md)              | :white\_check\_mark: |
 |        [DeformConv](operators/neural-network/nn.deform\_conv_.md)             | :white\_check\_mark: |
 |                    [Sinh](operators/tensor/tensor.sinh.md)                    | :white\_check\_mark: |

diff --git a/docs/framework/operators/neural-network/README.md b/docs/framework/operators/neural-network/README.md
@@ -39,4 +39,5 @@ 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.conv_integer`](nn.conv\_integer.md) | Performs integer convolution  |
 
diff --git a/docs/framework/operators/neural-network/nn.conv_integer.md b/docs/framework/operators/neural-network/nn.conv_integer.md
@@ -0,0 +1,122 @@
+# tensor.conv_integer
+
+```rust
+
+fn conv_integer(
+    X: @Tensor<T>,
+    W: @Tensor<T>,
+    X_zero_point: Option<@Tensor<T>>,
+    W_zero_point: Option<@Tensor<T>>,
+    auto_pad: Option<AUTO_PAD>,
+    dilations: Option<Span<usize>>,
+    group: Option<usize>,
+    kernel_shape: Option<Span<usize>>,
+    pads: Option<Span<usize>>,
+    strides: Option<Span<usize>>,
+) -> Tensor<T>;
+```
+
+Performs integer convolution 
+
+The integer convolution operator consumes an input tensor, its zero-point, a filter, and its zero-point, and computes the output.
+
+## Args
+
+* `X`(`@Tensor<i8>`) - Input data tensor, has size (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 width. Note that this is for the 2D image. Otherwise the size is (N x C x D1 x D2 ... x Dn)..
+* `W`(`@Tensor<i8>`) - Weight tensor that will be used in the convolutions; has size (M x C/group x kH x kW), where C is the number of channels, and kH and kW are the height and width of the kernel, and M is the number of feature maps. For more than 2 dimensions, the kernel shape will be (M x C/group x k1 x k2 x ... x kn), where (k1 x k2 x ... kn) is the dimension of the kernel. 
+* `X_zero_point`(`@Tensor<T>`) - Zero point for input `X`
+* `W_zero_point`(`@Tensor<T>`) - Zero point for input `W`. 
+* `auto_pad`(`Option<AUTO_PAD>`) - Default is NOTSET, auto_pad must be either NOTSET, SAME_UPPER, SAME_LOWER or VALID. NOTSET means explicit padding is used. SAME_UPPER or SAME_LOWER mean pad the input so that `output_shape[i] = ceil(input_shape[i] / strides[i])` for each axis `i`.
+* `dilations`(`Option<Span<usize>>`) - Dilation value along each spatial axis of the filter. If not present, the dilation defaults to 1 along each spatial axis.
+* `group`(`Option<usize>`) - Default is 1, number of groups input channels and output channels are divided into.
+* `kernel_shape`(`Option<Span<usize>>`) - The shape of the convolution kernel. If not present, should be inferred from input W.
+* `pads`(`Option<Span<usize>>`) - Padding for the beginning and ending along each spatial axis, it can take any value greater than or equal to 0. The value represent the number of pixels added to the beginning and end part of the corresponding axis. `pads` format should be as follow [x1_begin, x2_begin...x1_end, x2_end,...], where xi_begin the number of pixels added at the beginning of axis `i` and xi_end, the number of pixels added at the end of axis `i`. This attribute cannot be used simultaneously with auto_pad attribute. If not present, the padding defaults to 0 along start and end of each spatial axis.
+* `strides`(`Option<Span<usize>>`) - Stride along each spatial axis. If not present, the stride defaults to 1 along each spatial axis. 
+
+## Returns
+
+A new `Tensor<usize>`, containing the result of the convolution of the inputs.
+
+
+## Example
+
+```rust
+use core::array::{ArrayTrait, SpanTrait};
+
+use orion::operators::tensor::{TensorTrait, Tensor, I8Tensor, U32Tensor};
+use orion::operators::nn::NNTrait;
+use orion::operators::nn::U32NN;
+
+fn example_conv_integer() -> Tensor<usize> {
+    let mut shape = ArrayTrait::<usize>::new();
+    shape.append(1);
+    shape.append(1);
+    shape.append(3);
+    shape.append(3);
+
+    let mut data = ArrayTrait::new();
+    data.append(2);
+    data.append(3);
+    data.append(4);
+    data.append(5);
+    data.append(6);
+    data.append(7);
+    data.append(8);
+    data.append(9);
+    data.append(10);
+    let mut X = TensorTrait::<i8>::new(shape.span(), data.span());
+
+    let mut shape = ArrayTrait::<usize>::new();
+    shape.append(1);
+    shape.append(1);
+    shape.append(2);
+    shape.append(2);
+
+    let mut data = ArrayTrait::new();
+    data.append(1);
+    data.append(1);
+    data.append(1);
+    data.append(1);
+    let mut W = TensorTrait::new(shape.span(), data.span());
+
+    let mut shape = ArrayTrait::<usize>::new();
+    shape.append(1);
+    shape.append(1);
+    shape.append(1);
+    shape.append(1);
+
+    let mut data = ArrayTrait::new();
+    data.append(1);
+    let X_zero_point = TensorTrait::<i8>::new(shape.span(), data.span());
+
+
+    let mut shape = ArrayTrait::<usize>::new();
+    shape.append(1);
+    shape.append(1);
+    shape.append(2);
+    shape.append(2);
+
+    let mut data = ArrayTrait::new();
+    data.append(12);
+    data.append(16);
+    data.append(24);
+    data.append(28);
+    let expected_output = TensorTrait::new(shape.span(), data.span());
+
+    'data ok'.print();
+
+    return NNTrait::conv_integer(
+        @X,
+        @W,
+        Option::Some(@X_zero_point),
+        Option::None,
+        Option::None,
+        Option::None,
+        Option::None,
+        Option::None,
+        Option::None,
+        Option::None,
+    );
+}
+>>> [[12, 16], [24, 28]]
+```
diff --git a/nodegen/node/conv_integer.py b/nodegen/node/conv_integer.py
@@ -0,0 +1,116 @@
+import numpy as np
+from nodegen.node import RunAll
+from ..helpers import make_test, Tensor, Dtype, Trait
+from .conv import conv
+
+import numpy as np
+
+def conv_integer( 
+    X,
+    W,
+    x_zero_point=None,
+    w_zero_point=None,
+    auto_pad=None,
+    dilations=None,
+    group=None,
+    kernel_shape=None,
+    pads=None,
+    strides=None,
+):
+    if len(X.shape) < 3:
+        raise ValueError(
+            f"X must have at least 3 dimensions but its shape is {X.shape}."
+        )
+    X = X.astype(np.int32)
+    if x_zero_point:
+        X -= x_zero_point
+    W = W.astype(np.int32)
+    if w_zero_point:
+        W -= w_zero_point
+    return (
+        conv(
+            X, W, None, auto_pad, dilations, group, kernel_shape, pads, strides
+        ).astype(np.int32),
+    )
+
+
+class Conv_integer(RunAll):
+
+    @staticmethod
+    def export_without_padding() -> None:
+        x = (
+            np.array([2, 3, 4, 5, 6, 7, 8, 9, 10])
+            .astype(np.uint8)
+            .reshape((1, 1, 3, 3))
+        )
+        x_zero_point = (
+            np.array([1])
+            .astype(np.uint8)
+            .reshape((1, 1, 1, 1))
+        )
+        w = np.array([1, 1, 1, 1]).astype(np.uint8).reshape((1, 1, 2, 2))
+
+
+        y = conv_integer(x,w, x_zero_point)
+        y = np.array(y[0])
+
+        x = Tensor(Dtype.I8, x.shape, x.flatten())
+        x_zero_point = Tensor(Dtype.I8, x_zero_point.shape, x_zero_point.flatten())
+        w = Tensor(Dtype.I8, w.shape, w.flatten())
+        y = Tensor(Dtype.U32, y.shape, y.flatten())
+
+        name = "conv_interger_no_padding"
+        func_sig = "NNTrait::conv_integer("
+        func_sig += "@input_0,"
+        func_sig += "@input_1,"
+        func_sig += "Option::Some(@input_2),"
+        func_sig += "Option::None," 
+        func_sig += "Option::None," 
+        func_sig += "Option::None,"
+        func_sig += "Option::None,"
+        func_sig += "Option::None,"
+        func_sig += "Option::None,"
+        func_sig += "Option::None)"
+        make_test(
+            [x, w, x_zero_point], y, func_sig, name, Trait.NN)  
+
+    @staticmethod
+    def export_with_padding() -> None:
+        x = (
+            np.array([2, 3, 4, 5, 6, 7, 8, 9, 10])
+            .astype(np.uint8)
+            .reshape((1, 1, 3, 3))
+        )
+        x_zero_point = (
+            np.array([1])
+            .astype(np.uint8)
+            .reshape((1, 1, 1, 1))
+        )
+        w = np.array([1, 1, 1, 1]).astype(np.uint8).reshape((1, 1, 2, 2))
+
+        y = conv_integer(x,w, x_zero_point=x_zero_point, pads=[1, 1, 1, 1])
+        y = np.array(y[0])
+
+        x = Tensor(Dtype.I8, x.shape, x.flatten())
+        x_zero_point = Tensor(Dtype.I8, x_zero_point.shape, x_zero_point.flatten())
+        w = Tensor(Dtype.I8, w.shape, w.flatten())
+        y = Tensor(Dtype.U32, y.shape, y.flatten())
+
+        name = "conv_interger_with_padding"
+        func_sig = "NNTrait::conv_integer("
+        func_sig += "@input_0,"
+        func_sig += "@input_1,"
+        func_sig += "Option::Some(@input_2),"
+        func_sig += "Option::None," 
+        func_sig += "Option::None," 
+        func_sig += "Option::None,"
+        func_sig += "Option::None,"
+        func_sig += "Option::None,"
+        func_sig += "Option::Some(array![1, 1, 1, 1].span()),"
+        func_sig += "Option::None)"
+        make_test(
+            [x, w, x_zero_point], y, func_sig, name, Trait.NN)  
+
+
+
+