1import operator_benchmark as op_bench 2 3import torch 4from torch import nn 5 6 7""" 8Microbenchmarks for RNNs. 9""" 10 11qrnn_configs = op_bench.config_list( 12 attrs=[ 13 [1, 3, 1], 14 [5, 7, 4], 15 ], 16 # names: input_size, hidden_size, num_layers 17 attr_names=["I", "H", "NL"], 18 cross_product_configs={ 19 "B": (True,), # Bias always True for quantized 20 "D": (False, True), # Bidirectional 21 "dtype": (torch.qint8,), # Only qint8 dtype works for now 22 }, 23 tags=["short"], 24) 25 26 27class LSTMBenchmark(op_bench.TorchBenchmarkBase): 28 def init(self, I, H, NL, B, D, dtype): 29 sequence_len = 128 30 batch_size = 16 31 32 # The quantized.dynamic.LSTM has a bug. That's why we create a regular 33 # LSTM, and quantize it later. See issue #31192. 34 scale = 1.0 / 256 35 zero_point = 0 36 cell_nn = nn.LSTM( 37 input_size=I, 38 hidden_size=H, 39 num_layers=NL, 40 bias=B, 41 batch_first=False, 42 dropout=0.0, 43 bidirectional=D, 44 ) 45 cell_temp = nn.Sequential(cell_nn) 46 self.cell = torch.ao.quantization.quantize_dynamic( 47 cell_temp, {nn.LSTM, nn.Linear}, dtype=dtype 48 )[0] 49 50 x = torch.randn( 51 sequence_len, batch_size, I # sequence length # batch size 52 ) # Number of features in X 53 h = torch.randn( 54 NL * (D + 1), batch_size, H # layer_num * dir_num # batch size 55 ) # hidden size 56 c = torch.randn( 57 NL * (D + 1), batch_size, H # layer_num * dir_num # batch size 58 ) # hidden size 59 60 self.inputs = {"x": x, "h": h, "c": c} 61 self.set_module_name("QLSTM") 62 63 def forward(self, x, h, c): 64 return self.cell(x, (h, c))[0] 65 66 67op_bench.generate_pt_test(qrnn_configs, LSTMBenchmark) 68 69if __name__ == "__main__": 70 op_bench.benchmark_runner.main() 71