use std::env;
#[cfg(feature = "rerun")]
use std::{cmp::Ordering, path::PathBuf, sync::Arc};
#[path = "common/mod.rs"]
mod common;
#[cfg(feature = "rerun")]
use burn::train::MetricSummary;
use burn::{
backend::{
cuda::{Cuda, CudaDevice},
Autodiff,
},
train::logger::InMemoryMetricLogger,
};
#[cfg(feature = "otel")]
use burn_telemetry::otel::{OpenTelemetryConfig, OpenTelemetryLogger};
use burn_telemetry::MultiLogger;
#[cfg(feature = "rerun")]
use burn_telemetry::{
factory,
rerun::{RerunConfig, RerunLogger},
};
#[cfg(feature = "rerun")]
use common::MnistBatch;
use common::{run_training, MetricLoggerPair, MnistTrainingConfig, MnistTrainingRun};
#[cfg(feature = "rerun")]
use pollster::block_on;
#[cfg(feature = "rerun")]
use rerun::ColorModel;
type BaseBackend = Cuda;
type TrainingBackend = Autodiff<BaseBackend>;
type InferenceBackend = BaseBackend;
type Device = CudaDevice;
#[tokio::main]
async fn main() -> anyhow::Result<()> {
let args = ExampleArgs::parse();
let device = Device::default();
let mut telemetry = TelemetryTargets::new(&args).await?;
let metric_pair = telemetry.metric_pair.take();
let training_config: MnistTrainingConfig = args.training.clone().into();
let run = tokio::task::spawn_blocking(move || {
run_training::<TrainingBackend>(training_config, device, metric_pair)
})
.await
.expect("blocking training task panicked")?;
#[cfg(feature = "rerun")]
if let Some(rerun) = telemetry.rerun_ctx.as_ref() {
log_rerun_metric_curves(rerun, &run);
log_rerun_gallery(rerun, &run, args.preview_samples).await?;
}
Ok(())
}
#[cfg(feature = "rerun")]
fn log_rerun_metric_curves(ctx: &RerunContext, run: &MnistTrainingRun<InferenceBackend>) {
let Some(summary) = &run.summary else {
return;
};
log_metric_split(ctx, "train", &summary.metrics.train);
log_metric_split(ctx, "valid", &summary.metrics.valid);
}
#[cfg(feature = "rerun")]
fn log_metric_split(ctx: &RerunContext, split: &str, metrics: &[MetricSummary]) {
for metric in metrics {
log_metric_series(ctx, split, metric);
}
}
#[cfg(feature = "rerun")]
fn log_metric_series(ctx: &RerunContext, split: &str, metric: &MetricSummary) {
if metric.entries.is_empty() {
return;
}
let entity = metric_summary_entity_path(split, &metric.name);
for entry in &metric.entries {
ctx.logger.log_scalar(&entity, entry.value, entry.step);
}
if let Some((step, value, label)) = summarize_metric(metric) {
ctx.logger.log_text(
&format!("{entity}/summary"),
&format!("{label} {value:.4} at epoch {step}"),
step,
);
}
}
#[cfg(feature = "rerun")]
fn metric_summary_entity_path(split: &str, name: &str) -> String {
format!("metrics_summary/{split}/{}", metric_entity_leaf(name))
}
#[cfg(feature = "rerun")]
fn metric_entity_leaf(name: &str) -> String {
let mut normalized = String::with_capacity(name.len());
for ch in name.chars() {
if ch.is_ascii_alphanumeric() {
normalized.push(ch.to_ascii_lowercase());
} else if ch == '_' {
normalized.push(ch);
} else {
normalized.push('_');
}
}
if normalized.is_empty() {
normalized.push_str("metric");
}
normalized
}
#[cfg(feature = "rerun")]
fn summarize_metric(metric: &MetricSummary) -> Option<(usize, f64, &'static str)> {
let entries = &metric.entries;
if entries.is_empty() {
return None;
}
let name = metric.name.to_lowercase();
let minimize = name.contains("loss") || name.contains("error");
let best = if minimize {
entries.iter().min_by(|a, b| cmp_f64(a.value, b.value))?
} else {
entries.iter().max_by(|a, b| cmp_f64(a.value, b.value))?
};
Some((best.step, best.value, if minimize { "min" } else { "max" }))
}
#[cfg(feature = "rerun")]
fn cmp_f64(a: f64, b: f64) -> Ordering {
a.partial_cmp(&b).unwrap_or(Ordering::Equal)
}
#[cfg(feature = "rerun")]
async fn log_rerun_gallery(
ctx: &RerunContext,
run: &MnistTrainingRun<InferenceBackend>,
preview: usize,
) -> anyhow::Result<()> {
let sample = match run.sample_provider.sample(0, preview) {
Some(batch) => batch,
None => return Ok(()),
};
let outputs = run.model.forward(sample.images.clone());
let preds = outputs.argmax(1);
let preds_data = preds.into_data_async().await.convert::<i64>();
let frames = images_from_batch(&sample);
let labels_data = sample.targets.into_data_async().await.convert::<i64>();
let preds = preds_data.into_vec::<i64>().unwrap_or_default();
let labels = labels_data.into_vec::<i64>().unwrap_or_default();
for (idx, frame) in frames.into_iter().enumerate() {
let entity = format!("mnist/gallery/sample_{idx:02}");
ctx.logger.log_image(&entity, frame, idx);
let pred = preds.get(idx).copied().unwrap_or(-1);
let label = labels.get(idx).copied().unwrap_or(-1);
ctx.logger
.log_scalar(&format!("{entity}/prediction"), pred as f64, idx);
ctx.logger
.log_scalar(&format!("{entity}/label"), label as f64, idx);
ctx.logger.log_text(
&format!("{entity}/text"),
&format!("pred={pred}, label={label}"),
idx,
);
}
Ok(())
}
#[cfg(feature = "rerun")]
fn images_from_batch(batch: &MnistBatch<InferenceBackend>) -> Vec<rerun::Image> {
let data = block_on(batch.images.clone().into_data_async());
let values = data.into_vec::<f32>().unwrap_or_default();
let per_image = 28 * 28;
let mut frames = Vec::new();
for idx in 0..(values.len() / per_image) {
let start = idx * per_image;
let end = start + per_image;
let rgb = values[start..end]
.iter()
.flat_map(|&v| {
let pixel = ((v * 0.3081) + 0.1307) * 255.0;
let clamped = pixel.clamp(0.0, 255.0) as u8;
[clamped, clamped, clamped]
})
.collect::<Vec<_>>();
frames.push(rerun::Image::from_color_model_and_bytes(
rgb,
[28, 28],
ColorModel::RGB,
rerun::ChannelDatatype::U8,
));
}
frames
}
struct TelemetryTargets {
metric_pair: Option<MetricLoggerPair>,
#[cfg(feature = "rerun")]
rerun_ctx: Option<RerunContext>,
}
impl TelemetryTargets {
async fn new(args: &ExampleArgs) -> anyhow::Result<Self> {
let mut train_multi = MultiLogger::new().add_logger(InMemoryMetricLogger::default());
let mut valid_multi = MultiLogger::new().add_logger(InMemoryMetricLogger::default());
#[cfg(feature = "rerun")]
let rerun_ctx = if args.rerun.enabled {
let mut config = RerunConfig::new("mnist-demo")
.with_spawn_viewer(args.rerun.spawn_viewer)
.with_default_timeline("step");
if let Some(path) = &args.rerun.record_to {
config = config.with_record_to_path(path.clone());
}
let rerun =
Arc::new(RerunLogger::new(config, "mnist-demo").map_err(anyhow::Error::msg)?);
let (train, valid) = factory::rerun_metric_logger_pair(rerun.clone());
train_multi = train_multi.add_logger(train);
valid_multi = valid_multi.add_logger(valid);
Some(RerunContext { logger: rerun })
} else {
None
};
#[cfg(feature = "otel")]
{
if args.otel.enabled {
let otel_logger_train = OpenTelemetryLogger::new(
args.otel.config(),
"mnist-demo",
Some(args.otel.device_id.as_str()),
)
.expect("OTel logger");
let otel_logger_valid = OpenTelemetryLogger::new(
args.otel.config(),
"mnist-demo",
Some(args.otel.device_id.as_str()),
)
.expect("OTel logger");
train_multi = train_multi.add_logger(otel_logger_train);
valid_multi = valid_multi.add_logger(otel_logger_valid);
}
}
Ok(Self {
metric_pair: Some(MetricLoggerPair::new(train_multi, valid_multi)),
#[cfg(feature = "rerun")]
rerun_ctx,
})
}
}
#[cfg(feature = "rerun")]
struct RerunContext {
logger: Arc<RerunLogger>,
}
#[derive(Clone)]
struct TrainingCliArgs {
epochs: usize,
batch_size: usize,
valid_batch_size: usize,
num_workers: usize,
seed: u64,
limit_train_samples: Option<usize>,
limit_valid_samples: Option<usize>,
augment: bool,
}
impl From<TrainingCliArgs> for MnistTrainingConfig {
fn from(args: TrainingCliArgs) -> Self {
MnistTrainingConfig {
num_epochs: args.epochs,
batch_size: args.batch_size,
valid_batch_size: args.valid_batch_size,
num_workers: args.num_workers,
seed: args.seed,
limit_train_samples: args.limit_train_samples,
limit_valid_samples: args.limit_valid_samples,
augment: args.augment,
..MnistTrainingConfig::default()
}
}
}
#[cfg(feature = "rerun")]
#[derive(Clone)]
struct RerunCliArgs {
enabled: bool,
spawn_viewer: bool,
record_to: Option<PathBuf>,
}
#[cfg(feature = "otel")]
#[derive(Clone)]
struct OtelCliArgs {
enabled: bool,
endpoint: Option<String>,
device_id: String,
}
struct ExampleArgs {
training: TrainingCliArgs,
preview_samples: usize,
#[cfg(feature = "rerun")]
rerun: RerunCliArgs,
#[cfg(feature = "otel")]
otel: OtelCliArgs,
}
impl ExampleArgs {
fn parse() -> Self {
let mut args = env::args().skip(1);
let mut training = TrainingCliArgs {
epochs: 3,
batch_size: 128,
valid_batch_size: 128,
num_workers: 2,
seed: 42,
limit_train_samples: Some(10_000),
limit_valid_samples: Some(2_000),
augment: true,
};
let mut preview_samples = 8;
#[cfg(feature = "rerun")]
let mut rerun = RerunCliArgs {
enabled: true,
spawn_viewer: true,
record_to: None,
};
#[cfg(feature = "otel")]
let mut otel = OtelCliArgs {
enabled: false,
endpoint: None,
device_id: "cuda:0".into(),
};
while let Some(arg) = args.next() {
match arg.as_str() {
"--epochs" => training.epochs = parse_or_default(&mut args, training.epochs),
"--batch-size" => {
training.batch_size = parse_or_default(&mut args, training.batch_size)
}
"--valid-batch-size" => {
training.valid_batch_size =
parse_or_default(&mut args, training.valid_batch_size)
}
"--num-workers" => {
training.num_workers = parse_or_default(&mut args, training.num_workers)
}
"--seed" => training.seed = parse_or_default(&mut args, training.seed),
"--limit-train" => {
training.limit_train_samples = args.next().and_then(|v| v.parse().ok())
}
"--limit-valid" => {
training.limit_valid_samples = args.next().and_then(|v| v.parse().ok())
}
"--full-dataset" => {
training.limit_train_samples = None;
training.limit_valid_samples = None;
}
"--preview-samples" => {
preview_samples = parse_or_default(&mut args, preview_samples)
}
"--no-augment" => training.augment = false,
#[cfg(feature = "rerun")]
"--no-rerun" => rerun.enabled = false,
#[cfg(feature = "rerun")]
"--record-to" => rerun.record_to = args.next().map(PathBuf::from),
#[cfg(feature = "rerun")]
"--no-viewer" => rerun.spawn_viewer = false,
#[cfg(feature = "rerun")]
"--viewer" => rerun.spawn_viewer = true,
#[cfg(feature = "otel")]
"--otel" => otel.enabled = true,
#[cfg(feature = "otel")]
"--no-otel" => otel.enabled = false,
#[cfg(feature = "otel")]
"--endpoint" => {
otel.endpoint = args.next();
otel.enabled = true;
}
#[cfg(feature = "otel")]
"--device-id" => {
if let Some(id) = args.next() {
otel.device_id = id;
}
}
"--quick" => {
training.epochs = training.epochs.min(1);
training.limit_train_samples =
Some(training.limit_train_samples.unwrap_or(1_000).min(1_000));
training.limit_valid_samples =
Some(training.limit_valid_samples.unwrap_or(256).min(256));
preview_samples = preview_samples.min(4);
}
"--help" | "-h" => {
print_usage();
std::process::exit(0);
}
other => {
eprintln!("Unknown argument: {other}");
print_usage();
std::process::exit(1);
}
}
}
Self {
training,
preview_samples,
#[cfg(feature = "rerun")]
rerun,
#[cfg(feature = "otel")]
otel,
}
}
}
#[cfg(feature = "otel")]
impl OtelCliArgs {
fn config(&self) -> OpenTelemetryConfig {
let mut config = OpenTelemetryConfig::new("mnist-demo")
.with_resource_attribute("device.id", &self.device_id);
if let Some(endpoint) = &self.endpoint {
config = config.with_endpoint(endpoint);
}
config
}
}
fn parse_or_default<T: std::str::FromStr>(
args: &mut impl Iterator<Item = String>,
default: T,
) -> T {
args.next().and_then(|v| v.parse().ok()).unwrap_or(default)
}
fn print_usage() {
eprintln!(
"Usage: cargo run --example mnist [--epochs N] [--batch-size N] \
[--limit-train N] [--limit-valid N] [--full-dataset] [--preview-samples N] [--quick] \
[--record-to path] [--no-rerun] [--otel] [--endpoint URL] [--no-otel]"
);
}