Merge pull request #4 from twitchax/twitchax/model_fiddling

Twitchax/model fiddling

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "kord"
-version = "0.5.2"
+version = "0.5.3"
 edition = "2021"
 authors = ["Aaron Roney <[email protected]>"]
 license = "MIT"
@@ -29,7 +29,7 @@ analyze_file_alac = ["symphonia/alac", "symphonia/isomp4"]
 
 ml = ["ml_train", "ml_infer"]
 ml_base = ["serde", "byteorder", "bincode"]
-ml_train = ["ml_base", "rand", "burn-autodiff", "burn/train", "burn/std", "burn-ndarray/std"]
+ml_train = ["ml_base", "rand", "rayon", "burn-autodiff", "burn/train", "burn/std", "burn-ndarray/std"]
 ml_infer = ["ml_base", "burn", "burn-ndarray"]
 ml_gpu = ["ml_train", "burn-tch"]
 
@@ -76,10 +76,11 @@ serde = { version = "1.0.152", features = ["derive"], optional = true }
 rand = { version = "0.8.4", optional = true }
 byteorder = { version = "1.4.3", optional = true }
 bincode = { version = "2.0.0-rc.2", git = "https://github.com/bincode-org/bincode.git", default-features = false, optional = true, features = ["alloc", "serde"] }
-burn = { git = "https://github.com/burn-rs/burn", default-features = false, optional = true }
-burn-autodiff = { git = "https://github.com/burn-rs/burn", optional = true }
-burn-tch = { git = "https://github.com/burn-rs/burn", optional = true }
-burn-ndarray = { git = "https://github.com/burn-rs/burn", default-features = false, optional = true }
+rayon = { version = "1.7.0", optional = true }
+burn = { version = "0.6.0", default-features = false, optional = true }
+burn-autodiff = { version = "0.6.0", optional = true }
+burn-tch = { version = "0.6.0", optional = true }
+burn-ndarray = { version = "0.6.0", default-features = false, optional = true }
 
 # plot
 plotters = { version = "0.3.4", optional = true }

model/model_config.json

@@ -1,19 +1,23 @@
 {
-  "source": ".hidden/samples",
+  "source": "samples",
   "destination": "model",
   "log": ".hidden/train_log",
   "simulation_size": 100,
-  "mlp_layers": 0,
+  "simulation_peak_radius": 1.0,
+  "simulation_harmonic_decay": 0.1,
+  "simulation_frequency_wobble": 0.4,
+  "mlp_layers": 3,
   "mlp_size": 1024,
-  "mlp_dropout": 0.3,
+  "mlp_dropout": 0.1,
   "model_epochs": 32,
   "model_batch_size": 100,
-  "model_workers": 32,
+  "model_workers": 64,
   "model_seed": 76980,
-  "adam_learning_rate": 0.0001,
-  "adam_weight_decay": 0.00005,
+  "adam_learning_rate": 0.00001,
+  "adam_weight_decay": 0.0005,
   "adam_beta1": 0.9,
   "adam_beta2": 0.999,
   "adam_epsilon": 1.1920929e-7,
-  "sigmoid_strength": 1.0
+  "sigmoid_strength": 1.0,
+  "no_plots": false
 }

src/analyze/base.rs

@@ -83,6 +83,56 @@ pub fn get_time_space(data: &[f32]) -> Vec<(f32, f32)> {
     buffer.into_iter().enumerate().map(|(k, d)| (k as f32, d.abs())).collect::<Vec<_>>()
 }
 
+pub fn compute_cqt(frequency_space: &[f32]) -> Vec<f32> {
+    const Q_FACTOR: f32 = 24.7; // Q-factor for the CQT
+    const MIN_FREQ: f32 = 65.41; // minimum frequency for the CQT
+    const MAX_FREQ: f32 = 2093.0; // maximum frequency for the CQT
+    const N_BINS: usize = 60; // number of frequency bins for the CQT
+
+    let mut cqt_output = vec![vec![0.0; frequency_space.len()]; N_BINS];
+
+    let log_min_freq = MIN_FREQ.log2();
+    let log_max_freq = MAX_FREQ.log2();
+    let log_freq_step = (log_max_freq - log_min_freq) / (N_BINS as f32 - 1.0);
+
+    for i in 0..N_BINS {
+        let log_freq_center = log_min_freq + i as f32 * log_freq_step;
+        let freq_center = 2.0f32.powf(log_freq_center);
+        let freq_bw = freq_center / Q_FACTOR;
+        let fft_freq_step = 1.0;
+
+        let start_bin = (freq_center - freq_bw / 2.0) / fft_freq_step;
+        let end_bin = (freq_center + freq_bw / 2.0) / fft_freq_step;
+
+        let mut cqt_bin = vec![rustfft::num_complex::Complex::new(0.0, 0.0); frequency_space.len()];
+
+        for j in start_bin as usize..=end_bin as usize {
+            let weight = (j as f32 - freq_center / fft_freq_step) / freq_bw;
+            let weight = weight * std::f32::consts::PI * 2.0;
+            let fft_bin = frequency_space[j];
+            cqt_bin[j] = rustfft::num_complex::Complex::new(fft_bin * weight.sin(), 0.0);
+        }
+
+        let ifft = rustfft::FftPlanner::<f32>::new().plan_fft_inverse(cqt_bin.len());
+        ifft.process(&mut cqt_bin);
+
+        for j in 0..frequency_space.len() {
+            cqt_output[i][j] = cqt_bin[j].abs();
+        }
+    }
+
+    let mut result = vec![];
+    for k in 0..N_BINS {
+        let mut sum = 0.0;
+        for j in 0..frequency_space.len() {
+            sum += cqt_output[k][j];
+        }
+        result.push(sum);
+    }
+
+    result
+}
+
 /// Calculates the "smoothed" frequency space by normalizing to 1.0 seconds of playback.
 pub fn get_smoothed_frequency_space(frequency_space: &[(f32, f32)], length_in_seconds: u8) -> Vec<(f32, f32)> {
     let mut smoothed_frequency_space = Vec::new();
@@ -232,6 +282,7 @@ fn reduce_notes_by_harmonic_series(notes: &[(Note, f32)], cutoff: f32) -> Vec<No
 /// the one before, and the next one.
 ///
 /// Returns a vector of tuples, where the first element is the note, and the second element is the frequency window as a (low, high) tuple.
+/// The first and the last note supplied are ignored, so this method returns `notes.len() - 2` elements.
 pub fn get_frequency_bins(notes: &[Note]) -> Vec<(Note, (f32, f32))> {
     let mut bins = Vec::new();
 
@@ -344,7 +395,7 @@ pub(crate) mod tests {
 
     #[test]
     fn test_get_frequency_bins() {
-        let bins = get_frequency_bins(&ALL_PITCH_NOTES.iter().skip(23).take(62).cloned().collect::<Vec<_>>());
+        let bins = get_frequency_bins(&ALL_PITCH_NOTES.iter().skip(24).take(62).cloned().collect::<Vec<_>>());
 
         assert_eq!(bins.len(), 60);
     }

src/bin.rs

@@ -1,12 +1,12 @@
 use std::path::PathBuf;
 
 use clap::{ArgAction, Parser, Subcommand};
-use klib::core::{
+use klib::{core::{
     base::{Parsable, Res, Void},
     chord::{Chord, Chordable},
     note::Note,
     octave::Octave,
-};
+}};
 
 #[derive(Parser, Debug)]
 #[command(author, version, about, long_about = None)]
@@ -145,7 +145,7 @@ enum MlCommand {
     #[cfg(feature = "ml_train")]
     Train {
         /// The source directory for the gathered samples.
-        #[arg(long, default_value = ".hidden/samples")]
+        #[arg(long, default_value = "samples")]
         source: String,
 
         /// The destination directory for the trained model.
@@ -156,24 +156,36 @@ enum MlCommand {
         #[arg(long, default_value = ".hidden/train_log")]
         log: String,
 
+        /// The device to use for training.
+        #[arg(long, default_value = "gpu")]
+        device: String,
+
         /// Simulation data set size.
         #[arg(long, default_value_t = 100)]
         simulation_size: usize,
 
-        /// The device to use for training.
-        #[arg(long, default_value = "gpu")]
-        device: String,
+        /// Simulation peak radius.
+        #[arg(long, default_value_t = 1.0)]
+        simulation_peak_radius: f32,
+
+        /// Simulation harmonic decay.
+        #[arg(long, default_value_t = 0.1)]
+        simulation_harmonic_decay: f32,
+
+        /// Simulation frequency wobble.
+        #[arg(long, default_value_t = 0.4)]
+        simulation_frequency_wobble: f32,
 
         /// The number of Multi Layer Perceptron (MLP) layers.
-        #[arg(long, default_value_t = 0)]
+        #[arg(long, default_value_t = 3)]
         mlp_layers: usize,
 
         /// The number of neurons in each Multi Layer Perceptron (MLP) layer.
         #[arg(long, default_value_t = 1024)]
         mlp_size: usize,
 
         /// The Multi Layer Perceptron (MLP) dropout rate.
-        #[arg(long, default_value_t = 0.3)]
+        #[arg(long, default_value_t = 0.1)]
         mlp_dropout: f64,
 
         /// The number of epochs to train for.
@@ -185,19 +197,19 @@ enum MlCommand {
         model_batch_size: usize,
 
         /// The number of workers to use for training.
-        #[arg(long, default_value_t = 32)]
+        #[arg(long, default_value_t = 64)]
         model_workers: usize,
 
         /// The seed used for training.
         #[arg(long, default_value_t = 76980)]
         model_seed: u64,
 
         /// The Adam optimizer learning rate.
-        #[arg(long, default_value_t = 1e-4)]
+        #[arg(long, default_value_t = 1e-5)]
         adam_learning_rate: f64,
 
         /// The Adam optimizer weight decay.
-        #[arg(long, default_value_t = 5e-5)]
+        #[arg(long, default_value_t = 5e-4)]
         adam_weight_decay: f64,
 
         /// The Adam optimizer beta1.
@@ -215,6 +227,10 @@ enum MlCommand {
         /// The "sigmoid strength" of the final pass.
         #[arg(long, default_value_t = 1.0)]
         sigmoid_strength: f32,
+
+        /// Suppresses the training plots.
+        #[arg(long, action=ArgAction::SetTrue, default_value_t = false)]
+        no_plots: bool,
     },
 
     /// Records audio from the microphone, and using the trained model, guesses the chord.
@@ -238,6 +254,26 @@ enum MlCommand {
         #[arg(long, default_value_t = 8192.0)]
         x_max: f32,
     },
+
+    /// Runs the ML trainer across various hyperparameters, and outputs the results.
+    #[cfg(feature = "ml_train")]
+    Hpt {
+        /// The source directory for the gathered samples.
+        #[arg(long, default_value = "samples")]
+        source: String,
+
+        /// The destination directory for the trained model.
+        #[arg(long, default_value = "model")]
+        destination: String,
+
+        /// The log directory for training.
+        #[arg(long, default_value = ".hidden/train_log")]
+        log: String,
+
+        /// The device to use for training.
+        #[arg(long, default_value = "gpu")]
+        device: String,
+    }
 }
 
 #[derive(Subcommand, Debug)]
@@ -360,6 +396,9 @@ fn start(args: Args) -> Void {
                 log,
                 simulation_size,
                 device,
+                simulation_peak_radius,
+                simulation_harmonic_decay,
+                simulation_frequency_wobble,
                 mlp_layers,
                 mlp_size,
                 mlp_dropout,
@@ -373,6 +412,7 @@ fn start(args: Args) -> Void {
                 adam_beta2,
                 adam_epsilon,
                 sigmoid_strength,
+                no_plots
             }) => {
                 use burn_autodiff::ADBackendDecorator;
                 use klib::ml::base::TrainConfig;
@@ -382,6 +422,9 @@ fn start(args: Args) -> Void {
                     destination,
                     log,
                     simulation_size,
+                    simulation_peak_radius,
+                    simulation_harmonic_decay,
+                    simulation_frequency_wobble,
                     mlp_layers,
                     mlp_size,
                     mlp_dropout,
@@ -395,23 +438,27 @@ fn start(args: Args) -> Void {
                     adam_beta2,
                     adam_epsilon,
                     sigmoid_strength,
+                    no_plots,
                 };
 
                 match device.as_str() {
                     #[cfg(feature = "ml_gpu")]
                     "gpu" => {
                         use burn_tch::{TchBackend, TchDevice};
 
+                        #[cfg(not(target_os = "macos"))]
                         let device = TchDevice::Cuda(0);
+                        #[cfg(target_os = "macos")]
+                        let device = TchDevice::Mps;
 
-                        klib::ml::train::run_training::<ADBackendDecorator<TchBackend<f32>>>(device, &config, true)?;
+                        klib::ml::train::run_training::<ADBackendDecorator<TchBackend<f32>>>(device, &config, true, true)?;
                     }
                     "cpu" => {
                         use burn_ndarray::{NdArrayBackend, NdArrayDevice};
 
                         let device = NdArrayDevice::Cpu;
 
-                        klib::ml::train::run_training::<ADBackendDecorator<NdArrayBackend<f32>>>(device, &config, true)?;
+                        klib::ml::train::run_training::<ADBackendDecorator<NdArrayBackend<f32>>>(device, &config, true, true)?;
                     }
                     _ => {
                         return Err(anyhow::Error::msg("Invalid device (must choose either `gpu` [requires `ml_gpu` feature] or `cpu`)."));
@@ -464,11 +511,10 @@ fn start(args: Args) -> Void {
             Some(MlCommand::Plot { source, x_min, x_max }) => {
                 use anyhow::Context;
                 use klib::{
-                    analyze::base::translate_frequency_space_to_peak_space,
+                    analyze::base::{translate_frequency_space_to_peak_space, compute_cqt},
                     helpers::plot_frequency_space,
                     ml::{
-                        base::MEL_SPACE_SIZE,
-                        train::helpers::{load_kord_item, mel_filter_banks_from},
+                        base::{MEL_SPACE_SIZE, helpers::{load_kord_item, mel_filter_banks_from, harmonic_convolution}},
                     },
                 };
 
@@ -482,6 +528,16 @@ fn start(args: Args) -> Void {
                 let frequency_space = kord_item.frequency_space.into_iter().enumerate().map(|(k, v)| (k as f32, v)).collect::<Vec<_>>();
                 plot_frequency_space(&frequency_space, "KordItem Frequency Space", &frequency_file_name, x_min, x_max);
 
+                // Plot harmonic convolution.
+                let harmonic_file_name = format!("{}_harmonic", name);
+                let harmonic_space = harmonic_convolution(&kord_item.frequency_space).into_iter().enumerate().map(|(k, v)| (k as f32, v)).collect::<Vec<_>>();
+                plot_frequency_space(&harmonic_space, "KordItem Harmonic Space", &harmonic_file_name, x_min, x_max);
+
+                // Plot CQT space.
+                let cqt_file_name = format!("{}_cqt", name);
+                let cqt_space = compute_cqt(&kord_item.frequency_space).into_iter().enumerate().map(|(k, v)| (k as f32, v)).collect::<Vec<_>>();
+                plot_frequency_space(&cqt_space, "KordItem CQT Space", &cqt_file_name, 0.0, 256.0);
+
                 // Plot mel space.
                 let mel_file_name = format!("{}_mel", name);
                 let mel_space = mel_filter_banks_from(&kord_item.frequency_space)
@@ -521,6 +577,12 @@ fn start(args: Args) -> Void {
                 let time_space = klib::analyze::base::get_time_space(&peak_space);
                 plot_frequency_space(&time_space, "KordItem Time Space", &harmonic_file_name, x_min, x_max);
             }
+            #[cfg(feature = "ml_train")]
+            Some(MlCommand::Hpt { source, destination, log, device }) => {
+                use klib::ml::train::execute::hyper_parameter_tuning;
+
+                hyper_parameter_tuning(source, destination, log, device)?;
+            }
             None => {
                 return Err(anyhow::Error::msg("No subcommand given for `ml`."));
             }