Minor Finetuning Scheduler Tutorial Update (#176)

* update fts link, add advanced feature ref, cleanup depth logging
* recheck with updated fts deps
* Apply suggestions from code review

Co-authored-by: Jirka Borovec <[email protected]>
Co-authored-by: Rohit Gupta <[email protected]>

Author: 3 people

lightning_examples/finetuning-scheduler/.meta.yml

@@ -1,11 +1,11 @@
 title: Finetuning Scheduler
 author: "[Dan Dale](https://github.com/speediedan)"
 created: 2021-11-29
-updated: 2022-05-10
+updated: 2022-06-10
 license: CC BY-SA
-build: 3
+build: 0
 tags:
-  - finetuning
+  - Finetuning
 description: |
   This notebook introduces the [Finetuning Scheduler](https://finetuning-scheduler.readthedocs.io/en/stable/index.html) extension
   and demonstrates the use of it to finetune a small foundational model on the

lightning_examples/finetuning-scheduler/finetuning-scheduler.py

@@ -153,6 +153,8 @@
 # - ``DDP_SHARDED``
 # - ``DDP_SHARDED_SPAWN``
 #
+# Custom or officially unsupported strategies can be used by setting [FinetuningScheduler.allow_untested](https://finetuning-scheduler.readthedocs.io/en/stable/api/finetuning_scheduler.fts.html?highlight=allow_untested#finetuning_scheduler.fts.FinetuningScheduler.params.allow_untested) to ``True``.
+# Note that most currently unsupported strategies are so because they require varying degrees of modification to be compatible (e.g. ``deepspeed`` requires an ``add_param_group`` method, ``tpu_spawn`` an override of the current broadcast method to include python objects)
 # </div>
 
 # %% [markdown]
@@ -387,9 +389,12 @@ def training_step(self, batch, batch_idx):
         self.log("train_loss", loss)
         return loss
 
-    def training_epoch_end(self, outputs: List[Any]) -> None:
+    def on_train_epoch_start(self) -> None:
         if self.finetuningscheduler_callback:
-            self.log("finetuning_schedule_depth", float(self.finetuningscheduler_callback.curr_depth))
+            self.logger.log_metrics(
+                metrics={"finetuning_schedule_depth": float(self.finetuningscheduler_callback.curr_depth)},
+                step=self.global_step,
+            )
 
     def validation_step(self, batch, batch_idx, dataloader_idx=0):
         outputs = self(**batch)
@@ -524,6 +529,8 @@ def configure_optimizers(self):
 # used in other pytorch-lightning tutorials) also work with FinetuningScheduler. Though the LR scheduler is theoretically
 # justified [(Loshchilov & Hutter, 2016)](#f4), the particular values provided here are primarily empircally driven.
 #
+# [FinetuningScheduler](https://finetuning-scheduler.readthedocs.io/en/stable/api/finetuning_scheduler.fts.html#finetuning_scheduler.fts.FinetuningScheduler) also supports LR scheduler
+# reinitialization in both explicit and implicit finetuning schedule modes. See the [advanced usage documentation](https://finetuning-scheduler.readthedocs.io/en/stable/advanced/lr_scheduler_reinitialization.html) for explanations and demonstration of the extension's support for more complex requirements.
 # </div>