classy vision example

Summary:
WIP

This is an example of a simple classy vision model using torchx.

Reviewed By: kiukchung

Differential Revision: D28498759

fbshipit-source-id: 41e034d9d08c59a661181aaec7fda9f6b895b269

Author: d4l3k

torchx/examples/lightning_classy_vision/lightning_classy_vision.torchx

@@ -0,0 +1,49 @@
+"""
+description: Runs the example lightning_classy_vision app.
+arguments:
+  - name: --image
+    type: str
+    help: image to run (e.g. foobar:latest)
+  - name: --resource
+    type: str
+    help: resource spec
+    default: T1
+  - name: --output_path
+    type: str
+    help: output path for model checkpoints (e.g. file:///foo/bar)
+    required: true
+  - name: --load_path
+    type: str
+    help: path to load pretrained model from
+    default: ""
+  - name: --log_dir
+    type: str
+    help: path to save tensorboard logs to
+    default: "/logs"
+"""
+
+import torchx.specs.api as torchx
+import torchx.schedulers.fb.resource as resource
+
+container = torchx.Container(image=args.image).require(resources=resource.get(args.resource))
+entrypoint = "main"
+
+trainer_role = (
+    torchx.Role(
+        name="trainer"
+    )
+    .runs(
+      "main",
+      "--output_path",
+      args.output_path,
+      "--load_path",
+      args.load_path,
+      "--log_dir",
+      args.log_dir,
+    )
+    .on(container)
+    .replicas(1)
+)
+
+app = torchx.Application("examples-lightning_classy_vision").of(trainer_role)
+export(app)

torchx/examples/lightning_classy_vision/main.py

@@ -0,0 +1,128 @@
+#!/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.
+
+# pyre-unsafe
+
+import argparse
+import sys
+from typing import List
+
+import pytorch_lightning as pl
+import torch
+from classy_vision.dataset.classy_dataset import ClassyDataset
+from classy_vision.dataset.core.random_image_datasets import (
+    RandomImageDataset,
+    SampleType,
+)
+from pytorch_lightning.callbacks import ModelCheckpoint
+from pytorch_lightning.loggers import TensorBoardLogger
+from torch.nn import functional as F
+from torch.utils.data import DataLoader
+from torchvision import transforms
+
+
+class SyntheticMNIST(ClassyDataset):
+    def __init__(self, transform):
+        batchsize_per_replica = 16
+        shuffle = True
+        num_samples = 1000
+        dataset = RandomImageDataset(
+            crop_size=28,
+            num_channels=3,
+            num_samples=num_samples,
+            num_classes=10,
+            seed=1234,
+            sample_type=SampleType.TUPLE,
+        )
+        super().__init__(
+            dataset, batchsize_per_replica, shuffle, transform, num_samples
+        )
+
+
+class MNISTModel(pl.LightningModule):
+    def __init__(self):
+        super(MNISTModel, self).__init__()
+        self.l1 = torch.nn.Linear(28 * 28, 10)
+
+    def forward(self, x):
+        return torch.relu(self.l1(x.view(x.size(0), -1)))
+
+    def training_step(self, batch, batch_nb):
+        x, y = batch
+        loss = F.cross_entropy(self(x), y)
+        return loss
+
+    def configure_optimizers(self):
+        return torch.optim.Adam(self.parameters(), lr=0.02)
+
+
+def parse_args(argv: List[str]) -> argparse.Namespace:
+    parser = argparse.ArgumentParser(
+        description="pytorch lightning + classy vision TorchX example app"
+    )
+    parser.add_argument(
+        "--epochs", type=int, default=3, help="number of epochs to train"
+    )
+    parser.add_argument(
+        "--batch_size", type=int, default=32, help="batch size to use for traiing"
+    )
+    parser.add_argument("--load_path", type=str, help="checkpoint path to load from")
+    parser.add_argument(
+        "--output_path",
+        type=str,
+        help="path to place checkpoints and model outputs",
+        required=True,
+    )
+    parser.add_argument(
+        "--log_dir", type=str, help="directory to place the logs", default="/tmp"
+    )
+
+    return parser.parse_args(argv)
+
+
+def main(argv):
+    args = parse_args(argv)
+
+    # Init our model
+    mnist_model = MNISTModel()
+
+    # Init DataLoader from MNIST Dataset
+    img_transform = transforms.Compose(
+        [
+            transforms.Grayscale(),
+            transforms.ToTensor(),
+        ]
+    )
+    train_ds = SyntheticMNIST(
+        transform=lambda x: (img_transform(x[0]), x[1]),
+    )
+    train_loader = DataLoader(train_ds, batch_size=args.batch_size)
+
+    checkpoint_callback = ModelCheckpoint(
+        monitor="train_loss",
+        dirpath=args.output_path,
+        save_last=True,
+    )
+    if args.load_path:
+        print(f"loading checkpoint: {args.load_path}...")
+        mnist_model.load_from_checkpoint(checkpoint_path=args.load_path)
+
+    logger = TensorBoardLogger(save_dir=args.log_dir, version=1, name="lightning_logs")
+
+    # Initialize a trainer
+    trainer = pl.Trainer(
+        logger=logger,
+        max_epochs=args.epochs,
+        callbacks=[checkpoint_callback],
+    )
+
+    # Train the model ⚡
+    trainer.fit(mnist_model, train_loader)
+
+
+if __name__ == "__main__":
+    main(sys.argv[1:])