Extracting deep learning image embeddings in Rust

State of Rust Machine Learning Rust may not be the first choice to develop a deep learning model. Right now the site that tracks its progress in the domain is about right that it is not yet ready for this task – http://www.arewelearningyet.com/. While the language is more than capable of handling this task, the ecosystem is […]
Blog
Machine learning

Table of Contents

State of Rust Machine Learning

Rust may not be the first choice to develop a deep learning model. Right now the site that tracks its progress in the domain is about right that it is not yet ready for this task – http://www.arewelearningyet.com/.

While the language is more than capable of handling this task, the ecosystem is not ready yet for general machine learning development. The problem is that in my opinion there are no obvious ways how to deal with dataframes just to name the first thing. I’m pretty sure that in 2-3 years Rust will be a language that can serve for Machine Learning experimentation and development of models but right now its main usage is in the model deployments.

Vision models deployment using Rust

However Rust can offer some value in Machine Learning deployments. In building RecoAI (our fast, accurate and fair recommendation engine) we had the idea to generate image embeddings using deep learning models.

It turns out that probably the best way to do it currently is not to use a generic deep learning framework like https://github.com/LaurentMazare/tch-rs but instead rely on an intermediate ONNX format developed by Microsoft.

Let’s create a command line app in Rust that can inspect the ONNX model (print out the layer names) and extract embeddings for a selected image.

use tract_onnx::prelude::*;
use clap::{AppSettings, Clap};
use std::error::Error;

#[derive(Clap, Debug)]
struct Image {
    #[clap(long, default_value = "cat.jpeg")]
    image_path: String,
    #[clap(long, default_value = "Reshape_103")]
    layer_name: String,
    #[clap(long)]
    normalize: bool,
    #[clap(long, default_value = "224")]
    image_size: usize,
}

#[derive(Clap, Debug)]
enum SubCommand {
    Inspect,
    Embed(Image),
}

#[derive(Clap, Debug)]
#[clap(version = "0.1", author = "Paweł Jankiewicz")]
#[clap(setting = AppSettings::ColoredHelp)]
struct Opts {
    #[clap(long, default_value = "mobilenetv2-7.onnx")]
    model_path: String,
    #[clap(subcommand)]
    subcmd: SubCommand,
}

We are using an amazing clap crate to create a command-line interface. Subcommand struct creates 2 tasks: inspect, embed.

fn inspect_model(opts: Opts) -> Result<(), Box<dyn error="">>  {
    let model = tract_onnx::onnx()
        .model_for_path(opts.model_path)?;

    for name in model.node_names() {
        println!("{:?}", name);
    }

    Ok(())
}</dyn>

inspect_model function takes the options as a parameter and loads the ONNX model using the provided path. We enumerate the node names. After running

cargo run -- --model-path mobilenetv2-7.onnx inspect

It should print out:

(...redacted for brevity)
"Clip_91"
"Conv_92"
"Clip_93"
"Conv_94"
"Conv_95"
"Clip_96"
"GlobalAveragePool_97"
"Shape_98"
"Constant_99"
"Gather_100"
"Unsqueeze_101"
"Concat_102"
"Reshape_103"
"Gemm_104"

In this case we are interested in Reshape_103 node which has a vector 1280 floats representing last but one layer before classification. The code to extract this embedding is a little bit more involved mostly because we need not provide the input image as a vector of floats, and even worse we need to normalize the input according to the mean and the standard deviation used when the model was trained originally.

fn embed(opts: &Opts, image_opt: &Image) -> Result<(), Box<dyn error="">> {
    let image_size = image_opt.image_size;
    let model = tract_onnx::onnx()
        .model_for_path(opts.model_path.clone())?
        .with_input_fact(0, InferenceFact::dt_shape(f32::datum_type(), tvec!(1, 3, image_size, image_size)))?
        .with_output_names(vec!(image_opt.layer_name.clone()))?
        .into_optimized()?
        .into_runnable()?;


    let image = image::open(&image_opt.image_path).unwrap().to_rgb8();
    let resized =
        image::imageops::resize(&image, image_size as u32, image_size as u32, ::image::imageops::FilterType::Triangle);

    let image: Tensor = if image_opt.normalize {
        tract_ndarray::Array4::from_shape_fn((1, 3, image_size, image_size), |(_, c, y, x)| {
            let mean = [0.485, 0.456, 0.406][c];
            let std = [0.229, 0.224, 0.225][c];
            (resized[(x as _, y as _)][c] as f32 / 255.0 - mean) / std})
            .into()
    } else {
        tract_ndarray::Array4::from_shape_fn((1, 3, image_size, image_size), |(_, c, y, x)| {
            resized[(x as _, y as _)][c] as f32
        }).into()
    };

    // run the model on the input
    let result = model.run(tvec!(image))?;
    let best: Vec<_> = result[0]
        .to_array_view::<f32>()?
        .iter()
        .cloned()
        .collect();

    println!("{:?}", best);
    Ok(())
}
</f32></dyn>

To run image embedding on a cat image, run

cargo run -- --model-path mobilenetv2-7.onnx embed --normalize --image-size 224 --image-path cat.jpeg

This command should print out a list of 1280 floats. You can see the code for the whole project in the repository.

Performance

There are other considerations than just writing Rust for fun. Rust and ONNX is meant to increase the performance on Deep Learning model deployments. Let’s compile our small project and see how fast it loads the model and extracts the embeddings

cargo build --release
time target/release/image-embedding-rust --model-path mobilenetv2-7.onnx embed --normalize --image-size 224 --image-path cat.jpeg

On my old Thinkpad P52 with i7 processor it takes 0.2 seconds to load the model and perform feature extraction.

Even as a cold cached model it is not bad at all. This result could be improved by loading the model to memory.

What’s next

In the next episode we will wrap this code and create a web service that can index the images and return most similar image if queried. Stay tuned.

If you are interested in problems like this and using Rust in production development of the fastest real-time recommendation engine check out RecoAI website or drop us an e-mail at hello@logicai.io.

Topics: Onnx | Tutorial

CTO / Data Scientist / Kaggle Grandmaster

CTO / Data Scientist / Kaggle Grandmaster

Other stories in category

BlogKaggle Days
4 – Nature never goes out of style!

4 – Nature never goes out of style!

4 – Nature never goes out of ...

Five continents, twelve events, one grand finale, and a community of more than 10 million - that's Kaggle Days, a nonprofit event for data science enthusiasts and Kagglers. Beginning in November 2021, hundreds of participants attending each meetup face a daunting task to be on the podium and win one of three invitations to the finals in Barcelona and prizes from Kaggle Days and Z by HPZ by HP.

Paras Varshney

16 Aug 2022

BlogKaggle Days
3 – Now you are playing with power

3 – Now you are playing with power

3 – Now you are playing with ...

"It was amazing," commented attendees of the third Kaggle Days X Z by HP World Championship meetup, and we fully agree. The Moscow event brought together as many as 280 data science enthusiasts in one place to take on the challenge and compete for three spots in the grand finale of Kaggle Days in Barcelona. Of course, we already know the winning teams that best handled the contest task. In addition to the excitement of the competition, in Moscow were also inspiring lectures, speeches, and fascinating presentations of modern equipment. As always, at Kaggle Days, a lot was going on.

Paras Varshney

16 Aug 2022

BlogKaggle Days
2 – Water Water everywhere, not a drop to drink

2 – Water Water everywhere, not a drop to drink

2 – Water Water everywhere, n...

"Happy to be part of shaping the future." "It's the Way of The Future." That is how the participants summed up another meetup organized as part of Kaggle Days, a non-profit event for data science enthusiasts who want to grow and compete for prizes under the watchful eye of top Kaggle mentors and grandmasters. The second meetup in New Delhi is behind us. Three hundred participants, more than one hundred teams, and only three invitations to the finals in Barcelona mean that the excitement could not be lacking.

Paras Varshney

16 Aug 2022