• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1from pt import configs
2
3import operator_benchmark as op_bench
4
5import torch
6import torch.ao.nn.quantized as nnq
7
8
9"""
10Microbenchmarks for qConv operators.
11"""
12
13
14class QConv1dBenchmark(op_bench.TorchBenchmarkBase):
15    # def init(self, N, IC, OC, L, G, kernel, stride, pad):
16    def init(self, IC, OC, kernel, stride, N, L, device):
17        G = 1
18        pad = 0
19        self.scale = 1.0 / 255
20        self.zero_point = 0
21        X = torch.randn(N, IC, L, dtype=torch.float32)
22        qX = torch.quantize_per_tensor(
23            X, scale=self.scale, zero_point=self.zero_point, dtype=torch.quint8
24        )
25        # Convert the tensor to NHWC format
26        W = torch.randn(OC, IC // G, kernel, dtype=torch.float32)
27        self.qW = torch.quantize_per_tensor(
28            W, scale=self.scale, zero_point=0, dtype=torch.qint8
29        )
30
31        self.inputs = {"input": qX}
32
33        self.qconv1d = nnq.Conv1d(IC, OC, kernel, stride=stride, padding=pad, groups=G)
34        self.qconv1d.set_weight_bias(self.qW, None)
35        self.qconv1d.scale = torch.tensor(self.scale, dtype=torch.double)
36        self.qconv1d.zero_point = torch.tensor(self.zero_point, dtype=torch.int)
37        self.set_module_name("QConv1d")
38
39    def forward(self, input):
40        return self.qconv1d(input)
41
42
43class QConv2dBenchmark(op_bench.TorchBenchmarkBase):
44    # def init(self, N, IC, OC, H, W, G, kernel, stride, pad):
45    def init(self, IC, OC, kernel, stride, N, H, W, G, pad, device):
46        # super().init(N, IC, OC, (H, W), G, (kernel, kernel), stride, pad)
47
48        self.scale = 1.0 / 255
49        self.zero_point = 0
50        X = torch.randn(N, IC, H, W, dtype=torch.float32)
51        qX = torch.quantize_per_tensor(
52            X, scale=self.scale, zero_point=self.zero_point, dtype=torch.quint8
53        )
54        # Convert the tensor to NHWC format
55        W = torch.randn(OC, IC // G, kernel, kernel, dtype=torch.float32)
56        self.qW = torch.quantize_per_tensor(
57            W, scale=self.scale, zero_point=0, dtype=torch.qint8
58        )
59
60        self.inputs = {"input": qX}
61
62        self.qconv2d = nnq.Conv2d(IC, OC, kernel, stride=stride, padding=pad, groups=G)
63        self.qconv2d.set_weight_bias(self.qW, None)
64        self.qconv2d.scale = torch.tensor(self.scale, dtype=torch.double)
65        self.qconv2d.zero_point = torch.tensor(self.zero_point, dtype=torch.int)
66        self.set_module_name("QConv2d")
67
68    def forward(self, input):
69        return self.qconv2d(input)
70
71
72op_bench.generate_pt_test(
73    configs.remove_cuda(configs.conv_1d_configs_short + configs.conv_1d_configs_long),
74    QConv1dBenchmark,
75)
76op_bench.generate_pt_test(
77    configs.remove_cuda(configs.conv_2d_configs_short + configs.conv_2d_configs_long),
78    QConv2dBenchmark,
79)
80
81
82if __name__ == "__main__":
83    op_bench.benchmark_runner.main()
84