We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Handling for aliases of PyTorch Lightning
## 🐛 Bug

Starting from PyTorch Lightning 1.8, this package has been renamed to `Lightning`. While both `lightning` and `pytorch_lightning` are in use as a package (`lightning` is a superset package for `pytorch_lightning` now).

This has caused a problem that `aim`'s code only support for PyTorch Lightning when `pytorch_lightning` is installed. When `lightning` is installed for some other users/developers, they will find that `aim` will wrongly suggest that they have not installed `pytorch_lightning` package.

Since parallel distribution for both `lightning` and `pytorch_lightning` would [stays for a long time](https://github.com/Lightning-AI/lightning/discussions/17095). It may be necessary to support both packages in `aim`'s code.

### To reproduce

Use `pip install lightning` instead of `pip install pytorch_lightning` to install `PyTorch Lightning`.
### Expected behavior

pytorch lightning addons should work for both those who installed `lightning` package and `pytorch-lightning` package.

### Environment

- Aim Version (e.g., 3.0.1): latest
- Python version: 3.10.1
- pip version
- OS (e.g., Linux): Linux
- PyTorch Lightning >= 1.8.0

### Additional context

None


--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `aim/sdk/adapters/pytorch_lightning.py`
Content:
```
1 import os
2 from typing import Any, Dict, Optional, Union
3 from argparse import Namespace
4 
5 import packaging.version
6 
7 try:
8     import pytorch_lightning as pl
9 
10     if packaging.version.parse(pl.__version__) < packaging.version.parse("1.7"):
11         from pytorch_lightning.loggers.base import (
12             LightningLoggerBase as Logger,
13             rank_zero_experiment,
14         )
15     else:
16         from pytorch_lightning.loggers.logger import (
17             Logger,
18             rank_zero_experiment,
19         )
20 
21     from pytorch_lightning.utilities import rank_zero_only
22 except ImportError:
23     raise RuntimeError(
24         'This contrib module requires PyTorch Lightning to be installed. '
25         'Please install it with command: \n pip install pytorch-lightning'
26     )
27 
28 from aim.sdk.run import Run
29 from aim.sdk.repo import Repo
30 from aim.sdk.utils import clean_repo_path, get_aim_repo_name
31 from aim.ext.resource.configs import DEFAULT_SYSTEM_TRACKING_INT
32 
33 
34 class AimLogger(Logger):
35     def __init__(self,
36                  repo: Optional[str] = None,
37                  experiment: Optional[str] = None,
38                  train_metric_prefix: Optional[str] = 'train_',
39                  val_metric_prefix: Optional[str] = 'val_',
40                  test_metric_prefix: Optional[str] = 'test_',
41                  system_tracking_interval: Optional[int]
42                  = DEFAULT_SYSTEM_TRACKING_INT,
43                  log_system_params: Optional[bool] = True,
44                  capture_terminal_logs: Optional[bool] = True,
45                  run_name: Optional[str] = None,
46                  run_hash: Optional[str] = None,
47                  ):
48         super().__init__()
49 
50         self._experiment_name = experiment
51         self._run_name = run_name
52         self._repo_path = repo
53 
54         self._train_metric_prefix = train_metric_prefix
55         self._val_metric_prefix = val_metric_prefix
56         self._test_metric_prefix = test_metric_prefix
57         self._system_tracking_interval = system_tracking_interval
58         self._log_system_params = log_system_params
59         self._capture_terminal_logs = capture_terminal_logs
60 
61         self._run = None
62         self._run_hash = run_hash
63 
64     @staticmethod
65     def _convert_params(params: Union[Dict[str, Any], Namespace]) -> Dict[str, Any]:
66         # in case converting from namespace
67         if isinstance(params, Namespace):
68             params = vars(params)
69 
70         if params is None:
71             params = {}
72 
73         return params
74 
75     @property
76     @rank_zero_experiment
77     def experiment(self) -> Run:
78         if self._run is None:
79             if self._run_hash:
80                 self._run = Run(
81                     self._run_hash,
82                     repo=self._repo_path,
83                     system_tracking_interval=self._system_tracking_interval,
84                     capture_terminal_logs=self._capture_terminal_logs
85                 )
86                 if self._run_name is not None:
87                     self._run.name = self._run_name
88             else:
89                 self._run = Run(
90                     repo=self._repo_path,
91                     experiment=self._experiment_name,
92                     system_tracking_interval=self._system_tracking_interval,
93                     log_system_params=self._log_system_params,
94                     capture_terminal_logs=self._capture_terminal_logs,
95                 )
96                 self._run_hash = self._run.hash
97         return self._run
98 
99     @rank_zero_only
100     def log_hyperparams(self, params: Union[Dict[str, Any], Namespace]):
101         params = self._convert_params(params)
102 
103         # Handle OmegaConf object
104         try:
105             from omegaconf import OmegaConf
106         except ModuleNotFoundError:
107             pass
108         else:
109             # Convert to primitives
110             if OmegaConf.is_config(params):
111                 params = OmegaConf.to_container(params, resolve=True)
112 
113         for key, value in params.items():
114             self.experiment.set(('hparams', key), value, strict=False)
115 
116     @rank_zero_only
117     def log_metrics(self, metrics: Dict[str, float],
118                     step: Optional[int] = None):
119         assert rank_zero_only.rank == 0, \
120             'experiment tried to log from global_rank != 0'
121 
122         metric_items: Dict[str: Any] = {k: v for k, v in metrics.items()}
123 
124         if 'epoch' in metric_items:
125             epoch: int = metric_items.pop('epoch')
126         else:
127             epoch = None
128 
129         for k, v in metric_items.items():
130             name = k
131             context = {}
132             if self._train_metric_prefix \
133                     and name.startswith(self._train_metric_prefix):
134                 name = name[len(self._train_metric_prefix):]
135                 context['subset'] = 'train'
136             elif self._test_metric_prefix \
137                     and name.startswith(self._test_metric_prefix):
138                 name = name[len(self._test_metric_prefix):]
139                 context['subset'] = 'test'
140             elif self._val_metric_prefix \
141                     and name.startswith(self._val_metric_prefix):
142                 name = name[len(self._val_metric_prefix):]
143                 context['subset'] = 'val'
144             self.experiment.track(v, name=name, step=step, epoch=epoch, context=context)
145 
146     @rank_zero_only
147     def finalize(self, status: str = '') -> None:
148         super().finalize(status)
149         if self._run:
150             self._run.close()
151             del self._run
152             self._run = None
153 
154     def __del__(self):
155         self.finalize()
156 
157     @property
158     def save_dir(self) -> str:
159         repo_path = clean_repo_path(self._repo_path) or Repo.default_repo_path()
160         return os.path.join(repo_path, get_aim_repo_name())
161 
162     @property
163     def name(self) -> str:
164         return self._experiment_name
165 
166     @property
167     def version(self) -> str:
168         return self.experiment.hash
169
```

Path: `examples/pytorch_lightning_track.py`
Content:
```
1 from aim.pytorch_lightning import AimLogger
2 
3 from argparse import ArgumentParser
4 
5 import torch
6 import pytorch_lightning as pl
7 from torch.nn import functional as F
8 from torch.utils.data import DataLoader, random_split
9 
10 try:
11     from torchvision.datasets.mnist import MNIST
12     from torchvision import transforms
13 except Exception as e:
14     from tests.base.datasets import MNIST
15 
16 
17 class LitClassifier(pl.LightningModule):
18     def __init__(self, hidden_dim=128, learning_rate=1e-3):
19         super().__init__()
20         self.save_hyperparameters()
21 
22         self.l1 = torch.nn.Linear(28 * 28, self.hparams.hidden_dim)
23         self.l2 = torch.nn.Linear(self.hparams.hidden_dim, 10)
24 
25     def forward(self, x):
26         x = x.view(x.size(0), -1)
27         x = torch.relu(self.l1(x))
28         x = torch.relu(self.l2(x))
29         return x
30 
31     def training_step(self, batch, batch_idx):
32         x, y = batch
33         y_hat = self(x)
34         loss = F.cross_entropy(y_hat, y)
35         self.log('train_loss', loss)
36         return loss
37 
38     def validation_step(self, batch, batch_idx):
39         x, y = batch
40         y_hat = self(x)
41         loss = F.cross_entropy(y_hat, y)
42         self.log('val_loss', loss)
43 
44     def test_step(self, batch, batch_idx):
45         x, y = batch
46         y_hat = self(x)
47         loss = F.cross_entropy(y_hat, y)
48         self.log('test_loss', loss)
49         # Track metrics manually
50         self.logger.experiment.track(1, name='manually_tracked_metric')
51 
52     def configure_optimizers(self):
53         return torch.optim.Adam(self.parameters(), lr=self.hparams.learning_rate)
54 
55     @staticmethod
56     def add_model_specific_args(parent_parser):
57         parser = ArgumentParser(parents=[parent_parser], add_help=False)
58         parser.add_argument('--hidden_dim', type=int, default=128)
59         parser.add_argument('--learning_rate', type=float, default=0.0001)
60         return parser
61 
62 
63 def cli_main():
64     pl.seed_everything(1234)
65 
66     # ------------
67     # args
68     # ------------
69     parser = ArgumentParser()
70     parser.add_argument('--batch_size', default=32, type=int)
71     parser = pl.Trainer.add_argparse_args(parser)
72     parser = LitClassifier.add_model_specific_args(parser)
73     args = parser.parse_args()
74 
75     # ------------
76     # data
77     # ------------
78     dataset = MNIST('', train=True, download=True, transform=transforms.ToTensor())
79     mnist_test = MNIST('', train=False, download=True, transform=transforms.ToTensor())
80     mnist_train, mnist_val = random_split(dataset, [55000, 5000])
81 
82     train_loader = DataLoader(mnist_train, batch_size=args.batch_size)
83     val_loader = DataLoader(mnist_val, batch_size=args.batch_size)
84     test_loader = DataLoader(mnist_test, batch_size=args.batch_size)
85 
86     # ------------
87     # model
88     # ------------
89     model = LitClassifier(args.hidden_dim, args.learning_rate)
90 
91     # ------------
92     # training
93     # ------------
94     aim_logger = AimLogger(
95         experiment='pt_lightning_exp',
96         train_metric_prefix='train_',
97         test_metric_prefix='test_',
98         val_metric_prefix='val_',
99     )
100     trainer = pl.Trainer(logger=aim_logger)
101     trainer.fit(model, train_loader, val_loader)
102 
103     # ------------
104     # testing
105     # ------------
106     trainer.test(dataloaders=test_loader)
107 
108 
109 if __name__ == '__main__':
110     cli_main()
111
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/aim/sdk/adapters/pytorch_lightning.py b/aim/sdk/adapters/pytorch_lightning.py
--- a/aim/sdk/adapters/pytorch_lightning.py
+++ b/aim/sdk/adapters/pytorch_lightning.py
@@ -1,10 +1,19 @@
 import os
+import importlib.util
 from typing import Any, Dict, Optional, Union
 from argparse import Namespace
 
 import packaging.version
 
-try:
+if importlib.util.find_spec("lightning"):
+    import lightning.pytorch as pl
+
+    from lightning.pytorch.loggers.logger import (
+        Logger, rank_zero_experiment
+    )
+
+    from lightning.pytorch.utilities import rank_zero_only
+elif importlib.util.find_spec("pytorch_lightning"):
     import pytorch_lightning as pl
 
     if packaging.version.parse(pl.__version__) < packaging.version.parse("1.7"):
@@ -19,10 +28,11 @@
         )
 
     from pytorch_lightning.utilities import rank_zero_only
-except ImportError:
+else:
     raise RuntimeError(
         'This contrib module requires PyTorch Lightning to be installed. '
         'Please install it with command: \n pip install pytorch-lightning'
+        'or \n pip install lightning'
     )
 
 from aim.sdk.run import Run
diff --git a/examples/pytorch_lightning_track.py b/examples/pytorch_lightning_track.py
--- a/examples/pytorch_lightning_track.py
+++ b/examples/pytorch_lightning_track.py
@@ -1,9 +1,19 @@
+import importlib.util
 from aim.pytorch_lightning import AimLogger
 
 from argparse import ArgumentParser
 
 import torch
-import pytorch_lightning as pl
+if importlib.util.find_spec("lightning"):
+    import lightning.pytorch as pl
+elif importlib.util.find_spec("pytorch_lightning"): # noqa F401
+    import pytorch_lightning as pl
+else:
+    raise RuntimeError(
+        'This contrib module requires PyTorch Lightning to be installed. '
+        'Please install it with command: \n pip install pytorch-lightning \n'
+        'or \n pip install lightning'
+    )
 from torch.nn import functional as F
 from torch.utils.data import DataLoader, random_split

{"golden_diff": "diff --git a/aim/sdk/adapters/pytorch_lightning.py b/aim/sdk/adapters/pytorch_lightning.py\n--- a/aim/sdk/adapters/pytorch_lightning.py\n+++ b/aim/sdk/adapters/pytorch_lightning.py\n@@ -1,10 +1,19 @@\n import os\n+import importlib.util\n from typing import Any, Dict, Optional, Union\n from argparse import Namespace\n \n import packaging.version\n \n-try:\n+if importlib.util.find_spec(\"lightning\"):\n+    import lightning.pytorch as pl\n+\n+    from lightning.pytorch.loggers.logger import (\n+        Logger, rank_zero_experiment\n+    )\n+\n+    from lightning.pytorch.utilities import rank_zero_only\n+elif importlib.util.find_spec(\"pytorch_lightning\"):\n     import pytorch_lightning as pl\n \n     if packaging.version.parse(pl.__version__) < packaging.version.parse(\"1.7\"):\n@@ -19,10 +28,11 @@\n         )\n \n     from pytorch_lightning.utilities import rank_zero_only\n-except ImportError:\n+else:\n     raise RuntimeError(\n         'This contrib module requires PyTorch Lightning to be installed. '\n         'Please install it with command: \\n pip install pytorch-lightning'\n+        'or \\n pip install lightning'\n     )\n \n from aim.sdk.run import Run\ndiff --git a/examples/pytorch_lightning_track.py b/examples/pytorch_lightning_track.py\n--- a/examples/pytorch_lightning_track.py\n+++ b/examples/pytorch_lightning_track.py\n@@ -1,9 +1,19 @@\n+import importlib.util\n from aim.pytorch_lightning import AimLogger\n \n from argparse import ArgumentParser\n \n import torch\n-import pytorch_lightning as pl\n+if importlib.util.find_spec(\"lightning\"):\n+    import lightning.pytorch as pl\n+elif importlib.util.find_spec(\"pytorch_lightning\"): # noqa F401\n+    import pytorch_lightning as pl\n+else:\n+    raise RuntimeError(\n+        'This contrib module requires PyTorch Lightning to be installed. '\n+        'Please install it with command: \\n pip install pytorch-lightning \\n'\n+        'or \\n pip install lightning'\n+    )\n from torch.nn import functional as F\n from torch.utils.data import DataLoader, random_split\n", "issue": "Handling for aliases of PyTorch Lightning\n## \ud83d\udc1b Bug\r\n\r\nStarting from PyTorch Lightning 1.8, this package has been renamed to `Lightning`. While both `lightning` and `pytorch_lightning` are in use as a package (`lightning` is a superset package for `pytorch_lightning` now).\r\n\r\nThis has caused a problem that `aim`'s code only support for PyTorch Lightning when `pytorch_lightning` is installed. When `lightning` is installed for some other users/developers, they will find that `aim` will wrongly suggest that they have not installed `pytorch_lightning` package.\r\n\r\nSince parallel distribution for both `lightning` and `pytorch_lightning` would [stays for a long time](https://github.com/Lightning-AI/lightning/discussions/17095). It may be necessary to support both packages in `aim`'s code.\r\n\r\n### To reproduce\r\n\r\nUse `pip install lightning` instead of `pip install pytorch_lightning` to install `PyTorch Lightning`.\r\n### Expected behavior\r\n\r\npytorch lightning addons should work for both those who installed `lightning` package and `pytorch-lightning` package.\r\n\r\n### Environment\r\n\r\n- Aim Version (e.g., 3.0.1): latest\r\n- Python version: 3.10.1\r\n- pip version\r\n- OS (e.g., Linux): Linux\r\n- PyTorch Lightning >= 1.8.0\r\n\r\n### Additional context\r\n\r\nNone\r\n\n", "before_files": [{"content": "import os\nfrom typing import Any, Dict, Optional, Union\nfrom argparse import Namespace\n\nimport packaging.version\n\ntry:\n    import pytorch_lightning as pl\n\n    if packaging.version.parse(pl.__version__) < packaging.version.parse(\"1.7\"):\n        from pytorch_lightning.loggers.base import (\n            LightningLoggerBase as Logger,\n            rank_zero_experiment,\n        )\n    else:\n        from pytorch_lightning.loggers.logger import (\n            Logger,\n            rank_zero_experiment,\n        )\n\n    from pytorch_lightning.utilities import rank_zero_only\nexcept ImportError:\n    raise RuntimeError(\n        'This contrib module requires PyTorch Lightning to be installed. '\n        'Please install it with command: \\n pip install pytorch-lightning'\n    )\n\nfrom aim.sdk.run import Run\nfrom aim.sdk.repo import Repo\nfrom aim.sdk.utils import clean_repo_path, get_aim_repo_name\nfrom aim.ext.resource.configs import DEFAULT_SYSTEM_TRACKING_INT\n\n\nclass AimLogger(Logger):\n    def __init__(self,\n                 repo: Optional[str] = None,\n                 experiment: Optional[str] = None,\n                 train_metric_prefix: Optional[str] = 'train_',\n                 val_metric_prefix: Optional[str] = 'val_',\n                 test_metric_prefix: Optional[str] = 'test_',\n                 system_tracking_interval: Optional[int]\n                 = DEFAULT_SYSTEM_TRACKING_INT,\n                 log_system_params: Optional[bool] = True,\n                 capture_terminal_logs: Optional[bool] = True,\n                 run_name: Optional[str] = None,\n                 run_hash: Optional[str] = None,\n                 ):\n        super().__init__()\n\n        self._experiment_name = experiment\n        self._run_name = run_name\n        self._repo_path = repo\n\n        self._train_metric_prefix = train_metric_prefix\n        self._val_metric_prefix = val_metric_prefix\n        self._test_metric_prefix = test_metric_prefix\n        self._system_tracking_interval = system_tracking_interval\n        self._log_system_params = log_system_params\n        self._capture_terminal_logs = capture_terminal_logs\n\n        self._run = None\n        self._run_hash = run_hash\n\n    @staticmethod\n    def _convert_params(params: Union[Dict[str, Any], Namespace]) -> Dict[str, Any]:\n        # in case converting from namespace\n        if isinstance(params, Namespace):\n            params = vars(params)\n\n        if params is None:\n            params = {}\n\n        return params\n\n    @property\n    @rank_zero_experiment\n    def experiment(self) -> Run:\n        if self._run is None:\n            if self._run_hash:\n                self._run = Run(\n                    self._run_hash,\n                    repo=self._repo_path,\n                    system_tracking_interval=self._system_tracking_interval,\n                    capture_terminal_logs=self._capture_terminal_logs\n                )\n                if self._run_name is not None:\n                    self._run.name = self._run_name\n            else:\n                self._run = Run(\n                    repo=self._repo_path,\n                    experiment=self._experiment_name,\n                    system_tracking_interval=self._system_tracking_interval,\n                    log_system_params=self._log_system_params,\n                    capture_terminal_logs=self._capture_terminal_logs,\n                )\n                self._run_hash = self._run.hash\n        return self._run\n\n    @rank_zero_only\n    def log_hyperparams(self, params: Union[Dict[str, Any], Namespace]):\n        params = self._convert_params(params)\n\n        # Handle OmegaConf object\n        try:\n            from omegaconf import OmegaConf\n        except ModuleNotFoundError:\n            pass\n        else:\n            # Convert to primitives\n            if OmegaConf.is_config(params):\n                params = OmegaConf.to_container(params, resolve=True)\n\n        for key, value in params.items():\n            self.experiment.set(('hparams', key), value, strict=False)\n\n    @rank_zero_only\n    def log_metrics(self, metrics: Dict[str, float],\n                    step: Optional[int] = None):\n        assert rank_zero_only.rank == 0, \\\n            'experiment tried to log from global_rank != 0'\n\n        metric_items: Dict[str: Any] = {k: v for k, v in metrics.items()}\n\n        if 'epoch' in metric_items:\n            epoch: int = metric_items.pop('epoch')\n        else:\n            epoch = None\n\n        for k, v in metric_items.items():\n            name = k\n            context = {}\n            if self._train_metric_prefix \\\n                    and name.startswith(self._train_metric_prefix):\n                name = name[len(self._train_metric_prefix):]\n                context['subset'] = 'train'\n            elif self._test_metric_prefix \\\n                    and name.startswith(self._test_metric_prefix):\n                name = name[len(self._test_metric_prefix):]\n                context['subset'] = 'test'\n            elif self._val_metric_prefix \\\n                    and name.startswith(self._val_metric_prefix):\n                name = name[len(self._val_metric_prefix):]\n                context['subset'] = 'val'\n            self.experiment.track(v, name=name, step=step, epoch=epoch, context=context)\n\n    @rank_zero_only\n    def finalize(self, status: str = '') -> None:\n        super().finalize(status)\n        if self._run:\n            self._run.close()\n            del self._run\n            self._run = None\n\n    def __del__(self):\n        self.finalize()\n\n    @property\n    def save_dir(self) -> str:\n        repo_path = clean_repo_path(self._repo_path) or Repo.default_repo_path()\n        return os.path.join(repo_path, get_aim_repo_name())\n\n    @property\n    def name(self) -> str:\n        return self._experiment_name\n\n    @property\n    def version(self) -> str:\n        return self.experiment.hash\n", "path": "aim/sdk/adapters/pytorch_lightning.py"}, {"content": "from aim.pytorch_lightning import AimLogger\n\nfrom argparse import ArgumentParser\n\nimport torch\nimport pytorch_lightning as pl\nfrom torch.nn import functional as F\nfrom torch.utils.data import DataLoader, random_split\n\ntry:\n    from torchvision.datasets.mnist import MNIST\n    from torchvision import transforms\nexcept Exception as e:\n    from tests.base.datasets import MNIST\n\n\nclass LitClassifier(pl.LightningModule):\n    def __init__(self, hidden_dim=128, learning_rate=1e-3):\n        super().__init__()\n        self.save_hyperparameters()\n\n        self.l1 = torch.nn.Linear(28 * 28, self.hparams.hidden_dim)\n        self.l2 = torch.nn.Linear(self.hparams.hidden_dim, 10)\n\n    def forward(self, x):\n        x = x.view(x.size(0), -1)\n        x = torch.relu(self.l1(x))\n        x = torch.relu(self.l2(x))\n        return x\n\n    def training_step(self, batch, batch_idx):\n        x, y = batch\n        y_hat = self(x)\n        loss = F.cross_entropy(y_hat, y)\n        self.log('train_loss', loss)\n        return loss\n\n    def validation_step(self, batch, batch_idx):\n        x, y = batch\n        y_hat = self(x)\n        loss = F.cross_entropy(y_hat, y)\n        self.log('val_loss', loss)\n\n    def test_step(self, batch, batch_idx):\n        x, y = batch\n        y_hat = self(x)\n        loss = F.cross_entropy(y_hat, y)\n        self.log('test_loss', loss)\n        # Track metrics manually\n        self.logger.experiment.track(1, name='manually_tracked_metric')\n\n    def configure_optimizers(self):\n        return torch.optim.Adam(self.parameters(), lr=self.hparams.learning_rate)\n\n    @staticmethod\n    def add_model_specific_args(parent_parser):\n        parser = ArgumentParser(parents=[parent_parser], add_help=False)\n        parser.add_argument('--hidden_dim', type=int, default=128)\n        parser.add_argument('--learning_rate', type=float, default=0.0001)\n        return parser\n\n\ndef cli_main():\n    pl.seed_everything(1234)\n\n    # ------------\n    # args\n    # ------------\n    parser = ArgumentParser()\n    parser.add_argument('--batch_size', default=32, type=int)\n    parser = pl.Trainer.add_argparse_args(parser)\n    parser = LitClassifier.add_model_specific_args(parser)\n    args = parser.parse_args()\n\n    # ------------\n    # data\n    # ------------\n    dataset = MNIST('', train=True, download=True, transform=transforms.ToTensor())\n    mnist_test = MNIST('', train=False, download=True, transform=transforms.ToTensor())\n    mnist_train, mnist_val = random_split(dataset, [55000, 5000])\n\n    train_loader = DataLoader(mnist_train, batch_size=args.batch_size)\n    val_loader = DataLoader(mnist_val, batch_size=args.batch_size)\n    test_loader = DataLoader(mnist_test, batch_size=args.batch_size)\n\n    # ------------\n    # model\n    # ------------\n    model = LitClassifier(args.hidden_dim, args.learning_rate)\n\n    # ------------\n    # training\n    # ------------\n    aim_logger = AimLogger(\n        experiment='pt_lightning_exp',\n        train_metric_prefix='train_',\n        test_metric_prefix='test_',\n        val_metric_prefix='val_',\n    )\n    trainer = pl.Trainer(logger=aim_logger)\n    trainer.fit(model, train_loader, val_loader)\n\n    # ------------\n    # testing\n    # ------------\n    trainer.test(dataloaders=test_loader)\n\n\nif __name__ == '__main__':\n    cli_main()\n", "path": "examples/pytorch_lightning_track.py"}], "after_files": [{"content": "import os\nimport importlib.util\nfrom typing import Any, Dict, Optional, Union\nfrom argparse import Namespace\n\nimport packaging.version\n\nif importlib.util.find_spec(\"lightning\"):\n    import lightning.pytorch as pl\n\n    from lightning.pytorch.loggers.logger import (\n        Logger, rank_zero_experiment\n    )\n\n    from lightning.pytorch.utilities import rank_zero_only\nelif importlib.util.find_spec(\"pytorch_lightning\"):\n    import pytorch_lightning as pl\n\n    if packaging.version.parse(pl.__version__) < packaging.version.parse(\"1.7\"):\n        from pytorch_lightning.loggers.base import (\n            LightningLoggerBase as Logger,\n            rank_zero_experiment,\n        )\n    else:\n        from pytorch_lightning.loggers.logger import (\n            Logger,\n            rank_zero_experiment,\n        )\n\n    from pytorch_lightning.utilities import rank_zero_only\nelse:\n    raise RuntimeError(\n        'This contrib module requires PyTorch Lightning to be installed. '\n        'Please install it with command: \\n pip install pytorch-lightning'\n        'or \\n pip install lightning'\n    )\n\nfrom aim.sdk.run import Run\nfrom aim.sdk.repo import Repo\nfrom aim.sdk.utils import clean_repo_path, get_aim_repo_name\nfrom aim.ext.resource.configs import DEFAULT_SYSTEM_TRACKING_INT\n\n\nclass AimLogger(Logger):\n    def __init__(self,\n                 repo: Optional[str] = None,\n                 experiment: Optional[str] = None,\n                 train_metric_prefix: Optional[str] = 'train_',\n                 val_metric_prefix: Optional[str] = 'val_',\n                 test_metric_prefix: Optional[str] = 'test_',\n                 system_tracking_interval: Optional[int]\n                 = DEFAULT_SYSTEM_TRACKING_INT,\n                 log_system_params: Optional[bool] = True,\n                 capture_terminal_logs: Optional[bool] = True,\n                 run_name: Optional[str] = None,\n                 run_hash: Optional[str] = None,\n                 ):\n        super().__init__()\n\n        self._experiment_name = experiment\n        self._run_name = run_name\n        self._repo_path = repo\n\n        self._train_metric_prefix = train_metric_prefix\n        self._val_metric_prefix = val_metric_prefix\n        self._test_metric_prefix = test_metric_prefix\n        self._system_tracking_interval = system_tracking_interval\n        self._log_system_params = log_system_params\n        self._capture_terminal_logs = capture_terminal_logs\n\n        self._run = None\n        self._run_hash = run_hash\n\n    @staticmethod\n    def _convert_params(params: Union[Dict[str, Any], Namespace]) -> Dict[str, Any]:\n        # in case converting from namespace\n        if isinstance(params, Namespace):\n            params = vars(params)\n\n        if params is None:\n            params = {}\n\n        return params\n\n    @property\n    @rank_zero_experiment\n    def experiment(self) -> Run:\n        if self._run is None:\n            if self._run_hash:\n                self._run = Run(\n                    self._run_hash,\n                    repo=self._repo_path,\n                    system_tracking_interval=self._system_tracking_interval,\n                    capture_terminal_logs=self._capture_terminal_logs\n                )\n                if self._run_name is not None:\n                    self._run.name = self._run_name\n            else:\n                self._run = Run(\n                    repo=self._repo_path,\n                    experiment=self._experiment_name,\n                    system_tracking_interval=self._system_tracking_interval,\n                    log_system_params=self._log_system_params,\n                    capture_terminal_logs=self._capture_terminal_logs,\n                )\n                self._run_hash = self._run.hash\n        return self._run\n\n    @rank_zero_only\n    def log_hyperparams(self, params: Union[Dict[str, Any], Namespace]):\n        params = self._convert_params(params)\n\n        # Handle OmegaConf object\n        try:\n            from omegaconf import OmegaConf\n        except ModuleNotFoundError:\n            pass\n        else:\n            # Convert to primitives\n            if OmegaConf.is_config(params):\n                params = OmegaConf.to_container(params, resolve=True)\n\n        for key, value in params.items():\n            self.experiment.set(('hparams', key), value, strict=False)\n\n    @rank_zero_only\n    def log_metrics(self, metrics: Dict[str, float],\n                    step: Optional[int] = None):\n        assert rank_zero_only.rank == 0, \\\n            'experiment tried to log from global_rank != 0'\n\n        metric_items: Dict[str: Any] = {k: v for k, v in metrics.items()}\n\n        if 'epoch' in metric_items:\n            epoch: int = metric_items.pop('epoch')\n        else:\n            epoch = None\n\n        for k, v in metric_items.items():\n            name = k\n            context = {}\n            if self._train_metric_prefix \\\n                    and name.startswith(self._train_metric_prefix):\n                name = name[len(self._train_metric_prefix):]\n                context['subset'] = 'train'\n            elif self._test_metric_prefix \\\n                    and name.startswith(self._test_metric_prefix):\n                name = name[len(self._test_metric_prefix):]\n                context['subset'] = 'test'\n            elif self._val_metric_prefix \\\n                    and name.startswith(self._val_metric_prefix):\n                name = name[len(self._val_metric_prefix):]\n                context['subset'] = 'val'\n            self.experiment.track(v, name=name, step=step, epoch=epoch, context=context)\n\n    @rank_zero_only\n    def finalize(self, status: str = '') -> None:\n        super().finalize(status)\n        if self._run:\n            self._run.close()\n            del self._run\n            self._run = None\n\n    def __del__(self):\n        self.finalize()\n\n    @property\n    def save_dir(self) -> str:\n        repo_path = clean_repo_path(self._repo_path) or Repo.default_repo_path()\n        return os.path.join(repo_path, get_aim_repo_name())\n\n    @property\n    def name(self) -> str:\n        return self._experiment_name\n\n    @property\n    def version(self) -> str:\n        return self.experiment.hash\n", "path": "aim/sdk/adapters/pytorch_lightning.py"}, {"content": "import importlib.util\nfrom aim.pytorch_lightning import AimLogger\n\nfrom argparse import ArgumentParser\n\nimport torch\nif importlib.util.find_spec(\"lightning\"):\n    import lightning.pytorch as pl\nelif importlib.util.find_spec(\"pytorch_lightning\"): # noqa F401\n    import pytorch_lightning as pl\nelse:\n    raise RuntimeError(\n        'This contrib module requires PyTorch Lightning to be installed. '\n        'Please install it with command: \\n pip install pytorch-lightning \\n'\n        'or \\n pip install lightning'\n    )\nfrom torch.nn import functional as F\nfrom torch.utils.data import DataLoader, random_split\n\ntry:\n    from torchvision.datasets.mnist import MNIST\n    from torchvision import transforms\nexcept Exception as e:\n    from tests.base.datasets import MNIST\n\n\nclass LitClassifier(pl.LightningModule):\n    def __init__(self, hidden_dim=128, learning_rate=1e-3):\n        super().__init__()\n        self.save_hyperparameters()\n\n        self.l1 = torch.nn.Linear(28 * 28, self.hparams.hidden_dim)\n        self.l2 = torch.nn.Linear(self.hparams.hidden_dim, 10)\n\n    def forward(self, x):\n        x = x.view(x.size(0), -1)\n        x = torch.relu(self.l1(x))\n        x = torch.relu(self.l2(x))\n        return x\n\n    def training_step(self, batch, batch_idx):\n        x, y = batch\n        y_hat = self(x)\n        loss = F.cross_entropy(y_hat, y)\n        self.log('train_loss', loss)\n        return loss\n\n    def validation_step(self, batch, batch_idx):\n        x, y = batch\n        y_hat = self(x)\n        loss = F.cross_entropy(y_hat, y)\n        self.log('val_loss', loss)\n\n    def test_step(self, batch, batch_idx):\n        x, y = batch\n        y_hat = self(x)\n        loss = F.cross_entropy(y_hat, y)\n        self.log('test_loss', loss)\n        # Track metrics manually\n        self.logger.experiment.track(1, name='manually_tracked_metric')\n\n    def configure_optimizers(self):\n        return torch.optim.Adam(self.parameters(), lr=self.hparams.learning_rate)\n\n    @staticmethod\n    def add_model_specific_args(parent_parser):\n        parser = ArgumentParser(parents=[parent_parser], add_help=False)\n        parser.add_argument('--hidden_dim', type=int, default=128)\n        parser.add_argument('--learning_rate', type=float, default=0.0001)\n        return parser\n\n\ndef cli_main():\n    pl.seed_everything(1234)\n\n    # ------------\n    # args\n    # ------------\n    parser = ArgumentParser()\n    parser.add_argument('--batch_size', default=32, type=int)\n    parser = pl.Trainer.add_argparse_args(parser)\n    parser = LitClassifier.add_model_specific_args(parser)\n    args = parser.parse_args()\n\n    # ------------\n    # data\n    # ------------\n    dataset = MNIST('', train=True, download=True, transform=transforms.ToTensor())\n    mnist_test = MNIST('', train=False, download=True, transform=transforms.ToTensor())\n    mnist_train, mnist_val = random_split(dataset, [55000, 5000])\n\n    train_loader = DataLoader(mnist_train, batch_size=args.batch_size)\n    val_loader = DataLoader(mnist_val, batch_size=args.batch_size)\n    test_loader = DataLoader(mnist_test, batch_size=args.batch_size)\n\n    # ------------\n    # model\n    # ------------\n    model = LitClassifier(args.hidden_dim, args.learning_rate)\n\n    # ------------\n    # training\n    # ------------\n    aim_logger = AimLogger(\n        experiment='pt_lightning_exp',\n        train_metric_prefix='train_',\n        test_metric_prefix='test_',\n        val_metric_prefix='val_',\n    )\n    trainer = pl.Trainer(logger=aim_logger)\n    trainer.fit(model, train_loader, val_loader)\n\n    # ------------\n    # testing\n    # ------------\n    trainer.test(dataloaders=test_loader)\n\n\nif __name__ == '__main__':\n    cli_main()\n", "path": "examples/pytorch_lightning_track.py"}]}