perf: tweak hyper-params

model/src/inference.rs

@@ -116,10 +116,41 @@ pub fn infer<B: Backend, D: TextClassificationDataset + 'static>(
         let class_index = prediction.argmax(1).squeeze::<1>(1).into_scalar(); // Get class index with the highest value
         let class = D::class_name(class_index.elem::<i32>() as usize); // Get class name
 
+        // Apply confidence threshold
+        let confidence_threshold = 0.6; // Define a confidence threshold
+        let max_logit = Tensor::<B, 2, Float>::from_data(logits.clone(), &device)
+            .max()
+            .to_data()
+            .iter()
+            .collect::<Vec<f64>>()[0];
+
         // Print sample text, predicted logits and predicted class
-        println!(
-            "\n=== Item {i} ===\n- Text: {text}\n- Logits: {logits}\n- Prediction: \
-             {class}\n================"
-        );
+        println!("\n=== Item {i} ===\n- Text: {text}");
+
+        println!("- Prediction:");
+        if max_logit < confidence_threshold {
+            println!(
+                "🤔 Model is unsure about the sentiment (confidence {:.2}).",
+                max_logit
+            );
+            println!(
+                "Top prediction would have been class {} with low confidence.",
+                class
+            );
+        } else {
+            println!(
+                "Predicted sentiment = {} (confidence {:.1}%)",
+                class,
+                max_logit * 100.0
+            );
+        }
+
+        // Print logits for each class alongside their labels
+        print!("- Logits: [");
+        for (j, logit) in logits.iter::<f64>().enumerate() {
+            let class_label = D::class_name(j);
+            print!(" ({}: {:.2}) ", class_label, logit);
+        }
+        print!("]\n====================");
     }
 }

model/src/lib.rs

@@ -1,3 +1,3 @@
 pub mod data;
 pub mod inference;
-pub mod model;
+pub mod model;

model/src/model.rs

@@ -8,6 +8,7 @@ use crate::data::{
     TextClassificationInferenceBatch, TextClassificationTrainingBatch,
 };
 use crate::inference::ExperimentConfig;
+use burn::nn::{Dropout, DropoutConfig};
 use burn::{
     nn::{
         Embedding, EmbeddingConfig, Linear, LinearConfig,
@@ -24,6 +25,8 @@ use std::sync::Arc;
 #[derive(Config)]
 pub struct TextClassificationModelConfig {
     transformer: TransformerEncoderConfig,
+    #[config(default = 0.1)]
+    dropout_rate: f64,
     n_classes: usize,
     vocab_size: usize,
     max_seq_length: usize,
@@ -35,6 +38,8 @@ pub struct TextClassificationModel<B: Backend> {
     transformer: TransformerEncoder<B>,
     embedding_token: Embedding<B>,
     embedding_pos: Embedding<B>,
+    embed_dropout: Dropout,
+    output_dropout: Dropout,
     output: Linear<B>,
     n_classes: usize,
     max_seq_length: usize,
@@ -44,17 +49,28 @@ pub struct TextClassificationModel<B: Backend> {
 impl TextClassificationModelConfig {
     /// Initializes a model with default weights
     pub fn init<B: Backend>(&self, device: &B::Device) -> TextClassificationModel<B> {
-        let output = LinearConfig::new(self.transformer.d_model, self.n_classes).init(device);
-        let transformer = self.transformer.init(device);
+        let transformer = self
+            .transformer
+            .clone()
+            .with_dropout(self.dropout_rate)
+            .init(device);
+
         let embedding_token =
             EmbeddingConfig::new(self.vocab_size, self.transformer.d_model).init(device);
         let embedding_pos =
             EmbeddingConfig::new(self.max_seq_length, self.transformer.d_model).init(device);
 
+        let embed_dropout = DropoutConfig::new(self.dropout_rate).init();
+        let output_dropout = DropoutConfig::new(self.dropout_rate).init();
+
+        let output = LinearConfig::new(self.transformer.d_model, self.n_classes).init(device);
+
         TextClassificationModel {
             transformer,
             embedding_token,
             embedding_pos,
+            embed_dropout,
+            output_dropout,
             output,
             n_classes: self.n_classes,
             max_seq_length: self.max_seq_length,
@@ -83,24 +99,33 @@ impl<B: Backend> TextClassificationModel<B> {
         let embedding_tokens = self.embedding_token.forward(tokens);
         let embedding = (embedding_positions + embedding_tokens) / 2;
 
-        // Perform transformer encoding, calculate output and loss
+        // Apply dropout to the embeddings
+        let embedding = self.embed_dropout.forward(embedding);
+
+        // Perform transformer encoding
         let encoded = self
             .transformer
             .forward(TransformerEncoderInput::new(embedding).mask_pad(mask_pad));
+
+        // Apply dropout to the output of the transformer
+        let encoded = self.output_dropout.forward(encoded);
+
+        // Calculate the output using the linear layer
         let output = self.output.forward(encoded);
 
-        let output_classification = output
+        let logits = output
             .slice([0..batch_size, 0..1])
             .reshape([batch_size, self.n_classes]);
 
+        // Compute the loss using Cross-Entropy
         let loss = CrossEntropyLossConfig::new()
-            .init(&output_classification.device())
-            .forward(output_classification.clone(), labels.clone());
+            .init(&logits.device())
+            .forward(logits.clone(), labels.clone());
 
         // Return the output and loss
         ClassificationOutput {
             loss,
-            output: output_classification,
+            output: logits,
             targets: labels,
         }
     }

trainer/src/main.rs

@@ -15,6 +15,7 @@ fn main() {
 mod training_runner {
     use crate::training;
     use burn::backend::{Autodiff, Wgpu};
+    use burn::grad_clipping::GradientClippingConfig;
     use burn::nn::transformer::TransformerEncoderConfig;
     use burn::optim::AdamConfig;
     use burn::optim::decay::WeightDecayConfig;
@@ -35,7 +36,12 @@ mod training_runner {
             TransformerEncoderConfig::new(256, 1024, 8, 4)
                 .with_norm_first(true)
                 .with_quiet_softmax(true),
-            AdamConfig::new().with_weight_decay(Some(WeightDecayConfig::new(5e-5))),
+            AdamConfig::new()
+                .with_weight_decay(Some(WeightDecayConfig::new(0.01)))
+                .with_grad_clipping(Some(GradientClippingConfig::Norm(1.0))) // clip gradients by L2 norm (max 1.0)
+                .with_beta_1(0.9)
+                .with_beta_2(0.999)
+                .with_epsilon(1e-8),
         );
 
         training::train::<B, TweetSentimentDataset>(

trainer/src/training.rs

@@ -5,6 +5,8 @@
 // to build a learner, which is used to train the model. The trained model and the configuration are
 // then saved to the specified directory.
 
+use burn::train::metric::store::{Aggregate, Direction, Split};
+use burn::train::{MetricEarlyStoppingStrategy, StoppingCondition};
 use burn::{
     data::{dataloader::DataLoaderBuilder, dataset::transform::SamplerDataset},
     lr_scheduler::noam::NoamLrSchedulerConfig,
@@ -49,21 +51,21 @@ pub fn train<B: AutodiffBackend, D: TextClassificationDataset + 'static>(
     .init::<B>(&devices[0]);
 
     // Initialize data loaders for training and testing data
-    let dataloader_train = DataLoaderBuilder::new(batcher.clone())
+    let train_loader = DataLoaderBuilder::new(batcher.clone())
         .batch_size(config.batch_size)
         .num_workers(1)
         .build(SamplerDataset::new(dataset_train, 50_000));
-    let dataloader_test = DataLoaderBuilder::new(batcher)
+    let valid_loader = DataLoaderBuilder::new(batcher)
         .batch_size(config.batch_size)
         .num_workers(1)
         .build(SamplerDataset::new(dataset_test, 5_000));
-
+    
     // Initialize optimizer
-    let optim = config.optimizer.init();
+    let optimizer = config.optimizer.init();
 
     // Initialize learning rate scheduler
-    let lr_scheduler = NoamLrSchedulerConfig::new(1e-2)
-        .with_warmup_steps(1000)
+    let lr_scheduler = NoamLrSchedulerConfig::new(1e-4)
+        .with_warmup_steps(8_000)
         .with_model_size(config.transformer.d_model)
         .init()
         .unwrap();
@@ -79,13 +81,20 @@ pub fn train<B: AutodiffBackend, D: TextClassificationDataset + 'static>(
         .metric_valid_numeric(AccuracyMetric::new())
         .metric_train_numeric(LearningRateMetric::new())
         .with_file_checkpointer(CompactRecorder::new())
+        .early_stopping(MetricEarlyStoppingStrategy::new::<LossMetric<B>>(
+            &LossMetric::new(),
+            Aggregate::Mean,
+            Direction::Lowest,
+            Split::Valid,
+            StoppingCondition::NoImprovementSince { n_epochs: 2 },
+        )) // stop if no val loss improvement for 2 epochs
         .devices(devices)
         .num_epochs(config.num_epochs)
         .summary()
-        .build(model, optim, lr_scheduler);
+        .build(model, optimizer, lr_scheduler);
 
     // Train the model
-    let model_trained = learner.fit(dataloader_train, dataloader_test);
+    let model_trained = learner.fit(train_loader, valid_loader);
 
     // Save the configuration and the trained model
     config.save(format!("{artifact_dir}/config.json")).unwrap();