TurnkeyML Code Structure
January 28, 2025 ยท View on GitHub
Repo Organization
The TurnkeyML source code has a few major top-level directories:
docs: documentation for the entire project.examples: example scripts for use with the TurnkeyML tools.examples/turnkey/cli: tutorial series starting inexamples/turnkey/cli/readme.mdto help learn theturnkeyCLI.examples/turnkey/cli/scripts: example scripts that can be fed as input into theturnkeyCLI. These scripts each have a docstring that recommends one or moreturnkeyCLI commands to try out.
examples/turnkey/api: examples scripts that invokeToolsvia APIs.
models: the corpora of models that makes up the TurnkeyML models (see the models readme).- Each subdirectory under
modelsrepresents a corpus of models pulled from somewhere on the internet. For example,models/torch_hubis a corpus of models from Torch Hub.
- Each subdirectory under
src/turnkeyml: source code for the TurnkeyML package.src/turnkeyml/tools: implementsTooland defines the tools built in toturnkey.src/turnkeyml/sequence: implementsSequenceand defines the plugin API forTools.src/turnkeyml/run: implementsBaseRT, an abstract base class that defines TurnkeyML's vendor-agnostic benchmarking functionality. This module also includes the runtime and device plugin APIs and the built-in runtimes and devices.src/turnkeyml/cli: implements theturnkeyCLI.src/turnkeyml/common: functions common to the other modules.src/turnkeyml/version.py: defines the package version number.src/turnkeyml/state.py: implements theStateclass.src/turnkeyml/files_api.py: implements theevaluate_files()API, which is the top-level API called by the CLI.
test: tests for the TurnkeyML tools.test/turnkey/analysis.py: tests focusing on thediscoverTool.test/turnkey/cli.py: tests focusing on top-level CLI features.
Tool Classes
All of the logic for actually building models is contained in Tool classes. Generally, a FirstTool class obtains a model, and each subsequent Tool is a model-to-model transformation. For example:
- the
Discover(FirstTool)(akadiscoverin the CLI) obtains a PyTorch model instance from a python script. - the
ExportPytorchModel(Tool)(akaexport-pytorchin the CLI) transforms a PyTorch model instance into an ONNX model file.
Composability
Tools are designed to be composable, for example, there are already a few ONNX-to-ONNX Tools defined in src/turnkeyml/tools/onnx.py that could sequenced in any order.
This composability is facilitated by the State class, which is how Tools communicate with each other. Every Tool takes an instance of State as input and then returns an instance of State. For example:
Discover(FirstTool)takes a freshly initialized instance ofStateas input, and modifies it so thatstate.resultpoints to a PyTorch model.ExportPytorchModel(Tool)takes a PyTorch model instate.resultthen modifiesStatesuch thatstate.resultpoints to the exported ONNX file.
Implementation
See tools.py for a definition of each method of Tool that must be implemented to create a new Tool subclass.