def configure_optimizers(self):
optimizer_1 = torch.optim.SGD(self.layer.parameters(), lr=0.1)
optimizer_2 = torch.optim.Adam(self.layer.parameters(), lr=0.1)
optimizer_1 = LightningOptimizer(optimizer_1, 4)
lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer_1, step_size=1)
return [optimizer_1, optimizer_2], [lr_scheduler]
def optimizer_step(self, optimizer, opt_idx, batch_idx,
train_step_and_backward_closure):
model_ref = self.trainer.lightning_module
is_lbfgs = isinstance(optimizer, torch.optim.LBFGS)
using_native_amp = self.trainer.amp_backend == AMPType.NATIVE
# native amp + lbfgs is a no go right now
if using_native_amp and is_lbfgs:
raise MisconfigurationException(
'native PyTorch amp and lbfgs are not compatible.'
' To request, please file a Github issue in PyTorch and tag @mcarilli'
)
# wraps into LightningOptimizer only for running step
optimizer = LightningOptimizer._to_lightning_optimizer(
optimizer, self.trainer, opt_idx)
# model hook
model_ref.optimizer_step(
self.trainer.current_epoch,
batch_idx,
optimizer,
opt_idx,
train_step_and_backward_closure,
on_tpu=self.trainer._device_type == DeviceType.TPU
and _TPU_AVAILABLE,
using_native_amp=using_native_amp,
using_lbfgs=is_lbfgs,
)
def test_lightning_optimizer_keeps_hooks(tmpdir):
model = BoringModel()
optimizer = OptimizerWithHooks(model)
lightning_optimizer = LightningOptimizer(optimizer)
assert len(optimizer._fwd_handles) == 1
del lightning_optimizer
assert len(optimizer._fwd_handles) == 1
def test_state(tmpdir):
model = torch.nn.Linear(3, 4)
optimizer = torch.optim.Adam(model.parameters())
lightning_optimizer = LightningOptimizer(optimizer)
assert isinstance(lightning_optimizer, LightningOptimizer)
assert isinstance(lightning_optimizer, Adam)
assert isinstance(lightning_optimizer, Optimizer)
lightning_dict = {}
special_attrs = ["_accumulate_grad_batches", "_optimizer", "_optimizer_idx",
"_trainer", "_use_accumulate_grad_batches_from_trainer", "_lightning_step"]
for k, v in lightning_optimizer.__dict__.items():
if k not in special_attrs:
lightning_dict[k] = v
assert lightning_dict == optimizer.__dict__
assert optimizer.state_dict() == lightning_optimizer.state_dict()
assert optimizer.state == lightning_optimizer.state
def configure_optimizers(self):
optimizer = torch.optim.SGD(self.layer.parameters(), lr=0.1)
if not auto:
# note: this is not recommended, only done for coverage
optimizer = LightningOptimizer(optimizer)
lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer,
step_size=1)
return [optimizer], [lr_scheduler]
def _convert_to_lightning_optimizer(optimizer: Optimizer) -> LightningOptimizer:
if not isinstance(optimizer, LightningOptimizer):
optimizer = LightningOptimizer(optimizer) # type: ignore [assignment]
optimizer._trainer = self
for opt_idx, opt in enumerate(self.optimizers):
if opt == optimizer._optimizer:
optimizer._optimizer_idx = opt_idx
break
return optimizer # type: ignore [return-value]
def training_step(self, batch, batch_idx):
opt = self.optimizers()
if not isinstance(opt, LightningOptimizer):
opt = LightningOptimizer.to_lightning_optimizer(opt, self.trainer)
output = self.layer(batch)
loss = self.loss(batch, output)
self.losses.append(loss.detach().item())
self.manual_backward(loss, opt)
opt.step()
def test_state(tmpdir):
model = torch.nn.Linear(3, 4)
optimizer = torch.optim.Adam(model.parameters())
lightning_optimizer = LightningOptimizer(optimizer)
# test state
assert optimizer.state == lightning_optimizer.state
lightning_optimizer.state = optimizer.state
assert optimizer.state == lightning_optimizer.state
# test param_groups
assert optimizer.param_groups == lightning_optimizer.param_groups
lightning_optimizer.param_groups = optimizer.param_groups
assert optimizer.param_groups == lightning_optimizer.param_groups
# test defaults
assert optimizer.defaults == lightning_optimizer.defaults
lightning_optimizer.defaults = optimizer.defaults
assert optimizer.defaults == lightning_optimizer.defaults
assert isinstance(lightning_optimizer, LightningOptimizer)
assert isinstance(lightning_optimizer, Adam)
assert isinstance(lightning_optimizer, Optimizer)
lightning_dict = {
k: v
for k, v in lightning_optimizer.__dict__.items()
if k not in {"_optimizer", "_optimizer_idx", "_trainer"}
}
assert lightning_dict == optimizer.__dict__
assert optimizer.state_dict() == lightning_optimizer.state_dict()
assert optimizer.state == lightning_optimizer.state
def test_state(tmpdir):
model = torch.nn.Linear(3, 4)
optimizer = torch.optim.Adam(model.parameters())
lightning_optimizer = LightningOptimizer(optimizer)
# test state
assert optimizer.state == lightning_optimizer.state
lightning_optimizer.state = optimizer.state
assert optimizer.state == lightning_optimizer.state
# test param_groups
assert optimizer.param_groups == lightning_optimizer.param_groups
lightning_optimizer.param_groups = optimizer.param_groups
assert optimizer.param_groups == lightning_optimizer.param_groups
# test defaults
assert optimizer.defaults == lightning_optimizer.defaults
lightning_optimizer.defaults = optimizer.defaults
assert optimizer.defaults == lightning_optimizer.defaults
assert isinstance(lightning_optimizer, LightningOptimizer)
assert isinstance(lightning_optimizer, Adam)
assert isinstance(lightning_optimizer, Optimizer)
lightning_dict = {}
special_attrs = ["_accumulate_grad_batches", "_optimizer", "_optimizer_idx", "_support_closure",
"_trainer", "__getstate__", "__setstate__", "state_dict", "load_state_dict",
"zero_grad", "__setstate__", "add_param_group"]
for k, v in lightning_optimizer.__dict__.items():
if k not in special_attrs:
lightning_dict[k] = v
assert lightning_dict == optimizer.__dict__
assert optimizer.state_dict() == lightning_optimizer.state_dict()
assert optimizer.state == lightning_optimizer.state
def _optimizer_step(
self,
optimizer: Union[Optimizer, LightningOptimizer],
opt_idx: int,
batch_idx: int,
train_step_and_backward_closure: Callable[[], Optional[Tensor]],
) -> None:
"""Performs the optimizer step and some sanity checking.
Args:
optimizer: the optimizer to perform the step with
opt_idx: the index of the current :param:`optimizer`
batch_idx: the index of the current batch
train_step_and_backward_closure: the closure function performing the train step and computing the
gradients. By default, called by the optimizer (if possible)
"""
is_lbfgs = isinstance(optimizer, torch.optim.LBFGS)
# wraps into LightningOptimizer only for running step
if self.trainer.amp_backend == AMPType.APEX:
# apex overrides .step function and need to be wrapped on each step
optimizer = LightningOptimizer._to_lightning_optimizer(
optimizer, self.trainer.strategy, opt_idx)
else:
optimizer = self.trainer.strategy._lightning_optimizers[opt_idx]
# if `strategy.handles_gradient_accumulation`, this method will be called to route into the strategy, but we
# need to check again if `should_accumulate` before increasing the counters
should_accumulate = self.trainer.fit_loop._should_accumulate()
if not should_accumulate:
self.optim_progress.optimizer.step.increment_ready()
# model hook
self.trainer._call_lightning_module_hook(
"optimizer_step",
self.trainer.current_epoch,
batch_idx,
optimizer,
opt_idx,
train_step_and_backward_closure,
on_tpu=isinstance(self.trainer.accelerator, TPUAccelerator),
using_native_amp=(self.trainer.amp_backend == AMPType.NATIVE),
using_lbfgs=is_lbfgs,
)
if not should_accumulate:
self.optim_progress.optimizer.step.increment_completed()
def _optimizer_step(self, optimizer: torch.optim.Optimizer, opt_idx: int,
batch_idx: int,
train_step_and_backward_closure: Callable) -> None:
"""Performs the optimizer step and some sanity checking.
Args:
optimizer: the optimizer to perform the step with
opt_idx: the index of the current :param:`optimizer`
batch_idx: the index of the current batch
train_step_and_backward_closure: the closure function performing the train step and computing the
gradients. By default called by the optimizer (if possible)
"""
model_ref = self.trainer.lightning_module
is_lbfgs = isinstance(optimizer, torch.optim.LBFGS)
using_native_amp = self.trainer.amp_backend == AMPType.NATIVE
# native amp + lbfgs is a no go right now
if using_native_amp and is_lbfgs:
raise MisconfigurationException(
'native PyTorch amp and lbfgs are not compatible.'
' To request, please file a Github issue in PyTorch and tag @mcarilli'
)
# wraps into LightningOptimizer only for running step
optimizer = LightningOptimizer._to_lightning_optimizer(
optimizer, self.trainer, opt_idx)
self.optim_progress.optimizer.step.increment_ready()
# model hook
model_ref.optimizer_step(
self.trainer.current_epoch,
batch_idx,
optimizer,
opt_idx,
train_step_and_backward_closure,
on_tpu=(self.trainer._device_type == DeviceType.TPU
and _TPU_AVAILABLE),
using_native_amp=using_native_amp,
using_lbfgs=is_lbfgs,
)
self.optim_progress.optimizer.step.increment_completed()
def _optimizer_step(
self,
optimizer: Union[Optimizer, LightningOptimizer],
opt_idx: int,
batch_idx: int,
train_step_and_backward_closure: Callable[[], Optional[Tensor]],
) -> None:
"""Performs the optimizer step and some sanity checking.
Args:
optimizer: the optimizer to perform the step with
opt_idx: the index of the current :param:`optimizer`
batch_idx: the index of the current batch
train_step_and_backward_closure: the closure function performing the train step and computing the
gradients. By default called by the optimizer (if possible)
"""
is_lbfgs = isinstance(optimizer, torch.optim.LBFGS)
# wraps into LightningOptimizer only for running step
if self.trainer.amp_backend == AMPType.APEX:
# apex overrides .step function and need to be wrapped on each step
optimizer = LightningOptimizer._to_lightning_optimizer(
optimizer, self.trainer.strategy, opt_idx)
else:
optimizer = self.trainer.strategy._lightning_optimizers[opt_idx]
self.optim_progress.optimizer.step.increment_ready()
# model hook
self.trainer._call_lightning_module_hook(
"optimizer_step",
self.trainer.current_epoch,
batch_idx,
optimizer,
opt_idx,
train_step_and_backward_closure,
on_tpu=(self.trainer._device_type == _AcceleratorType.TPU
and _TPU_AVAILABLE),
using_native_amp=(self.trainer.amp_backend == AMPType.NATIVE),
using_lbfgs=is_lbfgs,
)
self.optim_progress.optimizer.step.increment_completed()
def optimizer_step(
self,
epoch: int = None,
batch_idx: int = None,
optimizer: Optimizer = None,
optimizer_idx: int = None,
optimizer_closure: Optional[Callable] = None,
on_tpu: bool = None,
using_native_amp: bool = None,
using_lbfgs: bool = None,
) -> None:
# warm-up + decay schedule placed here since LARSWrapper is not optimizer class
# adjust LR of optim contained within LARSWrapper
for param_group in optimizer.param_groups:
param_group["lr"] = self.lr_schedule[self.trainer.global_step]
# from lightning
if not isinstance(optimizer, LightningOptimizer):
# wraps into LightingOptimizer only for running step
optimizer = LightningOptimizer.to_lightning_optimizer(optimizer, self.trainer)
optimizer.step(closure=optimizer_closure)
def _optimizer_step(
self,
optimizer: Optimizer,
opt_idx: int,
batch_idx: int,
train_step_and_backward_closure: Callable[[], Optional[Tensor]],
) -> None:
"""Performs the optimizer step and some sanity checking.
Args:
optimizer: the optimizer to perform the step with
opt_idx: the index of the current :param:`optimizer`
batch_idx: the index of the current batch
train_step_and_backward_closure: the closure function performing the train step and computing the
gradients. By default called by the optimizer (if possible)
"""
lightning_module = self.trainer.lightning_module
is_lbfgs = isinstance(optimizer, torch.optim.LBFGS)
# wraps into LightningOptimizer only for running step
optimizer = LightningOptimizer._to_lightning_optimizer(optimizer, self.trainer, opt_idx)
self.optim_progress.optimizer.step.increment_ready()
# model hook
lightning_module.optimizer_step(
self.trainer.current_epoch,
batch_idx,
optimizer,
opt_idx,
train_step_and_backward_closure,
on_tpu=(self.trainer._device_type == DeviceType.TPU and _TPU_AVAILABLE),
using_native_amp=(self.trainer.amp_backend is not None and self.trainer.amp_backend == AMPType.NATIVE),
using_lbfgs=is_lbfgs,
)
self.optim_progress.optimizer.step.increment_completed()
def test_get_optimizer_name_with_lightning_optimizer():
from pytorch_lightning.core.optimizer import LightningOptimizer
adam = torch.optim.Adam(torch.nn.Linear(1, 1).parameters())
assert _get_optimizer_name(LightningOptimizer(adam)) == "Adam"
def _convert_to_lightning_optimizer(trainer, optimizer):
if not isinstance(optimizer, LightningOptimizer):
optimizer = LightningOptimizer(optimizer)
optimizer._on_trainer_init(trainer)
return optimizer
def optimizers(self, optimizers: List[Optimizer]) -> None:
self._optimizers = optimizers
self._lightning_optimizers = {
idx: LightningOptimizer._to_lightning_optimizer(opt, self, idx) for idx, opt in enumerate(self.optimizers)
}