How to quantize model have more than one output?

username: m.safari
cooked: <p>I want to quantize segmentation models like u2net(<a href="https://github.com/xuebinqin/U-2-Net/blob/master/model/u2net.py" rel="noopener nofollow ugc">u2net model github</a><br>
u2net structure is like below:</p>
<pre><code class="lang-auto">   ...
   ...   
   ...
    def forward(self,x):

        hx = x

        #stage 1
        hx1 = self.stage1(hx)
        hx = self.pool12(hx1)

        #stage 2
        hx2 = self.stage2(hx)
        hx = self.pool23(hx2)

        #stage 3
        hx3 = self.stage3(hx)
        hx = self.pool34(hx3)

        #stage 4
        hx4 = self.stage4(hx)
        hx = self.pool45(hx4)

        #stage 5
        hx5 = self.stage5(hx)
        hx = self.pool56(hx5)

        #stage 6
        hx6 = self.stage6(hx)
        hx6up = _upsample_like(hx6,hx5)

        #-------------------- decoder --------------------
        hx5d = self.stage5d(torch.cat((hx6up,hx5),1))
        hx5dup = _upsample_like(hx5d,hx4)

        hx4d = self.stage4d(torch.cat((hx5dup,hx4),1))
        hx4dup = _upsample_like(hx4d,hx3)

        hx3d = self.stage3d(torch.cat((hx4dup,hx3),1))
        hx3dup = _upsample_like(hx3d,hx2)

        hx2d = self.stage2d(torch.cat((hx3dup,hx2),1))
        hx2dup = _upsample_like(hx2d,hx1)

        hx1d = self.stage1d(torch.cat((hx2dup,hx1),1))


        #side output
        d1 = self.side1(hx1d)

        d2 = self.side2(hx2d)
        d2 = _upsample_like(d2,d1)

        d3 = self.side3(hx3d)
        d3 = _upsample_like(d3,d1)

        d4 = self.side4(hx4d)
        d4 = _upsample_like(d4,d1)

        d5 = self.side5(hx5d)
        d5 = _upsample_like(d5,d1)

        d6 = self.side6(hx6)
        d6 = _upsample_like(d6,d1)

        d0 = self.outconv(torch.cat((d1,d2,d3,d4,d5,d6),1))

        return F.sigmoid(d0), F.sigmoid(d1), F.sigmoid(d2), F.sigmoid(d3), F.sigmoid(d4), F.sigmoid(d5), F.sigmoid(d6)
</code></pre>
<p>has 7 output !!!<br>
I define Quantizedu2net class like below:</p>
<pre><code class="lang-auto">class Quantizedu2net(nn.Module):
    def __init__(self, model_fp32):
        super(Quantizedu2net, self).__init__()
        # QuantStub converts tensors from floating point to quantized.
        # This will only be used for inputs.
        self.quant = torch.quantization.QuantStub()
        # DeQuantStub converts tensors from quantized to floating point.
        # This will only be used for outputs.
        self.dequant = torch.quantization.DeQuantStub()
        # FP32 model
        self.model_fp32 = model_fp32

    def forward(self, x):
        # manually specify where tensors will be converted from floating
        # point to quantized in the quantized model
        x = self.quant(x)
        x = self.model_fp32(x)
        # manually specify where tensors will be converted from quantized
        # to floating point in the quantized model
        x = self.dequant(x)
        return x
</code></pre>
<p>but when I want to save a quantized model, I get error:</p>
<pre><code class="lang-auto">forward(__torch__.torch.nn.quantized.modules.DeQuantize self, Tensor Xq) -&gt; (Tensor):
Expected a value of type 'Tensor (inferred)' for argument 'Xq' but instead found type 'Tuple[Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor]'.
Inferred 'Xq' to be of type 'Tensor' because it was not annotated with an explicit type.
:
  File "/home/segmentation_u2net/func_u2_net_v1.py", line 1086
        # manually specify where tensors will be converted from quantized
        # to floating point in the quantized model
        x = self.dequant(x)
            ~~~~~~~~~~~~ &lt;--- HERE
        return x
</code></pre>
<p>can anyone help me?</p>

username: jerryzh168
cooked: <aside class="quote no-group" data-username="m.safari" data-post="1" data-topic="143137">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/m.safari/48/49969_2.png" class="avatar"> m.safari:</div>
<blockquote>
<pre><code class="lang-auto">            ~~~~~~~~~~~~ &lt;--- HERE
</code></pre>
</blockquote>
</aside>
<p>yeah you would need to modify the original model instead of wrapping it inside <code>self.model_fp32</code> here I think. Or if you don’t have access to the original model, you could unpack the output of <code>self.model_fp32(x)</code> and Add a dequant stub for the ones that needs to be dequantized.</p>