[Nightly] Packed params no longer returned via state_dict() method

masahi: <p>Hi, I’ve installed the nightly build 1.6.0.dev20200607 today and ran my scripts that exercises quantization and jit.</p>
<p>Until v1.5, I was able to get all packed params of a top level quantized module, say quantized resnet in torchvision, via state_dict() method. But now with nightly, I only get <code>quant.scale</code> and <code>quant.zero_point</code> from the same module.</p>
<p>I also noticed that a packed param is now an instance of <code>torch._C.ScriptObject</code>, instead of QTensor as was the case until v1.5</p>
<p>How do I get all parameters from quantized + jitted model now? Can you point me to github issues/PRs that introduced relevant changes?</p>
<p><a class="mention" href="/u/jerryzh168">@jerryzh168</a> <a class="mention" href="/u/raghuramank100">@raghuramank100</a></p>

jerryzh168: <p>yeah, <a href="https://github.com/pytorch/pytorch/pull/35923" rel="nofollow noopener">https://github.com/pytorch/pytorch/pull/35923</a> and <a href="https://github.com/pytorch/pytorch/pull/34140" rel="nofollow noopener">https://github.com/pytorch/pytorch/pull/34140</a> are relevant changes.</p>
<p>We are using TorchBind object for the packed params now.</p>

masahi: <p>Thanks, I’ll take a look. Seems like a big change.</p>

masahi: <p>Hi <a class="mention" href="/u/jerryzh168">@jerryzh168</a>,</p>
<p>Now that v1.6 is out, I came back to this issue. As I mentioned,  <code>state_dict()</code> method on traced quantized networks like qresnet from torchvision no longer returns all quantized parameters.</p>
<p>After some digging and thanks to the onnx export implementation below, I found that I can use <code>torch._C._jit_pass_lower_graph(graph, model._c)</code> to get at quantized paremeters I’ve been looking for. Is this the recommend way for third party pkg like TVM to get quantized parameters? Having to pass <code>model._c</code> seems like a very internal API…</p>
<aside class="onebox githubblob">
  <header class="source">
      <a href="https://github.com/pytorch/pytorch/blob/master/torch/onnx/utils.py#L357" target="_blank" rel="nofollow noopener">github.com</a>
  </header>
  <article class="onebox-body">
    <h4><a href="https://github.com/pytorch/pytorch/blob/master/torch/onnx/utils.py#L357" target="_blank" rel="nofollow noopener">pytorch/pytorch/blob/master/torch/onnx/utils.py#L357</a></h4>
<pre class="onebox"><code class="lang-py"><ol class="start lines" start="347" style="counter-reset: li-counter 346 ;">
<li>
</li><li>if isinstance(example_outputs, torch.Tensor):</li>
<li>    example_outputs = [example_outputs]</li>
<li>
</li><li>torch_out = None</li>
<li>if isinstance(model, torch.jit.ScriptModule):</li>
<li>    assert example_outputs is not None, "example_outputs must be provided when exporting a ScriptModule"</li>
<li>    try:</li>
<li>        graph = model.forward.graph</li>
<li>        torch._C._jit_pass_onnx_function_substitution(graph)</li>
<li class="selected">        method_graph, params = torch._C._jit_pass_lower_graph(graph, model._c)</li>
<li>        in_vars, in_desc = torch.jit._flatten(tuple(args) + tuple(params))</li>
<li>        graph = _propagate_and_assign_input_shapes(</li>
<li>            method_graph, tuple(in_vars), False, propagate)</li>
<li>    except AttributeError:</li>
<li>        raise RuntimeError('\'forward\' method must be a script method')</li>
<li>elif isinstance(model, torch.jit.ScriptFunction):</li>
<li>    assert example_outputs is not None, "example_outputs must be provided when exporting a TorchScript ScriptFunction"</li>
<li>    method = model</li>
<li>    params = ()</li>
<li>    in_vars, in_desc = torch.jit._flatten(tuple(args))</li>
</ol></code></pre>


  </article>
  <div class="onebox-metadata">
    
    
  </div>
  <div style="clear: both"></div>
</aside>


jerryzh168: <p>cc <a class="mention" href="/u/james_reed">@James_Reed</a> can you take a look?</p>

jerryzh168: <p>does the conv packed params/linear packed params appear in state_dict? you can call unpack on these object to get the parameters I think: <a href="https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/native/quantized/cpu/conv_packed_params.h#L17" rel="nofollow noopener">https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/native/quantized/cpu/conv_packed_params.h#L17</a></p>

masahi: <p>Ok here is the test script.</p>
<pre><code class="lang-auto">import torch
from torchvision.models.quantization import mobilenet as qmobilenet


def quantize_model(model, inp):
    model.fuse_model()
    model.qconfig = torch.quantization.get_default_qconfig('fbgemm')
    torch.quantization.prepare(model, inplace=True)
    model(inp)
    torch.quantization.convert(model, inplace=True)


qmodel = qmobilenet.mobilenet_v2(pretrained=True).eval()

pt_inp = torch.rand((1, 3, 224, 224))
quantize_model(qmodel, pt_inp)
script_module = torch.jit.trace(qmodel, pt_inp).eval()

graph = script_module.graph
print(script_module.state_dict())
_, params = torch._C._jit_pass_lower_graph(graph, script_module._c)
</code></pre>
<p>The model is quantized mobilenet v2 from torchvision. The output of <code>state_dict()</code> from above script is different between v1.6 and v1.5.1:</p>
<ul>
<li>With v1.5.1, all packed quantized parameters (conv and linear) are returned. Unpacking is also no problem.</li>
<li>With v1.6, I only get <code>OrderedDict([('quant.scale', tensor([0.0079])), ('quant.zero_point', tensor([0]))]) </code>, so there is nothing that can be unpacked.</li>
</ul>
<p>In both version, the last line in the above script, <code>torch._C._jit_pass_lower_graph(graph, script_module._c)</code>, returns all quantized parameters. So technically my original problem is solved. My question is if this is an expected behavior.</p>

jerryzh168: <aside class="quote no-group" data-username="masahi" data-post="7" data-topic="84588">
<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/masahi/48/19725_2.png" class="avatar"> masahi:</div>
<blockquote>
<p>In both version, the last line in the above script, <code>torch._C._jit_pass_lower_graph(graph, script_module._c)</code> , returns all quantized parameters. So technically my original problem is solved. My question is if this is an expected behavior.</p>
</blockquote>
</aside>
<p>probably not, I’ll create an issue for this, this for reporting</p>

jerryzh168: <p>actually I think you are not supposed to call state_dict on a torchscript model, could you call state_dict before script/trace instead?</p>

masahi: <p>We (TVM) take jitted models as input, so we don’t get to see the original models. Fortunately, I found a workaround for this problem without the use of <code>state_dict</code>, so it is no longer a problem for us. Thanks.</p>