-
Notifications
You must be signed in to change notification settings - Fork 3.5k
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
implement conv3d op #4400
implement conv3d op #4400
Conversation
python/tvm/contrib/cudnn.py
Outdated
filter_d, | ||
filter_h, | ||
filter_w): | ||
"""Get weight shape for a 2D convolution |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
3D
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
will fix next commit
python/tvm/contrib/cudnn.py
Outdated
filter height | ||
filter_w: int | ||
filter width | ||
|
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
filter_d
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
will fix next commit
python/tvm/contrib/cudnn.py
Outdated
0: CUDNN_CONVOLUTION | ||
1: CUDNN_CROSS_CORRELATION | ||
tensor_format: int | ||
0: CUDNN_TENSOR_NCHW |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Is this value compatible with 3D convolution?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Yes. cudnn doesn't have specific tensor format for 3D and up dimentions convolution, the leading format is same "NCHW" and the additional dimentions would be append to the tail for 'channel first' format. Please see below citation from nvidia DL SDK
format
Input.Type of the filter layout format. If this input is set to CUDNN_TENSOR_NCHW, which is one of the enumerant values allowed by cudnnTensorFormat_t descriptor, then the layout of the filter is as follows:
For N=4, a 4D filter descriptor, the filter layout is in the form of KCRS:
K represents the number of output feature maps
C is the number of input feature maps
R is the number of rows per filter
S is the number of columns per filter
For N=3, a 3D filter descriptor, the number S (number of columns per filter) is omitted.
For N=5 and greater, the layout of the higher dimensions immediately follow RS.
On the other hand, if this input is set to CUDNN_TENSOR_NHWC, then the layout of the filter is as follows:
For N=4, a 4D filter descriptor, the filter layout is in the form of KRSC.
For N=3, a 3D filter descriptor, the number S (number of columns per filter) is omitted and the layout of C immediately follows R.
For N=5 and greater, the layout of the higher dimensions are inserted between S and C. For more information, see cudnnTensorFormat_t.
- **weight**: (channels, in_channels, kernel_size[0], kernel_size[1]) | ||
- **out**: This depends on the `layout` parameter. Output is 4D array of shape | ||
(batch_size, channels, out_height, out_width) if `layout` is `NCHW`. | ||
|
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
need to update document above for 3D
@optima2005 thank you very much for working on this. Can we use cuDNN's ND API like cudnnSetConvolutionNdDescriptor for 2D convolution? It would be great if we could unify 2D and 3D implementation. |
@optima2005 can you add test cases to topi/tests/python too? |
@masahi, glad to contribute! Many thanks for the review! For "unify 2D and 3D" cudnn convolution, I am wondering whether it is better to keep separation. The APIs are 2 groups in cudnn lib(2D vs ND), I guess whether there would be some specific optimazation for 2D group. What do you think about that? For other comments, I have revised in the latest commit, please check again, thanks a lot! |
So I think we can go with the ND API for both 2D and 3D convolution. If you are worried about performance you can always do benchmarks. Can you first send a PR that refactors our cuDNN convolution to use the ND API? It will make reviewing this PR easier. |
83e05ac
to
828127e
Compare
@masahi, I have rebased to the master. Please go on to review. Thanks! |
Thanks @optima2005 this is merged. |
* implement conv3d op * add back missed conv2d_output_shape by mistake * fix typo and docs, add topi test * rebase to master and merge 2d/3d unification * use cudnn.conv_forward
* implement conv3d op * add back missed conv2d_output_shape by mistake * fix typo and docs, add topi test * rebase to master and merge 2d/3d unification * use cudnn.conv_forward
Hi @optima2005 , @masahi , |
This is a start attempt to implement #4009
The implemation in this PR is a very basic version of conv3d( NCDHW layout only). I am proposing this version to confirm that I am in the correct direction. And if so, I think I can go on or others can keep working it from this on.
I can run through the below testing on a x86 cpu and a nvidia gpu server, both linux OS.
TVM_FFI=ctypes python -m pytest -v tests/python/relay/test_op_level2.py -k test_conv3d_run
TVM_FFI=ctypes python -m pytest -v tests/python/contrib/test_cudnn.py -k test_conv3d