Lightning-AI / Lightning-AI/pytorch-lightning
Trainer validates `gradient_clip_algorithm` although `configure_gradient_clipping` is defined
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 31.4k
- Forks
- 3.8k
- Avg merge
- 6d 7h
- Merged PRs (30d)
- 6
Description
🐛 Bug
I can't pass a custom gradient clipping algorithm, although I implemented configure_gradient_clipping hook, Documentation and release notes hints that you can use configure_gradient_clipping to implement your custom gradient clipping algorithm (.release notes: ... This means you can now implement state-of-the-art clipping algorithms with Lightning! ...) .
Please, allow to pass custom algorithm names in gradient_clip_algorithm when the model implements configure_gradient_clipping.
To Reproduce
from pytorch_lightning import LightningModule, Trainer
class BoringModel(LightningModule):
def __init__(self):
super().__init__()
self.layer = torch.nn.Linear(32, 2)
def forward(self, x):
return self.layer(x)
def training_step(self, batch, batch_idx):
loss = self(batch).sum()
self.log("train_loss", loss)
return {"loss": loss}
def validation_step(self, batch, batch_idx):
loss = self(batch).sum()
self.log("valid_loss", loss)
def test_step(self, batch, batch_idx):
loss = self(batch).sum()
self.log("test_loss", loss)
def configure_optimizers(self):
return torch.optim.SGD(self.layer.parameters(), lr=0.1)
def configure_gradient_clipping(
self,
optimizer: Optimizer,
optimizer_idx: int,
gradient_clip_val: Optional[Union[int, float]] = None,
gradient_clip_algorithm: Optional[str] = None,
):
if gradient_clip_algorithm == "my_custom_clipping_algorithm":
my_custom_clipping(optimizer, gradient_clip_val, gradient_clip_algorithm)
else: # Lightning will handle the gradient clipping
self.clip_gradients(
optimizer,
gradient_clip_val=gradient_clip_val,
gradient_clip_algorithm=gradient_clip_algorithm
)
trainer = Trainer(
default_root_dir=os.getcwd(),
limit_train_batches=1,
limit_val_batches=1,
limit_test_batches=1,
num_sanity_val_steps=0,
max_epochs=1,
enable_model_summary=False,
gradient_clip_algorithm="my_custom_clipping_algorithm"
)
Expected behavior
Environment
- PyTorch Lightning Version: 1.5.4
Additional context
I want to train NFNets with Adaptive Gradient Clipping and compare with standard L2 gradient clipping.
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Assessment
This issue has not been assessed yet.