Feat/cosine_annealing_lr

src/cosine_annealing.rs

@@ -0,0 +1,75 @@
+use burn::{lr_scheduler::LRScheduler, LearningRate};
+#[derive(Clone, Debug)]
+pub struct CosineAnnealingLR {
+ t_max: f64,
+ eta_min: f64,
+ init_lr: f64,
+ step_count: f64,
+ current_lr: LearningRate,
+}
+
+impl CosineAnnealingLR {
+ pub fn init(t_max: f64, init_lr: f64) -> CosineAnnealingLR {
+ CosineAnnealingLR {
+ t_max,
+ eta_min: 0.0,
+ init_lr,
+ step_count: 0.0,
+ current_lr: init_lr,
+ }
+ }
+}
+
+impl LRScheduler for CosineAnnealingLR {
+ type Record = usize;
+
+ fn step(&mut self) -> LearningRate {
+ self.step_count += 1.0;
+ use std::f64::consts::PI;
+ fn cosine_annealing_lr(
+ init_lr: f64,
+ lr: f64,
+ step_count: f64,
+ t_max: f64,
+ eta_min: f64,
+ ) -> f64 {
+ let cosine_arg = PI * step_count / t_max;
+ if (step_count - 1.0 - t_max) % (2.0 * t_max) == 0.0 {
+ (init_lr - eta_min) * (1.0 - f64::cos(PI / t_max)) / 2.0
+ } else {
+ (1.0 + f64::cos(cosine_arg)) / (1.0 + f64::cos(PI * (step_count - 1.0) / t_max))
+ * (lr - eta_min)
+ + eta_min
+ }
+ }
+ self.current_lr = cosine_annealing_lr(
+ self.init_lr,
+ self.current_lr,
+ self.step_count,
+ self.t_max,
+ self.eta_min,
+ );
+ self.current_lr
+ }
+
+ fn to_record(&self) -> Self::Record {
+ self.step_count as usize
+ }
+
+ fn load_record(mut self, record: Self::Record) -> Self {
+ self.step_count = record as f64;
+ self
+ }
+}
+
+#[test]
+fn test_lr_scheduler() {
+ let mut lr_scheduler = CosineAnnealingLR::init(100000.0, 1.0e-1);
+ for i in 0..400000 {
+ if i % 5000 == 0 {
+ println!("{}", lr_scheduler.current_lr);
+ }
+ lr_scheduler.step();
+ }
+ println!("{}", lr_scheduler.current_lr);
+}

src/dataset.rs

@@ -139,6 +139,10 @@ impl FSRSDataset {
  Self::new()
  }
 
+ pub fn len(&self) -> usize {
+ self.dataset.len()
+ }
+
  fn new() -> Self {
  let dataset = InMemDataset::<FSRSItem>::new(anki_to_fsrs());
  Self { dataset }

src/lib.rs

@@ -1,4 +1,5 @@
 pub mod convertor;
+mod cosine_annealing;
 pub mod dataset;
 pub mod model;
 pub mod training;

src/training.rs

@@ -1,3 +1,4 @@
+use crate::cosine_annealing::CosineAnnealingLR;
 use crate::dataset::{FSRSBatch, FSRSBatcher, FSRSDataset};
 use crate::model::{Model, ModelConfig};
 use crate::weight_clipper::weight_clipper;
@@ -75,13 +76,13 @@ pub struct TrainingConfig {
  pub optimizer: AdamConfig,
  #[config(default = 10)]
  pub num_epochs: usize,
- #[config(default = 2)]
+ #[config(default = 1)]
  pub batch_size: usize,
  #[config(default = 4)]
  pub num_workers: usize,
  #[config(default = 42)]
  pub seed: u64,
- #[config(default = 1.0e-4)]
+ #[config(default = 1.0e-3)]
  pub learning_rate: f64,
 }
 
@@ -113,18 +114,23 @@ pub fn train<B: ADBackend<FloatElem = f32>>(
  .num_workers(config.num_workers)
  .build(FSRSDataset::test());
 
+ let lr_scheduler = CosineAnnealingLR::init(
+ (FSRSDataset::train().len() * config.num_epochs) as f64,
+ config.learning_rate,
+ );
+
  let learner = LearnerBuilder::new(artifact_dir)
  // .metric_train_plot(AccuracyMetric::new())
  // .metric_valid_plot(AccuracyMetric::new())
  // .metric_train_plot(LossMetric::new())
  // .metric_valid_plot(LossMetric::new())
- .with_file_checkpointer(1, PrettyJsonFileRecorder::<FullPrecisionSettings>::new())
+ .with_file_checkpointer(10, PrettyJsonFileRecorder::<FullPrecisionSettings>::new())
  .devices(vec![device])
  .num_epochs(config.num_epochs)
  .build(
  config.model.init::<B>(),
  config.optimizer.init(),
- config.learning_rate,
+ lr_scheduler,
  );
 
  let mut model_trained = learner.fit(dataloader_train, dataloader_test);