Skip to content

Commit

Permalink
model_dump: Add a helper to produce html with a single call
Browse files Browse the repository at this point in the history
Differential Revision: [D31281091](https://our.internmc.facebook.com/intern/diff/D31281091/)

**NOTE FOR REVIEWERS**: This PR has internal Facebook specific changes or comments, please review them on [Phabricator](https://our.internmc.facebook.com/intern/diff/D31281091/)!

[ghstack-poisoned]
  • Loading branch information
dreiss committed Oct 1, 2021
1 parent b777d79 commit 01e002e
Show file tree
Hide file tree
Showing 2 changed files with 8 additions and 3 deletions.
4 changes: 1 addition & 3 deletions test/test_model_dump.py
Expand Up @@ -101,9 +101,7 @@ def do_dump_model(self, model, extra_files=None):
def open_html_model(self, wd, model, extra_files=None):
buf = io.BytesIO()
torch.jit.save(model, buf, _extra_files=extra_files)
info = torch.utils.model_dump.get_model_info(buf)
skeleton = torch.utils.model_dump.get_inline_skeleton()
page = torch.utils.model_dump.burn_in_info(skeleton, info)
page = torch.utils.model_dump.get_info_and_burn_skeleton(buf)
wd.get("data:text/html;charset=utf-8," + urllib.parse.quote(page))

def open_section_and_get_body(self, wd, name):
Expand Down
7 changes: 7 additions & 0 deletions torch/utils/model_dump/__init__.py
Expand Up @@ -367,6 +367,13 @@ def burn_in_info(skeleton, info):
"BURNED_IN_MODEL_INFO = " + json.dumps(info, sort_keys=True).replace("/", "\\/"))


def get_info_and_burn_skeleton(path_or_bytesio, **kwargs):
model_info = get_model_info(path_or_bytesio, **kwargs)
skeleton = get_inline_skeleton()
page = burn_in_info(skeleton, model_info)
return page


def main(argv, *, stdout=None):
parser = argparse.ArgumentParser()
parser.add_argument("--style", choices=["json", "html"])
Expand Down

0 comments on commit 01e002e

Please sign in to comment.