lightning_classy_vision: add model interpretability example (#54)

Summary:
This is a model interpretability example. It loads the model from a training step and then runs integrated gradients on it via captum (https://captum.ai/tutorials/CIFAR_TorchVision_Interpret) and then writes out 5 samples to the output path.

Pull Request resolved: #54

Test Plan:
```
cd examples/apps && torchx run --scheduler local lightning_classy_vision/component.py:interpret --image (pwd) --output_path /tmp/output --data_path /tmp/data.tar.gz --load_path /tmp/model-out/last.ckpt
scripts/kfpint.py
```

http://5ab6bab9-istiosystem-istio-2af2-1926929629.us-west-2.elb.amazonaws.com/_/pipeline/?ns=torchx-dev#/runs/details/0a4a26f2-e1b2-4b99-be4e-3c0d74875770

Reviewed By: kiukchung

Differential Revision: D29085948

Pulled By: d4l3k

fbshipit-source-id: 3c3922e4f8eb6cdb97251eb20f3f403b4cb2bbbb

Author: d4l3k

.pyre_configuration

@@ -1,5 +1,6 @@
 {
   "source_directories": [
+    "examples/apps/lightning_classy_vision",
     "."
   ],
   "strict": true,

dev-requirements.txt

@@ -11,3 +11,4 @@ classy-vision>=0.5.0
 flake8==3.9.0
 ts>=0.5.1
 torchserve>=0.4.0
+captum>=0.3.1

docs/Makefile

@@ -19,6 +19,7 @@ help:
 clean:
 	@echo "Deleting build directory"
 	rm -rf "$(BUILDDIR)"
+	rm -rf "$(SOURCEDIR)/examples_apps" "$(SOURCEDIR)/examples_pipelines"
 
 .PHONY: help Makefile clean livehtml
 

examples/apps/Dockerfile

@@ -1,6 +1,6 @@
 FROM pytorch/pytorch
 
-RUN pip install classy_vision pytorch-lightning fsspec[s3] torch-model-archiver
+RUN pip install classy_vision pytorch-lightning fsspec[s3] torch-model-archiver captum
 
 WORKDIR /app
 

examples/apps/lightning_classy_vision/component.py

@@ -11,44 +11,81 @@
 This is a component definition that runs the example lightning_classy_vision app.
 """
 
+from typing import Optional
+
 import torchx.specs.api as torchx
+from torchx.components.base import named_resource
+from torchx.components.base.binary_component import binary_component
 
 
-def classy_vision(
+def trainer(
     image: str,
     output_path: str,
+    data_path: str,
     load_path: str = "",
     log_dir: str = "/logs",
+    resource: Optional[str] = None,
 ) -> torchx.AppDef:
     """Runs the example lightning_classy_vision app.
 
-    Runs the example lightning_classy_vision app.
-
     Args:
         image: image to run (e.g. foobar:latest)
-        resource: resource spec
         output_path: output path for model checkpoints (e.g. file:///foo/bar)
         load_path: path to load pretrained model from
+        data_path: path to the data to load
         log_dir: path to save tensorboard logs to
+        resource: the resources to use
     """
-    entrypoint = "main"
-
-    trainer_role = (
-        torchx.Role(
-            name="trainer",
-            image=image,
-            resource=torchx.Resource(cpu=1, gpu=1, memMB=1024),
-        )
-        .runs(
-            "main",
+    return binary_component(
+        name="examples-lightning_classy_vision-trainer",
+        entrypoint="lightning_classy_vision/train.py",
+        args=[
             "--output_path",
             output_path,
             "--load_path",
             load_path,
             "--log_dir",
             log_dir,
-        )
-        .replicas(1)
+            "--data_path",
+            data_path,
+        ],
+        image=image,
+        resource=named_resource(resource)
+        if resource
+        else torchx.Resource(cpu=1, gpu=0, memMB=1024),
     )
 
-    return torchx.AppDef("examples-lightning_classy_vision").of(trainer_role)
+
+def interpret(
+    image: str,
+    load_path: str,
+    data_path: str,
+    output_path: str,
+    resource: Optional[str] = None,
+) -> torchx.AppDef:
+    """Runs the model intepretability 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)
+        resource: the resources to use
+    """
+    return binary_component(
+        name="examples-lightning_classy_vision-intepret",
+        entrypoint="lightning_classy_vision/interpret.py",
+        args=[
+            "--load_path",
+            load_path,
+            "--data_path",
+            data_path,
+            "--output_path",
+            output_path,
+        ],
+        image=image,
+        resource=named_resource(resource)
+        if resource
+        else torchx.Resource(cpu=1, gpu=0, memMB=1024),
+    )

examples/apps/lightning_classy_vision/data.py

@@ -0,0 +1,130 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+# 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.
+
+"""
+Trainer Datasets Example
+========================
+
+This is the datasets used for the training example. It's using stock Pytorch
+Lightning + Classy Vision libraries.
+"""
+
+import os.path
+import tarfile
+from typing import Optional, Callable
+
+import fsspec
+import pytorch_lightning as pl
+from classy_vision.dataset.classy_dataset import ClassyDataset
+from torch.utils.data import DataLoader
+from torchvision import datasets, transforms
+
+# %%
+# This uses classy vision to define a dataset that we will then later use in our
+# Pytorch Lightning data module.
+
+
+class TinyImageNetDataset(ClassyDataset):
+    """
+    TinyImageNetDataset is a ClassyDataset for the tiny imagenet dataset.
+    """
+
+    def __init__(self, data_path: str, transform: Callable[[object], object]) -> None:
+        batchsize_per_replica = 16
+        shuffle = False
+        num_samples = 1000
+        dataset = datasets.ImageFolder(data_path)
+        super().__init__(
+            # pyre-fixme[6]
+            dataset,
+            batchsize_per_replica,
+            shuffle,
+            transform,
+            num_samples,
+        )
+
+
+# %%
+# For easy of use, we define a lightning data module so we can reuse it across
+# our trainer and other components that need to load data.
+
+# pyre-fixme[13]: Attribute `test_ds` is never initialized.
+# pyre-fixme[13]: Attribute `train_ds` is never initialized.
+# pyre-fixme[13]: Attribute `val_ds` is never initialized.
+class TinyImageNetDataModule(pl.LightningDataModule):
+    """
+    TinyImageNetDataModule is a pytorch LightningDataModule for the tiny
+    imagenet dataset.
+    """
+
+    train_ds: TinyImageNetDataset
+    val_ds: TinyImageNetDataset
+    test_ds: TinyImageNetDataset
+
+    def __init__(self, data_dir: str, batch_size: int = 16) -> None:
+        super().__init__()
+        self.data_dir = data_dir
+        self.batch_size = batch_size
+
+    def setup(self, stage: Optional[str] = None) -> None:
+        # Setup data loader and transforms
+        img_transform = transforms.Compose(
+            [
+                transforms.Grayscale(),
+                transforms.ToTensor(),
+            ]
+        )
+        self.train_ds = TinyImageNetDataset(
+            data_path=os.path.join(self.data_dir, "train"),
+            transform=lambda x: (img_transform(x[0]), x[1]),
+        )
+        self.val_ds = TinyImageNetDataset(
+            data_path=os.path.join(self.data_dir, "val"),
+            transform=lambda x: (img_transform(x[0]), x[1]),
+        )
+        self.test_ds = TinyImageNetDataset(
+            data_path=os.path.join(self.data_dir, "test"),
+            transform=lambda x: (img_transform(x[0]), x[1]),
+        )
+
+    def train_dataloader(self) -> DataLoader:
+        # pyre-fixme[6]
+        return DataLoader(self.train_ds, batch_size=self.batch_size)
+
+    def val_dataloader(self) -> DataLoader:
+        # pyre-fixme[6]:
+        return DataLoader(self.val_ds, batch_size=self.batch_size)
+
+    def test_dataloader(self) -> DataLoader:
+        # pyre-fixme[6]
+        return DataLoader(self.test_ds, batch_size=self.batch_size)
+
+    def teardown(self, stage: Optional[str] = None) -> None:
+        pass
+
+
+# %%
+# To pass data between the different components we use fsspec which allows us to
+# read/write to cloud or local file storage.
+
+
+def download_data(remote_path: str, tmpdir: str) -> str:
+    """
+    download_data downloads the training data from the specified remote path via
+    fsspec and places it in the tmpdir unextracted.
+    """
+    tar_path = os.path.join(tmpdir, "data.tar.gz")
+    print(f"downloading dataset from {remote_path} to {tar_path}...")
+    fs, _, rpaths = fsspec.get_fs_token_paths(remote_path)
+    assert len(rpaths) == 1, "must have single path"
+    fs.get(rpaths[0], tar_path)
+
+    data_path = os.path.join(tmpdir, "data")
+    print(f"extracting {tar_path} to {data_path}...")
+    with tarfile.open(tar_path, mode="r") as f:
+        f.extractall(data_path)
+
+    return data_path

examples/apps/lightning_classy_vision/interpret.py

@@ -0,0 +1,130 @@
+#!/usr/bin/env python3
+# Copyright (c) Facebook, Inc. and its affiliates.
+# 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.
+
+"""
+Model Interpretability App Example
+=============================================
+
+This is an example TorchX app that uses captum to analyze inputs to for model
+intepretability purposes. It consumes the trained model from the trainer app
+example and the preprocessed examples from the datapreproc app example. The
+output is a series of images with integrated gradient attributions overlayed on
+them.
+
+See https://captum.ai/tutorials/CIFAR_TorchVision_Interpret for more info on
+using captum.
+"""
+
+import argparse
+import itertools
+import os.path
+import sys
+import tempfile
+from typing import List
+
+import fsspec
+import torch
+from data import TinyImageNetDataModule, download_data
+from model import TinyImageNetModel
+
+
+# FIXME: captum must be imported after torch otherwise it causes python to crash
+if True:
+    import numpy as np
+    from captum.attr import IntegratedGradients
+    from captum.attr import visualization as viz
+
+
+def parse_args(argv: List[str]) -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="example TorchX captum app")
+    parser.add_argument(
+        "--load_path",
+        type=str,
+        help="checkpoint path to load model weights from",
+        required=True,
+    )
+    parser.add_argument(
+        "--data_path",
+        type=str,
+        help="path to load the training data from",
+        required=True,
+    )
+    parser.add_argument(
+        "--output_path",
+        type=str,
+        help="path to place analysis results",
+        required=True,
+    )
+
+    return parser.parse_args(argv)
+
+
+def convert_to_rgb(arr: torch.Tensor) -> np.ndarray:
+    """
+    This converts the image from a torch tensor with size (1, 1, 64, 64) to
+    numpy array with size (64, 64, 3).
+    """
+    squeezed = arr.squeeze()
+    stacked = torch.stack([squeezed, squeezed, squeezed], dim=2)
+    print("img stats: ", stacked.count_nonzero(), stacked.mean())
+    return stacked.numpy()
+
+
+def main(argv: List[str]) -> None:
+    with tempfile.TemporaryDirectory() as tmpdir:
+        args = parse_args(argv)
+
+        # Init our model
+        model = TinyImageNetModel()
+
+        print(f"loading checkpoint: {args.load_path}...")
+        model.load_from_checkpoint(checkpoint_path=args.load_path)
+
+        # Download and setup the data module
+        data_path = download_data(args.data_path, tmpdir)
+        data = TinyImageNetDataModule(
+            data_dir=data_path,
+            batch_size=1,
+        )
+
+        ig = IntegratedGradients(model)
+
+        data.setup("test")
+        dataloader = data.test_dataloader()
+
+        # process first 5 images
+        for i, (input, label) in enumerate(itertools.islice(dataloader, 5)):
+            print(f"analyzing example {i}")
+            model.zero_grad()
+            attr_ig, delta = ig.attribute(
+                input,
+                target=label,
+                baselines=input * 0,
+                return_convergence_delta=True,
+            )
+
+            if attr_ig.count_nonzero() == 0:
+                # Our toy model sometimes has no IG results.
+                print("skipping due to zero gradients")
+                continue
+
+            fig, axis = viz.visualize_image_attr(
+                convert_to_rgb(attr_ig),
+                convert_to_rgb(input),
+                method="blended_heat_map",
+                sign="all",
+                show_colorbar=True,
+                title="Overlayed Integrated Gradients",
+            )
+            out_path = os.path.join(args.output_path, f"ig_{i}.png")
+            print(f"saving heatmap to {out_path}")
+            with fsspec.open(out_path, "wb") as f:
+                fig.savefig(f)
+
+
+if __name__ == "__main__":
+    main(sys.argv[1:])