Refactor LearningRateScheduler (#711)

Refactor LearningRateScheduler

Co-authored-by: Dave Abrahams <[email protected]>

Author: xihui-wu

Examples/BERT-CoLA/main.swift

@@ -123,18 +123,8 @@ func clipGradByGlobalNorm<L: TrainingLoopProtocol>(_ loop: inout L, event: Train
   }
 }
 
-/// A function that returns a LinearlyDecayedParameter but with first 10 steps linearly warmed up;
-/// for remaining steps it decays at slope of -(peakLearningRate / `totalStepCount`).
-let scheduledParameterGetter = { (_ totalStepCount: Float) -> LinearlyDecayedParameter in
-  LinearlyDecayedParameter(
-    baseParameter: LinearlyWarmedUpParameter(
-      baseParameter: FixedParameter<Float>(peakLearningRate),
-      warmUpStepCount: 10,
-      warmUpOffset: 0),
-    slope: -(peakLearningRate / totalStepCount),  // The LR decays linearly to zero.
-    startStep: 10
-  )
-}
+/// A linear shape to the learning rate in both warmup and decay phases.
+let linear = Shape({ $0 })
 
 var trainingLoop: TrainingLoop = TrainingLoop(
   training: cola.trainingEpochs,
@@ -144,10 +134,15 @@ var trainingLoop: TrainingLoop = TrainingLoop(
   metrics: [.matthewsCorrelationCoefficient],
   callbacks: [
     clipGradByGlobalNorm,
-    LearningRateScheduler(
-      scheduledParameterGetter: scheduledParameterGetter,
-      biasCorrectionBeta1: beta1,
-      biasCorrectionBeta2: beta2).schedule
+    learningRateScheduler(
+      schedule: makeSchedule(
+        [
+          ScheduleSegment(shape: linear, startRate: 0, endRate: peakLearningRate, stepCount: 10),
+          ScheduleSegment(shape: linear, endRate: 0)
+        ]
+      ),
+      biasCorrectionBeta: (beta1, beta2)
+    ),
   ])
 
 print("Training \(bertPretrained.name) for the CoLA task!")

TrainingLoop/CMakeLists.txt

@@ -5,7 +5,8 @@ add_library(TrainingLoop
   Callbacks/StatisticsRecorder.swift
   Callbacks/ProgressPrinter.swift
   Callbacks/CSVLogger.swift
-  Callbacks/LearningRateScheduler.swift)
+  Callbacks/LearningRateScheduler/LearningRateScheduler.swift
+  Callbacks/LearningRateScheduler/LearningRateSchedule.swift)
 target_link_libraries(TrainingLoop PUBLIC
   ModelSupport)
 set_target_properties(TrainingLoop PROPERTIES