feat: deform conv 2D

chachaleo · Feb 27, 2024 · 0882868 · 0882868
1 parent 7d84a42
commit 0882868
Show file tree

Hide file tree

Showing 46 changed files with 2,322 additions and 26 deletions.
diff --git a/docs/SUMMARY.md b/docs/SUMMARY.md
@@ -172,10 +172,11 @@
     * [nn.gemm](framework/operators/neural-network/nn.gemm.md)
     * [nn.grid\_sample](framework/operators/neural-network/nn.grid\_sample.md)
     * [nn.col2im](framework/operators/neural-network/nn.col2im.md)
-    * [nn.conv_transpose](framework/operators/neural-network/nn.conv\_transpose.md)
+    * [nn.conv\_transpose](framework/operators/neural-network/nn.conv\_transpose.md)
     * [nn.conv](framework/operators/neural-network/nn.conv.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.deform\_conv](framework/operators/neural-network/nn.deform\_conv_.md)
   * [Machine Learning](framework/operators/machine-learning/README.md)
     * [Tree Ensemble Classifier](framework/operators/machine-learning/tree-ensemble-classifier/README.md)
       * [tree\_ensemble\_classifier.predict](framework/operators/machine-learning/tree-ensemble-classifier/tree\_ensemble\_classifier.predict.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: |
+|        [DeformConv](operators/neural-network/nn.deform\_conv_.md)             | :white\_check\_mark: |
 |                    [Sinh](operators/tensor/tensor.sinh.md)                    | :white\_check\_mark: |
 |                   [Asinh](operators/tensor/tensor.asinh.md)                   | :white\_check\_mark: |
 |                   [Atanh](operators/tensor/tensor.atanh.md)                   | :white\_check\_mark: |

diff --git a/...ework/operators/machine-learning/linear-classifier/linear_classifier.predict.md b/...ework/operators/machine-learning/linear-classifier/linear_classifier.predict.md
@@ -1,7 +1,7 @@
 # LinearClassifierTrait::predict
 
 ```rust 
-   fn predict(ref self: LinearClassifier<T>, X: Tensor<T>) -> Tensor<T>;
+   fn predict(classifier: LinearClassifier<T>, X: Tensor<T>) -> Tensor<T>;
 ```
 
 Linear Classifier. Performs the linear classification.
@@ -85,7 +85,7 @@ fn linear_classifier_helper(
 fn linear_classifier_multi_softmax() -> (Span<usize>, Tensor<FP16x16>) {
     let (mut classifier, X) = linear_classifier_helper(POST_TRANSFORM::SOFTMAX);
 
-    let (labels, mut scores) = LinearClassifierTrait::predict(ref classifier, X);
+    let (labels, mut scores) = LinearClassifierTrait::predict(classifier, X);
 
     (labels, scores)
 }

diff --git a/...amework/operators/machine-learning/linear-regressor/linear_regressor.predict.md b/...amework/operators/machine-learning/linear-regressor/linear_regressor.predict.md
@@ -1,14 +1,14 @@
 # LinearRegressorTrait::predict
 
 ```rust 
-   fn predict(ref self: LinearRegressor<T>, X: Tensor<T>) -> Tensor<T>;
+   fn predict(regressor: LinearRegressor<T>, X: Tensor<T>) -> Tensor<T>;
 ```
 
 Linear Regressor. Performs the generalized linear regression evaluation.
 
 ## Args
 
-* `self`: LinearRegressor<T> - A LinearRegressor object.
+* `regressor`: LinearRegressor<T> - A LinearRegressor object.
 * `X`:  Input 2D tensor.
 
 ## Returns
@@ -68,7 +68,7 @@ fn example_linear_regressor() -> Tensor<FP16x16> {
         post_transform
     };
 
-    let scores = LinearRegressorTrait::predict(ref regressor, X);
+    let scores = LinearRegressorTrait::predict(regressor, X);
 
     scores
 }
@@ -120,7 +120,7 @@ fn example_linear_regressor_2() -> Tensor<FP16x16> {
         post_transform
     };
 
-    let scores = LinearRegressorTrait::predict(ref regressor, X);
+    let scores = LinearRegressorTrait::predict(regressor, X);
 
     scores
 }

diff --git a/...s/machine-learning/tree-ensemble-classifier/tree_ensemble_classifier.predict.md b/...s/machine-learning/tree-ensemble-classifier/tree_ensemble_classifier.predict.md
@@ -1,7 +1,7 @@
 # TreeEnsembleClassifier::predict
 
 ```rust 
-   fn predict(ref self: TreeEnsembleClassifier<T>, X: Tensor<T>) -> (Span<usize>, MutMatrix::<T>);
+   fn predict(classifier: TreeEnsembleClassifier<T>, X: Tensor<T>) -> (Span<usize>, MutMatrix::<T>);
 ```
 
 Tree Ensemble classifier. Returns the top class for each of N inputs.
@@ -185,7 +185,7 @@ fn tree_ensemble_classifier_helper(
 fn test_tree_ensemble_classifier_multi_pt_softmax() -> (Span<usize>, MutMatrix::<FP16x16>) {
     let (mut classifier, X) = tree_ensemble_classifier_helper(POST_TRANSFORM::SOFTMAX);
 
-    let (labels, scores) = TreeEnsembleClassifierTrait::predict(ref classifier, X);
+    let (labels, scores) = TreeEnsembleClassifierTrait::predict(classifier, X);
     (labels, scores)
 }   
 

diff --git a/...ors/machine-learning/tree-ensemble-regressor/tree_ensemble_regressor.predict.md b/...ors/machine-learning/tree-ensemble-regressor/tree_ensemble_regressor.predict.md
@@ -1,7 +1,7 @@
 # TreeEnsembleRegressor::predict
 
 ```rust 
-   fn predict(ref self: TreeEnsembleRegressor<T>, X: Tensor<T>) -> (Span<usize>, MutMatrix::<T>);
+   fn predict(regressor: TreeEnsembleRegressor<T>, X: Tensor<T>) -> (Span<usize>, MutMatrix::<T>);
 ```
 
 Tree Ensemble regressor. Returns the regressed values for each input in N.
@@ -160,7 +160,7 @@ fn tree_ensemble_regressor_helper(
 
 fn test_tree_ensemble_regressor_SUM() -> MutMatrix::<FP16x16> {
     let (mut regressor, X) = tree_ensemble_regressor_helper(AGGREGATE_FUNCTION::SUM);
-    let mut res = TreeEnsembleRegressorTrait::predict(ref regressor, X);
+    let mut res = TreeEnsembleRegressorTrait::predict(regressor, X);
     res
 }
 >>> 

diff --git a/docs/framework/operators/neural-network/nn.col2im.md b/docs/framework/operators/neural-network/nn.col2im.md
@@ -1,4 +1,3 @@
-
 # NNTrait::col2im
 
 ```rust

diff --git a/docs/framework/operators/neural-network/nn.deform_conv.md b/docs/framework/operators/neural-network/nn.deform_conv.md
@@ -0,0 +1,152 @@
+# NNTrait::deform_conv
+
+```rust
+   fn deform_conv(
+       X: @Tensor<T>,
+       W: @Tensor<T>,
+       offset: @Tensor<T>,
+       B: Option<Span<T>>,
+       mask: Option<Tensor<T>>,
+       dilations: Option<Span<usize>>,
+       group: Option<usize>,
+       kernel_shape: Option<Span<usize>>,
+       offset_group: Option<usize>,
+       pads: Option<Span<usize>>,
+       strides: Option<Span<usize>>,
+   ) -> Tensor<T>
+```
+
+Performs deformable convolution as described in https://arxiv.org/abs/1703.06211 and https://arxiv.org/abs/1811.11168. This operator specification supports the 2-D case.
+
+## Args
+
+   X: @Tensor<T>,
+   W: @Tensor<T>,
+   offset: @Tensor<T>,
+   B: Option<Span<T>>,
+   mask: Option<Tensor<T>>,
+   dilations: Option<Span<usize>>,
+   group: Option<usize>,
+   kernel_shape: Option<Span<usize>>,
+   offset_group: Option<usize>,
+   pads: Option<Span<usize>>,
+   strides: Option<Span<usize>>,
+
+* `X`(`@Tensor<T>`) - Input data tensor. For 2D image data, it has shape (N, C, H, W) where N is the batch size, C is the number of input channels, and H and W are the height and width. 
+* `W`(`@Tensor<T>`) - Weight tensor that will be used in the convolutions. It has shape (oC, C/group, kH, kW), where oC is the number of output channels and kH and kW are the kernel height and width.
+* `offset`(`@Tensor<T>`) - Offset tensor denoting the offset for the sampling locations in the convolution kernel. It has shape (N, offset_group * kH * kW * 2, oH, oW) for 2D data
+* `B`(`Option<Span<T>>`) - Default is a tensor of zeros, optional 1D bias of length oC to be added to the convolution.
+* `mask`(`Option<Tensor<T>>`) -  Default is a tensor of ones, the mask tensor to be applied to each position in the convolution kernel. It has shape (N, offset_group * kH * kW, oH, oW) for 2D data.
+* `dilations`(`Option<Span<usize>>`) - Default is 1 along each axis, dilation value along each spatial axis of the kernel.
+* `group`(`usize`) - Default is 1, number of groups the input and output channels, C and oC, are divided into.
+* `kernel_shape`(`Option<Span<usize>>`) - Shape of the convolution kernel. If not present, it is inferred from the shape of input W.
+* `offset_group`(`Option<usize>`) - Default is 1, number of groups of offset. C must be divisible by offset_group.
+* `pads`(`Option<Span<usize>>`) - Default is 0 along each axis, padding for the beginning and end along each spatial axis. The values represent the number of pixels added to the beginning and end of the corresponding axis and can take any nonnegative value.
+* `strides`(`Option<Span<usize>>`) - Default is 1 along each axis, stride along each spatial axis.
+
+## Returns
+
+A `Tensor<T>` output tensor that contains the result of convolution.
+
+## Examples
+
+```rust
+fn example_deform_conv() -> Tensor<FP16x16> {
+    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(FP16x16 { mag: 0, sign: false });
+    data.append(FP16x16 { mag: 65536, sign: false });
+    data.append(FP16x16 { mag: 131072, sign: false });
+    data.append(FP16x16 { mag: 196608, sign: false });
+    data.append(FP16x16 { mag: 262144, sign: false });
+    data.append(FP16x16 { mag: 327680, sign: false });
+    data.append(FP16x16 { mag: 393216, sign: false });
+    data.append(FP16x16 { mag: 458752, sign: false });
+    data.append(FP16x16 { mag: 524288, sign: false });
+    let mut X = TensorTrait::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(FP16x16 { mag: 65536, sign: false });
+    data.append(FP16x16 { mag: 65536, sign: false });
+    data.append(FP16x16 { mag: 65536, sign: false });
+    data.append(FP16x16 { mag: 65536, sign: false });
+    let mut W = TensorTrait::new(shape.span(), data.span());
+
+    let mut shape = ArrayTrait::<usize>::new();
+    shape.append(1);
+    shape.append(8);
+    shape.append(2);
+    shape.append(2);
+
+    let mut data = ArrayTrait::new();
+    data.append(FP16x16 { mag: 32768, sign: false });
+    data.append(FP16x16 { mag: 0, sign: false });
+    data.append(FP16x16 { mag: 0, sign: false });
+    data.append(FP16x16 { mag: 0, sign: false });
+    data.append(FP16x16 { mag: 0, sign: false });
+    data.append(FP16x16 { mag: 0, sign: false });
+    data.append(FP16x16 { mag: 0, sign: false });
+    data.append(FP16x16 { mag: 0, sign: false });
+    data.append(FP16x16 { mag: 0, sign: false });
+    data.append(FP16x16 { mag: 0, sign: false });
+    data.append(FP16x16 { mag: 0, sign: false });
+    data.append(FP16x16 { mag: 0, sign: false });
+    data.append(FP16x16 { mag: 0, sign: false });
+    data.append(FP16x16 { mag: 0, sign: false });
+    data.append(FP16x16 { mag: 0, sign: false });
+    data.append(FP16x16 { mag: 0, sign: false });
+    data.append(FP16x16 { mag: 0, sign: false });
+    data.append(FP16x16 { mag: 0, sign: false });
+    data.append(FP16x16 { mag: 0, sign: false });
+    data.append(FP16x16 { mag: 0, sign: false });
+    data.append(FP16x16 { mag: 0, sign: false });
+    data.append(FP16x16 { mag: 6553, sign: true });
+    data.append(FP16x16 { mag: 0, sign: false });
+    data.append(FP16x16 { mag: 0, sign: false });
+    data.append(FP16x16 { mag: 0, sign: false });
+    data.append(FP16x16 { mag: 0, sign: false });
+    data.append(FP16x16 { mag: 0, sign: false });
+    data.append(FP16x16 { mag: 0, sign: false });
+    data.append(FP16x16 { mag: 0, sign: false });
+    data.append(FP16x16 { mag: 0, sign: false });
+    data.append(FP16x16 { mag: 0, sign: false });
+    data.append(FP16x16 { mag: 0, sign: false });
+    let mut offset = TensorTrait::new(shape.span(), data.span());
+
+
+    return NNTrait::deform_conv(
+        @X,
+        @W,
+        @offset,
+        Option::None,
+        Option::None,
+        Option::None,
+        Option::None,
+        Option::Some(array![2, 2].span()),
+        Option::None,
+        Option::Some(array![0, 0, 0, 0].span()),
+        Option::None,
+    );
+}
+
+>>> [
+        [
+            [
+                [9.5, 11.9],  
+                [20.0, 24.0],
+            ]
+        ]
+    ]
+
+````