WIP: burn-train in the browser

Cargo.toml

@@ -24,10 +24,11 @@ members = [
     "burn-train",
     "xtask",
     "examples/*",
+    "examples/train-web/train",
     "backend-comparison",
 ]
 
-exclude = ["examples/notebook"]
+exclude = ["examples/notebook", "examples/train-web"]
 
 [workspace.dependencies]
 async-trait = "0.1.73"
@@ -66,7 +67,7 @@ thiserror = "1.0.49"
 tracing-appender = "0.2.2"
 tracing-core = "0.1.31"
 tracing-subscriber = "0.3.17"
-wasm-bindgen = "0.2.87"
+wasm-bindgen = "=0.2.88"
 wasm-bindgen-futures = "0.4.37"
 wasm-logger = "0.2.0"
 

NOTICES.md

@@ -245,4 +245,28 @@ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 See the License for the specific language governing permissions and
 limitations under the License.
 
+## Rust threads in the browser
 
+**Source**: https://github.com/tweag/rust-wasm-threads/tree/main/shared-memory
+
+MIT License
+
+Copyright (c) 2023 Matthew Toohey and Modus Create LLC
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

burn-train/Cargo.toml

@@ -12,20 +12,14 @@ version = "0.11.0"
 
 [features]
 default = ["metrics", "tui"]
-metrics = [
-    "nvml-wrapper",
-    "sysinfo",
-    "systemstat"
-]
-tui = [
-    "ratatui",
-    "crossterm"
-]
+metrics = ["nvml-wrapper", "sysinfo", "systemstat"]
+tui = ["ratatui", "crossterm"]
+browser = ["js-sys", "web-sys", "wasm-bindgen"]
 
 [dependencies]
-burn-core = {path = "../burn-core", version = "0.11.0" }
+burn-core = { path = "../burn-core", version = "0.11.0" }
 
-log = {workspace = true}
+log = { workspace = true }
 tracing-subscriber.workspace = true
 tracing-appender.workspace = true
 tracing-core.workspace = true
@@ -40,8 +34,16 @@ ratatui = { version = "0.23", optional = true, features = ["all-widgets"] }
 crossterm = { version = "0.27", optional = true }
 
 # Utilities
-derive-new = {workspace = true}
-serde = {workspace = true, features = ["std", "derive"]}
+derive-new = { workspace = true }
+serde = { workspace = true, features = ["std", "derive"] }
+
+js-sys = { version = "0.3.64", optional = true }
+web-sys = { version = "0.3.65", optional = true, features = [
+  "Worker",
+  "WorkerOptions",
+  "WorkerType",
+] }
+wasm-bindgen = { workspace = true, optional = true }
 
 [dev-dependencies]
-burn-ndarray = {path = "../burn-ndarray", version = "0.11.0" }
+burn-ndarray = { path = "../burn-ndarray", version = "0.11.0" }

burn-train/src/checkpoint/strategy/metric.rs

@@ -76,7 +76,7 @@ mod tests {
         },
         TestBackend,
     };
-    use std::sync::Arc;
+    use std::rc::Rc;
 
     use super::*;
 
@@ -93,7 +93,7 @@ mod tests {
         store.register_logger_train(InMemoryMetricLogger::default());
         // Register the loss metric.
         metrics.register_train_metric_numeric(LossMetric::<TestBackend>::new());
-        let store = Arc::new(EventStoreClient::new(store));
+        let store = Rc::new(EventStoreClient::new(store));
         let mut processor = MinimalEventProcessor::new(metrics, store.clone());
 
         // Two points for the first epoch. Mean 0.75

burn-train/src/learner/base.rs

@@ -6,6 +6,7 @@ use burn_core::lr_scheduler::LrScheduler;
 use burn_core::module::Module;
 use burn_core::optim::Optimizer;
 use burn_core::tensor::backend::Backend;
+use std::rc::Rc;
 use std::sync::atomic::{AtomicBool, Ordering};
 use std::sync::Arc;
 
@@ -24,7 +25,7 @@ pub struct Learner<LC: LearnerComponents> {
     pub(crate) interrupter: TrainingInterrupter,
     pub(crate) early_stopping: Option<Box<dyn EarlyStoppingStrategy>>,
     pub(crate) event_processor: LC::EventProcessor,
-    pub(crate) event_store: Arc<EventStoreClient>,
+    pub(crate) event_store: Rc<EventStoreClient>,
 }
 
 #[derive(new)]

burn-train/src/learner/builder.rs

@@ -1,4 +1,4 @@
-use std::sync::Arc;
+use std::rc::Rc;
 
 use super::log::install_file_logger;
 use super::Learner;
@@ -312,7 +312,7 @@
                 ));
         }
 
-        let event_store = Arc::new(EventStoreClient::new(self.event_store));
+        let event_store = Rc::new(EventStoreClient::new(self.event_store));
         let event_processor = FullEventProcessor::new(self.metrics, renderer, event_store.clone());
 
         let checkpointer = self.checkpointers.map(|(model, optim, scheduler)| {

burn-train/src/learner/early_stopping.rs

@@ -104,7 +104,7 @@ impl MetricEarlyStoppingStrategy {
 
 #[cfg(test)]
 mod tests {
-    use std::sync::Arc;
+    use std::{rc::Rc, sync::Arc};
 
     use crate::{
         logger::InMemoryMetricLogger,
@@ -188,7 +188,7 @@ mod tests {
         store.register_logger_train(InMemoryMetricLogger::default());
         metrics.register_train_metric_numeric(LossMetric::<TestBackend>::new());
 
-        let store = Arc::new(EventStoreClient::new(store));
+        let store = Rc::new(EventStoreClient::new(store));
         let mut processor = MinimalEventProcessor::new(metrics, store.clone());
 
         let mut epoch = 1;

burn-train/src/lib.rs

@@ -2,6 +2,8 @@
 
 //! A library for training neural networks using the burn crate.
 
+pub mod util;
+
 #[macro_use]
 extern crate derive_new;
 

burn-train/src/metric/processor/full.rs

@@ -1,22 +1,22 @@
 use super::{Event, EventProcessor, Metrics};
 use crate::metric::store::EventStoreClient;
 use crate::renderer::{MetricState, MetricsRenderer};
-use std::sync::Arc;
+use std::rc::Rc;
 
 /// An [event processor](EventProcessor) that handles:
 ///   - Computing and storing metrics in an [event store](crate::metric::store::EventStore).
 ///   - Render metrics using a [metrics renderer](MetricsRenderer).
 pub struct FullEventProcessor<T, V> {
     metrics: Metrics<T, V>,
     renderer: Box<dyn MetricsRenderer>,
-    store: Arc<EventStoreClient>,
+    store: Rc<EventStoreClient>,
 }
 
 impl<T, V> FullEventProcessor<T, V> {
     pub(crate) fn new(
         metrics: Metrics<T, V>,
         renderer: Box<dyn MetricsRenderer>,
-        store: Arc<EventStoreClient>,
+        store: Rc<EventStoreClient>,
     ) -> Self {
         Self {
             metrics,

burn-train/src/metric/processor/minimal.rs

@@ -1,13 +1,13 @@
 use super::{Event, EventProcessor, Metrics};
 use crate::metric::store::EventStoreClient;
-use std::sync::Arc;
+use std::rc::Rc;
 
 /// An [event processor](EventProcessor) that handles:
 ///   - Computing and storing metrics in an [event store](crate::metric::store::EventStore).
 #[derive(new)]
 pub(crate) struct MinimalEventProcessor<T, V> {
     metrics: Metrics<T, V>,
-    store: Arc<EventStoreClient>,
+    store: Rc<EventStoreClient>,
 }
 
 impl<T, V> EventProcessor for MinimalEventProcessor<T, V> {

burn-train/src/metric/store/client.rs

@@ -1,11 +1,12 @@
 use super::EventStore;
 use super::{Aggregate, Direction, Event, Split};
-use std::{sync::mpsc, thread::JoinHandle};
+use crate::util;
+use std::sync::mpsc;
 
 /// Type that allows to communicate with an [event store](EventStore).
 pub struct EventStoreClient {
     sender: mpsc::Sender<Message>,
-    handler: Option<JoinHandle<()>>,
+    handler: Option<Box<dyn FnOnce() -> Result<(), ()>>>,
 }
 
 impl EventStoreClient {
@@ -17,7 +18,7 @@ impl EventStoreClient {
         let (sender, receiver) = mpsc::channel();
         let thread = WorkerThread::new(store, receiver);
 
-        let handler = std::thread::spawn(move || thread.run());
+        let handler = util::spawn(move || thread.run());
         let handler = Some(handler);
 
         Self { sender, handler }
@@ -153,7 +154,7 @@ impl Drop for EventStoreClient {
         let handler = self.handler.take();
 
         if let Some(handler) = handler {
-            handler.join().expect("The event store thread should stop.");
+            handler().expect("The event store thread should stop.");
         }
     }
 }

burn-train/src/util.rs

@@ -0,0 +1,64 @@
+#![allow(missing_docs)]
+
+#[cfg(feature = "browser")]
+use wasm_bindgen::prelude::*;
+
+#[cfg(not(feature = "browser"))]
+pub fn spawn<F>(f: F) -> Box<dyn FnOnce() -> Result<(), ()>>
+where
+    F: FnOnce(),
+    F: Send + 'static,
+{
+    let handle = std::thread::spawn(f);
+    Box::new(move || handle.join().map_err(|_| ()))
+}
+
+// High level description at https://www.tweag.io/blog/2022-11-24-wasm-threads-and-messages/
+// Mostly copied from https://github.com/tweag/rust-wasm-threads/blob/main/shared-memory/src/lib.rs
+#[cfg(feature = "browser")]
+pub fn spawn<F>(f: F) -> Box<dyn FnOnce() -> Result<(), ()>>
+where
+    F: FnOnce(),
+    F: Send + 'static,
+{
+    let mut worker_options = web_sys::WorkerOptions::new();
+    worker_options.type_(web_sys::WorkerType::Module);
+    // Double-boxing because `dyn FnOnce` is unsized and so `Box<dyn FnOnce()>` has
+    let w = web_sys::Worker::new_with_options(
+        WORKER_URL
+            .get()
+            .expect("You must first call `init` with the worker's url."),
+        &worker_options,
+    )
+    .unwrap_or_else(|_| panic!("Error initializing worker at {:?}", WORKER_URL));
+    // an undefined layout (although I think in practice its a pointer and a length?).
+    let ptr = Box::into_raw(Box::new(Box::new(f) as Box<dyn FnOnce()>));
+
+    // See `worker.js` for the format of this message.
+    let msg: js_sys::Array = [
+        &wasm_bindgen::module(),
+        &wasm_bindgen::memory(),
+        &JsValue::from(ptr as u32),
+    ]
+    .into_iter()
+    .collect();
+    if let Err(e) = w.post_message(&msg) {
+        // We expect the worker to deallocate the box, but if there was an error then
+        // we'll do it ourselves.
+        let _ = unsafe { Box::from_raw(ptr) };
+        panic!("Error initializing worker during post_message: {:?}", e)
+    } else {
+        Box::new(move || {
+            w.terminate();
+            Ok(())
+        })
+    }
+}
+
+#[cfg(feature = "browser")]
+static WORKER_URL: std::sync::OnceLock<String> = std::sync::OnceLock::new();
+
+#[cfg(feature = "browser")]
+pub fn init(worker_url: String) -> Result<(), String> {
+    WORKER_URL.set(worker_url)
+}

burn/Cargo.toml

@@ -21,6 +21,8 @@ train = ["burn-train/default", "autodiff", "dataset"]
 # Useful when targeting WASM and not using WGPU.
 wasm-sync = ["burn-core/wasm-sync"]
 
+browser = ["burn-train/browser"]
+
 ## Include nothing
 train-minimal = ["burn-train"]
 
@@ -62,4 +64,12 @@ burn-core = { path = "../burn-core", version = "0.11.0", default-features = fals
 burn-train = { path = "../burn-train", version = "0.11.0", optional = true, default-features = false }
 
 [package.metadata.docs.rs]
-features = ["dataset", "default", "std", "train", "train-tui", "train-metrics", "dataset-sqlite"]
+features = [
+  "dataset",
+  "default",
+  "std",
+  "train",
+  "train-tui",
+  "train-metrics",
+  "dataset-sqlite",
+]

examples/train-web/readme.md

@@ -0,0 +1,14 @@
+## Getting Started
+
+Your `wasm-bindgen-cli` version must *exactly* match the `wasm-bindgen` version in [Cargo.toml](../../Cargo.toml) since `wasm-bindgen-cli` is implicitly used by `wasm-pack`.
+
+For example, run `cargo install --version 0.2.88 wasm-bindgen-cli --force`. The version in this example command is not guaranteed to be up to date!
+
+Install [PNPM](https://pnpm.io/).
+
+Then in separate terminals:
+
+1. `cd train && dev.sh`
+2. `cd web && pnpm i && pnpm dev`
+
+Any changes to `/train` or `burn` should trigger a recompilation. When a new binary is generated, `web` will automatically refresh the page.

examples/train-web/train/.cargo/config.toml

@@ -0,0 +1,6 @@
+[unstable]
+build-std = ['std', 'panic_abort']
+
+[build]
+target = "wasm32-unknown-unknown"
+rustflags = '-Ctarget-feature=+atomics,+bulk-memory,+mutable-globals'

examples/train-web/train/.gitignore

@@ -0,0 +1 @@
+pkg

examples/train-web/train/Cargo.toml

@@ -0,0 +1,21 @@
+[package]
+name = "train"
+version = "0.1.0"
+edition = "2021"
+
+[lib]
+crate-type = ["cdylib"]
+
+[dependencies]
+wasm-bindgen = { workspace = true }
+log = { workspace = true }
+console_error_panic_hook = "0.1.7"
+console_log = { version = "1", features = ["color"] }
+burn = { path = "../../../burn", default-features = false, features = [
+  "autodiff",
+  "ndarray-no-std",
+  "train-minimal",
+  "wasm-sync",
+  "browser",
+] }
+serde = { workspace = true }