[ROCM] added a cudnn switch of conv2d for rocm platform #31836
Conversation
|
Thanks for your contribution! |
python/paddle/nn/layer/conv.py
Outdated
| @@ -629,6 +633,11 @@ def __init__(self, | |||
| weight_attr=None, | |||
| bias_attr=None, | |||
| data_format="NCHW"): | |||
| use_cudnn = True | |||
There was a problem hiding this comment.
这里写成类似这样的形式
use_cudnn = True if (core.is_compiled_with_cuda() and
cudnn_version is not None) else False
python/paddle/nn/layer/conv.py
Outdated
| @@ -153,6 +156,7 @@ def _get_default_param_initializer(): | |||
| out_channels % in_channels == 0): | |||
| self._op_type = 'depthwise_conv2d' | |||
| self._use_cudnn = False | |||
| self._use_cudnn = self._use_cudnn and use_cudnn | |||
There was a problem hiding this comment.
加一行注释，解释一下这里这么修改的原因，例如
// add environment of ROCM_CONV2D_DISABLE_CUDNN to disable use_cudnn in Conv2D
// If use_cudnn is not defined, default value is True, then self._use_cudnn = self._use_cudnn
python/paddle/fluid/layers/nn.py
Outdated
| @@ -1498,6 +1498,9 @@ def conv2d(input, | |||
| conv2d = paddle.static.nn.conv2d(input=data, num_filters=2, filter_size=3, act="relu") | |||
| print(conv2d.shape) # [-1, 2, 30, 30] | |||
| """ | |||
| use_cudnn = False if ( | |||
| core.is_compiled_with_rocm() and | |||
| bool(os.getenv("ROCM_CONV2D_DISABLE_CUDNN"))) else use_cudnn | |||
There was a problem hiding this comment.
If ROCM_CONV2D_DISABLE_CUDNN is not a built-in variable in ROCM but created by Paddle, I suggest following the naming rule like other environments like "FLAGS_xxx_xxx".
There was a problem hiding this comment.
Plus, it's better to get the value by fluid.get_flags instead of os.getenv
ronny1996
force-pushed
the
rocm-dev
branch
2 times, most recently
from
e1cf412
to
ae72bcd
Compare
PR types
Others
PR changes
Others
Describe
added a cudnn switch of conv2d for rocm platform