You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
TurnkeyML currently supports only torch.nn.modules within the turnkey benchmark command. However, there is interest in using TurnkeyML with PyTorch Lightning modules as well.
Scope of Work
At least the following improvements would be needed to support PyTorch Lightning:
In the analyze module, detect PyTorch Lightning modules and pass them to explore_invocation(). Currently, we only pass torch.nn.module to explore_invocation().
Add a new build Stage to export the Lightning module to ONNX. This is hopefully as simple as calling lightning_module.to_onnx().
Add this new build stage to the ExportPlaceholder Stage so that it can be automatically selected whenever a Lightning module is encountered.
Contribution
I am setting a priority of P2 until we have an e2e example available that we can use for testing.
I believe that this would make a good first issue if someone wanted to contribute it.
Problem Statement
TurnkeyML currently supports only
torch.nn.modules
within theturnkey benchmark
command. However, there is interest in using TurnkeyML with PyTorch Lightning modules as well.Scope of Work
At least the following improvements would be needed to support PyTorch Lightning:
analyze
module, detect PyTorch Lightning modules and pass them toexplore_invocation()
. Currently, we only passtorch.nn.module
toexplore_invocation()
.Stage
to export the Lightning module to ONNX. This is hopefully as simple as callinglightning_module.to_onnx()
.Contribution
P2
until we have an e2e example available that we can use for testing.cc @andife with regards to #69
The text was updated successfully, but these errors were encountered: