Extract constants from Torchscript module

jit
1

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?

Thanks a lot!

2

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 constants.pkl in the serialized zip file, which is in the same format as Python’s pickler (our reader is here).

3

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 PythonPrint. @ifeherva can you file a feature request on github for this?

4

I am using currently the .code attribute of the TracedModule class. For example a resnet50 backbone with a constant scaling factor looks like this:

def forward(self,
    input: Tensor) -> 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)

I did’t consider the zip file yet, I thought there is a way to get it ‘on-the-fly’.

5

Will do that, thanks!

6