Add interpret docs to example component, remove test arg from cv tr… (#255)

Summary:
…ainer

Pull Request resolved: #255

Reviewed By: kiukchung

Differential Revision: D31660518

Pulled By: aivanou

fbshipit-source-id: dc4e3c4698545a318875b5b8fe970cf7bb5b2cd4

Author: aivanou

torchx/components/interpret.py

@@ -16,6 +16,6 @@
 See the
 :ref:`examples_apps/lightning_classy_vision/interpret:Model Interpretability App Example`
 and the corresponding
-:ref:`Interpret component definition<examples_apps/lightning_classy_vision/component:Trainer Component Examples>`
+:ref:`Interpret component definition<examples_apps/lightning_classy_vision/component:Interpreting the Model>`
 for an example of how to use Captum with TorchX.
 """

torchx/examples/apps/lightning_classy_vision/component.py

@@ -46,7 +46,7 @@
 #
 #    torchx run --scheduler local_cwd \
 #    ./torchx/examples/apps/lightning_classy_vision/component.py:trainer \
-#    --output_path /tmp
+#    --output_path /tmp/$USER
 #
 # Single trainer component code:
 
@@ -258,44 +258,58 @@ def trainer_dist(
 
 
 # %%
-# Model Interpretability
+# Interpreting the Model
 # #######################
-# TODO(aivanou): add documentation
+# Defines a component that interprets the model
+#
+# Train a single trainer example: :ref:`examples_apps/lightning_classy_vision/component:Single Trainer Component`
+# And use the following cmd to try out:
+#
+# .. code:: bash
+#
+#    torchx run --scheduler local_cwd \
+#    ./torchx/examples/apps/lightning_classy_vision/component.py:interpret \
+#    --output_path /tmp/aivanou/interpret  --load_path /tmp/$USER/last.ckpt
 
 
 def interpret(
-    image: str,
     load_path: str,
-    data_path: str,
     output_path: str,
+    data_path: Optional[str] = None,
+    image: str = TORCHX_IMAGE,
     resource: Optional[str] = None,
 ) -> torchx.AppDef:
     """Runs the model interpretability app on the model outputted by the training
     component.
 
     Args:
-        image: image to run (e.g. foobar:latest)
         load_path: path to load pretrained model from
-        data_path: path to the data to load
         output_path: output path for model checkpoints (e.g. file:///foo/bar)
+        data_path: path to the data to load
+        image: image to run (e.g. foobar:latest)
         resource: the resources to use
     """
+    args = [
+        "-m",
+        "torchx.examples.apps.lightning_classy_vision.interpret",
+        "--load_path",
+        load_path,
+        "--output_path",
+        output_path,
+    ]
+    if data_path:
+        args += [
+            "--data_path",
+            data_path,
+        ]
+
     return torchx.AppDef(
         name="cv-interpret",
         roles=[
             torchx.Role(
                 name="worker",
                 entrypoint="python",
-                args=[
-                    "-m",
-                    "torchx.examples.apps.lightning_classy_vision.interpret",
-                    "--load_path",
-                    load_path,
-                    "--data_path",
-                    data_path,
-                    "--output_path",
-                    output_path,
-                ],
+                args=args,
                 image=image,
                 resource=named_resources[resource]
                 if resource

torchx/examples/apps/lightning_classy_vision/data.py

@@ -148,7 +148,7 @@ def download_data(remote_path: str, tmpdir: str) -> str:
     return data_path
 
 
-def create_random_data(output_path: str, num_images: int = 5) -> None:
+def create_random_data(output_path: str, num_images: int = 250) -> None:
     """
     Fills the given path with randomly generated 64x64 images.
     This can be used for quick testing of the workflow of the model.

torchx/examples/apps/lightning_classy_vision/interpret.py

@@ -35,10 +35,10 @@
 from torchx.examples.apps.lightning_classy_vision.data import (
     TinyImageNetDataModule,
     download_data,
+    create_random_data,
 )
 from torchx.examples.apps.lightning_classy_vision.model import TinyImageNetModel
 
-
 # FIXME: captum must be imported after torch otherwise it causes python to crash
 if True:
     import numpy as np
@@ -57,8 +57,7 @@ def parse_args(argv: List[str]) -> argparse.Namespace:
     parser.add_argument(
         "--data_path",
         type=str,
-        help="path to load the training data from",
-        required=True,
+        help="path to load the training data from, if not provided, random dataset will be created",
     )
     parser.add_argument(
         "--output_path",
@@ -91,7 +90,12 @@ def main(argv: List[str]) -> None:
         model.load_from_checkpoint(checkpoint_path=args.load_path)
 
         # Download and setup the data module
-        data_path = download_data(args.data_path, tmpdir)
+        if not args.data_path:
+            data_path = os.path.join(tmpdir, "data")
+            os.makedirs(data_path)
+            create_random_data(data_path)
+        else:
+            data_path = download_data(args.data_path, tmpdir)
         data = TinyImageNetDataModule(
             data_dir=data_path,
             batch_size=1,

torchx/examples/apps/lightning_classy_vision/train.py

@@ -56,12 +56,7 @@ def parse_args(argv: List[str]) -> argparse.Namespace:
     parser.add_argument(
         "--batch_size", type=int, default=32, help="batch size to use for training"
     )
-    parser.add_argument("--num_samples", type=int, default=None, help="num_samples")
-    parser.add_argument(
-        "--test",
-        help="Sets to test mode, training on a much smaller set of randomly generated images",
-        action="store_true",
-    )
+    parser.add_argument("--num_samples", type=int, default=10, help="num_samples")
     parser.add_argument(
         "--data_path",
         type=str,
@@ -122,7 +117,7 @@ def main(argv: List[str]) -> None:
         data = TinyImageNetDataModule(
             data_dir=data_path,
             batch_size=args.batch_size,
-            num_samples=5 if args.test else args.num_samples,
+            num_samples=args.num_samples,
         )
 
         # Setup model checkpointing