diff --git a/examples/qualcomm/CMakeLists.txt b/examples/qualcomm/CMakeLists.txt index 14cd7eb969e..cdc0b013149 100644 --- a/examples/qualcomm/CMakeLists.txt +++ b/examples/qualcomm/CMakeLists.txt @@ -82,6 +82,14 @@ set(_common_include_directories # build qnn_executor_runner add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/executor_runner) +# build qnn_toy_model_trainer for the joint-graph training example +option(BUILD_QNN_TRAINING_EXAMPLE "Build QNN/HTP joint-graph training example" + OFF +) +if(BUILD_QNN_TRAINING_EXAMPLE) + add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/training/toy_model/on-device) +endif() + # build qnn_llama_runner for llama add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/oss_scripts/llama) diff --git a/examples/qualcomm/README.md b/examples/qualcomm/README.md index 32a710eaef5..691341be23e 100644 --- a/examples/qualcomm/README.md +++ b/examples/qualcomm/README.md @@ -13,6 +13,8 @@ We have separated the example scripts into the following subfolders, please refe 4. util_scripts: This folder includes tutorial example scripts designed to showcase the utilities we've developed. For example, we provide a debugging tool [qnn_intermediate_debugger](./util_scripts/qnn_intermediate_debugger_demo.py) that allow users to compare the intermediate outputs of QNNs V.S. CPUs. By reviewing these scripts, we aim to help users smoothly integrate these utilities into their own projects. +5. training: A minimal, self-contained on-device training example using a QNN/HTP-delegated joint forward+backward graph. See [training/toy_model/README.md](training/toy_model/README.md) for details. + Please check helper of each examples for detailed arguments. diff --git a/examples/qualcomm/training/toy_model/README.md b/examples/qualcomm/training/toy_model/README.md new file mode 100644 index 00000000000..6395cd34f27 --- /dev/null +++ b/examples/qualcomm/training/toy_model/README.md @@ -0,0 +1,128 @@ +# QNN/HTP joint-graph training example: 4-layer toy model + +A minimal, self-contained example of on-device training on the QNN/HTP +(Hexagon) backend: export a forward+backward "joint graph" for a plain +4-layer toy model, lower it to the QNN delegate, and run a real host-driven +training loop against the compiled `.pte`. + +- **No quantization.** Compiled fp16 for the HTP. +- **No pretrained weights.** `ToyModel`'s weights are whatever `nn.Linear`'s + default init produces; pass `--seed` for reproducibility. + +This example builds the joint graph with the lower-level +`aot_export_joint_with_descriptors` API (see `joint_graph.py`) and lowers it +to the QNN/HTP delegate with `to_edge_transform_and_lower_to_qnn` (see +`export_toy_model.py`). Because the resulting joint `GraphModule` has no +`nn.Parameter` attributes -- every weight arrives as an anonymous `primals_N` +placeholder -- re-exporting it inside `to_edge_transform_and_lower_to_qnn` +classifies every weight as a plain input automatically. + +## Minimal decomposition table + +`joint_graph.py`'s `build_decompositions()` starts from +`core_aten_decompositions()`, removes everything QNN supports natively via +`get_skip_decomp_table()`, and registers only the one decomp a bias-free +ToyModel + MSE-loss graph actually needs: + +- `mean.default` -> `mean.dim` (`mse_loss`'s `reduction="mean"` produces a + full-tensor mean; QNN has no `mean.default` builder, only `mean.dim`) + +The `nn.Linear` layers have no bias, so they lower straight to `mm` (for which QNN builder is present). + +## One .pte per weight-sharing setting + +`export_toy_model.py` compiles two methods -- `forward` (inference) and +`forward_backward` (the joint graph, used for training) -- named from the +`name` config.json field. Both methods take weights as explicit `USER_INPUT`s, +so trained weights can be fed directly at inference time without re-exporting. +Which variant gets built is controlled by the `weight_sharing` config.json +field (or `--weight_sharing`): + +- `false` (default): QNN weight-sharing disabled. Compiles `forward` and + `forward_backward` into **separate** `.pte` files (`_fwd.pte` and + `_fwd_bwd.pte`) under `/no_ws/`. The inference runner + uses `_fwd.pte` and calls `"forward"` directly with trained weights. +- `true`: QNN weight-sharing enabled. Compiles both methods into a **single** + `.pte` (`.pte`) under `/ws/`, one shared compile context + for both methods. Weight-sharing is a hardware-only feature -- requires real Hexagon/HTP + hardware. The inference runner uses this single `.pte` and also calls + `"forward"` with trained weights. + +`manifest.json` (written next to the `.pte`(s), in the same `ws`/`no_ws` +subdir) records `pte_name`, `fwd_pte_name` (no_ws only), `weight_sharing`, +`param_names`, `x_shape`, and `target_shape`. `export_toy_model.py` also saves +the initial weights next to it as `init_weights.safetensors`, keyed by +`param_names`. + +## On-device C++ training runner + +`on-device/trainer.cpp` (built as `qnn_toy_model_trainer`, wired into +`examples/qualcomm/CMakeLists.txt` behind the `BUILD_QNN_TRAINING_EXAMPLE` +CMake option, OFF by default -- pass `-DBUILD_QNN_TRAINING_EXAMPLE=ON` when +configuring `examples/qualcomm` to build it), +for exercising the training loop mechanics on real Hexagon/HTP hardware -- +including the `weight_sharing=true` variant. It uses the plain +`executorch::extension::Module` API (`method_meta` + `execute`), joint graph's weights +are plain method I/O (anonymous `primals_N` placeholders), not +`nn.Parameter`-like mutable buffers. + +Usage: + +```bash +qnn_toy_model_trainer # train; writes trained_model.safetensors +qnn_toy_model_inference # infer from trained_model.safetensors +``` + +The `.pte` and `init_weights.safetensors` are expected alongside `manifest.json` +(`export_toy_model.py`'s layout) -- resolved from `manifest.json`'s own +directory, not a separate flag. Initial weights are reordered to +`manifest.json`'s `param_names`; a single `x`/`target` batch is generated +on-device once via `N(0,1)`, seeded from `config.json`'s `seed` if present, +and reused every step -- the model memorizes it, so loss decreases +monotonically. `steps`/`lr`/`log_every` come from `config.json`. + +To run on a connected Android device, use `run_toy_model.py` (requires +`export_toy_model.py` to have been run first): + +```bash +# push + train + infer (default) +python run_toy_model.py --device --host \ + --soc_model SM8850 --build_folder /path/to/build-android + +# push only +python run_toy_model.py ... --push_only + +# skip push, train only (re-run without re-pushing) +python run_toy_model.py ... --skip_push --skip_inference + +# skip push, infer only (after training already done) +python run_toy_model.py ... --skip_push --skip_train +``` + +Use `--weight_sharing true` to select the `ws/` variant (default: `no_ws/`). +`QNN_SDK_ROOT` must be set (same requirement as all other Qualcomm examples). + +This lands `qnn_toy_model_trainer`, `qnn_toy_model_inference`, and the QNN/HTP +shared libraries into `bin/`, and `config.json`/`manifest.json`/the `.pte`(s)/ +`init_weights.safetensors` into `artifacts/` -- both under +`/data/local/tmp/$(whoami)/qnn_toy_model_trainer/` on the device. + +## Configuration + +All scripts read defaults from `config.json` at the `training/toy_model/` root (override the path with +`--config`, or any individual field with its CLI flag). Key fields: + +- `name`: base filename for the exported artifacts (`_fwd.pte`, `_fwd_bwd.pte`). +- `output_dir`: base artifacts directory (default `outputs/artifacts`). + `export_toy_model.py` writes into `/ws/` or `/no_ws/` + depending on `weight_sharing`; +- `seed`: seeds `torch.manual_seed` at export time (reproducible initial + weights) and the on-device RNG used to generate the `x`/`target` batch. +- `steps`, `lr`, `log_every`: training-loop hyperparameters. + +```bash +cd examples/qualcomm/training/toy_model + +# Export: writes *.pte's + manifest.json to outputs/artifacts/no_ws/ +python export/export_toy_model.py +``` \ No newline at end of file diff --git a/examples/qualcomm/training/toy_model/config.json b/examples/qualcomm/training/toy_model/config.json new file mode 100644 index 00000000000..e5f54f9edde --- /dev/null +++ b/examples/qualcomm/training/toy_model/config.json @@ -0,0 +1,14 @@ +{ + "name": "toy_model", + "weight_sharing": false, + "soc_model": "SM8850", + "output_dir": "outputs/artifacts", + "seed": 0, + "batch_size": 4, + "in_dim": 8, + "hidden_dim": 32, + "out_dim": 4, + "steps": 500, + "lr": 0.1, + "log_every": 50 +} \ No newline at end of file diff --git a/examples/qualcomm/training/toy_model/export/config.py b/examples/qualcomm/training/toy_model/export/config.py new file mode 100644 index 00000000000..d6a42c7115a --- /dev/null +++ b/examples/qualcomm/training/toy_model/export/config.py @@ -0,0 +1,17 @@ +# Copyright (c) Qualcomm Innovation Center, Inc. +# All rights reserved +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""Shared config.json loader for the ToyModel training example scripts.""" + +import json +import os + + +def load_config(path): + if not path or not os.path.exists(path): + return {} + with open(path) as f: + return json.load(f) diff --git a/examples/qualcomm/training/toy_model/export/export_toy_model.py b/examples/qualcomm/training/toy_model/export/export_toy_model.py new file mode 100644 index 00000000000..61ac06045c0 --- /dev/null +++ b/examples/qualcomm/training/toy_model/export/export_toy_model.py @@ -0,0 +1,239 @@ +# Copyright (c) Qualcomm Innovation Center, Inc. +# All rights reserved +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""Export a 4-layer toy model as a joint (forward+backward) QNN/HTP training graph. + +For weight_sharing=false (default): compiles forward and forward_backward into +separate .pte files (_fwd.pte and _fwd_bwd.pte). The forward .pte +uses the functional_call trick so weights are USER_INPUTs -- trained weights can +be fed at inference time without re-exporting. + +For weight_sharing=true: compiles both methods into a single .pte (one shared +compile context). Requires real Hexagon/HTP hardware. + +Methods: + - "forward": inference graph. Weights are USER_INPUTs via functional_call -- + feed trained_model.safetensors at inference time, no re-export needed. + - "forward_backward": the joint graph. Weights are plain USER_INPUTs, fed + in and read back out (as gradients) by the host training loop each step. + +No checkpointing, no quantization, no pretrained weights -- everything here +is randomly initialized and self-contained. +""" + +import argparse +import json +import logging +import os + +import torch +import torch.nn as nn + +from config import load_config + +from executorch.backends.qualcomm.utils.utils import ( + generate_htp_compiler_spec, + generate_qnn_executorch_compiler_spec, + get_soc_to_chipset_map, + to_edge_transform_and_lower_to_qnn, +) +from executorch.extension.export_util.utils import save_pte_program +from joint_graph import JointGraphCreator +from model import ToyModel, ToyModelWithLoss +from safetensors.torch import save_file + + +def _make_fwd_wrapper(toy_model: ToyModel) -> nn.Module: + """Return an nn.Module whose forward(w1, w2, w3, w4, x) runs toy_model via + functional_call so the exported graph has weights as USER_INPUTs, not baked + constants. The wrapper holds no parameters of its own.""" + param_names = list(dict(toy_model.named_parameters()).keys()) + + class _FwdWrapper(nn.Module): + def forward(self, *args): + weights = args[: len(param_names)] + x = args[len(param_names)] + params = dict(zip(param_names, weights)) + return torch.func.functional_call(toy_model, params, (x,)) + + return _FwdWrapper() + + +def _compile_and_save( + fwd_wrapper, + joint_graph, + fwd_inputs, + joint_inputs, + soc_model, + output_folder, + use_weight_sharing, + pte_name, +): + backend_options = generate_htp_compiler_spec( + use_fp16=True, + use_weight_sharing=use_weight_sharing, + ) + compile_specs = generate_qnn_executorch_compiler_spec( + soc_model=get_soc_to_chipset_map()[soc_model], + backend_options=backend_options, + ) + + if use_weight_sharing: + # Weight-sharing requires both methods in one compile call. + edge_program = to_edge_transform_and_lower_to_qnn( + module={"forward": fwd_wrapper, "forward_backward": joint_graph}, + inputs={"forward": fwd_inputs, "forward_backward": joint_inputs}, + compiler_specs=compile_specs, + ) + save_pte_program(edge_program.to_executorch(), pte_name, output_folder) + return {"pte_name": f"{pte_name}.pte"} + else: + # Compile forward and forward_backward into separate .pte files so + # each can be loaded independently on-device. + fwd_name = f"{pte_name}_fwd" + fwd_bwd_name = f"{pte_name}_fwd_bwd" + + fwd_program = to_edge_transform_and_lower_to_qnn( + module={"forward": fwd_wrapper}, + inputs={"forward": fwd_inputs}, + compiler_specs=compile_specs, + ) + save_pte_program(fwd_program.to_executorch(), fwd_name, output_folder) + + fwd_bwd_program = to_edge_transform_and_lower_to_qnn( + module={"forward_backward": joint_graph}, + inputs={"forward_backward": joint_inputs}, + compiler_specs=compile_specs, + ) + save_pte_program(fwd_bwd_program.to_executorch(), fwd_bwd_name, output_folder) + return { + "pte_name": f"{fwd_bwd_name}.pte", + "fwd_pte_name": f"{fwd_name}.pte", + } + + +def export_toy_model( + toy_model, x, target, soc_model, output_folder, weight_sharing, pte_name +): + train_module = ToyModelWithLoss(toy_model) + creator = JointGraphCreator(train_module, (x, target)) + joint_graph = creator.create() + + placeholder_names = { + n.name for n in joint_graph.graph.nodes if n.op == "placeholder" + } + params_by_name = dict(train_module.named_parameters()) + param_names = [ + name + for i, name in enumerate(creator.joint_desc.params_spec or []) + if f"primals_{i + 1}" in placeholder_names + ] + weights = [params_by_name[name] for name in param_names] + + # aot_export_joint_with_descriptors's returned GraphModule takes exactly + # two positional args -- forward(primals, tangents) -- each a flat list, + # not one positional arg per primals_N/tangents_N placeholder node. This + # grouped shape is what torch.export.export (called internally by + # to_edge_transform_and_lower_to_qnn) must be given to retrace it. + tangent = torch.ones([1]) + joint_inputs = ((*weights, x, target), (tangent,)) + + # functional_call wrapper: same ToyModel computation, weights as USER_INPUTs. + fwd_wrapper = _make_fwd_wrapper(toy_model) + fwd_inputs = (*weights, x) + + pte_info = _compile_and_save( + fwd_wrapper, + joint_graph, + fwd_inputs, + joint_inputs, + soc_model, + output_folder, + weight_sharing, + pte_name, + ) + + save_file( + {name: w.contiguous() for name, w in zip(param_names, weights)}, + os.path.join(output_folder, "init_weights.safetensors"), + ) + + grad_names = [f"grad_{name}" for name in param_names] + manifest = { + "param_names": param_names, + "inputs": [*param_names, "x", "target", "tangent"], + "outputs": ["loss", "pred", *grad_names], + "x_shape": list(x.shape), + "target_shape": list(target.shape), + "weight_sharing": weight_sharing, + **pte_info, + } + with open(os.path.join(output_folder, "manifest.json"), "w") as f: + json.dump(manifest, f, indent=2) + + +def main(): + logging.basicConfig(level=logging.INFO) + toy_model_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) + parser = argparse.ArgumentParser( + description="Export a 4-layer toy model as joint QNN/HTP training .pte's", + ) + parser.add_argument( + "--config", type=str, default=os.path.join(toy_model_dir, "config.json") + ) + config_args, _ = parser.parse_known_args() + config = load_config(config_args.config) + + parser.add_argument( + "--soc_model", type=str, default=config.get("soc_model", "SM8850") + ) + parser.add_argument( + "-f", "--output_folder", type=str, default=config.get("output_dir", "") + ) + parser.add_argument("--seed", type=int, default=config.get("seed")) + parser.add_argument("--batch_size", type=int, default=config.get("batch_size", 4)) + parser.add_argument("--in_dim", type=int, default=config.get("in_dim", 8)) + parser.add_argument("--hidden_dim", type=int, default=config.get("hidden_dim", 32)) + parser.add_argument("--out_dim", type=int, default=config.get("out_dim", 4)) + parser.add_argument( + "--weight_sharing", + action=argparse.BooleanOptionalAction, + default=config.get("weight_sharing", False), + ) + parser.add_argument("--name", type=str, default=config.get("name", "toy_model")) + args = parser.parse_args() + + if args.seed is not None: + torch.manual_seed(args.seed) + + output_folder = args.output_folder + if output_folder: + if not os.path.isabs(output_folder): + output_folder = os.path.join(toy_model_dir, output_folder) + output_folder = os.path.join( + output_folder, "ws" if args.weight_sharing else "no_ws" + ) + os.makedirs(output_folder, exist_ok=True) + + toy_model = ToyModel( + in_dim=args.in_dim, hidden_dim=args.hidden_dim, out_dim=args.out_dim + ) + x = torch.randn(args.batch_size, args.in_dim) + target = torch.randn(args.batch_size, args.out_dim) + + export_toy_model( + toy_model, + x, + target, + args.soc_model, + output_folder, + args.weight_sharing, + args.name, + ) + + +if __name__ == "__main__": + main() diff --git a/examples/qualcomm/training/toy_model/export/joint_graph.py b/examples/qualcomm/training/toy_model/export/joint_graph.py new file mode 100644 index 00000000000..8e66788c130 --- /dev/null +++ b/examples/qualcomm/training/toy_model/export/joint_graph.py @@ -0,0 +1,172 @@ +# Copyright (c) Qualcomm Innovation Center, Inc. +# All rights reserved +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""Joint (forward+backward) graph creation for QNN HTP training. + +Trimmed for a bias-free toy model + MSE loss: only the one decomposition such a +graph actually needs is registered (full-tensor ``mean``); QNN-native ops are +left undecomposed via ``get_skip_decomp_table()``. +""" + +import logging +from contextlib import ExitStack + +import torch + +from executorch.backends.qualcomm.partition.qnn_partitioner import get_skip_decomp_table +from torch import nn +from torch._decomp import core_aten_decompositions, remove_decompositions +from torch._functorch.aot_autograd import aot_export_joint_with_descriptors + +# _PyTreeCodeGen/_PyTreeInfo are private torch internals; verified against +# PyTorch 2.13. Graph.set_codegen() (public) is used for the write path; +# ._codegen is still needed to read orig_args/in_spec (no public getter exists). +from torch.fx.graph import _PyTreeCodeGen, _PyTreeInfo +from torch.utils._pytree import tree_flatten + +logger = logging.getLogger(__name__) + + +class JointGraphCreator: + """Builds the forward+backward joint graph via the AOT descriptor API. + + The returned ``GraphModule`` has no ``nn.Parameter`` attributes -- every + weight is an anonymous ``primals_N`` placeholder input. That is what lets + a later ``torch.export.export`` re-trace of this graph (done internally + by ``to_edge_transform_and_lower_to_qnn``) classify every weight as a + plain ``USER_INPUT``, with no manual ``ExportedProgram`` construction. + """ + + def __init__(self, module: nn.Module, forward_inputs: tuple, decompositions=None): + self.module = module + self.forward_inputs = forward_inputs + self.decompositions = decompositions or self.build_decompositions() + self._joint_desc = None + self._gm = None + + @staticmethod + def build_decompositions() -> dict: + decomps = core_aten_decompositions() + + # QNN-native ops stay undecomposed. + remove_decompositions(decomps, get_skip_decomp_table()) + + # mse_loss's reduction="mean" produces a full-tensor mean; QNN has no + # mean.default builder, only mean.dim. + decomps[torch.ops.aten.mean.default] = _mean_default_decomp + return decomps + + def create(self) -> torch.fx.GraphModule: + with ExitStack() as stack: + self._joint_desc = aot_export_joint_with_descriptors( + stack, + self.module, + self.forward_inputs, + decompositions=self.decompositions, + ) + gm = self._joint_desc.graph_module + + gm = self._flatten_outputs_and_remove_none(gm) + gm = self._remove_unused_placeholders(gm) + self._gm = gm + return gm + + @property + def joint_desc(self): + return self._joint_desc + + @property + def graph_module(self): + return self._gm + + def _remove_unused_placeholders(self, gm): + """Remove unused parameter/buffer placeholders. + + In the joint graph, aot_export_joint assigns primals_N to all params + then all buffers in sequence. User inputs (forward primals beyond + that range and tangents) must never be removed. + """ + params_spec = list(self._joint_desc.params_spec or []) + buffers_spec = list(self._joint_desc.buffers_spec or []) + n_params = len(params_spec) + n_buffers = len(buffers_spec) + + def _is_user_input(name): + if name.startswith("tangents_"): + return True + if name.startswith("primals_"): + try: + return int(name[len("primals_") :]) - 1 >= n_params + n_buffers + except ValueError: + pass + return False + + changed = True + while changed: + changed = False + for node in list(gm.graph.nodes): + if ( + node.op == "placeholder" + and not _is_user_input(node.name) + and len(node.users) == 0 + ): + logger.info("Removing unused placeholder: %s", node.name) + gm.graph.erase_node(node) + changed = True + gm.graph.lint() + gm.recompile() + return gm + + def _flatten_outputs_and_remove_none(self, gm): + """Flatten output node, remove None entries, update meta['val'], and + reset the codegen pytree out_spec so torch.fx.Interpreter can process + the outputs without a leaf-count mismatch. + """ + new_outputs = [] + for node in gm.graph.nodes: + if node.op == "output": + flat_outputs, _ = tree_flatten(node.args[0]) + new_outputs = [o for o in flat_outputs if o is not None] + node.args = (new_outputs,) + if "val" in node.meta and isinstance(node.meta["val"], (tuple, list)): + node.meta["val"] = tuple( + v for v in node.meta["val"] if v is not None + ) + break + + gm.graph.lint() + + if new_outputs: + codegen = gm.graph._codegen # pylint: disable=protected-access + if isinstance(codegen, _PyTreeCodeGen) and codegen.pytree_info is not None: + dummy_out = tuple(torch.empty(1) for _ in new_outputs) + _, new_out_spec = tree_flatten(dummy_out) + try: + old_info = codegen.pytree_info + gm.graph.set_codegen( + _PyTreeCodeGen( + _PyTreeInfo( + old_info.orig_args, + old_info.in_spec, + new_out_spec, + ) + ) + ) + except (AttributeError, TypeError) as exc: + logger.warning( + "Could not reset pytree out_spec (%s); export will " + "likely fail with a leaf-count mismatch.", + exc, + ) + + gm.recompile() + return gm + + +def _mean_default_decomp(self): + """mean.default (full-tensor mean) -> mean.dim, which QNN supports.""" + dims = list(range(self.dim())) + return torch.ops.aten.mean.dim(self, dims) diff --git a/examples/qualcomm/training/toy_model/export/model.py b/examples/qualcomm/training/toy_model/export/model.py new file mode 100644 index 00000000000..d10a39f7d5a --- /dev/null +++ b/examples/qualcomm/training/toy_model/export/model.py @@ -0,0 +1,42 @@ +# Copyright (c) Qualcomm Innovation Center, Inc. +# All rights reserved +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +import torch +import torch.nn as nn +import torch.nn.functional as F + + +class ToyModel(nn.Module): + """Plain 4-layer toy model, randomly initialized (no checkpoint loading).""" + + def __init__(self, in_dim: int = 8, hidden_dim: int = 32, out_dim: int = 4): + super().__init__() + self.fc1 = nn.Linear(in_dim, hidden_dim, bias=False) + self.fc2 = nn.Linear(hidden_dim, hidden_dim, bias=False) + self.fc3 = nn.Linear(hidden_dim, hidden_dim, bias=False) + self.fc4 = nn.Linear(hidden_dim, out_dim, bias=False) + + def forward(self, x): + x = self.fc1(x) + x = self.fc2(x) + x = self.fc3(x) + return torch.tanh(self.fc4(x)) + + +class ToyModelWithLoss(nn.Module): + """Wraps ToyModel with the loss embedded, as required for on-device training. + + MSE (not cross-entropy) keeps the joint-graph decomposition table minimal + -- no ``nll_loss_backward`` to handle. + """ + + def __init__(self, toy_model: ToyModel): + super().__init__() + self.toy_model = toy_model + + def forward(self, x, target): + pred = self.toy_model(x) + return F.mse_loss(pred, target).reshape(1), pred.detach() diff --git a/examples/qualcomm/training/toy_model/on-device/CMakeLists.txt b/examples/qualcomm/training/toy_model/on-device/CMakeLists.txt new file mode 100644 index 00000000000..742be075ace --- /dev/null +++ b/examples/qualcomm/training/toy_model/on-device/CMakeLists.txt @@ -0,0 +1,32 @@ +# Copyright (c) Qualcomm Innovation Center, Inc. +# All rights reserved +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +# Shared utilities (safetensors loader, json helpers, batch helpers). +add_library(qnn_toy_model_common STATIC common.cpp) +target_include_directories( + qnn_toy_model_common PUBLIC ${_common_include_directories} +) +target_link_libraries(qnn_toy_model_common executorch_core extension_tensor) +target_compile_options(qnn_toy_model_common PUBLIC ${_common_compile_options}) + +function(_add_toy_model_binary target source) + add_executable(${target} ${source}) + target_include_directories(${target} PUBLIC ${_common_include_directories}) + target_link_libraries( + ${target} + qnn_toy_model_common + qnn_executorch_backend + executorch_core + extension_module_static + extension_tensor + full_portable_ops_lib + ) + target_compile_options(${target} PUBLIC ${_common_compile_options}) + set_target_properties(${target} PROPERTIES LINK_FLAGS "-Wl,-rpath,\\$ORIGIN") +endfunction() + +_add_toy_model_binary(qnn_toy_model_trainer trainer.cpp) +_add_toy_model_binary(qnn_toy_model_inference inference.cpp) diff --git a/examples/qualcomm/training/toy_model/on-device/common.cpp b/examples/qualcomm/training/toy_model/on-device/common.cpp new file mode 100644 index 00000000000..e2b38379f3a --- /dev/null +++ b/examples/qualcomm/training/toy_model/on-device/common.cpp @@ -0,0 +1,187 @@ +/* + * Copyright (c) Qualcomm Innovation Center, Inc. + * All rights reserved + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#include "common.h" + +#include +#include + +#include + +using executorch::aten::SizesType; +using executorch::extension::make_tensor_ptr; +using executorch::extension::TensorPtr; + +namespace qnn_toy_model_training { + +nlohmann::json load_json(const std::string& path) { + std::ifstream file(path); + ET_CHECK_MSG(file.is_open(), "Failed to open %s", path.c_str()); + nlohmann::json j; + file >> j; + return j; +} + +std::unordered_map load_safetensors( + const std::string& path) { + std::ifstream file(path, std::ios::binary); + ET_CHECK_MSG(file.is_open(), "Failed to open %s", path.c_str()); + + uint64_t header_size = 0; + file.read(reinterpret_cast(&header_size), sizeof(header_size)); + ET_CHECK_MSG(file.good(), "Failed to read header size from %s", path.c_str()); + + std::string header_bytes(header_size, '\0'); + file.read(header_bytes.data(), static_cast(header_size)); + ET_CHECK_MSG(file.good(), "Failed to read header from %s", path.c_str()); + const nlohmann::json header = nlohmann::json::parse(header_bytes); + + const auto data_start = + static_cast(sizeof(header_size) + header_size); + + std::unordered_map tensors; + for (const auto& [name, info] : header.items()) { + if (name == "__metadata__") { + continue; + } + const auto dtype = info.at("dtype").get(); + ET_CHECK_MSG( + dtype == "F32", + "Unsupported safetensors dtype %s for tensor %s (expected F32).", + dtype.c_str(), + name.c_str()); + + const auto shape = info.at("shape").get>(); + const std::vector sizes(shape.begin(), shape.end()); + + const auto offsets = info.at("data_offsets").get>(); + ET_CHECK_MSG( + offsets.size() == 2 && offsets.at(1) >= offsets.at(0), + "Invalid data_offsets for tensor %s in %s.", + name.c_str(), + path.c_str()); + const size_t nbytes = offsets.at(1) - offsets.at(0); + + size_t numel = 1; + for (const auto s : sizes) { + numel *= static_cast(s); + } + ET_CHECK_MSG( + nbytes == numel * sizeof(float), + "Tensor %s in %s: data_offsets span %zu bytes, expected %zu for its shape.", + name.c_str(), + path.c_str(), + nbytes, + numel * sizeof(float)); + + std::vector data(nbytes / sizeof(float)); + file.seekg(data_start + static_cast(offsets.at(0))); + file.read( + reinterpret_cast(data.data()), + static_cast(nbytes)); + ET_CHECK_MSG( + file.good(), + "Failed to read tensor data for %s from %s", + name.c_str(), + path.c_str()); + + tensors.emplace(name, make_tensor_ptr(sizes, std::move(data))); + } + return tensors; +} + +std::vector load_weights( + const nlohmann::json& manifest, + const std::string& safetensors_path) { + auto tensors = load_safetensors(safetensors_path); + std::vector weights; + for (const auto& name : manifest.at("param_names")) { + const auto key = name.get(); + const auto it = tensors.find(key); + ET_CHECK_MSG( + it != tensors.end(), + "Missing weight %s in safetensors file.", + key.c_str()); + weights.push_back(it->second); + } + return weights; +} + +std::vector shape_from_json(const nlohmann::json& shape_array) { + const auto shape = shape_array.get>(); + return std::vector(shape.begin(), shape.end()); +} + +TensorPtr make_random_tensor( + const std::vector& shape, + std::mt19937& rng) { + size_t n = 1; + for (const auto s : shape) { + n *= static_cast(s); + } + std::normal_distribution dist(0.0f, 1.0f); + std::vector data(n); + for (auto& v : data) { + v = dist(rng); + } + return make_tensor_ptr(shape, std::move(data)); +} + +void save_safetensors( + const std::string& path, + const std::vector& names, + const std::vector& tensors) { + ET_CHECK_MSG( + names.size() == tensors.size(), + "save_safetensors: %zu names but %zu tensors.", + names.size(), + tensors.size()); + + // Build header JSON: each entry has dtype, shape, data_offsets. + nlohmann::json header; + size_t offset = 0; + for (size_t i = 0; i < names.size(); i++) { + const auto& t = *tensors[i]; + const size_t nbytes = t.numel() * sizeof(float); + std::vector shape(t.sizes().begin(), t.sizes().end()); + header[names[i]] = { + {"dtype", "F32"}, + {"shape", shape}, + {"data_offsets", {offset, offset + nbytes}}, + }; + offset += nbytes; + } + + const std::string header_str = header.dump(); + // safetensors spec: header must be padded to a multiple of 8 bytes with + // spaces. + const size_t padded = ((header_str.size() + 7) / 8) * 8; + std::string header_padded = header_str; + header_padded.resize(padded, ' '); + + std::ofstream file(path, std::ios::binary); + ET_CHECK_MSG(file.is_open(), "Failed to open %s for writing.", path.c_str()); + + const uint64_t header_size = static_cast(padded); + file.write(reinterpret_cast(&header_size), sizeof(header_size)); + file.write(header_padded.data(), static_cast(padded)); + for (const auto& t : tensors) { + const size_t nbytes = t->numel() * sizeof(float); + file.write( + reinterpret_cast(t->const_data_ptr()), + static_cast(nbytes)); + } + ET_CHECK_MSG(file.good(), "Failed to write safetensors to %s.", path.c_str()); +} + +std::string manifest_dir(const std::string& manifest_path) { + const auto slash = manifest_path.find_last_of('/'); + return slash == std::string::npos ? "." : manifest_path.substr(0, slash); +} + +} // namespace qnn_toy_model_training diff --git a/examples/qualcomm/training/toy_model/on-device/common.h b/examples/qualcomm/training/toy_model/on-device/common.h new file mode 100644 index 00000000000..dda251ad77e --- /dev/null +++ b/examples/qualcomm/training/toy_model/on-device/common.h @@ -0,0 +1,64 @@ +/* + * Copyright (c) Qualcomm Innovation Center, Inc. + * All rights reserved + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +#include +#include +#include +#include + +#include +#include +#include + +namespace qnn_toy_model_training { + +// Fallbacks for fields missing from config.json -- mirror config.json's values. +constexpr int kDefaultSteps = 500; +constexpr float kDefaultLr = 0.1f; +constexpr int kDefaultLogEvery = 50; + +struct Batch { + executorch::extension::TensorPtr x; + executorch::extension::TensorPtr target; +}; + +// Loads and parses a JSON file (config.json or manifest.json). +nlohmann::json load_json(const std::string& path); + +// Parses a safetensors file into owned F32 tensors keyed by name. +std::unordered_map +load_safetensors(const std::string& path); + +// Loads init_weights.safetensors and reorders tensors to match manifest's +// "param_names" order (the joint graph's primals_N placeholder order). +std::vector load_weights( + const nlohmann::json& manifest, + const std::string& safetensors_path); + +// Converts a JSON int array (e.g. manifest's "x_shape") to tensor sizes. +std::vector shape_from_json( + const nlohmann::json& shape_array); + +// Generates one random F32 tensor of `shape` via N(0,1), matching the +// torch.randn semantics the Python export scripts used. +executorch::extension::TensorPtr make_random_tensor( + const std::vector& shape, + std::mt19937& rng); + +// Writes tensors (keyed by name, all F32) to a safetensors file. +void save_safetensors( + const std::string& path, + const std::vector& names, + const std::vector& tensors); + +// Resolves manifest_dir from manifest_path (directory part). +std::string manifest_dir(const std::string& manifest_path); + +} // namespace qnn_toy_model_training diff --git a/examples/qualcomm/training/toy_model/on-device/inference.cpp b/examples/qualcomm/training/toy_model/on-device/inference.cpp new file mode 100644 index 00000000000..2239df6d7d5 --- /dev/null +++ b/examples/qualcomm/training/toy_model/on-device/inference.cpp @@ -0,0 +1,93 @@ +/* + * Copyright (c) Qualcomm Innovation Center, Inc. + * All rights reserved + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +// On-device inference runner for the QNN/HTP-delegated ToyModel. +// +// Loads trained weights from trained_model.safetensors (written by trainer +// after the training loop), and feeds (*weights, x) into the "forward" method. + +// Usage: qnn_toy_model_inference +// trained_model.safetensors is expected alongside manifest.json. Each batch's +// x is generated on-device via N(0,1), seeded from config.json's "seed" if +// present. + +#include "common.h" + +#include + +using executorch::extension::Module; +using executorch::extension::TensorPtr; +using executorch::runtime::EValue; + +namespace qnn_toy_model_training { + +static void run_inference( + Module& module, + const nlohmann::json& config, + const nlohmann::json& manifest, + const std::string& trained_weights_path) { + const int log_every = config.value("log_every", kDefaultLogEvery); + + std::vector weights = load_weights(manifest, trained_weights_path); + + const auto x_shape = shape_from_json(manifest.at("x_shape")); + std::mt19937 rng( + config.contains("seed") ? config.at("seed").get() + : std::random_device{}()); + + const int num_batches = config.value("steps", kDefaultSteps); + + for (int i = 0; i < num_batches; i++) { + TensorPtr x = make_random_tensor(x_shape, rng); + + std::vector inputs; + inputs.reserve(weights.size() + 1); + for (auto& w : weights) { + inputs.push_back(*w); + } + inputs.push_back(*x); + + auto outputs = module.execute("forward", inputs); + ET_CHECK_MSG(outputs.ok(), "forward execution failed at batch %d.", i); + + if (i % log_every == 0 || i == num_batches - 1) { + const auto& pred = outputs->at(0).toTensor(); + ET_LOG( + Info, "batch %5d pred[0] %.6f", i, pred.const_data_ptr()[0]); + } + } +} + +} // namespace qnn_toy_model_training + +int main(int argc, char** argv) { + if (argc != 3) { + ET_LOG(Error, "Usage: %s ", argv[0]); + return 1; + } + const std::string config_path = argv[1]; + const std::string manifest_path = argv[2]; + + const nlohmann::json config = qnn_toy_model_training::load_json(config_path); + const nlohmann::json manifest = + qnn_toy_model_training::load_json(manifest_path); + + const std::string dir = qnn_toy_model_training::manifest_dir(manifest_path); + // no_ws: fwd_pte_name is the forward-only .pte; ws: pte_name holds both + // methods. + const std::string pte_name = manifest.count("fwd_pte_name") + ? manifest.at("fwd_pte_name").get() + : manifest.at("pte_name").get(); + const std::string pte_path = dir + "/" + pte_name; + const std::string trained_weights_path = dir + "/trained_model.safetensors"; + + executorch::extension::Module module(pte_path); + qnn_toy_model_training::run_inference( + module, config, manifest, trained_weights_path); + return 0; +} diff --git a/examples/qualcomm/training/toy_model/on-device/trainer.cpp b/examples/qualcomm/training/toy_model/on-device/trainer.cpp new file mode 100644 index 00000000000..7915649c999 --- /dev/null +++ b/examples/qualcomm/training/toy_model/on-device/trainer.cpp @@ -0,0 +1,137 @@ +/* + * Copyright (c) Qualcomm Innovation Center, Inc. + * All rights reserved + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +// On-device C++ trainer runs the QNN/HTP-delegated +// "forward_backward" method directly via executorch::extension::Module and +// applies SGD. +// +// Usage: qnn_toy_model_trainer +// The .pte and init_weights.safetensors are expected alongside manifest.json +// (export_toy_model.py's layout). A single (x, target) batch is generated +// on-device via N(0,1) once, seeded from config.json's "seed" if present, and +// reused every step -- the model memorizes it, so loss should decrease +// monotonically. + +#include "common.h" + +#include + +using executorch::extension::Module; +using executorch::extension::TensorPtr; +using executorch::runtime::EValue; + +namespace qnn_toy_model_training { + +static void train( + Module& module, + const nlohmann::json& config, + const nlohmann::json& manifest, + const std::string& safetensors_path, + const std::string& model_weights_path) { + const size_t num_weights = manifest.at("param_names").size(); + + auto meta_result = module.method_meta("forward_backward"); + ET_CHECK_MSG( + meta_result.ok(), "Failed to load forward_backward method metadata."); + const auto& meta = *meta_result; + ET_CHECK_MSG( + meta.num_outputs() == num_weights + 2, + "Expected %zu outputs (loss, pred, %zu grads), got %zu.", + num_weights + 2, + num_weights, + meta.num_outputs()); + + std::vector weights = load_weights(manifest, safetensors_path); + ET_CHECK_MSG( + weights.size() == num_weights, + "Loaded %zu weights from safetensors, expected %zu.", + weights.size(), + num_weights); + + const auto x_shape = shape_from_json(manifest.at("x_shape")); + const auto target_shape = shape_from_json(manifest.at("target_shape")); + std::mt19937 rng( + config.contains("seed") ? config.at("seed").get() + : std::random_device{}()); + + // export_toy_model.py's tangent is always a constant torch.ones([1]). + auto tangent = executorch::extension::make_tensor_ptr( + std::vector{1}, std::vector{1.0f}); + + const int steps = config.value("steps", kDefaultSteps); + const float lr = config.value("lr", kDefaultLr); + const int log_every = config.value("log_every", kDefaultLogEvery); + + TensorPtr x = make_random_tensor(x_shape, rng); + TensorPtr target = make_random_tensor(target_shape, rng); + + for (int step = 0; step < steps; step++) { + std::vector inputs; + inputs.reserve(num_weights + 3); + for (auto& w : weights) { + inputs.push_back(*w); + } + inputs.push_back(*x); + inputs.push_back(*target); + inputs.push_back(*tangent); + + auto outputs = module.execute("forward_backward", inputs); + ET_CHECK_MSG( + outputs.ok(), "forward_backward execution failed at step %d.", step); + + for (size_t i = 0; i < num_weights; i++) { + const auto& grad = outputs->at(2 + i).toTensor(); + float* w = weights[i]->mutable_data_ptr(); + const float* g = grad.const_data_ptr(); + const size_t n = weights[i]->numel(); + for (size_t j = 0; j < n; j++) { + w[j] -= lr * g[j]; + } + } + + if (step % log_every == 0 || step == steps - 1) { + const auto& loss = outputs->at(0).toTensor(); + ET_LOG( + Info, "step %5d loss %.6f", step, loss.const_data_ptr()[0]); + } + } + + // Save trained weights so inference.cpp can load them. + std::vector param_names; + for (const auto& name : manifest.at("param_names")) { + param_names.push_back(name.get()); + } + save_safetensors(model_weights_path, param_names, weights); + ET_LOG(Info, "Saved trained weights to %s", model_weights_path.c_str()); +} + +} // namespace qnn_toy_model_training + +int main(int argc, char** argv) { + if (argc != 3) { + ET_LOG(Error, "Usage: %s ", argv[0]); + return 1; + } + const std::string config_path = argv[1]; + const std::string manifest_path = argv[2]; + + const nlohmann::json config = qnn_toy_model_training::load_json(config_path); + const nlohmann::json manifest = + qnn_toy_model_training::load_json(manifest_path); + + const std::string dir = qnn_toy_model_training::manifest_dir(manifest_path); + const std::string pte_path = + dir + "/" + manifest.at("pte_name").get(); + const std::string safetensors_path = dir + "/init_weights.safetensors"; + const std::string model_weights_path = dir + "/trained_model.safetensors"; + + executorch::extension::Module module(pte_path); + qnn_toy_model_training::train( + module, config, manifest, safetensors_path, model_weights_path); + return 0; +} diff --git a/examples/qualcomm/training/toy_model/run_toy_model.py b/examples/qualcomm/training/toy_model/run_toy_model.py new file mode 100644 index 00000000000..7c543719f09 --- /dev/null +++ b/examples/qualcomm/training/toy_model/run_toy_model.py @@ -0,0 +1,214 @@ +# Copyright (c) Qualcomm Innovation Center, Inc. +# All rights reserved +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""Push, train, and run inference for the QNN/HTP joint-graph ToyModel. + +Extends SimpleADB to push the custom trainer/inference binaries and model +artifacts, then run training followed by on-device inference. + +Usage: + python run_toy_model.py --device --soc_model SM8850 \\ + --build_folder /path/to/build-android + + # push only + python run_toy_model.py ... --push_only + + # skip push, train only (re-run without re-pushing) + python run_toy_model.py ... --skip_push --skip_inference + + # skip push, infer only (after training already done on device) + python run_toy_model.py ... --skip_push --skip_train + +All config.json fields are read from --config (default: toy_model/config.json). +QNN_SDK_ROOT must be set (same requirement as all other Qualcomm examples). +""" + +import argparse +import logging +import os + +from executorch.backends.qualcomm.export_utils import ( + QnnConfig, + setup_common_args_and_variables, + SimpleADB, +) + +if __package__: + from .export.config import load_config +else: + import sys + + sys.path.insert(0, os.path.join(os.path.dirname(__file__), "export")) + from config import load_config + +logger = logging.getLogger(__name__) + +_DEVICE_ROOT_TMPL = "/data/local/tmp/{user}/qnn_toy_model_trainer" +_BIN_SUBDIR = "bin" +_ARTIFACTS_SUBDIR = "artifacts" + + +class TrainADB(SimpleADB): + """SimpleADB subclass for the ToyModel on-device training example. + + chmod_binaries() and push_artifacts() finish what SimpleADB.push() leaves + to do (chmod the pushed trainer/inference binaries; push the artifacts/ + files alongside them); train() / infer() then run the binaries from the + bin/ subdir, referencing artifacts/ for config and manifest. + """ + + def __init__(self, qnn_config: QnnConfig, workspace: str): + super().__init__( + qnn_config=qnn_config, + pte_path=[], # ptes live in artifacts/, pushed separately below. + workspace=os.path.join(workspace, _BIN_SUBDIR), + runner="examples/qualcomm/training/toy_model/on-device/qnn_toy_model_trainer", + ) + self.workspace_root = workspace + + def chmod_binaries(self): + """chmod+x the trainer/inference binaries already pushed by push().""" + if self.skip_push: + return + bin_dir = f"{self.workspace_root}/{_BIN_SUBDIR}" + for binary in ("qnn_toy_model_trainer", "qnn_toy_model_inference"): + self._adb(["shell", f"chmod +x {bin_dir}/{binary}"]) + + def push_artifacts(self, artifact_files, pte_files): + """Push config/manifest/weights/pte files to the artifacts/ subdir.""" + if self.skip_push: + return + art_dir = f"{self.workspace_root}/{_ARTIFACTS_SUBDIR}" + self._adb(["shell", f"rm -rf {art_dir}"]) + self._adb(["shell", f"mkdir -p {art_dir}"]) + for f in artifact_files + pte_files: + self._adb(["push", f, f"{art_dir}/"]) + + def train(self): + """Run qnn_toy_model_trainer on-device.""" + bin_dir = f"{self.workspace_root}/{_BIN_SUBDIR}" + art_dir = f"{self.workspace_root}/{_ARTIFACTS_SUBDIR}" + self.execute( + custom_runner_cmd=" ".join( + [ + f"cd {bin_dir} &&", + "export LD_LIBRARY_PATH=. && export ADSP_LIBRARY_PATH=. &&", + f"./qnn_toy_model_trainer {art_dir}/config.json {art_dir}/manifest.json", + ] + ) + ) + + def infer(self): + """Run qnn_toy_model_inference on-device.""" + bin_dir = f"{self.workspace_root}/{_BIN_SUBDIR}" + art_dir = f"{self.workspace_root}/{_ARTIFACTS_SUBDIR}" + self.execute( + custom_runner_cmd=" ".join( + [ + f"cd {bin_dir} &&", + "export LD_LIBRARY_PATH=. && export ADSP_LIBRARY_PATH=. &&", + f"./qnn_toy_model_inference {art_dir}/config.json {art_dir}/manifest.json", + ] + ) + ) + + +def main(args): + script_dir = os.path.dirname(os.path.abspath(__file__)) + cfg = load_config( + args.config if args.config else os.path.join(script_dir, "config.json") + ) + + weight_sharing = ( + args.weight_sharing + if args.weight_sharing is not None + else cfg.get("weight_sharing", False) + ) + output_dir = args.output_dir or cfg.get("output_dir", "outputs/artifacts") + if not os.path.isabs(output_dir): + output_dir = os.path.join(script_dir, output_dir) + artifact_subdir = os.path.join(output_dir, "ws" if weight_sharing else "no_ws") + + assert os.path.isdir( + artifact_subdir + ), f"Artifact dir not found: {artifact_subdir}; run export_toy_model.py first." + + pte_files = [ + os.path.join(artifact_subdir, f) + for f in os.listdir(artifact_subdir) + if f.endswith(".pte") + ] + assert ( + pte_files + ), f"No .pte found in {artifact_subdir}; run export_toy_model.py first." + + device_root = _DEVICE_ROOT_TMPL.format(user=os.environ.get("USER", "user")) + + qnn_config = QnnConfig.load_config(args) + adb = TrainADB(qnn_config=qnn_config, workspace=device_root) + + inference_bin = os.path.join( + args.build_folder, + "examples/qualcomm/training/toy_model/on-device/qnn_toy_model_inference", + ) + # Binaries + QNN libs go to bin/ (handled by SimpleADB.push via workspace=bin/). + # Artifacts (ptes already in pte_path, plus config/manifest/weights) go to artifacts/. + artifact_files = [ + os.path.join(script_dir, "config.json"), + os.path.join(artifact_subdir, "manifest.json"), + os.path.join(artifact_subdir, "init_weights.safetensors"), + ] + adb.push(files=[inference_bin]) + adb.chmod_binaries() + adb.push_artifacts(artifact_files, pte_files) + + if not args.push_only: + if not args.skip_train: + adb.train() + if not args.skip_inference: + adb.infer() + + +if __name__ == "__main__": + parser = setup_common_args_and_variables() + parser.add_argument( + "--config", + type=str, + default=None, + help="Path to config.json (default: config.json next to this script).", + ) + parser.add_argument( + "--output_dir", + type=str, + default=None, + help="Base artifacts directory (default: from config.json).", + ) + parser.add_argument( + "--weight_sharing", + action=argparse.BooleanOptionalAction, + default=None, + help="Use weight-sharing variant (ws/ subdir). Default: from config.json.", + ) + parser.add_argument( + "--skip_train", + action="store_true", + default=False, + help="Skip on-device training (run inference only).", + ) + parser.add_argument( + "--skip_inference", + action="store_true", + default=False, + help="Skip on-device inference (train only).", + ) + parser.add_argument( + "--push_only", + action="store_true", + default=False, + help="Push binaries/artifacts to device, but skip training and inference.", + ) + logging.basicConfig(level=logging.INFO) + main(parser.parse_args())