1. pass in lambd as scalar for CPU/CUDA_apply*; 2. removed tests for hardshrink at test_legacy_nn

weiyangfb · weiyangfb · commit 8a306d29cd4d · 2018-06-05T14:42:18.000-07:00
diff --git a/aten/src/ATen/native/Activation.cpp b/aten/src/ATen/native/Activation.cpp
@@ -34,36 +34,36 @@ Tensor & rrelu_(Tensor & self, Scalar lower, Scalar upper, bool training, Genera
 }
 
 Tensor hardshrink_cpu(const Tensor & self, Scalar lambd) {
-  auto lambd_tensor = at::zeros_like(self).fill_(lambd.toTensor());
-  auto out_tensor = self.clone();
+  auto lambd_tensor = lambd.toTensor().toType(self.type().scalarType()).toBackend(self.is_cuda() ? Backend::CUDA : Backend::CPU);
+  auto out_tensor = at::zeros_like(self);
   AT_DISPATCH_FLOATING_TYPES(self.type(), "hardshrink_cpu", [&] {
+    scalar_t* lambd_tensor_d = lambd_tensor.data<scalar_t>();
     at::CPU_tensor_apply2<scalar_t, scalar_t>(
-        out_tensor,
-        lambd_tensor,
-        [](scalar_t& out_tensor_val,
-           scalar_t& lambd_tensor_val) {
-             if (out_tensor_val >= -lambd_tensor_val && out_tensor_val <= lambd_tensor_val) {
-               out_tensor_val = convert<scalar_t, double>(0.0);
-             }
+      self,
+      out_tensor,
+      [lambd_tensor_d](
+        scalar_t& self_val,
+        scalar_t& out_tensor_val) {
+          out_tensor_val = (self_val >= -*lambd_tensor_d && self_val <= *lambd_tensor_d) ? convert<scalar_t, double>(0.0) : self_val;
     });
   });
   return out_tensor;
 }
 
 Tensor hardshrink_backward_cpu(const Tensor & grad, const Tensor & self, Scalar lambd) {
-  auto lambd_tensor = at::zeros_like(self).fill_(lambd);
-  auto out_tensor = grad.clone();
+  auto lambd_tensor = lambd.toTensor().toType(self.type().scalarType()).toBackend(self.is_cuda() ? Backend::CUDA : Backend::CPU);
+  auto out_tensor = at::zeros_like(self);
   AT_DISPATCH_FLOATING_TYPES(self.type(), "hardshrink_backward_cpu", [&] {
+    scalar_t* lambd_tensor_d = lambd_tensor.data<scalar_t>();
     at::CPU_tensor_apply3<scalar_t, scalar_t, scalar_t>(
-        out_tensor,
-        lambd_tensor,
-        self,
-        [](scalar_t& out_tensor_val,
-           scalar_t& lambd_tensor_val,
-           scalar_t& self_val) {
-             if (self_val >= -lambd_tensor_val && self_val <= lambd_tensor_val) {
-               out_tensor_val = convert<scalar_t, double>(0.0);
-             }
+      self,
+      grad,
+      out_tensor,
+      [lambd_tensor_d](
+        scalar_t& self_val,
+        scalar_t& grad_val,
+        scalar_t& out_tensor_val) {
+          out_tensor_val = (self_val >= -*lambd_tensor_d && self_val <= *lambd_tensor_d) ? convert<scalar_t, double>(0.0) : grad_val;
     });
   });
   return out_tensor;
diff --git a/aten/src/ATen/native/cuda/Activation.cu b/aten/src/ATen/native/cuda/Activation.cu
@@ -9,60 +9,48 @@
 namespace at { namespace native {
 
 template <typename scalar_t>
-void hardshrink_cuda_kernel(const Tensor& self, Tensor& out_tensor, Tensor& lambd_tensor) {
-  at::cuda::CUDA_tensor_apply3<scalar_t, scalar_t, scalar_t>(
-      self,
-      out_tensor,
-      lambd_tensor,
-      [] __device__ (scalar_t& self_val,
-        scalar_t& out_tensor_val,
-        scalar_t& lambd_tensor_val) {
-           if (self_val >= -lambd_tensor_val && self_val <= lambd_tensor_val) {
-             out_tensor_val = ScalarConvert<double, scalar_t>::to(0.0);
-           }
-           else {
-             out_tensor_val = self_val;
-           }
+void hardshrink_cuda_kernel(const Tensor& self, Tensor& out_tensor, scalar_t* lambd) {
+  at::cuda::CUDA_tensor_apply2<scalar_t, scalar_t>(
+    self,
+    out_tensor,
+    [lambd] __device__ (
+      scalar_t& self_val,
+      scalar_t& out_tensor_val,
+      bool early_exit) {
+        out_tensor_val = (self_val >= -*lambd && self_val <= *lambd) ? ScalarConvert<double, scalar_t>::to(0.0) : self_val;
   });
 }
 
 template <typename scalar_t>
-void hardshrink_backward_cuda_kernel(Tensor& out_tensor, Tensor& lambd_tensor, const Tensor& self, const Tensor& grad) {
-  at::cuda::CUDA_tensor_apply4<scalar_t, scalar_t, scalar_t, scalar_t>(
-      out_tensor,
-      lambd_tensor,
-      self,
-      grad,
-      [] __device__ (scalar_t& out_tensor_val,
-        scalar_t& lambd_tensor_val,
-        scalar_t& self_val,
-        scalar_t& grad_val) {
-           if (self_val >= -lambd_tensor_val && self_val <= lambd_tensor_val) {
-             out_tensor_val = ScalarConvert<double, scalar_t>::to(0.0);
-           }
-           else {
-             out_tensor_val = grad_val;
-           }
+void hardshrink_backward_cuda_kernel(Tensor& out_tensor, scalar_t* lambd, const Tensor& self, const Tensor& grad) {
+  at::cuda::CUDA_tensor_apply3<scalar_t, scalar_t, scalar_t>(
+    self,
+    grad,
+    out_tensor,
+    [lambd] __device__ (
+      scalar_t& self_val,
+      scalar_t& grad_val,
+      scalar_t& out_tensor_val) {
+        out_tensor_val = (self_val >= -*lambd && self_val <= *lambd) ? ScalarConvert<double, scalar_t>::to(0.0) : grad_val;
   });
 }
 
 Tensor hardshrink_cuda(const Tensor & self, Scalar lambd) {
-  auto lambd_tensor = at::zeros_like(self).fill_(lambd);
+  auto lambd_tensor = lambd.toTensor().toType(self.type().scalarType()).toBackend(self.is_cuda() ? Backend::CUDA : Backend::CPU);
   auto out_tensor = at::zeros_like(self);
   AT_DISPATCH_FLOATING_TYPES_AND_HALF(self.type(), "hardshrink_cuda", [&] {
     using cuda_scalar_t = cuda::into_type<scalar_t>;
-    hardshrink_cuda_kernel<cuda_scalar_t>(self, out_tensor, lambd_tensor);
+    hardshrink_cuda_kernel<cuda_scalar_t>(self, out_tensor, lambd_tensor.data<cuda_scalar_t>());
   });
   return out_tensor;
 }
 
 Tensor hardshrink_backward_cuda(const Tensor & grad, const Tensor & self, Scalar lambd) {
-  auto lambd_tensor = at::zeros_like(self).fill_(lambd);
-  // auto lambd_tensor = lambd.toTensor().toType(grad.type().scalarType()).toBackend(grad.is_cuda() ? Backend::CUDA : Backend::CPU);
+  auto lambd_tensor = lambd.toTensor().toType(self.type().scalarType()).toBackend(self.is_cuda() ? Backend::CUDA : Backend::CPU);
   auto out_tensor = at::zeros_like(grad);
   AT_DISPATCH_FLOATING_TYPES_AND_HALF(self.type(), "hardshrink_backward_cuda", [&] {
     using cuda_scalar_t = cuda::into_type<scalar_t>;
-    hardshrink_backward_cuda_kernel<cuda_scalar_t>(out_tensor, lambd_tensor, self, grad);
+    hardshrink_backward_cuda_kernel<cuda_scalar_t>(out_tensor, lambd_tensor.data<cuda_scalar_t>(), self, grad);
   });
   return out_tensor;
 }
diff --git a/test/test_legacy_nn.py b/test/test_legacy_nn.py
@@ -618,13 +618,20 @@ def add_test(test):
         test_params = deepcopy(test_params)
         name = test_params.pop('module_name')
         name = name_remap.get(name, name)
+        # hardshrink is deprecated in nn
+        if name == "HardShrink":
+            continue
+
         test_params['constructor'] = getattr(nn, name)
         test = OldModuleTest(**test_params)
         add_test(test)
     for test_params in criterion_tests:
         test_params = deepcopy(test_params)
         name = test_params.pop('module_name')
         name = name_remap.get(name, name.replace('Loss', 'Criterion'))
+        # hardshrink is deprecated in nn
+        if name == "HardShrink":
+            continue
 
         # nn.NLLLoss2d is deprecated, but there is a NLLLoss test for 2d
         if name == 'ClassNLLCriterion' and 'desc' in test_params.keys() and '2d' in test_params['desc']:
diff --git a/test/test_nn.py b/test/test_nn.py
@@ -5534,6 +5534,7 @@ def add(test_name, fn):
     else:
         add(cuda_test_name, lambda self, test=test: test.test_cuda(self))
 
+
 def wrap_functional(fn, **kwargs):
     class FunctionalModule(nn.Module):
         def forward(self, *args):
