Merge branch 'master' into bugfix/throughput-divisible

Author: awaelchli

.pre-commit-config.yaml

@@ -84,10 +84,12 @@ repos:
           - flake8-return
 
   - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: "v0.1.15"
+    rev: "v0.2.0"
     hooks:
       - id: ruff
         args: ["--fix", "--preview"]
+      - id: ruff-format
+        args: ["--preview"]
 
   - repo: https://github.com/executablebooks/mdformat
     rev: 0.7.17

README.md

@@ -92,12 +92,14 @@ pip install -iU https://test.pypi.org/simple/ pytorch-lightning
 
 ______________________________________________________________________
 
-## Lightning has 3 core packages
+## Lightning has 4 core packages
 
 [PyTorch Lightning: Train and deploy PyTorch at scale](#pytorch-lightning-train-and-deploy-pytorch-at-scale).
 <br/>
 [Lightning Fabric: Expert control](#lightning-fabric-expert-control).
 <br/>
+[Lightning Data: Blazing fast, distributed streaming of training data from cloud storage](https://github.com/Lightning-AI/pytorch-lightning/tree/master/src/lightning/data).
+<br/>
 [Lightning Apps: Build AI products and ML workflows](#lightning-apps-build-ai-products-and-ml-workflows).
 
 Lightning gives you granular control over how much abstraction you want to add over PyTorch.

docs/source-fabric/advanced/distributed_communication.rst

@@ -236,6 +236,11 @@ Full example:
     result = fabric.all_gather(data)
     print("Result of all-gather:", result)  # tensor([ 0, 10, 20, 30])
 
+.. warning::
+
+    For the special case where ``world_size`` is 1, no additional dimension is added to the tensor(s). This inconsistency
+    is kept for backward compatibility and you may need to handle this special case in your code to make it agnostic.
+
 
 ----
 

examples/app/dag/app.py

@@ -65,9 +65,9 @@ def __init__(self, models_paths: list):
         )
 
         # Step 3: Create the work to train the models_paths in parallel.
-        self.dict = Dict(
-            **{model_path.split(".")[-1]: ModelWork(model_path, parallel=True) for model_path in models_paths}
-        )
+        self.dict = Dict(**{
+            model_path.split(".")[-1]: ModelWork(model_path, parallel=True) for model_path in models_paths
+        })
 
         # Step 4: Some element to track components progress.
         self.has_completed = False

examples/app/server/app.py

@@ -20,13 +20,11 @@ def setup(self):
     def predict(self, request):
         image = base64.b64decode(request.image.encode("utf-8"))
         image = Image.open(io.BytesIO(image))
-        transforms = torchvision.transforms.Compose(
-            [
-                torchvision.transforms.Resize(224),
-                torchvision.transforms.ToTensor(),
-                torchvision.transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
-            ]
-        )
+        transforms = torchvision.transforms.Compose([
+            torchvision.transforms.Resize(224),
+            torchvision.transforms.ToTensor(),
+            torchvision.transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
+        ])
         image = transforms(image)
         image = image.to(self._device)
         prediction = self._model(image.unsqueeze(0))

examples/app/server_with_auto_scaler/app.py

@@ -34,13 +34,11 @@ def setup(self):
         self._model = torchvision.models.resnet18(pretrained=True).to(self._device)
 
     def predict(self, requests: BatchRequestModel):
-        transforms = torchvision.transforms.Compose(
-            [
-                torchvision.transforms.Resize(224),
-                torchvision.transforms.ToTensor(),
-                torchvision.transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
-            ]
-        )
+        transforms = torchvision.transforms.Compose([
+            torchvision.transforms.Resize(224),
+            torchvision.transforms.ToTensor(),
+            torchvision.transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
+        ])
         images = []
         for request in requests.inputs:
             image = app.components.serve.types.image.Image.deserialize(request.image)

examples/fabric/dcgan/train_fabric.py

@@ -4,6 +4,7 @@
 Code adapted from the official PyTorch DCGAN tutorial:
 https://pytorch.org/tutorials/beginner/dcgan_faces_tutorial.html
 """
+
 import os
 import time
 from pathlib import Path
@@ -55,14 +56,12 @@ def main():
         root=dataroot,
         split="all",
         download=True,
-        transform=transforms.Compose(
-            [
-                transforms.Resize(image_size),
-                transforms.CenterCrop(image_size),
-                transforms.ToTensor(),
-                transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
-            ]
-        ),
+        transform=transforms.Compose([
+            transforms.Resize(image_size),
+            transforms.CenterCrop(image_size),
+            transforms.ToTensor(),
+            transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
+        ]),
     )
 
     # Create the dataloader
@@ -227,7 +226,7 @@ def __init__(self):
             nn.ReLU(True),
             # state size. (ngf) x 32 x 32
             nn.ConvTranspose2d(ngf, nc, 4, 2, 1, bias=False),
-            nn.Tanh()
+            nn.Tanh(),
             # state size. (nc) x 64 x 64
         )
 

examples/fabric/dcgan/train_torch.py

@@ -4,6 +4,7 @@
 Code adapted from the official PyTorch DCGAN tutorial:
 https://pytorch.org/tutorials/beginner/dcgan_faces_tutorial.html
 """
+
 import os
 import random
 import time
@@ -55,14 +56,12 @@ def main():
         root=dataroot,
         split="all",
         download=True,
-        transform=transforms.Compose(
-            [
-                transforms.Resize(image_size),
-                transforms.CenterCrop(image_size),
-                transforms.ToTensor(),
-                transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
-            ]
-        ),
+        transform=transforms.Compose([
+            transforms.Resize(image_size),
+            transforms.CenterCrop(image_size),
+            transforms.ToTensor(),
+            transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
+        ]),
     )
 
     # Create the dataloader
@@ -236,7 +235,7 @@ def __init__(self):
             nn.ReLU(True),
             # state size. (ngf) x 32 x 32
             nn.ConvTranspose2d(ngf, nc, 4, 2, 1, bias=False),
-            nn.Tanh()
+            nn.Tanh(),
             # state size. (nc) x 64 x 64
         )
 

examples/fabric/meta_learning/train_fabric.py

@@ -14,6 +14,7 @@
 Run it with:
     lightning run model train_fabric.py --accelerator=cuda --devices=2 --strategy=ddp
 """
+
 import cherry
 import learn2learn as l2l
 import torch

examples/fabric/meta_learning/train_torch.py

@@ -15,6 +15,7 @@
 Run it with:
     torchrun --nproc_per_node=2 --standalone train_torch.py
 """
+
 import os
 import random
 

examples/fabric/reinforcement_learning/train_fabric.py

@@ -84,14 +84,10 @@ def main(args: argparse.Namespace):
     )
 
     # Environment setup
-    envs = gym.vector.SyncVectorEnv(
-        [
-            make_env(
-                args.env_id, args.seed + rank * args.num_envs + i, rank, args.capture_video, logger.log_dir, "train"
-            )
-            for i in range(args.num_envs)
-        ]
-    )
+    envs = gym.vector.SyncVectorEnv([
+        make_env(args.env_id, args.seed + rank * args.num_envs + i, rank, args.capture_video, logger.log_dir, "train")
+        for i in range(args.num_envs)
+    ])
     assert isinstance(envs.single_action_space, gym.spaces.Discrete), "only discrete action space is supported"
 
     # Define the agent and the optimizer and setup them with Fabric

examples/fabric/reinforcement_learning/train_fabric_decoupled.py

@@ -59,9 +59,9 @@ def player(args, world_collective: TorchCollective, player_trainer_collective: T
     )
 
     # Environment setup
-    envs = gym.vector.SyncVectorEnv(
-        [make_env(args.env_id, args.seed + i, 0, args.capture_video, log_dir, "train") for i in range(args.num_envs)]
-    )
+    envs = gym.vector.SyncVectorEnv([
+        make_env(args.env_id, args.seed + i, 0, args.capture_video, log_dir, "train") for i in range(args.num_envs)
+    ])
     assert isinstance(envs.single_action_space, gym.spaces.Discrete), "only discrete action space is supported"
 
     # Define the agent

examples/fabric/reinforcement_learning/train_torch.py

@@ -142,19 +142,17 @@ def main(args: argparse.Namespace):
         )
 
     # Environment setup
-    envs = gym.vector.SyncVectorEnv(
-        [
-            make_env(
-                args.env_id,
-                args.seed + global_rank * args.num_envs + i,
-                global_rank,
-                args.capture_video,
-                logger.log_dir if global_rank == 0 else None,
-                "train",
-            )
-            for i in range(args.num_envs)
-        ]
-    )
+    envs = gym.vector.SyncVectorEnv([
+        make_env(
+            args.env_id,
+            args.seed + global_rank * args.num_envs + i,
+            global_rank,
+            args.capture_video,
+            logger.log_dir if global_rank == 0 else None,
+            "train",
+        )
+        for i in range(args.num_envs)
+    ])
     assert isinstance(envs.single_action_space, gym.spaces.Discrete), "only discrete action space is supported"
 
     # Define the agent and the optimizer and setup them with DistributedDataParallel

examples/pytorch/basics/autoencoder.py

@@ -16,6 +16,7 @@
 To run: python autoencoder.py --trainer.max_epochs=50
 
 """
+
 from os import path
 from typing import Optional, Tuple
 

examples/pytorch/basics/backbone_image_classifier.py

@@ -16,6 +16,7 @@
 To run: python backbone_image_classifier.py --trainer.max_epochs=50
 
 """
+
 from os import path
 from typing import Optional
 

examples/pytorch/domain_templates/computer_vision_fine_tuning.py

@@ -119,14 +119,12 @@ def normalize_transform(self):
 
     @property
     def train_transform(self):
-        return transforms.Compose(
-            [
-                transforms.Resize((224, 224)),
-                transforms.RandomHorizontalFlip(),
-                transforms.ToTensor(),
-                self.normalize_transform,
-            ]
-        )
+        return transforms.Compose([
+            transforms.Resize((224, 224)),
+            transforms.RandomHorizontalFlip(),
+            transforms.ToTensor(),
+            self.normalize_transform,
+        ])
 
     @property
     def valid_transform(self):
@@ -269,13 +267,11 @@ def add_arguments_to_parser(self, parser):
         parser.link_arguments("data.batch_size", "model.batch_size")
         parser.link_arguments("finetuning.milestones", "model.milestones")
         parser.link_arguments("finetuning.train_bn", "model.train_bn")
-        parser.set_defaults(
-            {
-                "trainer.max_epochs": 15,
-                "trainer.enable_model_summary": False,
-                "trainer.num_sanity_val_steps": 0,
-            }
-        )
+        parser.set_defaults({
+            "trainer.max_epochs": 15,
+            "trainer.enable_model_summary": False,
+            "trainer.num_sanity_val_steps": 0,
+        })
 
 
 def cli_main():

examples/pytorch/domain_templates/generative_adversarial_net.py

@@ -18,6 +18,7 @@
 tensorboard --logdir default
 
 """
+
 from argparse import ArgumentParser, Namespace
 
 import numpy as np

examples/pytorch/domain_templates/imagenet.py

@@ -30,6 +30,7 @@
     python imagenet.py fit --help
 
 """
+
 import os
 from typing import Optional
 
@@ -139,14 +140,12 @@ def setup(self, stage: str):
             train_dir = os.path.join(self.data_path, "train")
             self.train_dataset = datasets.ImageFolder(
                 train_dir,
-                transforms.Compose(
-                    [
-                        transforms.RandomResizedCrop(224),
-                        transforms.RandomHorizontalFlip(),
-                        transforms.ToTensor(),
-                        normalize,
-                    ]
-                ),
+                transforms.Compose([
+                    transforms.RandomResizedCrop(224),
+                    transforms.RandomHorizontalFlip(),
+                    transforms.ToTensor(),
+                    normalize,
+                ]),
             )
         # all stages will use the eval dataset
         normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])

examples/pytorch/domain_templates/reinforce_learn_ppo.py

@@ -28,6 +28,7 @@
 [3] https://github.com/sid-sundrani/ppo_lightning
 
 """
+
 import argparse
 from typing import Callable, Iterator, List, Tuple
 

examples/pytorch/domain_templates/semantic_segmentation.py

@@ -333,14 +333,10 @@ def __init__(
         self.net = UNet(
             num_classes=19, num_layers=self.num_layers, features_start=self.features_start, bilinear=self.bilinear
         )
-        self.transform = transforms.Compose(
-            [
-                transforms.ToTensor(),
-                transforms.Normalize(
-                    mean=[0.35675976, 0.37380189, 0.3764753], std=[0.32064945, 0.32098866, 0.32325324]
-                ),
-            ]
-        )
+        self.transform = transforms.Compose([
+            transforms.ToTensor(),
+            transforms.Normalize(mean=[0.35675976, 0.37380189, 0.3764753], std=[0.32064945, 0.32098866, 0.32325324]),
+        ])
         self.trainset = KITTI(self.data_path, split="train", transform=self.transform)
         self.validset = KITTI(self.data_path, split="valid", transform=self.transform)