Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add axes argument to lr finder plot #15652

Merged
merged 4 commits into from
Nov 12, 2022
Merged

Add axes argument to lr finder plot #15652

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
1ff7d6f
add axes argument to lr finder plot
bipinKrishnan Nov 10, 2022
0f1e4f2
add changelog
awaelchli Nov 12, 2022
f39f853
Apply suggestions from code review
carmocca Nov 12, 2022
7a65884
Merge branch 'master' into 15416-axes-arg-for-plot
carmocca Nov 12, 2022
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
3 changes: 3 additions & 0 deletions src/pytorch_lightning/CHANGELOG.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,9 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
- Added support to upgrade all checkpoints in a folder using the `pl.utilities.upgrade_checkpoint` script ([#15333](https://github.com/Lightning-AI/lightning/pull/15333))


- Add an axes argument `ax` to the `.lr_find().plot()` to enable writing to a user-defined axes in a matplotlib figure ([#15652](https://github.com/Lightning-AI/lightning/pull/15652))


- Added a check to validate that wrapped FSDP models are used while initializing optimizers ([#15301](https://github.com/Lightning-AI/lightning/pull/15301))


Expand Down
13 changes: 9 additions & 4 deletions src/pytorch_lightning/tuner/lr_finder.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -38,9 +38,9 @@
from tqdm import tqdm

_MATPLOTLIB_AVAILABLE = RequirementCache("matplotlib")
if _MATPLOTLIB_AVAILABLE and TYPE_CHECKING:
if TYPE_CHECKING and _MATPLOTLIB_AVAILABLE:
import matplotlib.pyplot as plt

from matplotlib.axes import Axes
log = logging.getLogger(__name__)


Expand Down Expand Up @@ -130,12 +130,14 @@ def _exchange_scheduler(self, trainer: "pl.Trainer") -> None:
trainer.strategy.lr_scheduler_configs = [LRSchedulerConfig(scheduler, interval="step", opt_idx=0)]
_set_scheduler_opt_idx(trainer.optimizers, trainer.lr_scheduler_configs)

def plot(self, suggest: bool = False, show: bool = False) -> Optional["plt.Figure"]:
def plot(self, suggest: bool = False, show: bool = False, ax: Optional["Axes"] = None) -> Optional["plt.Figure"]:
"""Plot results from lr_find run
Args:
suggest: if True, will mark suggested lr to use with a red point

show: if True, will show figure

ax: Axes object to which the plot is to be drawn. If not provided, a new figure is created.
"""
if not _MATPLOTLIB_AVAILABLE:
raise MisconfigurationException(
Expand All @@ -147,7 +149,10 @@ def plot(self, suggest: bool = False, show: bool = False) -> Optional["plt.Figur
lrs = self.results["lr"]
losses = self.results["loss"]

fig, ax = plt.subplots()
if ax is None:
fig, ax = plt.subplots()
else:
fig = ax.figure

# Plot loss as a function of the learning rate
ax.plot(lrs, losses)
Expand Down