WIP: burn-train in the browser

Cargo.toml

@@ -1,5 +1,4 @@
 [workspace]
-# Try
 # require version 2 to avoid "feature" additiveness for dev-dependencies
 # https://doc.rust-lang.org/cargo/reference/resolver.html#feature-resolver-version-2
 resolver = "2"
@@ -28,10 +27,11 @@ members = [
     "xtask",
     "examples/*",
     "examples/pytorch-import/model",
+    "examples/train-web/train",
     "backend-comparison",
 ]
 
-exclude = ["examples/notebook"]
+exclude = ["examples/notebook", "examples/train-web"]
 
 [workspace.package]
 edition = "2021"
@@ -86,7 +86,7 @@ tokio = { version = "1.35.1", features = ["rt", "macros"] }
 tracing-appender = "0.2.3"
 tracing-core = "0.1.32"
 tracing-subscriber = "0.3.18"
-wasm-bindgen = "0.2.88"
+wasm-bindgen = "=0.2.90"
 wasm-bindgen-futures = "0.4.38"
 wasm-logger = "0.2.0"
 wasm-timer = "0.2.5"

burn-core/src/module/base.rs

@@ -177,6 +177,15 @@ pub trait Module<B: Backend>: Clone + Send + Sync + core::fmt::Debug {
         recorder.record(record, file_path.into())
     }
 
+    /// Return the module using [BytesRecorder](crate::record::BytesRecorder).
+    fn to_bytes<BR: crate::record::BytesRecorder<B>>(
+        self,
+        recorder: &BR,
+    ) -> Result<Vec<u8>, crate::record::RecorderError> {
+        let record = Self::into_record(self);
+        recorder.record(record, ())
+    }
+
     #[cfg(feature = "std")]
     /// Load the module from a file using the provided [file recorder](crate::record::FileRecorder).
     ///

burn-train/Cargo.toml

@@ -11,12 +11,20 @@ repository = "https://github.com/tracel-ai/burn/tree/main/burn-train"
 version.workspace = true
 
 [features]
-default = ["metrics", "tui"]
+default = ["metrics", "tui", "burn-core/default", "burn-core/dataset"]
 metrics = ["nvml-wrapper", "sysinfo", "systemstat"]
 tui = ["ratatui", "crossterm"]
+browser = [
+  "js-sys",
+  "web-sys",
+  "wasm-bindgen",
+  "wasm-bindgen-rayon",
+  "rayon",
+  "burn-core/std",
+]
 
 [dependencies]
-burn-core = { path = "../burn-core", version = "0.12.0", features = ["dataset"] }
+burn-core = { path = "../burn-core", version = "0.12.0", default-features = false }
 
 log = { workspace = true }
 tracing-subscriber = { workspace = true }
@@ -36,5 +44,17 @@ crossterm = { version = "0.27", optional = true }
 derive-new = { workspace = true }
 serde = { workspace = true, features = ["std", "derive"] }
 
+js-sys = { version = "0.3.64", optional = true }
+web-sys = { version = "0.3.64", optional = true, features = [
+  "Worker",
+  "WorkerOptions",
+  "WorkerType",
+  "MessageEvent",
+  "ErrorEvent",
+] }
+wasm-bindgen = { workspace = true, optional = true }
+wasm-bindgen-rayon = { version = "1.0.3", optional = true }
+rayon = { workspace = true, optional = true }
+
 [dev-dependencies]
 burn-ndarray = { path = "../burn-ndarray", version = "0.12.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

@@ -7,6 +7,7 @@ use burn_core::module::Module;
 use burn_core::optim::Optimizer;
 use burn_core::tensor::backend::Backend;
 use burn_core::tensor::Device;
+use std::rc::Rc;
 use std::sync::atomic::{AtomicBool, Ordering};
 use std::sync::Arc;
 
@@ -25,7 +26,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;
@@ -313,7 +313,7 @@ where
                 ));
         }
 
-        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

@@ -113,7 +113,7 @@ impl MetricEarlyStoppingStrategy {
 
 #[cfg(test)]
 mod tests {
-    use std::sync::Arc;
+    use std::rc::Rc;
 
     use crate::{
         logger::InMemoryMetricLogger,
@@ -197,7 +197,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,13 @@
 use super::EventStore;
 use super::{Aggregate, Direction, Event, Split};
-use std::{sync::mpsc, thread::JoinHandle};
+use crate::util::{self, AsyncTaskBoxed};
+use log::info;
+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<AsyncTaskBoxed>,
 }
 
 impl EventStoreClient {
@@ -17,7 +19,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 }

burn-train/src/util.rs

@@ -0,0 +1,46 @@
+#![allow(missing_docs)]
+
+pub trait AsyncTask {
+    fn join(self: Box<Self>) -> Result<(), ()>;
+}
+
+pub type AsyncTaskBoxed = Box<dyn AsyncTask>;
+
+struct Thread {
+    join: Box<dyn FnOnce() -> Result<(), ()>>,
+}
+
+impl AsyncTask for Thread {
+    fn join(self: Box<Self>) -> Result<(), ()> {
+        (self.join)()
+    }
+}
+
+impl Thread {
+    fn new(join: Box<dyn FnOnce() -> Result<(), ()>>) -> Self {
+        Thread { join }
+    }
+}
+
+#[cfg(not(feature = "browser"))]
+pub fn spawn<F>(f: F) -> AsyncTaskBoxed
+where
+    F: FnOnce(),
+    F: Send + 'static,
+{
+    let handle = std::thread::spawn(f);
+    Box::new(Thread::new(Box::new(move || handle.join().map_err(|_| ()))))
+}
+
+#[cfg(feature = "browser")]
+pub fn spawn<F>(f: F) -> AsyncTaskBoxed
+where
+    F: FnOnce(),
+    F: Send + 'static,
+{
+    rayon::spawn(f);
+    Box::new(Thread::new(Box::new(|| Ok(()))))
+}
+
+#[cfg(feature = "browser")]
+pub use wasm_bindgen_rayon::init_thread_pool;

burn/Cargo.toml

@@ -27,6 +27,8 @@ metrics = ["burn-train?/metrics"]
 # Useful when targeting WASM and not using WGPU.
 wasm-sync = ["burn-core/wasm-sync"]
 
+browser = ["burn-train/browser"]
+
 # Datasets
 dataset = ["burn-core/dataset"]
 

examples/train-web/readme.md

@@ -0,0 +1,18 @@
+## 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/).
+
+Install [cargo-watch](https://crates.io/crates/cargo-watch).
+
+The [`postinstall.sh`](./web/postinstall.sh) script expects the mnist database to be at `~/.cache/burn-dataset/mnist.db`. Running `burn/examples/guide` will generate this file. Alternatively, you can download it from [Hugging Face](https://huggingface.co/datasets/mnist).
+
+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,22 @@
+[package]
+name = "train-web"
+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",
+  "train",
+  "ndarray",
+  "wasm-sync",
+  "browser",
+] }
+serde = { workspace = true }
+image = { version = "0.24.7", features = ["png"] }

examples/train-web/train/build-for-web.sh

@@ -0,0 +1,24 @@
+#!/usr/bin/env bash
+
+set -eo pipefail # https://stackoverflow.com/a/2871034
+
+# Add wasm32 target for compiler.
+rustup target add wasm32-unknown-unknown
+
+if ! command -v wasm-pack &>/dev/null; then
+	echo "wasm-pack could not be found. Installing ..."
+	cargo install wasm-pack
+	exit
+fi
+
+# Set optimization flags
+if [[ $1 == "release" ]]; then
+	export RUSTFLAGS="-C lto=fat -C embed-bitcode=yes -C codegen-units=1 -C opt-level=3 --cfg web_sys_unstable_apis"
+else
+	# sets $1 to "dev"
+	set -- dev
+fi
+
+# Run wasm pack tool to build JS wrapper files and copy wasm to pkg directory.
+mkdir -p pkg
+wasm-pack build --out-dir pkg --$1 --target web --no-default-features

examples/train-web/train/dev.sh

@@ -0,0 +1,3 @@
+#!/usr/bin/env bash
+
+cargo watch -- ./build-for-web.sh

examples/train-web/train/rust-toolchain.toml

@@ -0,0 +1,2 @@
+[toolchain]
+channel = "nightly-2023-07-01"