Machine Learning in Rust

Hi everyone! In this article, we’re focusing on machine learning in Rust.

Luis Soares

11 Nov 2023 — 7 min read

Hi everyone! In this article, we’re focusing on machine learning in Rust.

As you all know, Rust is great for performance-sensitive tasks, and it turns out it’s pretty handy for machine learning too.

We’re going to look at some of the top Rust crates for ML. These tools are making it easier and more efficient to implement machine learning algorithms in Rust.

So, if you’re interested in how Rust can be used in the world of machine learning, you’re in the right place. Let’s get into it.

Popular Crates for Machine Learning in Rust

In Rust’s journey into machine learning, several crates stand out, each bringing unique strengths and functionalities. Let’s take a closer look at some of the most notable ones:

Tch-rs Crate

Bridging Rust and PyTorch, tch-rs allows Rust users to leverage PyTorch’s deep learning capabilities. It's particularly useful for those who appreciate PyTorch's dynamic computation graphs but want the memory safety and concurrency features of Rust. This crate is a game-changer for developing complex neural networks and executing them with the efficiency Rust is known for.

Key Features:

PyTorch Compatibility: Tch-rs allows Rust programs to use PyTorch, one of the leading deep learning frameworks. This compatibility means access to PyTorch's dynamic computation graphs, gradient calculations, and a vast array of pre-built layers and functions.
GPU Acceleration: It supports GPU acceleration for computations, which is essential for training complex neural networks efficiently.
Comprehensive Deep Learning: Offers a wide range of functionalities from basic tensor operations to advanced neural network architectures, making it suitable for both research and production.
C++ Interoperability: Since PyTorch is primarily written in C++, tch-rs bridges Rust with C++, enabling the use of C++ libraries and tools within Rust applications.

Example:

extern crate tch; 
use tch::{nn, nn::Module, nn::OptimizerConfig, Device, Tensor}; 
 
fn main() { 
    let vs = nn::VarStore::new(Device::cuda_if_available()); 
    let net = nn::seq() 
        .add(nn::linear(vs.root(), 784, 256, Default::default())) 
        .add_fn(|xs| xs.relu()) 
        .add(nn::linear(vs.root(), 256, 10, Default::default())); 
 
    // Assume input and labels are available 
    let mut opt = nn::Adam::default().build(&vs, 1e-3).unwrap(); 
    for epoch in 1..200 { 
        let loss = net.forward(&input).cross_entropy_for_logits(&labels); 
        opt.backward_step(&loss); 
    } 
}

Where:

Imports and Setup: The example imports necessary modules from tch and sets up the neural network.
VarStore: VarStore is used to store the variables (weights) of the neural network. Device::cuda_if_available() ensures that the network will use a GPU if one is available, falling back to the CPU otherwise.
Model Definition: A sequential model is defined with two layers:
A linear layer with 784 input features and 256 output features.
A ReLU activation function.
Another linear layer with 256 input features and 10 output features.
Optimizer: An Adam optimizer is used for training the model.
Training Loop: The network is trained for 200 epochs. In each epoch, it performs forward propagation, calculates loss using cross-entropy, and performs a backward step for optimization.

Documentation:

For more detailed information and usage guidelines, check out the Tch-rs Documentation.

Linfa Crate

Aiming to be Rust’s answer to Python’s scikit-learn, linfa provides a comprehensive collection of algorithms for machine learning. It covers various areas like clustering, classification, and regression. Linfa stands out for its ease of use, making it a go-to for those transitioning from Python to Rust, or for Rustaceans looking to implement standard ML algorithms without reinventing the wheel.

Example:

use linfa::prelude::*; 
use linfa_clustering::{KMeans, generate_blobs}; 
 
fn main() { 
    let dataset = generate_blobs(100, 2, Some(2)); 
    let model = KMeans::params(2).fit(&dataset).unwrap(); 
    let predictions = model.predict(&dataset); 
}

Where:

Dataset Generation: generate_blobs is used to create a synthetic dataset for clustering. It generates 100 samples with 2 features each and 2 centers.
Model Creation and Training: A KMeans model is instantiated to find 2 clusters (KMeans::params(2)) and then fitted to the dataset.
Prediction: The trained model is used to predict the cluster labels for the dataset.

Documentation:

Delve deeper into Linfa’s capabilities at the Linfa Documentation.

Download Now!

Rustlearn Crate

If you’re starting out in machine learning with Rust, rustlearn is a great place to begin. It offers a simple, straightforward approach to traditional machine learning algorithms like decision trees and random forests. This crate is perfect for those who want to grasp the basics of ML in Rust without getting overwhelmed by the complexities of deep learning.

Example:

extern crate rustlearn; 
use rustlearn::prelude::*; 
use rustlearn::trees::decision_tree; 
 
fn main() { 
    let mut dataset = Array::new(); 
    // Assume data is populated 
    let mut model = decision_tree::DecisionTree::new(); 
    model.fit(&dataset).unwrap(); 
    let prediction = model.predict(&dataset).unwrap(); 
}

Where:

Dataset: A dataset (Array) is created and assumed to be populated with features and labels.
Model Definition: A Decision Tree model is initialized.
Model Training: The model is trained on the dataset.
Prediction: The trained model is used to predict labels for the dataset.

Documentation:

Explore Rustlearn further at the Rustlearn Documentation.

Leaf Crate

For deep learning enthusiasts, Leaf is Rust's answer to popular frameworks like TensorFlow and Caffe. It is designed for building and training neural networks, offering a high degree of customization and control. Leaf stands out for its modular approach, allowing users to construct and tweak various layers and components of neural networks.

Example:

In this example, we’ll construct a simple neural network for a classification problem. Please note that the Leaf crate may have limited documentation and community support compared to some other frameworks, so this example is somewhat conceptual and may require adjustments for a specific setup.

extern crate leaf; 
extern crate collenchyma as co; 
 
use leaf::layers::*; 
use leaf::network::*; 
use leaf::optimizers::*; 
use co::backend::{Backend, BackendConfig}; 
use co::frameworks::Native; 
 
fn main() { 
    // Set up a backend (Native in this case) 
    let backend = Rc::new(Backend::<Native>::default().unwrap()); 
 
    // Define the network architecture 
    let mut net_cfg = SequentialConfig::default(); 
    net_cfg.add_input("Input", &[1, 28, 28]); // Example input shape for an MNIST image 
    net_cfg.add_layer(LayerConfig::new("Flatten", LayerType::Flatten)); 
    net_cfg.add_layer(LayerConfig::new("Dense", LayerType::Dense(128))); // 128 nodes 
    net_cfg.add_layer(LayerConfig::new("ReLU", LayerType::ReLU)); 
    net_cfg.add_layer(LayerConfig::new("Dense", LayerType::Dense(10)));  // 10 nodes for output (e.g., 10 classes) 
    net_cfg.add_layer(LayerConfig::new("Softmax", LayerType::Softmax)); 
 
    // Build the network 
    let mut network = Sequential::from_config(backend.clone(), &net_cfg); 
 
    // Define an optimizer 
    let mut optimizer = SGD::new(0.01, 0.9); // Learning rate and momentum 
 
    // Assume we have some training data and labels loaded... 
     
    // Training loop (pseudocode) 
    // for epoch in 0..num_epochs { 
    //     for (batch, label) in training_data.iter() { 
    //         network.forward(batch); 
    //         let loss = network.backward(label); 
    //         optimizer.update(&mut network); 
    //     } 
    // } 
     
    // The actual training process would involve loading data, 
    // feeding it through the network, and updating the weights using the optimizer. 
}

Where:

Backend Setup: We initialize a backend using Collenchyma, which Leaf uses for underlying computations. The Native backend is used here for simplicity.
Network Configuration: We configure a sequential neural network. The example demonstrates a network suitable for a task like MNIST digit classification.
Layers: The network consists of a Flatten layer, two Dense layers (one for hidden nodes and one for output classes), a ReLU activation layer, and a Softmax layer for classification.
Optimizer: We use Stochastic Gradient Descent (SGD) as the optimizer with a specified learning rate and momentum.
Training Loop: While the detailed training loop is not implemented in this example, it would involve feeding batches of training data into the network, performing backpropagation, and updating the model weights.

This example provides a basic framework for using Leaf to build a neural network. Adjustments may be needed based on the specific version of Leaf and Collenchyma, and the precise nature of the dataset and problem being tackled.

Documentation:

For more insights into using Leaf, visit the Leaf Documentation.

Candle Crate

Candle is a minimalist yet powerful ML framework that focuses on performance, including GPU support. It's ideal for those who need the computational efficiency of Rust in high-performance tasks. With Candle, you can develop machine learning models that are not only fast and reliable but also maintainable and easy to understand.

Example:

extern crate candle; 
use candle::{prelude::*, layers::*, models::Sequential, datasets}; 
 
fn main() { 
    let (train_data, train_labels) = datasets::iris::load(); 
    let mut model = Sequential::new(); 
    model.add(Dense::new(4, 10).activation(Activation::ReLU)); 
    model.add(Dense::new(10, 3).activation(Activation::Softmax)); 
    model.compile(Loss::CrossEntropy, Optimizer::SGD { lr: 0.05 }); 
    model.fit(&train_data, &train_labels, 32, 10); 
    let accuracy = model.evaluate(&train_data, &train_labels); 
}

Where:

Dataset Loading: The Iris dataset is loaded into train_data and train_labels.
Model Definition: A sequential model is defined with two dense layers.
The first dense layer has 4 input features (matching the Iris dataset) and 10 output nodes, with ReLU activation.
The second dense layer has 10 input features and 3 output nodes (for the 3 classes of Iris), with Softmax activation.
Compilation: The model is compiled with cross-entropy loss and SGD optimizer with a learning rate of 0.05.
Training: The model is trained on the dataset for 10 epochs with a batch size of 32.
Evaluation: The model’s accuracy is evaluated on the training dataset.

Documentation:

Learn more about Candle at the Candle Documentation.

Conclusion

And that’s a wrap! We’ve explored the landscape of machine learning in Rust, from the deep learning capabilities of tch-rs to the minimalist and performance-focused candle.

Each crate brings something special to the table, whether it's ease of use, a comprehensive set of algorithms, or GPU support for high-performance computing tasks.

The world of machine learning in Rust is constantly evolving, with new libraries and features appearing regularly. So, keep experimenting, keep learning, and who knows? Maybe you'll end up contributing to this exciting field too. Happy coding!

Check out some interesting hands-on Rust articles!

🌟 Developing a Fully Functional API Gateway in Rust — Discover how to set up a robust and scalable gateway that stands as the frontline for your microservices.

🌟 Implementing a Network Traffic Analyzer — Ever wondered about the data packets zooming through your network? Unravel their mysteries with this deep dive into network analysis.

🌟 Building an Application Container in Rust — Join us in creating a lightweight, performant, and secure container from scratch! Docker’s got nothing on this.

🌟 Implementing a P2P Database in Rust: Today, we’re going to roll up our sleeves and get our hands dirty building a Peer-to-Peer (P2P) key-value database.

🌟 Building a Function-as-a-Service (FaaS) in Rust: If you’ve been exploring cloud computing, you’ve likely come across FaaS platforms like AWS Lambda or Google Cloud Functions. In this article, we’ll be creating our own simple FaaS platform using Rust.

🌟 Building an Event Broker in Rust: We’ll explore essential concepts such as topics, event production, consumption, and even real-time event subscriptions.