Merge criterion and new criterion tests.

gchanan · gchanan · commit 1e3f728cf1df · 2020-08-28T11:07:04.000-07:00
They didn't actually do anything different.

[ghstack-poisoned]
diff --git a/test/test_nn.py b/test/test_nn.py
@@ -38,8 +38,8 @@
     get_function_arglist, load_tests, repeat_test_for_types, ALL_TENSORTYPES, \
     ALL_TENSORTYPES2, TemporaryFileName, TEST_WITH_UBSAN, IS_PPC
 from torch.testing._internal.common_cuda import TEST_CUDA, TEST_MULTIGPU, TEST_CUDNN, TEST_CUDNN_VERSION
-from torch.testing._internal.common_nn import NNTestCase, NewModuleTest, NewCriterionTest, \
-    module_tests, criterion_tests, new_criterion_tests, loss_reference_fns, \
+from torch.testing._internal.common_nn import NNTestCase, NewModuleTest, CriterionTest, \
+    module_tests, criterion_tests, loss_reference_fns, \
     ctcloss_reference, new_module_tests
 from torch.testing._internal.common_device_type import instantiate_device_type_tests, dtypes, \
     dtypesIfCUDA, skipCUDAIfNoCudnn, skipCUDAIfCudnnVersionLessThan, onlyCUDA, \
@@ -9010,10 +9010,10 @@ def reference_fn(i, p, m):
 
         add_test(test, decorator)
 
-for test_params in criterion_tests + new_criterion_tests:
+for test_params in criterion_tests:
     name = test_params.pop('module_name')
     test_params['constructor'] = getattr(nn, name)
-    test = NewCriterionTest(**test_params)
+    test = CriterionTest(**test_params)
     decorator = test_params.pop('decorator', None)
     add_test(test, decorator)
     if 'check_sum_reduction' in test_params:
@@ -9028,7 +9028,7 @@ def sum_reduction_constructor(*args, **kwargs):
             return sum_reduction_constructor
 
         test_params['constructor'] = gen_sum_reduction_constructor(test_params['constructor'])
-        test = NewCriterionTest(**test_params)
+        test = CriterionTest(**test_params)
         add_test(test, decorator)
 
 
diff --git a/torch/testing/_internal/common_nn.py b/torch/testing/_internal/common_nn.py
@@ -4086,9 +4086,6 @@ def padding3d_circular(input, pad):
         desc='margin',
         check_sum_reduction=True,
     ),
-]
-
-new_criterion_tests = [
     dict(
         module_name='BCEWithLogitsLoss',
         input_fn=lambda: torch.rand(15, 10).clamp_(1e-2, 1 - 1e-2),
@@ -4789,71 +4786,6 @@ def test_cuda(self, test_case):
                 raise
 
 
-class CriterionTest(TestBase):
-
-    _required_arg_names = TestBase._required_arg_names.union({'target'})
-
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
-        self.should_test_cuda = kwargs.get('test_cuda', True)
-        self.check_forward_only = kwargs.get('check_forward_only', True)
-
-    def _get_target(self):
-        return self._get_arg('target', True)
-
-    def __call__(self, test_case):
-        module = self.constructor(*self.constructor_args)
-        input = self._get_input()
-
-        # Check that these methods don't raise errors
-        module.__repr__()
-        str(module)
-
-        target = self._get_target()
-
-        if self.reference_fn is not None:
-            out = test_case._forward_criterion(module, input, target, extra_args=self.extra_args)
-            ref_args = (deepcopy(input), deepcopy(target)) + self.extra_args + (module,)
-            expected_out = self.reference_fn(*ref_args)
-            test_case.assertEqual(out, expected_out)
-
-        if self.check_forward_only:
-            return
-
-        test_case.check_criterion_jacobian(module, input, target)
-        self._do_extra_tests(test_case, module, input, target)
-
-    def test_cuda(self, test_case):
-        if not TEST_CUDA or not self.should_test_cuda:
-            raise unittest.SkipTest('Excluded from CUDA tests')
-        try:
-            cpu_input = self._get_input()
-            type_map = {
-                'torch.DoubleTensor': torch.cuda.FloatTensor,
-            }
-            gpu_input = to_gpu(cpu_input, type_map=type_map)
-
-            cpu_target = self._get_target()
-            gpu_target = to_gpu(cpu_target, type_map=type_map)
-
-            cpu_module = self.constructor(*self.constructor_args)
-            gpu_module = self.constructor(*self.constructor_args).float().cuda()
-
-            cpu_output = test_case._forward_criterion(cpu_module, cpu_input, cpu_target)
-            gpu_output = test_case._forward_criterion(gpu_module, gpu_input, gpu_target)
-            test_case.assertEqual(cpu_output, gpu_output, atol=4e-4, rtol=0)
-
-            gradOutput = torch.randn(())
-            cpu_gradInput = test_case._backward_criterion(cpu_module, cpu_input, cpu_target, gradOutput)
-            gpu_gradInput = test_case._backward_criterion(gpu_module, gpu_input, gpu_target, gradOutput)
-            test_case.assertEqual(cpu_gradInput, gpu_gradInput, atol=4e-4, rtol=0)
-        except NotImplementedError:
-            pass
-
-    def _do_extra_tests(self, test_case, module, input, target):
-        pass
-
-
 class InputVariableMixin(object):
     def _get_input(self):
         input = TestBase._get_input(self, False)
@@ -5024,17 +4956,43 @@ def constructor_args(self):
         return self._get_arg('constructor_args', False)
 
 
-class NewCriterionTest(InputVariableMixin, CriterionTest):
+class CriterionTest(InputVariableMixin, TestBase):
     # TODO: check that criterions don't ignore grad_output
 
+    _required_arg_names = TestBase._required_arg_names.union({'target'})
+
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
+        self.should_test_cuda = kwargs.get('test_cuda', True)
+        self.check_forward_only = kwargs.get('check_forward_only', True)
         self.check_gradgrad = kwargs.get('check_gradgrad', True)
         self.check_half = kwargs.get('check_half', True)
         self.check_bfloat16 = kwargs.get('check_bfloat16', False)
         self.convert_target = kwargs.get('convert_target', True)
         self.test_cpu = kwargs.get('test_cpu', True)
 
+    def __call__(self, test_case):
+        module = self.constructor(*self.constructor_args)
+        input = self._get_input()
+
+        # Check that these methods don't raise errors
+        module.__repr__()
+        str(module)
+
+        target = self._get_target()
+
+        if self.reference_fn is not None:
+            out = test_case._forward_criterion(module, input, target, extra_args=self.extra_args)
+            ref_args = (deepcopy(input), deepcopy(target)) + self.extra_args + (module,)
+            expected_out = self.reference_fn(*ref_args)
+            test_case.assertEqual(out, expected_out)
+
+        if self.check_forward_only:
+            return
+
+        test_case.check_criterion_jacobian(module, input, target)
+        self._do_extra_tests(test_case, module, input, target)
+
     def _do_extra_tests(self, test_case, module, input, target):
         if not self.check_gradgrad:
             return
