Skip to content

Commit

Permalink
Merge pull request #586 from chachaleo/feat/DeformConv
Browse files Browse the repository at this point in the history 
feat: deform conv 2D
  • Loading branch information
@raphaelDkhn
raphaelDkhn authored Apr 22, 2024
2 parents 508c48f + a0ce46e commit 62d55e9
Show file tree
Hide file tree
Showing 38 changed files with 2,239 additions and 10 deletions.
3 changes: 2 additions & 1 deletion docs/SUMMARY.md
View file
Original file line number Diff line number Diff line change
Expand Up @@ -175,11 +175,12 @@
* [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.max\_pool](framework/operators/neural-network/nn.max\_pool.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)
Expand Down
1 change: 1 addition & 0 deletions docs/framework/compatibility.md
View file
Original file line number Diff line number Diff line change
Expand Up @@ -48,6 +48,7 @@ You can see below the list of current supported ONNX Operators:
| [ConvTranspose](operators/neural-network/nn.conv\_transpose_.md) | :white\_check\_mark: |
| [Conv](operators/neural-network/nn.conv.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: |
| [Asinh](operators/tensor/tensor.asinh.md) | :white\_check\_mark: |
| [Atanh](operators/tensor/tensor.atanh.md) | :white\_check\_mark: |
Expand Down
152 changes: 152 additions & 0 deletions docs/framework/operators/neural-network/nn.deform_conv.md
View file
Original file line number Diff line number Diff line change
@@ -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],
]
]
]

````
Loading

0 comments on commit 62d55e9

Please sign in to comment.