mish is not support for cudnn for now

Signed-off-by: Adnios <[email protected]>
apache · Jun 3, 2021 · 6883541 · 6883541
1 parent 8ea2560
commit 6883541
Show file tree

Hide file tree

Showing 3 changed files with 12 additions and 8 deletions.
diff --git a/src/common/cuda/rtc/backward_functions-inl.h b/src/common/cuda/rtc/backward_functions-inl.h
@@ -53,9 +53,10 @@ backward_log_sigmoid(const DTypeGrad grad, const DType val) {
 template <typename DType, typename DTypeGrad>
 __device__ inline mixed_type<DTypeGrad, DType>
 backward_mish(const DTypeGrad grad, const DType val) {
-  const auto softrelu = op::log(1 + exp(val));
+  const mixed_type<DTypeGrad, DType> v = val;
+  const auto softrelu = op::log(1 + exp(v));
   const auto tanh = op::tanh(softrelu);
-  return grad * (tanh + val * sigmoid(val) * (1 - tanh * tanh));
+  return grad * (tanh + v * sigmoid(v) * (1 - tanh * tanh));
 }
 
 template <typename DType, typename DTypeGrad>

diff --git a/src/operator/nn/activation.cu b/src/operator/nn/activation.cu
@@ -56,13 +56,16 @@ void ActivationCompute<gpu>(const nnvm::NodeAttrs& attrs,
   const ActivationParam& param = nnvm::get<ActivationParam>(attrs.parsed);
   const int act_type = param.act_type;
 
-  // SoftReLU and kSoftSign are both not supported by CUDNN yet
+  // SoftReLU, kSoftSign and Mish are not supported by CUDNN yet
   if (act_type == activation::kSoftReLU) {
     ActivationForward<gpu, mshadow_op::softrelu, mshadow_op::softrelu_grad>(ctx,
       inputs[0], req[0], outputs[0]);
   } else if (act_type == activation::kSoftSign) {
     ActivationForward<gpu, mshadow_op::softsign, mshadow_op::softsign_grad>(ctx,
       inputs[0], req[0], outputs[0]);
+  } else if (act_type == activation::kMish) {
+    ActivationForward<gpu, mshadow_op::mish, mshadow_op::mish_grad>(ctx,
+      inputs[0], req[0], outputs[0]);
   } else {
     MSHADOW_REAL_TYPE_SWITCH(inputs[0].type_flag_, DType, {
       get_cudnn_op<DType>(param).Forward(ctx, inputs[0], req[0], outputs[0]);
@@ -84,10 +87,13 @@ void ActivationGradCompute<gpu>(const nnvm::NodeAttrs& attrs,
 
   bool do_memory_opt = dmlc::GetEnv("MXNET_MEMORY_OPT", 0);
 
-  // both SoftReLU and SoftSign not supported by CUDNN yet
+  // SoftReLU, SoftSign and Mish not supported by CUDNN yet
   if (act_type == activation::kSoftReLU) {
     ActivationBackward<gpu, mshadow_op::softrelu, mshadow_op::softrelu_grad>(
       ctx, inputs.at(0), inputs.at(1), req[0], outputs[0]);
+  } else if (act_type == activation::kMish) {
+    ActivationBackward<gpu, mshadow_op::mish, mshadow_op::mish_grad>(
+      ctx, inputs.at(0), inputs.at(1), req[0], outputs[0]);
   } else if (act_type == activation::kSoftSign) {
     if (do_memory_opt) {
       ActivationBackward<gpu, mshadow_op::softsign, mshadow_op::softsign_grad>(
@@ -118,9 +124,6 @@ void ActivationGradCompute<gpu>(const nnvm::NodeAttrs& attrs,
       } else if (act_type == activation::kLogSigmoid) {
         ActivationBackward<gpu, mshadow_op::log_sigmoid, mshadow_op::log_sigmoid_grad>(
           ctx, inputs.at(0), inputs.at(1), req[0], outputs[0]);
-      } else if (act_type == activation::kMish) {
-        ActivationBackward<gpu, mshadow_op::mish, mshadow_op::mish_grad>(
-          ctx, inputs.at(0), inputs.at(1), req[0], outputs[0]);
       } else {
         LOG(FATAL) << "unknown activation type";
       }

diff --git a/tests/python/unittest/test_operator.py b/tests/python/unittest/test_operator.py
@@ -710,7 +710,7 @@ def fmish_grad(a):
     check_symbolic_forward(y, [xa], [ya])
     check_symbolic_backward(y, [xa], [np.ones(shape)], [ya_grad])
 
-    A = mx.np.zeros(shape)
+    A = mx.np.random.uniform(low=-1.0, high=1.0, size=shape)
     A.attach_grad()
     np_out = fmish(A.asnumpy())
     with mx.autograd.record():