fast.ai: Handle resuming.

Found some bugs while working on https://github.com/iterative/dvc-get-started-cv/pull/59 .

- Use the ocasion to also remove the redundant logging of `epoch` as a metric.
- Expose `with_opt` option.
- Don't increase step when resuming.

Author: daavoo

src/dvclive/fastai.py

@@ -7,22 +7,36 @@
 
 
 class DVCLiveCallback(Callback):
-    def __init__(self, model_file=None, live: Optional[Live] = None, **kwargs):
+    def __init__(
+        self,
+        model_file: Optional[str] = None,
+        with_opt: bool = False,
+        live: Optional[Live] = None,
+        **kwargs
+    ):
         super().__init__()
         self.model_file = model_file
+        self.with_opt = with_opt
         self.live = live if live is not None else Live(**kwargs)
 
     def after_epoch(self):
+        logged_metrics = False
         for key, value in zip(
             self.learn.recorder.metric_names, self.learn.recorder.log
         ):
+            if key == "epoch":
+                continue
             self.live.log_metric(
                 standardize_metric_name(key, __name__), float(value)
             )
+            logged_metrics = True
 
-        if self.model_file:
-            self.learn.save(self.model_file)
-        self.live.next_step()
+        # When resuming (i.e. passing `start_epoch` to learner)
+        # fast.ai calls after_epoch but we don't want to increase the step.
+        if logged_metrics:
+            if self.model_file:
+                self.learn.save(self.model_file, with_opt=self.with_opt)
+            self.live.next_step()
 
     def after_fit(self):
         self.live.end()

tests/test_fastai.py

@@ -41,8 +41,7 @@ def data_loader():
 def test_fastai_callback(tmp_dir, data_loader, mocker):
     learn = tabular_learner(data_loader, metrics=accuracy)
     learn.remove_cb(ProgressCallback)
-    learn.model_dir = os.path.abspath("./")
-    callback = DVCLiveCallback("model")
+    callback = DVCLiveCallback()
     live = callback.live
 
     spy = mocker.spy(live, "end")
@@ -58,18 +57,38 @@ def test_fastai_callback(tmp_dir, data_loader, mocker):
     assert train_path.is_dir()
     assert valid_path.is_dir()
     assert (metrics_path / "accuracy.tsv").exists()
+    assert not (metrics_path / "epoch.tsv").exists()
 
 
-def test_fastai_model_file(tmp_dir, data_loader):
+def test_fastai_model_file(tmp_dir, data_loader, mocker):
     learn = tabular_learner(data_loader, metrics=accuracy)
     learn.remove_cb(ProgressCallback)
     learn.model_dir = os.path.abspath("./")
-    learn.fit_one_cycle(2, cbs=[DVCLiveCallback("model")])
+    save = mocker.spy(learn, "save")
+    learn.fit_one_cycle(2, cbs=[DVCLiveCallback("model", with_opt=True)])
     assert (tmp_dir / "model.pth").is_file()
+    save.assert_called_with("model", with_opt=True)
 
 
 def test_fastai_pass_logger():
     logger = Live("train_logs")
 
     assert DVCLiveCallback().live is not logger
     assert DVCLiveCallback(live=logger).live is logger
+
+
+def test_fast_ai_resume(tmp_dir, data_loader, mocker):
+    learn = tabular_learner(data_loader, metrics=accuracy)
+    learn.remove_cb(ProgressCallback)
+    callback = DVCLiveCallback()
+    live = callback.live
+
+    spy = mocker.spy(live, "next_step")
+    learn.fit_one_cycle(2, cbs=[callback])
+    assert spy.call_count == 2
+
+    callback = DVCLiveCallback(resume=True)
+    live = callback.live
+    spy = mocker.spy(live, "next_step")
+    learn.fit_one_cycle(3, cbs=[callback], start_epoch=live.step - 1)
+    assert spy.call_count == 1