File: qconv_test.py

package info (click to toggle)
pytorch-cuda 2.6.0%2Bdfsg-7
  • links: PTS, VCS
  • area: contrib
  • in suites: forky, sid, trixie
  • size: 161,620 kB
  • sloc: python: 1,278,832; cpp: 900,322; ansic: 82,710; asm: 7,754; java: 3,363; sh: 2,811; javascript: 2,443; makefile: 597; ruby: 195; xml: 84; objc: 68
file content (83 lines) | stat: -rw-r--r-- 2,807 bytes parent folder | download | duplicates (3)
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
from pt import configs

import operator_benchmark as op_bench

import torch
import torch.ao.nn.quantized as nnq


"""
Microbenchmarks for qConv operators.
"""


class QConv1dBenchmark(op_bench.TorchBenchmarkBase):
    # def init(self, N, IC, OC, L, G, kernel, stride, pad):
    def init(self, IC, OC, kernel, stride, N, L, device):
        G = 1
        pad = 0
        self.scale = 1.0 / 255
        self.zero_point = 0
        X = torch.randn(N, IC, L, dtype=torch.float32)
        qX = torch.quantize_per_tensor(
            X, scale=self.scale, zero_point=self.zero_point, dtype=torch.quint8
        )
        # Convert the tensor to NHWC format
        W = torch.randn(OC, IC // G, kernel, dtype=torch.float32)
        self.qW = torch.quantize_per_tensor(
            W, scale=self.scale, zero_point=0, dtype=torch.qint8
        )

        self.inputs = {"input": qX}

        self.qconv1d = nnq.Conv1d(IC, OC, kernel, stride=stride, padding=pad, groups=G)
        self.qconv1d.set_weight_bias(self.qW, None)
        self.qconv1d.scale = torch.tensor(self.scale, dtype=torch.double)
        self.qconv1d.zero_point = torch.tensor(self.zero_point, dtype=torch.int)
        self.set_module_name("QConv1d")

    def forward(self, input):
        return self.qconv1d(input)


class QConv2dBenchmark(op_bench.TorchBenchmarkBase):
    # def init(self, N, IC, OC, H, W, G, kernel, stride, pad):
    def init(self, IC, OC, kernel, stride, N, H, W, G, pad, device):
        # super().init(N, IC, OC, (H, W), G, (kernel, kernel), stride, pad)

        self.scale = 1.0 / 255
        self.zero_point = 0
        X = torch.randn(N, IC, H, W, dtype=torch.float32)
        qX = torch.quantize_per_tensor(
            X, scale=self.scale, zero_point=self.zero_point, dtype=torch.quint8
        )
        # Convert the tensor to NHWC format
        W = torch.randn(OC, IC // G, kernel, kernel, dtype=torch.float32)
        self.qW = torch.quantize_per_tensor(
            W, scale=self.scale, zero_point=0, dtype=torch.qint8
        )

        self.inputs = {"input": qX}

        self.qconv2d = nnq.Conv2d(IC, OC, kernel, stride=stride, padding=pad, groups=G)
        self.qconv2d.set_weight_bias(self.qW, None)
        self.qconv2d.scale = torch.tensor(self.scale, dtype=torch.double)
        self.qconv2d.zero_point = torch.tensor(self.zero_point, dtype=torch.int)
        self.set_module_name("QConv2d")

    def forward(self, input):
        return self.qconv2d(input)


op_bench.generate_pt_test(
    configs.remove_cuda(configs.conv_1d_configs_short + configs.conv_1d_configs_long),
    QConv1dBenchmark,
)
op_bench.generate_pt_test(
    configs.remove_cuda(configs.conv_2d_configs_short + configs.conv_2d_configs_long),
    QConv2dBenchmark,
)


if __name__ == "__main__":
    op_bench.benchmark_runner.main()