Model tweaks.

Author: twitchax

kord/Cargo.toml

@@ -63,7 +63,8 @@ ml_train = [
     "burn-ndarray/std",
 ]
 ml_infer = ["ml_base", "burn", "burn-ndarray", "burn-ndarray/std"]
-ml_gpu = ["ml_train", "burn-tch", "burn-wgpu"]
+ml_gpu = ["ml_train", "burn-tch"]
+ml_wgpu = ["ml_train", "burn-wgpu"]
 
 wasm = [
     "rodio/wasm-bindgen",

kord/src/bin.rs

@@ -453,7 +453,7 @@ fn start(args: Args) -> Void {
 
                         klib::ml::train::run_training::<Autodiff<LibTorch<f32>>>(device, &config, true, true)?;
                     }
-                    #[cfg(feature = "ml_gpu")]
+                    #[cfg(feature = "ml_wgpu")]
                     "wgpu" => {
                         use burn_wgpu::{Wgpu, WgpuDevice};
 

kord/src/ml/base/model.rs

@@ -68,24 +68,12 @@ impl<B: Backend> KordModel<B> {
         let output = self.forward(item.samples);
 
         let loss = BinaryCrossEntropyLossConfig::new().with_logits(false).init(&output.device());
-        let loss = loss.forward(output.clone(), targets.clone());
+        let mut loss = loss.forward(output.clone(), targets.clone());
 
-        // let loss = MeanSquareLoss::default();
-        // let loss = loss.forward(output.clone(), targets.clone());
-
-        // let loss = BinaryCrossEntropyLoss::default();
-        // let loss = loss.forward(output.clone(), targets.clone());
-
-        // let mut loss = FocalLoss::default();
-        // loss.gamma = 2.0;
-        // let loss = loss.forward(output.clone(), targets.clone());
-
-        //let loss = loss + l1_regularization(self, 1e-4);
-
-        // let harmonic_penalty_tensor = get_harmonic_penalty_tensor().to_device(&output.device());
-        // let harmonic_loss = output.clone().matmul(harmonic_penalty_tensor).sum_dim(0).mean().mul_scalar(0.0001);
-
-        // let loss = loss + harmonic_loss;
+        // Add L1 regularization
+        // let l1_reg_strength = 1e-4;
+        // let l1_penalty = self.output.weight.val().abs().sum() * l1_reg_strength;
+        // loss = loss + l1_penalty;
 
         MultiLabelClassificationOutput { loss, output, targets }
     }