Advanced Rust ML: Loading Pre-trained models with Tch-rs | ML Engineering

 

Hi All

Today we're looking at some Rust ML with Tch-rs. Now Tch-rs is a Rust bindings library for the C++ api of PyTorch. "The goal of the tch crate is to provide some thin wrappers around the C++ PyTorch api (a.k.a. libtorch)" - github . Let's get to it then.

Let's demonstrate tch-rs's core capability: creating tensors, moving them to devices, performing operations and getting results - all in Rust while using PyTorch's backend. See introductory example below. 


use tch::{Device, Tensor, Kind};

fn main() {
    let device = Device::cuda_if_available();
    let x = Tensor::of_slice(&[1.0, 2.0, 3.0]).to_device(device);
    let y = Tensor::of_slice(&[4.0, 5.0, 6.0]).to_device(device);
    let z = x + y;
    println!("{:?}", z);
}

Breakdown:

*    Imports: Device, Tensor and Kind  from Tch-rs, the core types for device management and memory operations.

*     Device Selection: Device::cuda_if_available(); automatically selects Cuda if GPU is available otherwise it fallback to CPU.

*    Tensor creation:

        *    Creates tensor x with values [1.0, 2.0, 3.0]

        *    Creates tensor y with values [4.0, 5.0, 6.0]

        *    Both are moved to the selected device (CPU or GPU) using .to_device()

*    Operation: x + y performs element wise addition (much like PyTorch vector addition). 

*    Output: print the result z, which will be [5.0, 7.0, 9.0]

Result: The program outputs [5, 7, 9] (or similar depending on formatting).

 An interesting feature Tch-rs has is the ability to load PyTorch models trained in Python. The below code defines a model and loads in Rust (Tch-rs).


import torch
model = torch.nn.Linear(10, 5)
torch.save(model.state_dict(), "model.pt")

  So far so good. Now let's call our PyTorch model in Rust, code below.


use tch::{nn, Device};

let vs = nn::VarStore::new(Device::Cpu);
let model: nn::Linear = nn::linear(&vs.root(), 10, 5, Default::default());

// Load from Python
vs.load("model.pt").unwrap();

 Just like that, we can load a pre-trained PyTorch model using Rust. I'll have more examples (hopefully some projects) soon on Tch-rs. Full gist source attached below.

 Intro Rust: https://gist.github.com/lightspeed001/edf2bce37ca135d1636f0a51250aa64a 

Python: https://gist.github.com/lightspeed001/1d45c8e7be0d9ad8fc44f635841b2646

Rust: https://gist.github.com/lightspeed001/377401a84e94218e1a882469ba954ee3

 

Later 

- Ed 

Comments

Post a Comment

Popular posts from this blog

PyTorch: Training TinyLlama 1.1B in Google Colab | ML Engineering

Image
Hi All Today we'll be loading a pre-trained model and train it on Colab. Here's a minimal example using Huggingface for transformers and accelerate for memory efficiency. View source below. Let's break down the above code. Block by block. I'm assuming you're using a notebook for this example. However in the full example in my GitHub I've included the Ops files too (Docker and k8s for GKE or AKS). I don't go into those here, I'll have a series of articles tacking MLOps in depth, this is about fine-tuning a pre-trained small LLM in Colab.  Prep: Setup Environment   !pip install -q transformers accelerate datasets peft bitsandbytes   *     transformers provides pre-trained models and training utilities *      peft enables LoRA (low-rank Adaptation), reducing memory usage *        bitsandbytes allows 8-bit precision to shrink model size *        accelerate helps manage GPU memory and multi-GPU setups ...
Read more