Prevent unexpected unpacking error when calling lr_finder.plot() with suggest_lr=True

tests/test_lr_finder.py

@@ -8,6 +8,7 @@
 import task as mod_task
 import dataset as mod_dataset
 
+import numpy as np
 import matplotlib.pyplot as plt
 
 # Check available backends for mixed precision training
@@ -404,17 +405,52 @@ def test_plot_with_skip_and_suggest_lr(suggest_lr, skip_start, skip_end):
         skip_start=skip_start, skip_end=skip_end, suggest_lr=suggest_lr, ax=ax
     )
 
-    if num_iter - skip_start - skip_end <= 1:
-        # handle data with one or zero lr
-        assert len(ax.lines) == 1
-        assert results is ax
+    # NOTE:
+    # - ax.lines[0]: the lr-loss curve. It should be always available once
+    #   `ax.plot(lrs, losses)` is called. But when there is no sufficent data
+    #   point (num_iter <= skip_start + skip_end), the coordinates will be
+    #   2 empty arrays.
+    # - ax.collections[0]: the point of suggested lr (type: <PathCollection>).
+    #   It's available only when there are sufficient data points to calculate
+    #   gradient of lr-loss curve.
+    assert len(ax.lines) == 1
+
+    if suggest_lr:
+        assert isinstance(results, tuple) and len(results) == 2
+
+        ret_ax, ret_lr = results
+        assert ret_ax is ax
+
+        # XXX: Currently suggested lr is selected according to gradient of
+        # lr-loss curve, so there should be at least 2 valid data points (after
+        # filtered by `skip_start` and `skip_end`). If not, the returned lr
+        # will be None.
+        # But we would need to rework on this if there are more suggestion
+        # methods is supported in the future.
+        if num_iter - skip_start - skip_end <= 1:
+            assert ret_lr is None
+            assert len(ax.collections) == 0
+        else:
+            assert len(ax.collections) == 1
     else:
-        # handle different suggest_lr
-        # for 'steepest': the point with steepest gradient (minimal gradient)
-        assert len(ax.lines) == 1
-        assert len(ax.collections) == int(suggest_lr)
-        if results is not ax:
-            assert len(results) == 2
+        # Not suggesting lr, so it just plots a lr-loss curve.
+        assert results is ax
+        assert len(ax.collections) == 0
+
+    # Check whether the data of plotted line is the same as the one filtered
+    # according to `skip_start` and `skip_end`.
+    lrs = np.array(lr_finder.history["lr"])
+    losses = np.array(lr_finder.history["loss"])
+    x, y = ax.lines[0].get_data()
+
+    # If skip_end is 0, we should replace it with None. Otherwise, it
+    # will create a slice as `x[0:-0]` which is an empty list.
+    _slice = slice(skip_start, -skip_end if skip_end != 0 else None, None)
+    assert np.allclose(x, lrs[_slice])
+    assert np.allclose(y, losses[_slice])
+
+    # Close figure to release memory
+    plt.close()
 
 
 def test_suggest_lr():

torch_lr_finder/lr_finder.py

@@ -538,7 +538,8 @@ def plot(
                 min_grad_idx = (np.gradient(np.array(losses))).argmin()
             except ValueError:
                 print(
-                    "Failed to compute the gradients, there might not be enough points."
+                    "Failed to compute the gradients, there might not be enough points. "
+                    "Please check whether num_iter >= (skip_start + skip_end + 2)."
                 )
             if min_grad_idx is not None:
                 print("Suggested LR: {:.2E}".format(lrs[min_grad_idx]))
@@ -565,8 +566,10 @@ def plot(
         if fig is not None:
             plt.show()
 
-        if suggest_lr and min_grad_idx is not None:
-            return ax, lrs[min_grad_idx]
+        if suggest_lr:
+            # If suggest_lr is set, then we should always return 2 values.
+            suggest_lr = None if min_grad_idx is None else lrs[min_grad_idx]
+            return ax, suggest_lr
         else:
             return ax