Handle flakines on autocast.

test/test_models.py

@@ -86,21 +86,50 @@ def _test_classification_model(self, name, input_shape, dev):
 
     def _test_segmentation_model(self, name, dev):
         set_rng_seed(0)
-        # passing num_class equal to a number other than 1000 helps in making the test
-        # more enforcing in nature
-        model = models.segmentation.__dict__[name](num_classes=2, pretrained_backbone=False)
+        # passing num_classes equal to a number other than 21 helps in making the test's
+        # expected file size smaller
+        model = models.segmentation.__dict__[name](num_classes=10, pretrained_backbone=False)
         model.eval().to(device=dev)
-        input_shape = (1, 3, 64, 64)
+        input_shape = (1, 3, 32, 32)
         # RNG always on CPU, to ensure x in cuda tests is bitwise identical to x in cpu tests
         x = torch.rand(input_shape).to(device=dev)
-        out = model(x)
-        self.assertExpected(out["out"].cpu(), prec=0.1, strip_suffix=f"_{dev}")
+        out = model(x)["out"]
+
+        def check_out(out):
+            prec = 0.01
+            strip_suffix = f"_{dev}"
+            try:
+                # We first try to assert the entire output if possible. This is not
+                # only the best way to assert results but also handles the cases
+                # where we need to create a new expected result.
+                self.assertExpected(out.cpu(), prec=prec, strip_suffix=strip_suffix)
+            except AssertionError:
+                # Unfortunately some segmentation models are flaky with autocast
+                # so instead of validating the probability scores, check that the class
+                # predictions match.
+                expected_file = self._get_expected_file(strip_suffix=strip_suffix)
+                expected = torch.load(expected_file)
+                self.assertEqual(out.argmax(dim=1), expected.argmax(dim=1), prec=prec)
+                return False  # Partial validation performed
+
+            return True  # Full validation performed
+
+        full_validation = check_out(out)
+
         self.check_jit_scriptable(model, (x,), unwrapper=script_model_unwrapper.get(name, None))
 
         if dev == torch.device("cuda"):
             with torch.cuda.amp.autocast():
-                out = model(x)
-                self.assertExpected(out["out"].cpu(), prec=0.1, strip_suffix=f"_{dev}")
+                out = model(x)["out"]
+                full_validation &= check_out(out)
+
+        if not full_validation:
+            msg = "The output of {} could only be partially validated. " \
+                  "This is likely due to unit-test flakiness, but you may " \
+                  "want to do additional manual checks if you made " \
+                  "significant changes to the codebase.".format(self._testMethodName)
+            warnings.warn(msg, RuntimeWarning)
+            raise unittest.SkipTest(msg)
 
     def _test_detection_model(self, name, dev):
         set_rng_seed(0)