Moved to lightning (#370)

docs/source/tutorials/offline_linear_eval.rst

@@ -8,9 +8,9 @@ As for pretraining, we start by importing the required packages:
 .. code-block:: python
 
     import torch
-    from pytorch_lightning import Trainer
-    from pytorch_lightning.loggers import WandbLogger
-    from pytorch_lightning.callbacks import LearningRateMonitor
+    from lightning.pytorch import Trainer
+    from lightning.pytorch.loggers import WandbLogger
+    from lightning.pytorch.callbacks import LearningRateMonitor
     from torchvision.models import resnet18
 
     from solo.methods.linear import LinearModel  # imports the linear eval class

docs/source/tutorials/overview.rst

@@ -25,10 +25,10 @@ We start by importing everything that we will need (we will be relying on Pytorc
 .. code-block:: python
 
     import torch
-    from pytorch_lightning import Trainer
-    from pytorch_lightning.callbacks import LearningRateMonitor
-    from pytorch_lightning.loggers import WandbLogger
-    from pytorch_lightning.plugins import DDPPlugin
+    from lightning.pytorch import Trainer
+    from lightning.pytorch.callbacks import LearningRateMonitor
+    from lightning.pytorch.loggers import WandbLogger
+    from lightning.pytorch.plugins import DDPPlugin
 
     from solo.methods import BarlowTwins  # imports the method class
     from solo.utils.checkpointer import Checkpointer

main_linear.py

@@ -24,11 +24,11 @@
 import hydra
 import torch
 import torch.nn as nn
+from lightning.pytorch import Trainer
+from lightning.pytorch.callbacks import LearningRateMonitor
+from lightning.pytorch.loggers import WandbLogger
+from lightning.pytorch.strategies.ddp import DDPStrategy
 from omegaconf import DictConfig, OmegaConf
-from pytorch_lightning import Trainer
-from pytorch_lightning.callbacks import LearningRateMonitor
-from pytorch_lightning.loggers import WandbLogger
-from pytorch_lightning.strategies.ddp import DDPStrategy
 from timm.data.mixup import Mixup
 from timm.loss import LabelSmoothingCrossEntropy, SoftTargetCrossEntropy
 
@@ -169,7 +169,6 @@ def main(cfg: DictConfig):
     callbacks = []
 
     if cfg.checkpoint.enabled:
-        # save checkpoint on last epoch only
         ckpt = Checkpointer(
             cfg,
             logdir=os.path.join(cfg.checkpoint.dir, "linear"),
@@ -211,23 +210,6 @@ def main(cfg: DictConfig):
     )
     trainer = Trainer(**trainer_kwargs)
 
-    # fix for incompatibility with nvidia-dali and pytorch lightning
-    # with dali 1.15 (this will be fixed on 1.16)
-    # https://github.com/Lightning-AI/lightning/issues/12956
-    try:
-        from pytorch_lightning.loops import FitLoop
-
-        class WorkaroundFitLoop(FitLoop):
-            @property
-            def prefetch_batches(self) -> int:
-                return 1
-
-        trainer.fit_loop = WorkaroundFitLoop(
-            trainer.fit_loop.min_epochs, trainer.fit_loop.max_epochs
-        )
-    except:
-        pass
-
     if cfg.data.format == "dali":
         trainer.fit(model, ckpt_path=ckpt_path, datamodule=dali_datamodule)
     else:

main_pretrain.py

@@ -22,12 +22,11 @@
 
 import hydra
 import torch
+from lightning.pytorch import Trainer, seed_everything
+from lightning.pytorch.callbacks import LearningRateMonitor
+from lightning.pytorch.loggers.wandb import WandbLogger
+from lightning.pytorch.strategies.ddp import DDPStrategy
 from omegaconf import DictConfig, OmegaConf
-from pytorch_lightning import Trainer, seed_everything
-from pytorch_lightning.callbacks import LearningRateMonitor
-from pytorch_lightning.loggers import WandbLogger
-from pytorch_lightning.strategies.ddp import DDPStrategy
-
 from solo.args.pretrain import parse_cfg
 from solo.data.classification_dataloader import prepare_data as prepare_data_classification
 from solo.data.pretrain_dataloader import (
@@ -177,7 +176,6 @@ def main(cfg: DictConfig):
     callbacks = []
 
     if cfg.checkpoint.enabled:
-        # save checkpoint on last epoch only
         ckpt = Checkpointer(
             cfg,
             logdir=os.path.join(cfg.checkpoint.dir, cfg.method),
@@ -230,23 +228,6 @@ def main(cfg: DictConfig):
     )
     trainer = Trainer(**trainer_kwargs)
 
-    # fix for incompatibility with nvidia-dali and pytorch lightning
-    # with dali 1.15 (this will be fixed on 1.16)
-    # https://github.com/Lightning-AI/lightning/issues/12956
-    try:
-        from pytorch_lightning.loops import FitLoop
-
-        class WorkaroundFitLoop(FitLoop):
-            @property
-            def prefetch_batches(self) -> int:
-                return 1
-
-        trainer.fit_loop = WorkaroundFitLoop(
-            trainer.fit_loop.min_epochs, trainer.fit_loop.max_epochs
-        )
-    except:
-        pass
-
     if cfg.data.format == "dali":
         trainer.fit(model, ckpt_path=ckpt_path, datamodule=dali_datamodule)
     else:

requirements.txt

@@ -1,7 +1,7 @@
 torch>=1.10.0
 torchvision>=0.11.1
 einops
-pytorch-lightning==2.0.2
+lightning==2.1.2
 torchmetrics>=0.6.0, <0.12.0
 tqdm
 wandb

scripts/pretrain/cifar-multicrop/swav.yaml

@@ -55,4 +55,4 @@ devices: [0]
 sync_batchnorm: True
 accelerator: "gpu"
 strategy: "ddp"
-precision: 16
+precision: 16-mixed

scripts/pretrain/cifar/barlow.yaml

@@ -50,4 +50,4 @@ devices: [0]
 sync_batchnorm: True
 accelerator: "gpu"
 strategy: "ddp"
-precision: 16
+precision: 16-mixed

scripts/pretrain/cifar/byol.yaml

@@ -53,4 +53,4 @@ devices: [0]
 sync_batchnorm: True
 accelerator: "gpu"
 strategy: "ddp"
-precision: 16
+precision: 16-mixed

scripts/pretrain/cifar/deepclusterv2.yaml

@@ -53,4 +53,4 @@ devices: [0]
 sync_batchnorm: True
 accelerator: "gpu"
 strategy: "ddp"
-precision: 16
+precision: 16-mixed

scripts/pretrain/cifar/dino.yaml

@@ -53,4 +53,4 @@ devices: [0]
 sync_batchnorm: True
 accelerator: "gpu"
 strategy: "ddp"
-precision: 16
+precision: 16-mixed

scripts/pretrain/cifar/mae.yaml

@@ -53,4 +53,4 @@ devices: [0]
 sync_batchnorm: True
 accelerator: "gpu"
 strategy: "ddp"
-precision: 16
+precision: 16-mixed

scripts/pretrain/cifar/mocov2plus.yaml

@@ -50,4 +50,4 @@ devices: [0]
 sync_batchnorm: True
 accelerator: "gpu"
 strategy: "ddp"
-precision: 16
+precision: 16-mixed

scripts/pretrain/cifar/mocov3.yaml

@@ -54,4 +54,4 @@ devices: [0]
 sync_batchnorm: True
 accelerator: "gpu"
 strategy: "ddp"
-precision: 16
+precision: 16-mixed

scripts/pretrain/cifar/nnbyol.yaml

@@ -54,4 +54,4 @@ devices: [0]
 sync_batchnorm: True
 accelerator: "gpu"
 strategy: "ddp"
-precision: 16
+precision: 16-mixed

scripts/pretrain/cifar/nnclr.yaml

@@ -52,4 +52,4 @@ devices: [0]
 sync_batchnorm: True
 accelerator: "gpu"
 strategy: "ddp"
-precision: 16
+precision: 16-mixed

scripts/pretrain/cifar/nnsiam.yaml

@@ -50,4 +50,4 @@ devices: [0]
 sync_batchnorm: True
 accelerator: "gpu"
 strategy: "ddp"
-precision: 16
+precision: 16-mixed

scripts/pretrain/cifar/ressl.yaml

@@ -53,4 +53,4 @@ devices: [0]
 sync_batchnorm: True
 accelerator: "gpu"
 strategy: "ddp"
-precision: 16
+precision: 16-mixed

scripts/pretrain/cifar/simclr.yaml

@@ -50,4 +50,4 @@ devices: [0]
 sync_batchnorm: True
 accelerator: "gpu"
 strategy: "ddp"
-precision: 16
+precision: 16-mixed

scripts/pretrain/cifar/simsiam.yaml

@@ -47,4 +47,4 @@ devices: [0]
 sync_batchnorm: True
 accelerator: "gpu"
 strategy: "ddp"
-precision: 16
+precision: 16-mixed

scripts/pretrain/cifar/supcon.yaml

@@ -46,4 +46,4 @@ devices: [0]
 sync_batchnorm: True
 accelerator: "gpu"
 strategy: "ddp"
-precision: 16
+precision: 16-mixed

scripts/pretrain/cifar/swav.yaml

@@ -54,4 +54,4 @@ devices: [0]
 sync_batchnorm: True
 accelerator: "gpu"
 strategy: "ddp"
-precision: 16
+precision: 16-mixed

scripts/pretrain/cifar/vibcreg.yaml

@@ -74,4 +74,4 @@ devices: [0]
 sync_batchnorm: True
 accelerator: "gpu"
 strategy: "ddp"
-precision: 16
+precision: 16-mixed

scripts/pretrain/cifar/vicreg.yaml

@@ -80,4 +80,4 @@ devices: [0]
 sync_batchnorm: True
 accelerator: "gpu"
 strategy: "ddp"
-precision: 16
+precision: 16-mixed

scripts/pretrain/cifar/wmse.yaml

@@ -70,4 +70,4 @@ devices: [0]
 sync_batchnorm: True
 accelerator: "gpu"
 strategy: "ddp"
-precision: 16
+precision: 16-mixed

scripts/pretrain/custom/byol.yaml

@@ -60,4 +60,4 @@ devices: [0, 1]
 sync_batchnorm: True
 accelerator: "gpu"
 strategy: "ddp"
-precision: 16
+precision: 16-mixed

scripts/pretrain/imagenet-100-multicrop/byol.yaml

@@ -53,4 +53,4 @@ devices: [0, 1]
 sync_batchnorm: True
 accelerator: "gpu"
 strategy: "ddp"
-precision: 16
+precision: 16-mixed

scripts/pretrain/imagenet-100-multicrop/simclr.yaml

@@ -50,4 +50,4 @@ devices: [0, 1]
 sync_batchnorm: True
 accelerator: "gpu"
 strategy: "ddp"
-precision: 16
+precision: 16-mixed

scripts/pretrain/imagenet-100-multicrop/supcon.yaml

@@ -46,4 +46,4 @@ devices: [0, 1]
 sync_batchnorm: True
 accelerator: "gpu"
 strategy: "ddp"
-precision: 16
+precision: 16-mixed

scripts/pretrain/imagenet-100/barlow.yaml

@@ -50,4 +50,4 @@ devices: [0, 1]
 sync_batchnorm: True
 accelerator: "gpu"
 strategy: "ddp"
-precision: 16
+precision: 16-mixed

scripts/pretrain/imagenet-100/byol.yaml

@@ -53,4 +53,4 @@ devices: [0, 1]
 sync_batchnorm: True
 accelerator: "gpu"
 strategy: "ddp"
-precision: 16
+precision: 16-mixed

scripts/pretrain/imagenet-100/deepclusterv2.yaml

@@ -55,4 +55,4 @@ devices: [0, 1]
 sync_batchnorm: True
 accelerator: "gpu"
 strategy: "ddp"
-precision: 16
+precision: 16-mixed

scripts/pretrain/imagenet-100/dino.yaml

@@ -54,4 +54,4 @@ devices: [0, 1]
 sync_batchnorm: True
 accelerator: "gpu"
 strategy: "ddp"
-precision: 16
+precision: 16-mixed

scripts/pretrain/imagenet-100/dino_vit.yaml

@@ -51,4 +51,4 @@ devices: [0, 1]
 sync_batchnorm: True
 accelerator: "gpu"
 strategy: "ddp"
-precision: 16
+precision: 16-mixed

scripts/pretrain/imagenet-100/mae.yaml

@@ -51,4 +51,4 @@ devices: [0, 1]
 sync_batchnorm: True
 accelerator: "gpu"
 strategy: "ddp"
-precision: 16
+precision: 16-mixed

scripts/pretrain/imagenet-100/mocov2plus.yaml

@@ -50,4 +50,4 @@ devices: [0, 1]
 sync_batchnorm: True
 accelerator: "gpu"
 strategy: "ddp"
-precision: 16
+precision: 16-mixed

scripts/pretrain/imagenet-100/mocov3.yaml

@@ -54,4 +54,4 @@ devices: [0, 1]
 sync_batchnorm: True
 accelerator: "gpu"
 strategy: "ddp"
-precision: 16
+precision: 16-mixed

scripts/pretrain/imagenet-100/mocov3_vit.yaml

@@ -50,4 +50,4 @@ devices: [0, 1, 2, 3, 4, 5, 6, 7]
 sync_batchnorm: True
 accelerator: "gpu"
 strategy: "ddp"
-precision: 16
+precision: 16-mixed

scripts/pretrain/imagenet-100/nnclr.yaml

@@ -52,4 +52,4 @@ devices: [0, 1]
 sync_batchnorm: True
 accelerator: "gpu"
 strategy: "ddp"
-precision: 16
+precision: 16-mixed

scripts/pretrain/imagenet-100/ressl.yaml

@@ -53,4 +53,4 @@ devices: [0, 1]
 sync_batchnorm: True
 accelerator: "gpu"
 strategy: "ddp"
-precision: 16
+precision: 16-mixed

scripts/pretrain/imagenet-100/simclr.yaml

@@ -50,4 +50,4 @@ devices: [0, 1]
 sync_batchnorm: True
 accelerator: "gpu"
 strategy: "ddp"
-precision: 16
+precision: 16-mixed

scripts/pretrain/imagenet-100/simsiam.yaml

@@ -48,4 +48,4 @@ devices: [0, 1]
 sync_batchnorm: True
 accelerator: "gpu"
 strategy: "ddp"
-precision: 16
+precision: 16-mixed

scripts/pretrain/imagenet-100/supcon.yaml

@@ -46,4 +46,4 @@ devices: [0, 1]
 sync_batchnorm: True
 accelerator: "gpu"
 strategy: "ddp"
-precision: 16
+precision: 16-mixed

scripts/pretrain/imagenet-100/swav.yaml

@@ -54,4 +54,4 @@ devices: [0, 1]
 sync_batchnorm: True
 accelerator: "gpu"
 strategy: "ddp"
-precision: 16
+precision: 16-mixed

scripts/pretrain/imagenet-100/vibcreg.yaml

@@ -53,4 +53,4 @@ devices: [0, 1]
 sync_batchnorm: True
 accelerator: "gpu"
 strategy: "ddp"
-precision: 16
+precision: 16-mixed

scripts/pretrain/imagenet-100/vicreg.yaml

@@ -81,4 +81,4 @@ devices: [0, 1]
 sync_batchnorm: True
 accelerator: "gpu"
 strategy: "ddp"
-precision: 16
+precision: 16-mixed

scripts/pretrain/imagenet-100/wmse.yaml

@@ -47,4 +47,4 @@ devices: [0, 1]
 sync_batchnorm: True
 accelerator: "gpu"
 strategy: "ddp"
-precision: 16
+precision: 16-mixed

scripts/pretrain/imagenet/barlow.yaml

@@ -51,4 +51,4 @@ devices: [0, 1, 2, 3]
 sync_batchnorm: True
 accelerator: "gpu"
 strategy: "ddp"
-precision: 16
+precision: 16-mixed

scripts/pretrain/imagenet/byol.yaml

@@ -53,5 +53,5 @@ devices: [0, 1]
 sync_batchnorm: True
 accelerator: "gpu"
 strategy: "ddp"
-precision: 16
+precision: 16-mixed
 accumulate_grad_batches: 16

scripts/pretrain/imagenet/mae.yaml

@@ -54,4 +54,4 @@ devices: [0, 1]
 sync_batchnorm: True
 accelerator: "gpu"
 strategy: "ddp"
-precision: 16
+precision: 16-mixed