Add torch.utils.show_pickle for showing pickle contents in saved models (pytorch#35168)

Summary:
Pull Request resolved: pytorch#35168

Sometimes when a saved model isn't working, it's nice to be able to look
at the contents of the pickle files.  Unfortunately, pickletools output
isn't particularly readable, and unpickling is often either not possible
or runs so much post-processing code that it's not possible to tell
exactly what is present in the pickled data.

This script uses a custom Unpickler to unpickle (almost) any data into
stub objects that have no dependency on torch or any other runtime types
and suppress (almost) any postprocessing code.

As a convenience, the wrapper can search through zip files, supporting
command lines like

`python -m torch.utils.show_pickle /path/to/model.pt1@*/data.pkl`

When the module is invoked as main, we also install a hack in pprint to
allow semi-resonable formatting of our stub objects.

Test Plan: Ran it on a data.pkl, constants.pkl, and a debug pkl

Differential Revision: D20842550

Pulled By: dreiss

fbshipit-source-id: ef662d8915fc5795039054d1f8fef2e1c51cf40a

Author: dreiss

test/run_test.py

@@ -59,6 +59,7 @@
     'quantization/test_quantize_script',
     'test_sparse',
     'test_serialization',
+    'test_show_pickle',
     'test_torch',
     'test_type_info',
     'test_type_hints',

test/test_show_pickle.py

@@ -0,0 +1,34 @@
+import unittest
+import io
+import tempfile
+import torch
+import torch.utils.show_pickle
+
+from torch.testing._internal.common_utils import IS_WINDOWS
+
+class TestShowPickle(unittest.TestCase):
+
+    @unittest.skipIf(IS_WINDOWS, "Can't re-open temp file on Windows")
+    def test_scripted_model(self):
+        class MyCoolModule(torch.nn.Module):
+            def __init__(self, weight):
+                super().__init__()
+                self.weight = weight
+
+            def forward(self, x):
+                return x * self.weight
+
+        m = torch.jit.script(MyCoolModule(torch.tensor([2.0])))
+
+        with tempfile.NamedTemporaryFile() as tmp:
+            torch.jit.save(m, tmp)
+            tmp.flush()
+            buf = io.StringIO()
+            torch.utils.show_pickle.main(["", tmp.name + "@*/data.pkl"], output_stream=buf)
+            output = buf.getvalue()
+            self.assertRegex(output, "MyCoolModule")
+            self.assertRegex(output, "weight")
+
+
+if __name__ == '__main__':
+    unittest.main()

torch/utils/show_pickle.py

@@ -0,0 +1,115 @@
+#!/usr/bin/env python3
+import sys
+import pickle
+import pprint
+import zipfile
+import fnmatch
+
+
+class FakeObject(object):
+    def __init__(self, module, name, args):
+        self.module = module
+        self.name = name
+        self.args = args
+        # NOTE: We don't distinguish between state never set and state set to None.
+        self.state = None
+
+    def __repr__(self):
+        state_str = "" if self.state is None else f"(state={self.state!r})"
+        return f"{self.module}.{self.name}{self.args!r}{state_str}"
+
+    def __setstate__(self, state):
+        self.state = state
+
+    @staticmethod
+    def pp_format(printer, obj, stream, indent, allowance, context, level):
+        if not obj.args and obj.state is None:
+            stream.write(repr(obj))
+            return
+        if obj.state is None:
+            stream.write(f"{obj.module}.{obj.name}")
+            printer._format(obj.args, stream, indent + 1, allowance + 1, context, level)
+            return
+        if not obj.args:
+            stream.write(f"{obj.module}.{obj.name}()(state=\n")
+            indent += printer._indent_per_level
+            stream.write(" " * indent)
+            printer._format(obj.state, stream, indent, allowance + 1, context, level + 1)
+            stream.write(")")
+            return
+        raise Exception("Need to implement")
+
+
+class FakeClass(object):
+    def __init__(self, module, name):
+        self.module = module
+        self.name = name
+        self.__new__ = self.fake_new
+
+    def __repr__(self):
+        return f"{self.module}.{self.name}"
+
+    def __call__(self, *args):
+        return FakeObject(self.module, self.name, args)
+
+    def fake_new(self, *args):
+        return FakeObject(self.module, self.name, args[1:])
+
+
+class DumpUnpickler(pickle._Unpickler):
+    def find_class(self, module, name):
+        return FakeClass(module, name)
+
+    def persistent_load(self, pid):
+        return FakeObject("pers", "obj", (pid,))
+
+    @classmethod
+    def dump(cls, in_stream, out_stream):
+        value = cls(in_stream).load()
+        pprint.pprint(value, stream=out_stream)
+
+
+def main(argv, output_stream=None):
+    if len(argv) != 2:
+        # Don't spam stderr if not using stdout.
+        if output_stream is not None:
+            raise Exception("Pass argv of length 2.")
+        sys.stderr.write("usage: show_pickle PICKLE_FILE\n")
+        sys.stderr.write("  PICKLE_FILE can be any of:\n")
+        sys.stderr.write("    path to a pickle file\n")
+        sys.stderr.write("    file.zip@member.pkl\n")
+        sys.stderr.write("    file.zip@*/pattern.*\n")
+        sys.stderr.write("      (shell glob pattern for members)\n")
+        sys.stderr.write("      (only first match will be shown)\n")
+        return 2
+
+    fname = argv[1]
+    if "@" not in fname:
+        with open(fname, "rb") as handle:
+            DumpUnpickler.dump(handle, output_stream)
+    else:
+        zfname, mname = fname.split("@", 1)
+        with zipfile.ZipFile(zfname) as zf:
+            if "*" not in mname:
+                with zf.open(mname) as handle:
+                    DumpUnpickler.dump(handle, output_stream)
+            else:
+                found = False
+                for info in zf.infolist():
+                    if fnmatch.fnmatch(info.filename, mname):
+                        with zf.open(info) as handle:
+                            DumpUnpickler.dump(handle, output_stream)
+                        found = True
+                        break
+                if not found:
+                    raise Exception(f"Could not find member matching {mname} in {zfname}")
+
+
+if __name__ == "__main__":
+    # This hack works on every version of Python I've tested.
+    # I've tested on the following versions:
+    #   3.7.4
+    if True:
+        pprint.PrettyPrinter._dispatch[FakeObject.__repr__] = FakeObject.pp_format
+
+    sys.exit(main(sys.argv))