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
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
|
import operator_benchmark as op_bench
import torch
import torch.nn as nn
from pt import configs
"""
Microbenchmarks for Conv1d and ConvTranspose1d operators.
"""
class Conv1dBenchmark(op_bench.TorchBenchmarkBase):
def init(self, IC, OC, kernel, stride, N, L, device):
self.inputs = {
"input": torch.rand(N, IC, L, device=device, requires_grad=self.auto_set())
}
self.conv1d = nn.Conv1d(IC, OC, kernel, stride=stride).to(device=device)
self.set_module_name('Conv1d')
def forward(self, input):
return self.conv1d(input)
class ConvTranspose1dBenchmark(op_bench.TorchBenchmarkBase):
def init(self, IC, OC, kernel, stride, N, L, device):
self.inputs = {
"input": torch.rand(N, IC, L, device=device)
}
self.convtranspose1d = nn.ConvTranspose1d(IC, OC, kernel, stride=stride).to(device=device)
self.set_module_name('ConvTranspose1d')
def forward(self, input):
return self.convtranspose1d(input)
op_bench.generate_pt_test(configs.conv_1d_configs_short + configs.conv_1d_configs_long,
Conv1dBenchmark)
op_bench.generate_pt_test(configs.conv_1d_configs_short + configs.conv_1d_configs_long,
ConvTranspose1dBenchmark)
"""
Microbenchmarks for Conv2d and ConvTranspose2d operators.
"""
class Conv2dBenchmark(op_bench.TorchBenchmarkBase):
def init(self, IC, OC, kernel, stride, N, H, W, G, pad, device):
self.inputs = {
"input": torch.rand(N, IC, H, W, device=device)
}
self.conv2d = nn.Conv2d(
IC, OC, kernel, stride=stride, groups=G, padding=pad).to(device=device)
self.set_module_name('Conv2d')
def forward(self, input):
return self.conv2d(input)
class ConvTranspose2dBenchmark(op_bench.TorchBenchmarkBase):
def init(self, IC, OC, kernel, stride, N, H, W, G, pad, device):
self.inputs = {
"input": torch.rand(N, IC, H, W, device=device)
}
self.convtranspose2d = nn.ConvTranspose2d(
IC, OC, kernel, stride=stride, groups=G, padding=pad).to(device=device)
self.set_module_name('ConvTranspose2d')
def forward(self, input):
return self.convtranspose2d(input)
op_bench.generate_pt_test(configs.conv_2d_configs_short + configs.conv_2d_configs_long,
Conv2dBenchmark)
op_bench.generate_pt_test(configs.conv_2d_configs_short + configs.conv_2d_configs_long,
ConvTranspose2dBenchmark)
"""
Microbenchmarks for Conv3d and ConvTranspose3d operators.
"""
class Conv3dBenchmark(op_bench.TorchBenchmarkBase):
def init(self, IC, OC, kernel, stride, N, D, H, W, device):
self.inputs = {
"input": torch.rand(N, IC, D, H, W, device=device)
}
self.conv3d = nn.Conv3d(IC, OC, kernel, stride=stride).to(device=device)
self.set_module_name('Conv3d')
def forward(self, input):
return self.conv3d(input)
class ConvTranspose3dBenchmark(op_bench.TorchBenchmarkBase):
def init(self, IC, OC, kernel, stride, N, D, H, W, device):
self.inputs = {
"input": torch.rand(N, IC, D, H, W, device=device)
}
self.convtranspose3d = nn.ConvTranspose3d(IC, OC, kernel, stride=stride).to(device=device)
self.set_module_name('ConvTranspose3d')
def forward(self, input):
return self.convtranspose3d(input)
op_bench.generate_pt_test(configs.conv_3d_configs_short, Conv3dBenchmark)
op_bench.generate_pt_test(configs.conv_3d_configs_short,
ConvTranspose3dBenchmark)
if __name__ == "__main__":
op_bench.benchmark_runner.main()
|
