• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1import numpy
2from pt import configs
3
4import operator_benchmark as op_bench
5
6import torch
7import torch.ao.nn.quantized as nnq
8
9
10"""
11Microbenchmarks for qEmbeddingBag operators.
12"""
13
14
15class QEmbeddingBagBenchmark(op_bench.TorchBenchmarkBase):
16    def init(
17        self,
18        embeddingbags,
19        dim,
20        mode,
21        input_size,
22        offset,
23        sparse,
24        include_last_offset,
25        device,
26    ):
27        self.embedding = nnq.EmbeddingBag(
28            num_embeddings=embeddingbags,
29            embedding_dim=dim,
30            mode=mode,
31            include_last_offset=include_last_offset,
32        ).to(device=device)
33        numpy.random.seed((1 << 32) - 1)
34        self.input = torch.tensor(
35            numpy.random.randint(0, embeddingbags, input_size), device=device
36        ).long()
37        offset = torch.LongTensor([offset], device=device)
38        self.offset = torch.cat(
39            (offset, torch.tensor([self.input.size(0)], dtype=torch.long)), 0
40        )
41        self.inputs = {"input": self.input, "offset": self.offset}
42        self.set_module_name("qEmbeddingBag")
43
44    def forward(self, input, offset):
45        return self.embedding(input, offset)
46
47
48op_bench.generate_pt_test(configs.embeddingbag_short_configs, QEmbeddingBagBenchmark)
49
50if __name__ == "__main__":
51    op_bench.benchmark_runner.main()
52