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
|
from . import benchmark
from . import tensor_engine
class NormalizationBench(benchmark.Benchmark):
def __init__(self, mode, device, dtype, N, C, H, W):
super().__init__(mode, device, dtype)
self.N = N
self.C = C
self.H = H
self.W = W
self.data = self.nchw_rand(
[self.N, self.C, self.H, self.W],
device=device, dtype=dtype,
requires_grad=self.requires_grad,
)
self.running_mean = self.rand([self.C], device=device, dtype=dtype)
self.running_var = self.rand([self.C], device=device, dtype=dtype)
self.training = self.mode == "both"
def config(self):
return [self.N, self.C, self.H, self.W]
def memory_workload(self):
if self.mode == "fwd":
sol_count = 1 + 1
algorithmic_count = 2 + 1
else:
sol_count = (1 + 1) + (1 + 1)
algorithmic_count = (2 + 1) + (3 + 1)
buffer_size = self.N * self.C * self.H * self.W * 4
return {
"sol": buffer_size * sol_count,
"algorithmic": buffer_size * algorithmic_count,
}
@staticmethod
def default_configs():
return [[128, 32, 128, 128]]
class BatchNormBench(NormalizationBench):
def forward(self):
y = self.batch_norm(
self.data, self.running_mean, self.running_var, training=self.training
)
return y
@staticmethod
def module():
return "batchnorm"
class InstanceNormBench(NormalizationBench):
def forward(self):
y = self.instance_norm(self.data)
return y
@staticmethod
def module():
return "instance_norm"
def is_supported(self):
return tensor_engine.is_supported(self.instance_norm)
class LayerNormBench(NormalizationBench):
def forward(self):
y = self.layer_norm(self.data, [self.H, self.W])
return y
@staticmethod
def module():
return "layernorm"
benchmark.register_benchmark_class(BatchNormBench)
benchmark.register_benchmark_class(InstanceNormBench)
benchmark.register_benchmark_class(LayerNormBench)
|
