1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
|
# Owner(s): ["module: dynamo"]
import unittest
import torch
import torch._dynamo as torchdynamo
from torch.testing._internal.common_utils import run_tests, TEST_CUDA, TestCase
try:
import tabulate # noqa: F401 # type: ignore[import]
from torch.utils.benchmark.utils.compile import bench_all
HAS_TABULATE = True
except ImportError:
HAS_TABULATE = False
@unittest.skipIf(not TEST_CUDA, "CUDA unavailable")
@unittest.skipIf(not HAS_TABULATE, "tabulate not available")
class TestCompileBenchmarkUtil(TestCase):
def test_training_and_inference(self):
class ToyModel(torch.nn.Module):
def __init__(self) -> None:
super().__init__()
self.weight = torch.nn.Parameter(torch.Tensor(2, 2))
def forward(self, x):
return x * self.weight
torchdynamo.reset()
model = ToyModel().cuda()
inference_table = bench_all(model, torch.ones(1024, 2, 2).cuda(), 5)
self.assertTrue(
"Inference" in inference_table
and "Eager" in inference_table
and "-" in inference_table
)
training_table = bench_all(
model,
torch.ones(1024, 2, 2).cuda(),
5,
optimizer=torch.optim.SGD(model.parameters(), lr=0.01),
)
self.assertTrue(
"Train" in training_table
and "Eager" in training_table
and "-" in training_table
)
if __name__ == "__main__":
run_tests()
|
