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) }