Extract constants from Torchscript module

ifeherva: <p>I am building a library that converts a torchscript module to mxnet gluon. Given the two language is very similar I have an easy time parsing the .code and recreate the model. However, I am not able to extract the constants that appear in the code (e.g. CONSTANTS.c0). Any idea how I can obtain the value of this from the jit compiled module?</p>
<p>Thanks a lot!</p>

driazati: <p>Are you trying to extract the constants from the serialized binary or in C++? If it’s the former you’ll have to parse the <code>constants.pkl</code> in the serialized zip file, which is in the same format as Python’s pickler (our reader is <a href="https://github.com/pytorch/pytorch/blob/master/torch/csrc/jit/serialization/pickler.cpp">here</a>).</p>

Michael_Suo: <p>Shouldn’t be too hard to add an endpoint that also passes back the constants table, given that we populate it with any call to <code>PythonPrint</code>. <a class="mention" href="/u/ifeherva">@ifeherva</a> can you file a feature request on github for this?</p>

ifeherva: <p>I am using currently the .code attribute of the TracedModule class. For example a resnet50 backbone with a constant scaling factor looks like this:</p>
<pre><code>def forward(self,
    input: Tensor) -&gt; Tensor:
_0 = self.model.relu
_1 = (self.model.bn1).forward((self.model.conv1).forward(input, ), )
_2 = self.model.layer1
_3 = (self.model.maxpool).forward((_0).forward(_1, ), )
_4 = self.model.layer3
_5 = (self.model.layer2).forward((_2).forward(_3, ), )
_6 = self._pool
_7 = (self.model.layer4).forward((_4).forward(_5, ), )
_8 = (_6).forward(_7, )
_9 = ops.prim.NumToTensor(torch.size(_8, 0))
x = torch.view(_8, [int(_9), -1])
return torch.mul(x, CONSTANTS.c0)
</code></pre>
<p>I did’t consider the zip file yet, I thought there is a way to get it ‘on-the-fly’.</p>

ifeherva: <p>Will do that, thanks!</p>

ifeherva: <aside class="quote no-group" data-username="ifeherva" data-post="4" data-topic="76666">
<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/ifeherva/48/57166_2.png" class="avatar"> ifeherva:</div>
<blockquote>
<p>TracedModule</p>
</blockquote>
</aside>
<aside class="onebox githubissue" data-onebox-src="https://github.com/pytorch/pytorch/issues/36625">
  <header class="source">

      <a href="https://github.com/pytorch/pytorch/issues/36625" target="_blank" rel="noopener nofollow ugc">github.com/pytorch/pytorch</a>
  </header>

  <article class="onebox-body">
    <div class="github-row">
  <div class="github-icon-container" title="Issue">
	  <svg width="60" height="60" class="github-icon" viewBox="0 0 14 16" aria-hidden="true"><path fill-rule="evenodd" d="M7 2.3c3.14 0 5.7 2.56 5.7 5.7s-2.56 5.7-5.7 5.7A5.71 5.71 0 0 1 1.3 8c0-3.14 2.56-5.7 5.7-5.7zM7 1C3.14 1 0 4.14 0 8s3.14 7 7 7 7-3.14 7-7-3.14-7-7-7zm1 3H6v5h2V4zm0 6H6v2h2v-2z"></path></svg>
  </div>

  <div class="github-info-container">
    <h4>
      <a href="https://github.com/pytorch/pytorch/issues/36625" target="_blank" rel="noopener nofollow ugc">Add TracedModule attribute for the constants table</a>
    </h4>

    <div class="github-info">
      <div class="date">
        opened <span class="discourse-local-date" data-format="ll" data-date="2020-04-14" data-time="22:33:29" data-timezone="UTC">10:33PM - 14 Apr 20 UTC</span>
      </div>

        <div class="date">
          closed <span class="discourse-local-date" data-format="ll" data-date="2020-05-01" data-time="03:46:24" data-timezone="UTC">03:46AM - 01 May 20 UTC</span>
        </div>

      <div class="user">
        <a href="https://github.com/ifeherva" target="_blank" rel="noopener nofollow ugc">
          <img alt="ifeherva" src="https://avatars.githubusercontent.com/u/3716849?v=4" class="onebox-avatar-inline" width="20" height="20">
          ifeherva
        </a>
      </div>
    </div>

    <div class="labels">
        <span style="display:inline-block;margin-top:2px;background-color: #B8B8B8;padding: 2px;border-radius: 4px;color: #fff;margin-left: 3px;">
          triage review
        </span>
        <span style="display:inline-block;margin-top:2px;background-color: #B8B8B8;padding: 2px;border-radius: 4px;color: #fff;margin-left: 3px;">
          oncall: jit
        </span>
        <span style="display:inline-block;margin-top:2px;background-color: #B8B8B8;padding: 2px;border-radius: 4px;color: #fff;margin-left: 3px;">
          module: bootcamp
        </span>
        <span style="display:inline-block;margin-top:2px;background-color: #B8B8B8;padding: 2px;border-radius: 4px;color: #fff;margin-left: 3px;">
          triaged
        </span>
        <span style="display:inline-block;margin-top:2px;background-color: #B8B8B8;padding: 2px;border-radius: 4px;color: #fff;margin-left: 3px;">
          enhancement
        </span>
    </div>
  </div>
</div>

  <div class="github-row">
    <p class="github-body-container">## 🚀 Feature

TracedModuel currently returns the compiled code via the .code()<span class="show-more-container"><a href="" rel="noopener" class="show-more">…</a></span><span class="excerpt hidden"> member function. This code can possible contain constants (e.g. CONSTANTS.c0), however the value of these are not accessible from python at the moment. It would really useful to surface the constant mapping to python so libraries that interpret the module can get the full picture.

## Motivation

I am currently writing a pytorch/torchscript -&gt; mxnet code converter and currently it cannot resolve the constants in the torchscript code.

## Pitch

tracedmodule.code -&gt; returns code (done)
tracedmoduel.constants -&gt; returns a mapping between constant name and value (dict)

## Alternatives

Any other function from jit that can return the constants of a tracedmodule

## Additional context

https://discuss.pytorch.org/t/extract-constants-from-torchscript-module/


cc @suo</span></p>
  </div>

  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>