1import operator_benchmark as op_bench 2 3import torch 4import torch.nn.functional as F 5 6 7"""Microbenchmarks for layernorm operator.""" 8 9layernorm_configs_short = op_bench.cross_product_configs( 10 dims=( 11 (1, 8, 16), 12 (8, 8, 16), 13 (32, 8, 16), 14 (64, 128, 56, 56), 15 ), 16 tags=["short"], 17) 18 19 20class LayerNormBenchmark(op_bench.TorchBenchmarkBase): 21 def init(self, dims): 22 input = (torch.rand(*dims) - 0.5) * 256 23 self.inputs = { 24 "input": input, 25 "weight": torch.rand(*input.size()[1:], dtype=torch.float), 26 "bias": torch.rand(*input.size()[1:], dtype=torch.float), 27 "eps": 1e-5, 28 } 29 30 def forward(self, input, weight, bias, eps: float): 31 return F.layer_norm(input, input.size()[1:], weight=weight, bias=bias, eps=eps) 32 33 34op_bench.generate_pt_test(layernorm_configs_short, LayerNormBenchmark) 35 36 37if __name__ == "__main__": 38 op_bench.benchmark_runner.main() 39