Repository: pytorch
Issue: lstm quantize issues
Labels: oncall: quantization,module: rnn,low priority,triaged,release notes: nn

### 🐛 Describe the bug

import numpy as np
import torch
import os

import torch
# Might not be necessary
os.environ.setdefault("KMP_DUPLICATE_LIB_OK", "True")

#import pytorch_GRU
#import quantGRUcell
device = torch.device("cpu")

class Net(torch.nn.Module):
    def __init__(self, seq_length):
        super(Net, self).__init__()

        self.hidden_size = 18
        self.input_size = 18

        self.seq_length = seq_length

        self.relu1 = torch.nn.Sigmoid()
        # Need to specify input sizes up front

        # batch_first specifies an input shape of (nBatches, nSeq, nFeatures),
        # otherwise this is (nSeq, nBatch, nFeatures)
        self.lstm = torch.nn.LSTM(input_size = self.input_size, hidden_size = self.hidden_size, batch_first = True)
        self.linear1 = torch.nn.Linear(self.hidden_size, self.hidden_size)
        self.linear2 = torch.nn.Linear(self.hidden_size, self.hidden_size)
        self.dropout = torch.nn.Dropout(0.5)        #self.squeeze = torch.squeeze
        self.linearOut = torch.nn.Linear(self.hidden_size, 20)
        self.sigmoidOut = torch.nn.Sigmoid()
        self.sqeeze1 = torch.Tensor.squeeze
        self.quant = torch.ao.quantization.QuantStub()
        self.dequant = torch.ao.quantization.DeQuantStub()
    def forward(self, x):
        print(type(x))
        x = self.quant(x)
        x, (h,c) = self.lstm(x)#, self.h0)


        # Get last output, x[:,l - 1,:], equivalent to (last) hidden state
        # Squeeze to remove length 1 dim
        x = h.reshape(1, h.shape[-1])

        x = self.dropout(x)
        x = self.linear2(x)
        x = self.relu1(x)
        x = self.linearOut(x)

        x = self.dequant(x)

        # Apply sigmoid either in the loss function, or in eval(...)
        return x
    def evaluate(self,x):
        return self.sigmoidOut(self.forward(x))
model_fp32 = Net(10)
model_fp32.eval()

model_fp32.qconfig = torch.ao.quantization.get_default_qat_qconfig('fbgemm')

model_fp32_prepared = torch.ao.quantization.prepare_qat(model_fp32.train(), inplace=True)

model_fp32_prepared.eval()
model_int8 = torch.ao.quantization.convert(model_fp32_prepared, [inplace=True)](url)
input_fp32 = torch.rand(1, 10, 18)
res = model_int8(input_fp32)

print(model_int8)
torch.onnx.export(
    model_int8,  # pytorch网络模型
    torch.rand(1, 10, 18),    # 随机的模拟输入
    "lstm_quant_fx.onnx",      # 导出的onnx文件位置
    export_params=True,     # 导出训练好的模型参数
    opset_version=17,       
    training=torch.onnx.TrainingMode.EVAL,  # 导出模型调整到推理状态，将dropout，BatchNorm等涉及的超参数固定
    input_names=["input_0"],    # 为静态网络图中的输入节点设置别名，在进行onnx推理时，将input_names字段与输入数据绑定
    output_names=["output_data"],  # 为输出节点设置别名
    # 如果不设置dynamic_axes，那么对于输入形状为[4, 3, 224, 224]，在以后使用onnx进行推理时也必须输入[4, 3, 224, 224]
    # 下面设置了输入的第0维是动态的，以后推理时batch_size的大小可以是其他动态值
    # dynamic_axes={
    #     # a dictionary to specify dynamic axes of input/output
    #     # each key must also be provided in input_names or output_names
    #     "input_0": {0: [2,3]},
    #     "output_0": {0: [2,3]}}
)

### Versions


Traceback (most recent call last):
  File "d:/笔记/量化/quant_lstm/torch_qat.py", line 46, in <module>
    torch.onnx.export(
  File "C:\Users\luohaifeng\Anaconda3\envs\tvm\lib\site-packages\torch\onnx\utils.py", line 551, in export
    _export(
  File "C:\Users\luohaifeng\Anaconda3\envs\tvm\lib\site-packages\torch\onnx\utils.py", line 1648, in _export
    graph, params_dict, torch_out = _model_to_graph(
  File "C:\Users\luohaifeng\Anaconda3\envs\tvm\lib\site-packages\torch\onnx\utils.py", line 1174, in _model_to_graph
    graph = _optimize_graph(
  File "C:\Users\luohaifeng\Anaconda3\envs\tvm\lib\site-packages\torch\onnx\utils.py", line 714, in _optimize_graph
    graph = _C._jit_pass_onnx(graph, operator_export_type)
  File "C:\Users\luohaifeng\Anaconda3\envs\tvm\lib\site-packages\torch\onnx\utils.py", line 1997, in _run_symbolic_function
    return symbolic_fn(graph_context, *inputs, **attrs)
  File "C:\Users\luohaifeng\Anaconda3\envs\tvm\lib\site-packages\torch\onnx\symbolic_opset10.py", line 892, in quantized_mul
    x, _, _, _ = symbolic_helper.dequantize_helper(g, x)
  File "C:\Users\luohaifeng\Anaconda3\envs\tvm\lib\site-packages\torch\onnx\symbolic_helper.py", line 1609, in dequantize_helper
    unpacked_qtensors = _unpack_quantized_tensor(qtensor)
  File "C:\Users\luohaifeng\Anaconda3\envs\tvm\lib\site-packages\torch\onnx\symbolic_helper.py", line 203, in _unpack_quantized_tensor
    raise errors.SymbolicValueError(
torch.onnx.errors.SymbolicValueError: ONNX symbolic expected the output of `%161 : Tensor = onnx::Sigmoid(%149), scope: lstm_qat_torch.Net::/torch.ao.nn.quantized.modules.rnn.LSTM::lstm/torch.ao.nn.quantizable.modules.rnn._LSTMLayer::layers.0/torch.ao.nn.quantizable.modules.rnn._LSTMSingleLayer::layer_fw/torch.ao.nn.quantizable.modules.rnn.LSTMCell::cell/torch.nn.modules.activation.Sigmoid::forget_gate # C:\Users\luohaifeng\Anaconda3\envs\tvm\lib\site-packages\torch\nn\modules\activation.py:301:0
` to be a quantized tensor. Is this likely due to missing support for quantized `onnx::Sigmoid`. Please create an issue on https://github.com/pytorch/pytorch/issues  [Caused by the value '161 defined in (%161 : Tensor = onnx::Sigmoid(%149), scope: lstm_qat_torch.Net::/torch.ao.nn.quantized.modules.rnn.LSTM::lstm/torch.ao.nn.quantizable.modules.rnn._LSTMLayer::layers.0/torch.ao.nn.quantizable.modules.rnn._LSTMSingleLayer::layer_fw/torch.ao.nn.quantizable.modules.rnn.LSTMCell::cell/torch.nn.modules.activation.Sigmoid::forget_gate # C:\Users\luohaifeng\Anaconda3\envs\tvm\lib\site-packages\torch\nn\modules\activation.py:301:0
)' (type 'Tensor') in the TorchScript graph. The containing node has kind 'onnx::Sigmoid'.]
    (node defined in C:\Users\luohaifeng\Anaconda3\envs\tvm\lib\site-packages\torch\nn\modules\activation.py(301): forward
C:\Users\luohaifeng\Anaconda3\envs\tvm\lib\site-packages\torch\nn\modules\module.py(1543): _slow_forward
C:\Users\luohaifeng\Anaconda3\envs\tvm\lib\site-packages\torch\nn\modules\module.py(1562): _call_impl
C:\Users\luohaifeng\Anaconda3\envs\tvm\lib\site-packages\torch\nn\modules\module.py(1553): _wrapped_call_impl
C:\Users\luohaifeng\Anaconda3\envs\tvm\lib\site-packages\torch\ao\nn\quantizable\modules\rnn.py(78): forward
C:\Users\luohaifeng\Anaconda3\envs\tvm\lib\site-packages\torch\nn\modules\module.py(1543): _slow_forward
C:\Users\luohaifeng\Anaconda3\envs\tvm\lib\site-packages\torch\nn\modules\module.py(1562): _call_impl
C:\Users\luohaifeng\Anaconda3\envs\tvm\lib\site-packages\torch\nn\modules\module.py(1553): _wrapped_call_impl
C:\Users\luohaifeng\Anaconda3\envs\tvm\lib\site-packages\torch\ao\nn\quantizable\modules\rnn.py(153): forward
C:\Users\luohaifeng\Anaconda3\envs\tvm\lib\site-packages\torch\nn\modules\module.py(1543): _slow_forward
C:\Users\luohaifeng\Anaconda3\envs\tvm\lib\site-packages\torch\nn\modules\module.py(1562): _call_impl
C:\Users\luohaifeng\Anaconda3\envs\tvm\lib\site-packages\torch\nn\modules\module.py(1553): _wrapped_call_impl
C:\Users\luohaifeng\Anaconda3\envs\tvm\lib\site-packages\torch\ao\nn\quantizable\modules\rnn.py(204): forward
C:\Users\luohaifeng\Anaconda3\envs\tvm\lib\site-packages\torch\nn\modules\module.py(1543): _slow_forward
C:\Users\luohaifeng\Anaconda3\envs\tvm\lib\site-packages\torch\nn\modules\module.py(1562): _call_impl
C:\Users\luohaifeng\Anaconda3\envs\tvm\lib\site-packages\torch\nn\modules\module.py(1553): _wrapped_call_impl
C:\Users\luohaifeng\Anaconda3\envs\tvm\lib\site-packages\torch\ao\nn\quantizable\modules\rnn.py(366): forward
C:\Users\luohaifeng\Anaconda3\envs\tvm\lib\site-packages\torch\nn\modules\module.py(1543): _slow_forward
C:\Users\luohaifeng\Anaconda3\envs\tvm\lib\site-packages\torch\nn\modules\module.py(1562): _call_impl
C:\Users\luohaifeng\Anaconda3\envs\tvm\lib\site-packages\torch\nn\modules\module.py(1553): _wrapped_call_impl
d:\笔记\量化\quant_lstm\lstm_qat_torch.py(39): forward
C:\Users\luohaifeng\Anaconda3\envs\tvm\lib\site-packages\torch\nn\modules\module.py(1543): _slow_forward
C:\Users\luohaifeng\Anaconda3\envs\tvm\lib\site-packages\torch\nn\modules\module.py(1562): _call_impl
C:\Users\luohaifeng\Anaconda3\envs\tvm\lib\site-packages\torch\nn\modules\module.py(1553): _wrapped_call_impl
C:\Users\luohaifeng\Anaconda3\envs\tvm\lib\site-packages\torch\jit\_trace.py(1275): trace_module
C:\Users\luohaifeng\Anaconda3\envs\tvm\lib\site-packages\torch\jit\_trace.py(695): _trace_impl
C:\Users\luohaifeng\Anaconda3\envs\tvm\lib\site-packages\torch\jit\_trace.py(1000): trace
C:\Users\luohaifeng\Anaconda3\envs\tvm\lib\site-packages\torch\onnx\utils.py(1124): _pre_trace_quant_model
C:\Users\luohaifeng\Anaconda3\envs\tvm\lib\site-packages\torch\onnx\utils.py(1169): _model_to_graph
C:\Users\luohaifeng\Anaconda3\envs\tvm\lib\site-packages\torch\onnx\utils.py(1648): _export
C:\Users\luohaifeng\Anaconda3\envs\tvm\lib\site-packages\torch\onnx\utils.py(551): export
d:/笔记/量化/quant_lstm/torch_qat.py(46): <module>
)

    Inputs:
        #0: 149 defined in (%149 : Tensor = onnx::Slice(%gates, %145, %148, %137), scope: lstm_qat_torch.Net::/torch.ao.nn.quantized.modules.rnn.LSTM::lstm/torch.ao.nn.quantizable.modules.rnn._LSTMLayer::layers.0/torch.ao.nn.quantizable.modules.rnn._LSTMSingleLayer::layer_fw/torch.ao.nn.quantizable.modules.rnn.LSTMCell::cell # C:\Users\luohaifeng\Anaconda3\envs\tvm\lib\site-packages\torch\ao\nn\quantizable\modules\rnn.py:75:0
    )  (type 'Tensor')
    Outputs:
        #0: 161 defined in (%161 : Tensor = onnx::Sigmoid(%149), scope: lstm_qat_torch.Net::/torch.ao.nn.quantized.modules.rnn.LSTM::lstm/torch.ao.nn.quantizable.modules.rnn._LSTMLayer::layers.0/torch.ao.nn.quantizable.modules.rnn._LSTMSingleLayer::layer_fw/torch.ao.nn.quantizable.modules.rnn.LSTMCell::cell/torch.nn.modules.activation.Sigmoid::forget_gate # C:\Users\luohaifeng\Anaconda3\envs\tvm\lib\site-packages\torch\nn\modules\activation.py:301:0
    )  (type 'Tensor')

cc @jerryzh168 @jianyuh @raghuramank100 @jamesr66a @vkuzo @jgong5 @Xia-Weiwen @leslie-fang-intel @msaroufim @mikaylagawarecki