Description

Burn Central is a new way of using Burn. It aims at providing a central platform for experiment tracking, model sharing, and deployment for all Burn users!

This repository contains the SDK associated with the project. It offers macros that help attach to your code and send training data to our application. To use this project you must first create an account on the application.

Also needed to use this is the new burn-cli.

Installation

Add Burn Central to your Cargo.toml:

[dependencies]
burn-central = "0.1.0"

Quick Start

Currently, we only support training. Here's how to integrate Burn Central into your training workflow:

1. Register your training function

Use the #[register] macro to register your training function:

use burn_central::{
    experiment::ExperimentRun,
    macros::register,
    runtime::{Args, ArtifactLoader, Model, MultiDevice},
};
use burn::prelude::*;

#[register(training, name = "mnist")]
pub fn training<B: AutodiffBackend>(
    client: &ExperimentRun,
    config: Args<YourExperimentConfig>,
    MultiDevice(devices): MultiDevice<B>,
    loader: ArtifactLoader<ModelArtifact<B>>,
) -> Result<Model<impl ModelArtifact<B::InnerBackend>>, String> {
    // Log your configuration
    client.log_config("Training Config", &training_config)
        .expect("Logging config failed");

    // Your training logic here...
    let model = train::<B>(client, artifact_dir, &training_config, devices[0].clone())?;

    Ok(Model(ModelArtifact {
        model_record: model.into_record(),
        config: training_config,
    }))
}

2. Integrate with your Learner

To enable experiment tracking, you need to add three key components to your LearnerBuilder:

use burn_central::{
    log::RemoteExperimentLoggerInstaller,
    metrics::RemoteMetricLogger,
    record::RemoteCheckpointRecorder,
};
use burn::train::{LearnerBuilder, metric::{AccuracyMetric, LossMetric}};

let learner = LearnerBuilder::new(artifact_dir)
    .metric_train_numeric(AccuracyMetric::new())
    .metric_valid_numeric(AccuracyMetric::new())
    .metric_train_numeric(LossMetric::new())
    .metric_valid_numeric(LossMetric::new())
    // Required: Remote metric logging
    .with_metric_logger(RemoteMetricLogger::new(client))
    // Required: Remote checkpoint saving
    .with_file_checkpointer(RemoteCheckpointRecorder::new(client))
    // Required: Remote application logging
    .with_application_logger(Some(Box::new(
        RemoteExperimentLoggerInstaller::new(client)
    )))
    .num_epochs(config.num_epochs)
    .summary()
    .build(
        model.init::<B>(&device),
        optimizer.init(),
        learning_rate,
        LearningStrategy::SingleDevice(device),
    );

3. Run your training

Once integrated, run your training using the burn-cli to automatically track metrics, checkpoints, and logs on Burn Central.

Requirements

• Rust 1.87.0 or higher
• A Burn Central account (create one at central.burn.dev)
• The burn-cli

Contribution

Contributions to this repository are welcome. You can also submit issues for features you would like to see in the near future.

License

Licensed under either of:

• Apache License, Version 2.0 (LICENSE-APACHE or http://www.apache.org/licenses/LICENSE-2.0)
• MIT license (LICENSE-MIT or http://opensource.org/licenses/MIT)

at your option.