Add metadata for torch jit TracedModules. (#17640)
authorPritam Damania <pritam.damania@fb.com>
Sun, 10 Mar 2019 05:31:42 +0000 (21:31 -0800)
committerFacebook Github Bot <facebook-github-bot@users.noreply.github.com>
Sun, 10 Mar 2019 05:37:15 +0000 (21:37 -0800)
commit24e7b824e0e8e82f9efa730da2e64025640ff3bc
tree14b3044bc8a1c6d35045c3e0fb925f8a6e5f0828
parent320c6977c25dd67dbbc3eadea506291bd083be4a
Add metadata for torch jit TracedModules. (#17640)

Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/17640

Pull Request resolved: https://github.com/pytorch/pytorch/pull/17311

I've extended our model metadata framework in this diff to support
traced modules as well. Re-used a lot of components from the previous
implementation of ScriptModule metadata.

Tracing is a little different from Scripting since you can't just create a
subclass of TopLevelTraceModule (type returned by torch.jit.trace) and attach
metadata the way we did for ScriptModule. As a result, I've introduced a
separate API torch.fb.jit_trace which returns an instance of
TracedModuleWithMetadata which is a subclass of TopLevelTracedModule. As a
result, we can now attach metadata to this instance.

Reviewed By: dzhulgakov

Differential Revision: D14117966

fbshipit-source-id: 3eee5eef733cb8d6a219c02e2f41d08698eca326
torch/jit/__init__.py