benchmark
3173d623 - Move the jit script code to __post_init__ (#761)

Commit
3 years ago
Move the jit script code to __post_init__ (#761) Summary: We need to remove current jit code from each model directory and use a unified entry for all the transformations. This is because if we do the jit script first, then change the precision to fp16, the CI test will fail with error: https://app.circleci.com/pipelines/github/pytorch/benchmark/3665/workflows/da928033-03fa-48d0-90a4-788d3ee794ed/jobs/3771 However, I noticed different models are using different `torch.jit` APIs: 1) `torch.jit.script(model)`, 2) `torch.jit.script(model, example_inputs)`, 3) `torch.jit.trace(model, example_inputs)` 4) an extra `torch.jit.optimize_for_inference()` for inference. Which one should I use if we are sharing the jit scripting code for all the models? Krovatkin The current design is to use `torch.jit.trace(model, example_inputs)` by default. For models that need to call `torch.jit.trace()` (like nvidia_deeprecommender), or models that need to script multiple `torch.nn.Module` instances, they should add a callback function, `jit_callback(self)`, to handle the JIT enablement. Pull Request resolved: https://github.com/pytorch/benchmark/pull/761 Reviewed By: davidberard98 Differential Revision: D34396461 Pulled By: xuzhao9 fbshipit-source-id: b51ef60b8ee28c0bd910404d549cfb4a75c0ae28
Author
Parents
Loading