Update pytorch/lightning to 2.2

conda/meta.yaml

@@ -21,7 +21,7 @@ requirements:
      - scipy~=1.11.0
      - jinja2~=3.0
      - torchvision
-     - pytorch>=1.12.0
+     - pytorch~=2.2.0
      - cudatoolkit
      - jsonschema
      - scikit-image~=0.21.0
@@ -30,9 +30,9 @@ requirements:
      - pyvips
      - coremltools
      - pyarrow
-     - lightning~=2.0
+     - lightning~=2.2
      - torchmetrics>=1.1.0
-     - conda-forge::threadpoolctl~=3.2.0
+     - conda-forge::threadpoolctl~=3.4.0
      - albumentations
      - rich
 about:

environment.yml

@@ -14,7 +14,7 @@ dependencies:
   - scipy~=1.10.0
   - jinja2~=3.0
   - conda-forge::torchvision-cpu>=0.5.0
-  - conda-forge::pytorch-cpu~=2.0.0
+  - conda-forge::pytorch-cpu~=2.2.0
   - jsonschema
   - scikit-learn~=1.2.1
   - scikit-image~=0.21.0
@@ -23,9 +23,9 @@ dependencies:
   - imagemagick>=7.1.0
   - pyarrow
   - importlib-resources>=1.3.0
-  - conda-forge::lightning~=2.0.0
+  - conda-forge::lightning~=2.2.0
   - conda-forge::torchmetrics>=1.1.0
-  - conda-forge::threadpoolctl~=3.2
+  - conda-forge::threadpoolctl~=3.4
   - pip
   - albumentations
   - rich

environment_cuda.yml

@@ -14,7 +14,7 @@ dependencies:
   - scipy~=1.10.0
   - jinja2~=3.0
   - conda-forge::torchvision>=0.5.0
-  - conda-forge::pytorch~=2.0.0
+  - conda-forge::pytorch~=2.2.0
   - cudatoolkit>=9.2
   - jsonschema
   - scikit-learn~=1.2.1
@@ -24,9 +24,9 @@ dependencies:
   - imagemagick>=7.1.0
   - pyarrow
   - importlib-resources>=1.3.0
-  - conda-forge::lightning~=2.0.0
+  - conda-forge::lightning~=2.2.0
   - conda-forge::torchmetrics>=1.1.0
-  - conda-forge::threadpoolctl~=3.2
+  - conda-forge::threadpoolctl~=3.4
   - pip
   - albumentations
   - rich

kraken/ketos/__init__.py

@@ -60,10 +60,10 @@
 def cli(ctx, verbose, seed, deterministic):
     ctx.meta['deterministic'] = False if not deterministic else 'warn'
     if seed:
-        from pytorch_lightning import seed_everything
+        from lightning.pytorch import seed_everything
         seed_everything(seed, workers=True)
     elif deterministic:
-        from pytorch_lightning import seed_everything
+        from lightning.pytorch import seed_everything
         seed_everything(42, workers=True)
 
     ctx.meta['verbose'] = verbose

kraken/lib/pretrain/model.py

@@ -36,11 +36,11 @@
 from typing import TYPE_CHECKING, Any, Dict, Optional, Sequence, Union
 
 import numpy as np
-import pytorch_lightning as pl
+import lightning as L
 import torch
 import torch.nn.functional as F
-from pytorch_lightning.callbacks import EarlyStopping
-from pytorch_lightning.utilities.memory import (garbage_collection_cuda,
+from lightning.pytorch.callbacks import EarlyStopping
+from lightning.pytorch.utilities.memory import (garbage_collection_cuda,
                                                 is_oom_error)
 from torch.optim import lr_scheduler
 from torch.utils.data import DataLoader, Subset, random_split
@@ -73,7 +73,7 @@ def _star_fun(fun, kwargs):
     return None
 
 
-class PretrainDataModule(pl.LightningDataModule):
+class PretrainDataModule(L.LightningDataModule):
     def __init__(self,
                  training_data: Union[Sequence[Union['PathLike', str]], Sequence[Dict[str, Any]]] = None,
                  evaluation_data: Optional[Union[Sequence[Union['PathLike', str]], Sequence[Dict[str, Any]]]] = None,
@@ -266,7 +266,7 @@ def setup(self, stage: Optional[str] = None):
         self.val_set.dataset.no_encode()
 
 
-class RecognitionPretrainModel(pl.LightningModule):
+class RecognitionPretrainModel(L.LightningModule):
     def __init__(self,
                  hyper_params: Dict[str, Any] = None,
                  output: str = 'model',

kraken/lib/progress.py

@@ -18,7 +18,7 @@
 from dataclasses import dataclass
 from typing import TYPE_CHECKING, Union
 
-from pytorch_lightning.callbacks.progress.rich_progress import (
+from lightning.pytorch.callbacks.progress.rich_progress import (
     CustomProgress, MetricsTextColumn, RichProgressBar)
 from rich import get_console, reconfigure
 from rich.default_styles import DEFAULT_STYLES

kraken/lib/ro/model.py

@@ -23,10 +23,10 @@
                     Sequence, Union)
 
 import numpy as np
-import pytorch_lightning as pl
+import lightning as L
 import torch
 import torch.nn.functional as F
-from pytorch_lightning.callbacks import EarlyStopping, LearningRateMonitor
+from lightning.pytorch.callbacks import EarlyStopping, LearningRateMonitor
 from torch.optim import lr_scheduler
 from torch.utils.data import DataLoader, Subset
 
@@ -64,7 +64,7 @@ def spearman_footrule_distance(s, t):
     return (s - t).abs().sum() / (0.5 * (len(s) ** 2 - (len(s) % 2)))
 
 
-class ROModel(pl.LightningModule):
+class ROModel(L.LightningModule):
     def __init__(self,
                  hyper_params: Dict[str, Any] = None,
                  output: str = 'model',