[Relay] Add space_to_batch_nd and batch_to_space_nd operators #6477
Conversation
|
A very quick review : I think as part of this we should update the tflite frontend at the same time. |
|
cc @kevinthesun @yongwww @srkreddy1238, please help review. Thanks. |
python/tvm/relay/op/nn/nn.py
Outdated
|
|
||
|
|
||
| def space_to_batch_nd(data, block_shape, paddings): | ||
| r"""Divide spatial dimensions of the data into a grid of blocks and interleave them into batch dim. |
There was a problem hiding this comment.
ci failed due to too long of this line
There was a problem hiding this comment.
I have corrected this.
@u99127 Modified tflite frontend to use batch_to_space_nd and space_to_batch_nd operators. |
|
also cc @masahi @FrozenGene please help manage this PR |
include/tvm/topi/nn.h
Outdated
| } | ||
|
|
||
| // pad the input with paddings provided | ||
| padded_t = pad(data, pad_before_int32, pad_after_int32, make_const(DataType::Int(32), 0)); |
There was a problem hiding this comment.
How about adding one pad_value parameter whose default value is 0 for space_to_batch? I imagine if we use it in the quantized model, its pad value will be zero_point, not 0.https:/apache/incubator-tvm/blob/master/python/tvm/relay/frontend/tflite.py#L2423
There was a problem hiding this comment.
@FrozenGene Added pad_value parameter to space_to_batch_nd operator with default value 0.
include/tvm/topi/nn.h
Outdated
| tvm::Array<Integer> axis; | ||
| tvm::Array<PrimExpr> o_shape; | ||
|
|
||
| size_t M = block_shape.size(); |
There was a problem hiding this comment.
I think we shouldn't use capital letter here
include/tvm/topi/nn.h
Outdated
| size_t M = block_shape.size(); | ||
| size_t N = in_shape.size(); |
| import numpy as np | ||
|
|
||
|
|
||
| def space_to_batch_nd_python(data, block_shape, pad_before, pad_after): |
There was a problem hiding this comment.
Add pad_value = 0 like previous comment said
…h_nd and correct variable names
|
@FrozenGene Could you please spare some time to review this. |
|
thanks @BhushanIMG |
…#6477) * [Relay] Add space_to_batch_nd and batch_to_space_nd operators * Correct python-format errors * correct lint errors * tflite frontend to use batch_to_space and space_to_batch operators * Add new pad_value parameter with default value is 0 for space_to_batch_nd and correct variable names * Fix cppdocs - add documentation for pad_value
…#6477) * [Relay] Add space_to_batch_nd and batch_to_space_nd operators * Correct python-format errors * correct lint errors * tflite frontend to use batch_to_space and space_to_batch operators * Add new pad_value parameter with default value is 0 for space_to_batch_nd and correct variable names * Fix cppdocs - add documentation for pad_value
…#6477) * [Relay] Add space_to_batch_nd and batch_to_space_nd operators * Correct python-format errors * correct lint errors * tflite frontend to use batch_to_space and space_to_batch operators * Add new pad_value parameter with default value is 0 for space_to_batch_nd and correct variable names * Fix cppdocs - add documentation for pad_value
Hi,
I am Bhushan from Imagination Technologies.
This PR adds space_to_batch_nd and batch_to_space_nd operators to Relay and TOPI.
I have added the tests for these operators and Tensorflow frontend is also modified to use these operators.
Please review @tqchen @masahi @Huyuwei and let me know your comments.