regression_model/
regression_model.rs

1//! Neural Architecture Search with Genetic Algorithms
2//!
3//! Demonstrates using fx_durable_ga to optimize neural network hyperparameters:
4//! hidden size, number of layers, activation function, bias usage, and learning rate.
5//!
6//! **IMPORTANT!**
7//! This library requires fx-durable-ga-example-simple-regression to be installed and available on the PATH. The crate is available here:
8//! https://github.com/jakob-lilliemarck/fx-durable-ga-simple-regression
9//!
10//! This workaround is required as the Autodiff backend of the Burn ML framework currently does not free memory between training run.
11//! As such, running multiple training runs will cause unbounded memory allocation.
12//!
13//! This example handles that by running each training run as a subprocess, in which case all memory allocations are freed after each run.
14
15use anyhow::Result;
16use fx_durable_ga::{
17    bootstrap,
18    models::{
19        Crossover, Distribution, Encodeable, Evaluator, FitnessGoal, GeneBounds, Mutagen,
20        MutationRate, Request, Schedule, Selector, Temperature, Terminated,
21    },
22    register_event_handlers, register_job_handlers,
23};
24use fx_mq_jobs::FX_MQ_JOBS_SCHEMA_NAME;
25use fx_mq_jobs::Queries;
26use serde::Deserialize;
27use sqlx::postgres::PgPoolOptions;
28use std::time::Duration;
29use std::{env, sync::Arc};
30use tracing::Level;
31use uuid::Uuid;
32
33const WORKERS: usize = 4;
34const FITNESS_TARGET: f64 = 0.1;
35
36#[derive(Debug, Clone, Copy, serde::Serialize, serde::Deserialize, PartialEq, Eq)]
37pub enum ActivationFunction {
38    Relu,
39    Gelu,
40    Sigmoid,
41}
42
43impl std::fmt::Display for ActivationFunction {
44    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
45        match self {
46            Self::Relu => write!(f, "relu"),
47            Self::Gelu => write!(f, "gelu"),
48            Self::Sigmoid => write!(f, "sigmoid"),
49        }
50    }
51}
52
53#[derive(Debug, Clone)]
54struct NeuralArchitecture {
55    pub hidden_size: usize,
56    pub num_hidden_layers: usize,
57    pub activation_fn: ActivationFunction,
58    pub use_bias: bool,
59    pub learning_rate: f64,
60}
61
62impl Encodeable for NeuralArchitecture {
63    const NAME: &'static str = "neural_architecture";
64
65    type Phenotype = NeuralArchitecture;
66
67    fn morphology() -> Vec<GeneBounds> {
68        vec![
69            GeneBounds::integer(0, 3, 4).unwrap(), // hidden_size: [32, 64, 128, 256]
70            GeneBounds::integer(0, 7, 8).unwrap(), // num_hidden_layers: [1, 2, 3, 4, 5, 6, 7, 8]
71            GeneBounds::integer(0, 2, 3).unwrap(), // activation_fn: [ReLU, GELU, Sigmoid]
72            GeneBounds::integer(0, 1, 2).unwrap(), // use_bias: [false, true]
73            GeneBounds::integer(0, 2, 3).unwrap(), // learning_rate: [1e-4, 1e-3, 1e-2]
74        ]
75    }
76
77    fn encode(&self) -> Vec<i64> {
78        let hidden_size_idx = match self.hidden_size {
79            32 => 0,
80            64 => 1,
81            128 => 2,
82            256 => 3,
83            _ => 1,
84        };
85
86        let layers_idx = (self.num_hidden_layers - 1).min(2) as i64;
87
88        let activation_idx = match self.activation_fn {
89            ActivationFunction::Relu => 0,
90            ActivationFunction::Gelu => 1,
91            ActivationFunction::Sigmoid => 2,
92        };
93
94        let bias_idx = if self.use_bias { 1 } else { 0 };
95
96        let lr_idx = if self.learning_rate <= 1e-4 {
97            0
98        } else if self.learning_rate <= 1e-3 {
99            1
100        } else {
101            2
102        };
103
104        vec![
105            hidden_size_idx,
106            layers_idx,
107            activation_idx,
108            bias_idx,
109            lr_idx,
110        ]
111    }
112
113    fn decode(genes: &[i64]) -> Self::Phenotype {
114        let hidden_size = match genes[0] {
115            0 => 32,
116            1 => 64,
117            2 => 128,
118            3 => 256,
119            _ => 64,
120        };
121
122        let num_hidden_layers = (genes[1] + 1).clamp(1, 8) as usize;
123
124        let activation_fn = match genes[2] {
125            0 => ActivationFunction::Relu,
126            1 => ActivationFunction::Gelu,
127            2 => ActivationFunction::Sigmoid,
128            _ => ActivationFunction::Relu,
129        };
130
131        let use_bias = genes[3] == 1;
132
133        let learning_rate = match genes[4] {
134            0 => 1e-4,
135            1 => 1e-3,
136            2 => 1e-2,
137            _ => 1e-3,
138        };
139
140        NeuralArchitecture {
141            hidden_size,
142            num_hidden_layers,
143            activation_fn,
144            use_bias,
145            learning_rate,
146        }
147    }
148}
149
150struct ArchitectureEvaluator;
151
152#[derive(Deserialize)]
153struct ResultOutput {
154    validation_loss: f64,
155}
156
157impl Evaluator<NeuralArchitecture> for ArchitectureEvaluator {
158    fn fitness<'a>(
159        &self,
160        _genotype_id: Uuid,
161        phenotype: NeuralArchitecture,
162        _: &Request,
163        _: &'a Box<dyn Terminated>,
164    ) -> futures::future::BoxFuture<'a, Result<f64, anyhow::Error>> {
165        Box::pin(async move {
166            // Spawn the binary
167            let output = tokio::process::Command::new("fx-example-regression")
168                .args([
169                    "--hidden-size",
170                    &phenotype.hidden_size.to_string(),
171                    "--num-hidden-layers",
172                    &phenotype.num_hidden_layers.to_string(),
173                    "--activation-fn",
174                    &phenotype.activation_fn.to_string(),
175                    "--learning-rate",
176                    &phenotype.learning_rate.to_string(),
177                ])
178                .output()
179                .await
180                .expect("Failed to run fx-example-regression");
181
182            // Print stderr logs (your tracing output)
183            eprintln!("{}", String::from_utf8_lossy(&output.stderr));
184
185            // Parse stdout as JSON (the ResultOutput struct)
186            let result: ResultOutput =
187                serde_json::from_slice(&output.stdout).expect("Invalid JSON from training binary");
188
189            Ok(result.validation_loss)
190        })
191    }
192}
193
194#[tokio::main]
195async fn main() -> Result<()> {
196    dotenv::from_filename(".env.local").ok();
197    tracing_subscriber::fmt()
198        .pretty()
199        .with_thread_ids(true)
200        .with_max_level(Level::INFO)
201        .init();
202
203    let database_url = env::var("DATABASE_URL").expect("DATABASE_URL must be set");
204    let pool = PgPoolOptions::new()
205        .max_connections(10)
206        .connect(&database_url)
207        .await?;
208
209    // Run all default migrations
210    fx_durable_ga::migrations::run_default_migrations(&pool).await?;
211
212    let service = Arc::new(
213        bootstrap(pool.clone())
214            .await?
215            .register::<NeuralArchitecture, _>(ArchitectureEvaluator)
216            .await?
217            .build(),
218    );
219
220    let mut registry = fx_event_bus::EventHandlerRegistry::new();
221    register_event_handlers(
222        Arc::new(Queries::new(FX_MQ_JOBS_SCHEMA_NAME)),
223        service.clone(),
224        &mut registry,
225    );
226    let mut listener = fx_event_bus::Listener::new(pool.clone(), registry);
227    tokio::spawn(async move { listener.listen(None).await });
228
229    let host_id = Uuid::parse_str("00000000-0000-0000-0000-123456789abc")?;
230    let mut jobs_listener = fx_mq_jobs::Listener::new(
231        pool.clone(),
232        register_job_handlers(&service, fx_mq_jobs::RegistryBuilder::new()),
233        WORKERS,
234        host_id,
235        Duration::from_secs(600),
236    )
237    .await?;
238    tokio::spawn(async move { jobs_listener.listen().await });
239
240    service
241        .new_optimization_request(
242            NeuralArchitecture::NAME,
243            NeuralArchitecture::HASH,
244            FitnessGoal::minimize(FITNESS_TARGET)?,
245            Schedule::generational(10, 10),
246            Selector::tournament(5, 15)?,
247            Mutagen::new(Temperature::constant(0.8)?, MutationRate::constant(0.4)?),
248            Crossover::uniform(0.5)?,
249            Distribution::latin_hypercube(15),
250            None::<()>,
251        )
252        .await?;
253
254    tokio::time::sleep(Duration::from_secs(3600)).await;
255    Ok(())
256}
regression_model/regression_model.rs

regression_model/
regression_model.rs
