stu1130 opened a new issue #18646:
URL: https://github.com/apache/incubator-mxnet/issues/18646
## Description
```
import mxnet as mx
from mxnet import autograd, np, npx, gluon, init
from mxnet.gluon import nn
import time
npx.set_np()
data = mx.np.random.uniform(size=(32, 100, 100), ctx=mx.gpu())
label = mx.np.ones((32, 100, 100), ctx=mx.gpu())
net = nn.Sequential()
net.add(nn.BatchNorm(axis=-1))
net.initialize(init.Xavier(), ctx=mx.gpu())
loss = gluon.loss.L2Loss()
t = time.time()
for _ in range(5000):
with autograd.record():
l = loss(net(data), label)
l.backward()
mx.nd.waitall()
print('spent: {}s'.format(time.time() - t))
```
I got around 5 sec with axis=1 and 30 sec with axis=-1.
## Solution
Thanks @ptrendx pointed to point it out, cudnn 7.4
(https://docs.nvidia.com/deeplearning/sdk/cudnn-release-notes/rel_7xx.html#rel_741)
added a new cudnnBatchNormalization*Ex API that gives much better speed for
axis = -1
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
