Description

Use dropout in CuDNN

Tested on p3.2x (V100). Test case:

import mxnet as mx
a = mx.nd.ones((10, 200, 300, 500), ctx=mx.gpu(0))
a.attach_grad()
mx.autograd.set_recording(True)
%timeit mx.nd.Dropout(a, 0.5, mode='always').wait_to_read()

Checklist

Essentials

Please feel free to remove inapplicable items for your PR.

• Changes are complete (i.e. I finished coding on this PR)
• All changes have test coverage:
• Unit tests are added for small changes to verify correctness (e.g. adding a new operator)
• Nightly tests are added for complicated/long-running ones (e.g. changing distributed kvstore)
• Build tests will be added for build configuration changes (e.g. adding a new build option with NCCL)
• Code is well-documented:
• For new C++ functions in header files, their functionalities and arguments are documented.
• Check the API doc at http://mxnet-ci-doc.s3-accelerate.dualstack.amazonaws.com/PR-$PR_ID/$BUILD_ID/index.html
• To the my best knowledge, examples are either not affected by this change, or have been fixed to be compatible with this change

Changes

• CuDNN Dropout.
• skip resource request for parallel random when using cudnn forward.
• extend cpptest to support stateful operators, resourcerequestex, ResourceRequest::kCuDNNDropoutDesc.
• fix bug where the backward pass on an inference-mode dropout forward didn't act as identity.
• add kCuDNNDropoutDesc resource for sharing dropout state across operators.
• replace dropout with identity when it's not dropping anything.

Comments

• cudnnSetDropoutDescriptor is an expensive call due to initialization on each of the stream multiprocessor on a GPU. Since cudnn 7, cudnnRestoreDropoutDescriptor becomes available so that the initialized space can be cached. This descriptor is currently used in both Dropout op and RNN op. We need a mechanism for caching this so that initialization on each stream happens only once, as the same desc can be shared among operators.