• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1import operator_benchmark as op_bench
2
3import torch
4
5
6# 2D pooling will have input matrix of rank 3 or 4
7qpool2d_long_configs = op_bench.config_list(
8    attrs=(
9        #  C    H    W   k       s       p
10        (1, 3, 3, (3, 3), (1, 1), (0, 0)),  # dummy        # noqa: E201,E241
11        (3, 64, 64, (3, 3), (2, 2), (1, 1)),  # dummy        # noqa: E201,E241
12        # VGG16 pools with original input shape: (-1, 3, 224, 224)
13        (64, 224, 224, (2, 2), (2, 2), (0, 0)),  # MaxPool2d-4  # noqa: E201
14        (256, 56, 56, (2, 2), (2, 2), (0, 0)),  # MaxPool2d-16 # noqa: E241
15    ),
16    attr_names=("C", "H", "W", "k", "s", "p"),  # Input layout  # Pooling parameters
17    cross_product_configs={
18        "N": (1, 4),
19        "contig": (False, True),
20        "dtype": (torch.quint8,),
21    },
22    tags=("long",),
23)
24
25qpool2d_short_configs = op_bench.config_list(
26    attrs=((1, 3, 3, (3, 3), (1, 1), (0, 0)),),  # dummy
27    attr_names=("C", "H", "W", "k", "s", "p"),  # Input layout  # Pooling parameters
28    cross_product_configs={
29        "N": (2,),
30        "contig": (True,),
31        "dtype": (torch.qint32, torch.qint8, torch.quint8),
32    },
33    tags=("short",),
34)
35
36qadaptive_avgpool2d_long_configs = op_bench.cross_product_configs(
37    input_size=(
38        # VGG16 pools with original input shape: (-1, 3, 224, 224)
39        (112, 112),  # MaxPool2d-9
40    ),
41    output_size=(
42        (448, 448),
43        # VGG16 pools with original input shape: (-1, 3, 224, 224)
44        (224, 224),  # MaxPool2d-4
45        (112, 112),  # MaxPool2d-9
46        (56, 56),  # MaxPool2d-16 # noqa: E201,E241
47        (14, 14),  # MaxPool2d-30 # noqa: E201,E241
48    ),
49    N=(1, 4),
50    C=(1, 3, 64, 128),
51    contig=(False, True),
52    dtype=(torch.quint8,),
53    tags=("long",),
54)
55
56qadaptive_avgpool2d_short_configs = op_bench.config_list(
57    attrs=((4, 3, (224, 224), (112, 112), True),),
58    attr_names=("N", "C", "input_size", "output_size", "contig"),
59    cross_product_configs={
60        "dtype": (torch.qint32, torch.qint8, torch.quint8),
61    },
62    tags=("short",),
63)
64
65
66class _QPool2dBenchmarkBase(op_bench.TorchBenchmarkBase):
67    def setup(self, N, C, H, W, dtype, contig):
68        # Input
69        if N == 0:
70            f_input = (torch.rand(C, H, W) - 0.5) * 256
71        else:
72            f_input = (torch.rand(N, C, H, W) - 0.5) * 256
73
74        scale = 1.0
75        zero_point = 0
76
77        # Quantize the tensor
78        self.q_input = torch.quantize_per_tensor(
79            f_input, scale=scale, zero_point=zero_point, dtype=dtype
80        )
81        if not contig:
82            # Permute into NHWC and back to make it non-contiguous
83            if N == 0:
84                self.q_input = self.q_input.permute(1, 2, 0).contiguous()
85                self.q_input = self.q_input.permute(2, 0, 1)
86            else:
87                self.q_input = self.q_input.permute(0, 2, 3, 1).contiguous()
88                self.q_input = self.q_input.permute(0, 3, 1, 2)
89
90        self.inputs = {"q_input": self.q_input}
91
92    def forward(self, q_input):
93        return self.pool_op(q_input)
94
95
96class QMaxPool2dBenchmark(_QPool2dBenchmarkBase):
97    def init(self, N, C, H, W, k, s, p, contig, dtype):
98        self.pool_op = torch.nn.MaxPool2d(
99            kernel_size=k,
100            stride=s,
101            padding=p,
102            dilation=(1, 1),
103            ceil_mode=False,
104            return_indices=False,
105        )
106        super().setup(N, C, H, W, dtype, contig)
107
108
109class QAvgPool2dBenchmark(_QPool2dBenchmarkBase):
110    def init(self, N, C, H, W, k, s, p, contig, dtype):
111        self.pool_op = torch.nn.AvgPool2d(
112            kernel_size=k, stride=s, padding=p, ceil_mode=False
113        )
114        super().setup(N, C, H, W, dtype, contig)
115
116
117class QAdaptiveAvgPool2dBenchmark(_QPool2dBenchmarkBase):
118    def init(self, N, C, input_size, output_size, contig, dtype):
119        self.pool_op = torch.nn.AdaptiveAvgPool2d(output_size=output_size)
120        super().setup(N, C, *input_size, dtype=dtype, contig=contig)
121
122
123op_bench.generate_pt_test(
124    qadaptive_avgpool2d_short_configs + qadaptive_avgpool2d_long_configs,
125    QAdaptiveAvgPool2dBenchmark,
126)
127op_bench.generate_pt_test(
128    qpool2d_short_configs + qpool2d_long_configs, QAvgPool2dBenchmark
129)
130op_bench.generate_pt_test(
131    qpool2d_short_configs + qpool2d_long_configs, QMaxPool2dBenchmark
132)
133
134
135if __name__ == "__main__":
136    op_bench.benchmark_runner.main()
137