Adding --one-shot argument to torchvision export (#1300)

Co-authored-by: Rahul Tuli <[email protected]>

Author: 2 people

src/sparseml/pytorch/torchvision/export_onnx.py

@@ -26,6 +26,7 @@
 import click
 from sparseml.pytorch.models.registry import ModelRegistry
 from sparseml.pytorch.opset import TORCH_DEFAULT_ONNX_OPSET
+from sparseml.pytorch.optim.manager import ScheduledModifierManager
 from sparseml.pytorch.torchvision import presets
 from sparseml.pytorch.utils import ModuleExporter
 from sparseml.pytorch.utils.model import load_model
@@ -60,6 +61,12 @@
     help="The root dir path where the dataset is stored or should "
     "be downloaded to if available",
 )
+@click.option(
+    "--one-shot",
+    default=None,
+    type=str,
+    help="Path to recipe to use to apply in a one-shot manner",
+)
 @click.option(
     "--labels-to-class-mapping",
     type=click.Path(dir_okay=False, file_okay=True, exists=True, path_type=Path),
@@ -118,6 +125,7 @@ def main(
     arch_key: str,
     checkpoint_path: str,
     dataset_path: Path,
+    one_shot: Optional[str],
     labels_to_class_mapping: Optional[Path],
     num_samples: int,
     onnx_opset: int,
@@ -159,6 +167,9 @@ def main(
 
     load_model(checkpoint_path, model, strict=True)
 
+    if one_shot is not None:
+        ScheduledModifierManager.from_yaml(one_shot).apply(model)
+
     if labels_to_class_mapping is not None:
         with open(labels_to_class_mapping) as fp:
             labels_to_class_mapping = json.load(fp)