Skip to content
This repository has been archived by the owner on Sep 18, 2024. It is now read-only.

adapt pytorch lighting 2.0 AKA lightning #5606

Open
wants to merge 9 commits into
base: master
Choose a base branch
from
Open

adapt pytorch lighting 2.0 AKA lightning #5606

Show file tree
Hide file tree
Changes from 3 commits
Commits
Show all changes
9 commits
Select commit Hold shift + click to select a range
b9ce086
adapt pytorch lighting 2.0 AKA lightning
qmpzzpmq Jun 11, 2023
9b08248
linting
qmpzzpmq Jun 13, 2023
b33ea6e
add document
qmpzzpmq Jun 14, 2023
daca954
Merge branch 'master' of https://github.com/microsoft/nni into pl2.0
qmpzzpmq Jul 14, 2023
aa0645f
cleaning
qmpzzpmq Jul 14, 2023
30ea9cc
add more
qmpzzpmq Jul 17, 2023
2fdebeb
cleaning
qmpzzpmq Jul 17, 2023
d03fc25
lint
qmpzzpmq Jul 17, 2023
88a5b5a
Merge branch 'master' into pl2.0
qmpzzpmq Aug 21, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
25 changes: 19 additions & 6 deletions nni/compression/pytorch/utils/evaluator.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,10 +14,16 @@
from torch.utils.hooks import RemovableHandle

try:
import pytorch_lightning as pl
from pytorch_lightning.callbacks import Callback
import lightning as pl
from lightning.callbacks import Callback
except ImportError:
LIGHTNING_INSTALLED = False
try:
import pytorch_lightning as pl
from pytorch_lightning.callbacks import Callback
except ImportError:
LIGHTNING_INSTALLED = False
else:
LIGHTNING_INSTALLED = True
else:
LIGHTNING_INSTALLED = True

Expand Down Expand Up @@ -315,10 +321,17 @@ class LightningEvaluator(Evaluator):
def __init__(self, trainer: pl.Trainer, data_module: pl.LightningDataModule,
dummy_input: Any | None = None):
assert LIGHTNING_INSTALLED, 'pytorch_lightning is not installed.'
err_msg_p = 'Only support traced {}, please use nni.trace({}) to initialize the trainer.'
err_msg = err_msg_p.format('pytorch_lightning.Trainer', 'pytorch_lightning.Trainer')
err_msg_p = (
'Only support traced {}, please use nni.trace({}) to initialize the trainer. '
'for pytorch_lightning version > 2.0, please using {}'
)
err_msg = err_msg_p.format('pytorch_lightning.Trainer', 'pytorch_lightning.Trainer', 'lightning.Trainer')
assert isinstance(trainer, pl.Trainer) and is_traceable(trainer), err_msg
err_msg = err_msg_p.format('pytorch_lightning.LightningDataModule', 'pytorch_lightning.LightningDataModule')
err_msg = err_msg_p.format(
'pytorch_lightning.LightningDataModule',
'pytorch_lightning.LightningDataModule',
'lightning.LightningDataModule',
)
assert isinstance(data_module, pl.LightningDataModule) and is_traceable(data_module), err_msg
self.trainer = trainer
self.data_module = data_module
Expand Down
22 changes: 17 additions & 5 deletions nni/contrib/compression/utils/evaluator.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,9 +16,14 @@
from torch.utils.hooks import RemovableHandle

try:
import pytorch_lightning as pl
import lightning as pl
except ImportError:
LIGHTNING_INSTALLED = False
try:
import pytorch_lightning as pl
except ImportError:
LIGHTNING_INSTALLED = False
else:
LIGHTNING_INSTALLED = True
else:
LIGHTNING_INSTALLED = True

Expand Down Expand Up @@ -370,10 +375,17 @@ class LightningEvaluator(Evaluator):
def __init__(self, trainer: pl.Trainer, data_module: pl.LightningDataModule,
dummy_input: Any | None = None):
assert LIGHTNING_INSTALLED, 'pytorch_lightning is not installed.'
err_msg_p = 'Only support traced {}, please use nni.trace({}) to initialize the trainer.'
err_msg = err_msg_p.format('pytorch_lightning.Trainer', 'pytorch_lightning.Trainer')
err_msg_p = (
'Only support traced {}, please use nni.trace({}) to initialize the trainer. '
'for pytorch_lightning version > 2.0, please using {}'
)
err_msg = err_msg_p.format('pytorch_lightning.Trainer', 'pytorch_lightning.Trainer', 'lightning.Trainer')
assert isinstance(trainer, pl.Trainer) and is_traceable(trainer), err_msg
err_msg = err_msg_p.format('pytorch_lightning.LightningDataModule', 'pytorch_lightning.LightningDataModule')
err_msg = err_msg_p.format(
'pytorch_lightning.LightningDataModule',
'pytorch_lightning.LightningDataModule',
'lightning.LightningDataModule',
)
assert isinstance(data_module, pl.LightningDataModule) and is_traceable(data_module), err_msg
self.trainer = trainer
self.data_module = data_module
Expand Down