Skip to content

Add axes argument to lr finder plot #15652

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

Merged
merged 4 commits into from
Nov 12, 2022
Merged

Add axes argument to lr finder plot #15652

Changes from 1 commit
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
11 changes: 8 additions & 3 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 @@ -40,7 +40,7 @@
_MATPLOTLIB_AVAILABLE = RequirementCache("matplotlib")
if _MATPLOTLIB_AVAILABLE and TYPE_CHECKING:
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 _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