xref: /aosp_15_r20/external/pytorch/test/test_functional_autograd_benchmark.py (revision da0073e96a02ea20f0ac840b70461e3646d07c45) 
1# Owner(s): ["module: autograd"]
2
3import os
4
5import subprocess
6import tempfile
7import unittest
8
9from torch.testing._internal.common_utils import (
10    IS_WINDOWS,
11    run_tests,
12    slowTest,
13    TestCase,
14)
15
16PYTORCH_COLLECT_COVERAGE = bool(os.environ.get("PYTORCH_COLLECT_COVERAGE"))
17
18
19# This is a very simple smoke test for the functional autograd benchmarking script.
20class TestFunctionalAutogradBenchmark(TestCase):
21    def _test_runner(self, model, disable_gpu=False):
22        # Note about windows:
23        # The temporary file is exclusively open by this process and the child process
24        # is not allowed to open it again. As this is a simple smoke test, we choose for now
25        # not to run this on windows and keep the code here simple.
26        with tempfile.NamedTemporaryFile() as out_file:
27            cmd = [
28                "python3",
29                "../benchmarks/functional_autograd_benchmark/functional_autograd_benchmark.py",
30            ]
31            # Only run the warmup
32            cmd += ["--num-iters", "0"]
33            # Only run the vjp task (fastest one)
34            cmd += ["--task-filter", "vjp"]
35            # Only run the specified model
36            cmd += ["--model-filter", model]
37            # Output file
38            cmd += ["--output", out_file.name]
39            if disable_gpu:
40                cmd += ["--gpu", "-1"]
41
42            res = subprocess.run(cmd)
43
44            self.assertTrue(res.returncode == 0)
45            # Check that something was written to the file
46            out_file.seek(0, os.SEEK_END)
47            self.assertTrue(out_file.tell() > 0)
48
49    @unittest.skipIf(
50        IS_WINDOWS,
51        "NamedTemporaryFile on windows does not have all the features we need.",
52    )
53    @unittest.skipIf(
54        PYTORCH_COLLECT_COVERAGE,
55        "Can deadlocks with gcov, see https://github.com/pytorch/pytorch/issues/49656",
56    )
57    def test_fast_tasks(self):
58        fast_tasks = [
59            "resnet18",
60            "ppl_simple_reg",
61            "ppl_robust_reg",
62            "wav2letter",
63            "transformer",
64            "multiheadattn",
65        ]
66
67        for task in fast_tasks:
68            self._test_runner(task)
69
70    @slowTest
71    @unittest.skipIf(
72        IS_WINDOWS,
73        "NamedTemporaryFile on windows does not have all the features we need.",
74    )
75    def test_slow_tasks(self):
76        slow_tasks = ["fcn_resnet", "detr"]
77        # deepspeech is voluntarily excluded as it takes too long to run without
78        # proper tuning of the number of threads it should use.
79
80        for task in slow_tasks:
81            # Disable GPU for slow test as the CI GPU don't have enough memory
82            self._test_runner(task, disable_gpu=True)
83
84
85if __name__ == "__main__":
86    run_tests()
87