Lightning Dataset (including optimized dataloading of s3 buckets)

.actions/assistant.py

@@ -43,6 +43,11 @@
         "requirements/fabric/base.txt",
         "requirements/fabric/strategies.txt",
     ),
+    "data": (
+        "requirements/data/data.txt",
+        "requirements/data/cloud.txt",
+        "requirements/data/examples.txt",
+    ),
 }
 REQUIREMENT_FILES_ALL = list(chain(*REQUIREMENT_FILES.values()))
 
@@ -404,6 +409,7 @@ def _replace_min(fname: str) -> None:
     def replace_oldest_ver(requirement_fnames: Sequence[str] = REQUIREMENT_FILES_ALL) -> None:
         """Replace the min package version by fixed one."""
         for fname in requirement_fnames:
+            print(fname)
             AssistantCLI._replace_min(fname)
 
     @staticmethod

.github/CODEOWNERS

@@ -41,6 +41,11 @@
 /src/lightning/pytorch/core/hooks.py         @williamfalcon @tchaton @awaelchli @carmocca
 /src/lightning/pytorch/core/module.py        @williamfalcon @tchaton @awaelchli @carmocca
 
+# Data Utilities
+/examples/data/           @nohalon @justusschock
+/src/lightning/data/      @nohalon @justusschock
+/tests/tests_data         @nohalon @justusschock
+
 # Lightning Fabric
 /src/lightning/fabric                       @awaelchli @carmocca @justusschock
 /src/lightning_fabric                       @awaelchli @carmocca @justusschock

.github/workflows/ci-tests-data.yml

@@ -0,0 +1,118 @@
+name: Test Data
+
+# see: https://help.github.com/en/actions/reference/events-that-trigger-workflows
+on:
+  push:
+    branches: [master, "release/*"]
+  pull_request:
+    branches: [master, "release/*"]
+    types: [opened, reopened, ready_for_review, synchronize]  # added `ready_for_review` since draft is skipped
+    paths:
+      - ".actions/**"
+      - "requirements/data/**"
+      - "src/lightning/data/**"
+      - "tests/tests_data/**"
+      - "pyproject.toml"  # includes pytest config
+      - ".github/workflows/ci-tests-data.yml"
+      - "!requirements/*/docs.txt"
+      - "!*.md"
+      - "!**/*.md"
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.ref }}-${{ github.head_ref }}
+  cancel-in-progress: ${{ ! (github.ref == 'refs/heads/master' || startsWith(github.ref, 'refs/heads/release/')) }}
+
+defaults:
+  run:
+    shell: bash
+
+jobs:
+  data-cpu:
+    runs-on: ${{ matrix.os }}
+    if: github.event.pull_request.draft == false
+    strategy:
+      fail-fast: false
+      matrix:
+        include:
+          - {os: "macOS-11", pkg-name: "lightning", python-version: "3.10", pytorch-version: "2.0"}
+          - {os: "ubuntu-20.04", pkg-name: "lightning", python-version: "3.10", pytorch-version: "2.0"}
+          - {os: "windows-2022", pkg-name: "lightning", python-version: "3.10", pytorch-version: "2.0"}
+          # "oldest" versions tests, only on minimum Python
+          - {os: "macOS-11", pkg-name: "lightning", python-version: "3.8", pytorch-version: "2.0", requires: "oldest"}
+          - {os: "ubuntu-20.04", pkg-name: "lightning", python-version: "3.8", pytorch-version: "2.0", requires: "oldest"}
+          - {os: "windows-2022", pkg-name: "lightning", python-version: "3.8", pytorch-version: "2.0", requires: "oldest"}
+    timeout-minutes: 25  # because of building grpcio on Mac
+    env:
+      PACKAGE_NAME: ${{ matrix.pkg-name }}
+      FREEZE_REQUIREMENTS: ${{ ! (github.ref == 'refs/heads/master' || startsWith(github.ref, 'refs/heads/release/')) }}
+      # PYPI_CACHE_DIR: "_pip-wheels"
+      TORCH_URL_STABLE: "https://download.pytorch.org/whl/cpu/torch_stable.html"
+      TORCH_URL_TEST: "https://download.pytorch.org/whl/test/cpu/torch_test.html"
+    steps:
+    - uses: actions/checkout@v3
+
+    - name: Set up Python ${{ matrix.python-version }}
+      uses: actions/setup-python@v4
+      with:
+        python-version: ${{ matrix.python-version }}
+
+    - name: basic setup
+      run: pip install -q -r .actions/requirements.txt
+
+    - name: Set min. dependencies
+      if: ${{ matrix.requires == 'oldest' }}
+      run: |
+        python .actions/assistant.py replace_oldest_ver
+
+    - name: Adjust PyTorch versions in requirements files
+      if: ${{ matrix.requires != 'oldest' && matrix.release != 'pre' }}
+      run: |
+        pip install -q wget packaging
+        python -m wget https://raw.githubusercontent.com/Lightning-AI/utilities/main/scripts/adjust-torch-versions.py
+        for fpath in `ls requirements/data/*.txt`; do \
+          python ./adjust-torch-versions.py $fpath ${{ matrix.pytorch-version }}; \
+        done
+        cat requirements/data/data.txt
+        cat requirements/data/cloud.txt
+
+    # - name: pip wheels cache
+    #   uses: actions/cache/restore@v3
+    #   with:
+    #     path: ${{ env.PYPI_CACHE_DIR }}
+    #     key: pypi_wheels
+    # - run: |
+    #     mkdir -p $PYPI_CACHE_DIR
+    #     ls -lh $PYPI_CACHE_DIR
+
+    # removing torch stable line:
+    # pip install -e ".[${extra}test]" "pytest-timeout" -U -f ${TORCH_URL} ${TORCH_PREINSTALL} -f ${PYPI_CACHE_DIR} --prefer-binary
+    - name: Install package & dependencies
+      run: |
+        python -m pip install -q pip -U
+        pip install -e ".[data-dev]" "pytest-timeout" -U -f ${TORCH_URL} --prefer-binary
+        pip list
+
+    - name: Testing Data
+      working-directory: tests/tests_data
+      # NOTE: do not include coverage report here, see: https://github.com/nedbat/coveragepy/issues/1003
+      run: |
+        python -m coverage run --source lightning \
+          -m pytest -v --timeout=30 --durations=50
+
+    - name: Statistics
+      if: success()
+      working-directory: tests/tests_data
+      run: |
+        coverage report
+        coverage xml
+
+    - name: Upload coverage to Codecov
+      uses: codecov/codecov-action@v3
+      # see: https://github.com/actions/toolkit/issues/399
+      continue-on-error: true
+      with:
+        token: ${{ secrets.CODECOV_TOKEN }}
+        file: tests/tests_data/coverage.xml
+        flags: lightning,cpu,pytest,python${{ matrix.python-version }}
+        name: CPU-coverage
+        fail_ci_if_error: false

.gitignore

@@ -189,6 +189,10 @@ our_model.tar
 test.png
 saved_models
 data/
+!src/lightning/data/
+!examples/data/
+!tests/tests_pytorch/utilities/data/
+!requirements/data/
 .shared
 .lightning
 node_modules/

examples/data/image/imagenet.py

@@ -0,0 +1,190 @@
+import os
+import traceback
+from argparse import ArgumentParser
+from typing import Callable, Literal, Optional
+
+import torch
+import torch.nn.functional as F
+import torch.optim as optim
+import torch.optim.lr_scheduler as lr_scheduler
+
+import lightning as L
+from lightning.pytorch.utilities.model_helpers import get_torchvision_model
+
+parser = ArgumentParser()
+parser.add_argument("--workers", default=4, type=int)
+parser.add_argument("--batchsize", default=56, type=int)
+parser.add_argument("-e", "--evaluate", dest="evaluate", action="store_true", help="evaluate model on validation set")
+args = parser.parse_args()
+
+# --------------------------------
+# Step 1: Define a LightningModule
+# --------------------------------
+
+
+class ImageNetLightningModel(L.LightningModule):
+    """
+    >>> ImageNetLightningModel(data_path='missing')  # doctest: +ELLIPSIS +NORMALIZE_WHITESPACE
+    ImageNetLightningModel(
+      (model): ResNet(...)
+    )
+    """
+
+    from torchvision.models.resnet import ResNet18_Weights
+
+    def __init__(
+        self,
+        data_path: str,
+        index_file_path: str = None,
+        arch: str = "resnet18",
+        weights=ResNet18_Weights.IMAGENET1K_V1,
+        lr: float = 1e-4,
+        momentum: float = 0.9,
+        weight_decay: float = 1e-4,
+        batch_size: int = 256,
+        workers: int = 4,
+    ):
+        super().__init__()
+        self.arch = arch
+        self.weights = weights
+        self.lr = lr
+        self.momentum = momentum
+        self.weight_decay = weight_decay
+        self.batch_size = batch_size
+        self.workers = workers
+        self.data_path = data_path
+        self.index_file_path = index_file_path
+        self.model = get_torchvision_model(self.arch, weights=self.weights)
+        self.train_dataset: Optional[Dataset] = None
+        self.eval_dataset: Optional[Dataset] = None
+
+    def forward(self, x):
+        return self.model(x)
+
+    def training_step(self, batch, batch_idx):
+        images, target = batch
+        output = self.model(images)
+        loss_train = F.cross_entropy(output, target)
+        self.log("train_loss", loss_train)
+        return loss_train
+
+    def eval_step(self, batch, batch_idx, prefix: str):
+        images, target = batch
+        output = self.model(images)
+        loss_val = F.cross_entropy(output, target)
+        self.log(f"{prefix}_loss", loss_val)
+        return loss_val
+
+    def validation_step(self, batch, batch_idx):
+        return self.eval_step(batch, batch_idx, "val")
+
+    def test_step(self, batch, batch_idx):
+        return self.eval_step(batch, batch_idx, "test")
+
+    def configure_optimizers(self):
+        optimizer = optim.SGD(self.parameters(), lr=self.lr, momentum=self.momentum, weight_decay=self.weight_decay)
+        scheduler = lr_scheduler.LambdaLR(optimizer, lambda epoch: 0.1 ** (epoch // 30))
+        return [optimizer], [scheduler]
+
+    def train_dataloader(self):
+        import torchvision as tv
+
+        transforms = tv.transforms.Compose([tv.transforms.RandomResizedCrop(224), tv.transforms.ToTensor()])
+
+        train_dataset = S3LightningImagenetDataset(
+            data_source=self.data_path, split="train", transforms=transforms, path_to_index_file=self.index_file_path
+        )
+
+        return torch.utils.data.DataLoader(
+            dataset=train_dataset, batch_size=self.batch_size, shuffle=True, num_workers=self.workers
+        )
+
+    def val_dataloader(self):
+        import torchvision as tv
+
+        transforms = tv.transforms.Compose([tv.transforms.RandomResizedCrop(224), tv.transforms.ToTensor()])
+
+        val_dataset = S3LightningImagenetDataset(
+            data_source=self.data_path, split="val", transforms=transforms, path_to_index_file=self.index_file_path
+        )
+
+        return torch.utils.data.DataLoader(
+            dataset=val_dataset, batch_size=self.batch_size, shuffle=True, num_workers=self.workers
+        )
+
+    def test_dataloader(self):
+        return self.val_dataloader()
+
+
+# -------------------
+# Step 2: Define data
+# -------------------
+
+
+class S3LightningImagenetDataset(L.LightningDataset):
+    def __init__(
+        self,
+        data_source: str,
+        split: Literal["train", "val"],
+        transforms: Optional[Callable] = None,
+        path_to_index_file: Optional[str] = None,
+    ):
+        from torchvision.models._meta import _IMAGENET_CATEGORIES
+
+        super().__init__(data_source=data_source, backend="s3", path_to_index_file=path_to_index_file)
+
+        # only get files for the split
+        self.files = tuple([x for x in self.files if split in x])
+
+        # get unique classes
+        self.classes = _IMAGENET_CATEGORIES
+
+        self.transforms = transforms
+
+    def load_sample(self, file_path, stream):
+        from PIL import Image
+
+        try:
+            img = Image.open(stream)
+
+            if self.transforms is not None:
+                img = self.transforms(img)
+
+            # Converting grey scale images to RGB
+            if img.shape[0] == 1:
+                img = img.repeat((3, 1, 1))
+
+            curr_cls = os.path.basename(os.path.dirname(file_path)).replace("_", " ")
+            cls_idx = self.classes.index(curr_cls)
+            return img, cls_idx
+        except Exception:
+            print(file_path, traceback.print_exc())
+            pass
+
+
+if __name__ == "__main__":
+    # os.environ["AWS_ACCESS_KEY"] = <your aws access key>
+    # os.environ["AWS_SECRET_ACCESS_KEY"] = <your aws secret key>
+
+    data_path = "s3://imagenet-tiny"
+    index_file_path = "imagenet/imagenet-index.txt"
+
+    # -------------------
+    # Step 3: Train
+    # -------------------
+
+    print("Instantiate Model")
+    model = ImageNetLightningModel(
+        weights=None,
+        data_path=data_path,
+        index_file_path=index_file_path,
+        batch_size=args.batchsize,
+        workers=args.workers,
+    )
+    trainer = L.Trainer()
+
+    print("Train Model")
+    if args.evaluate:
+        trainer.test(model)
+    else:
+        trainer.fit(model)

requirements/data/cloud.txt

@@ -0,0 +1,5 @@
+# NOTE: the upper bound for the package version is only set for CI stability, and it is dropped while installing this package
+#  in case you want to preserve/enforce restrictions on the latest compatible version, add "strict" as an in-line comment
+
+fsspec[http] >2021.06.0, <2023.5.0
+s3fs >=2022.5.0, <=2022.11.1

requirements/data/data.txt

@@ -0,0 +1,6 @@
+# NOTE: the upper bound for the package version is only set for CI stability, and it is dropped while installing this package
+#  in case you want to preserve/enforce restrictions on the latest compatible version, add "strict" as an in-line comment
+
+lightning-utilities >=0.8.0, <0.9.0
+torchdata >0.6.0, < 0.7.0
+torch >2.0.0, < 2.1.0

requirements/data/examples.txt

@@ -0,0 +1,3 @@
+Pillow >= 9.5.0
+# min version to match torch >= 2.0.1
+torchvision >=0.15.2, <=0.16

requirements/data/test.txt

@@ -0,0 +1,5 @@
+coverage ==7.2.5
+pytest ==7.3.1
+pytest-cov ==4.0.0
+pytest-rerunfailures ==10.3
+pytest-random-order ==1.1.0

requirements/fabric/base.txt

@@ -2,7 +2,7 @@
 #  in case you want to preserve/enforce restrictions on the latest compatible version, add "strict" as an in-line comment
 
 numpy >=1.17.2, <1.24.4
-torch >=1.11.0, <=2.0.0
+torch >=1.11.0, <2.1.0
 fsspec[http]>2021.06.0, <2023.5.0
 packaging >=17.1, <=23.0
 typing-extensions >=4.0.0, <=4.4.0

requirements/pytorch/base.txt

@@ -2,7 +2,7 @@
 #  in case you want to preserve/enforce restrictions on the latest compatible version, add "strict" as an in-line comment
 
 numpy >=1.17.2, <1.24.4
-torch >=1.11.0, <=2.0.0
+torch >=1.11.0, <2.1.0
 tqdm >=4.57.0, <4.66.0
 PyYAML >=5.4, <=6.0
 fsspec[http] >2021.06.0, <2023.5.0