Stimulator

機械学習とか好きな技術話とかエンジニア的な話とかを書く

f:id:vaaaaaanquish:20210810063410p:plain

- はじめに -

本記事は、近似最近傍探索（ANN: Approximate Nearest Neighbor）による画像検索をRustを用いて実装した際のメモである。

画像からの特徴量抽出にTensorFlow Rust bindings、ANNのインデックス管理にRustライブラリであるhoraを利用した。

RustとANNの現状および、実装について触れる。

- RustとANN -

Rustの機械学習関連クレート、事例をまとめたリポジトリがある。

github.com

この中でも、ANN関連のクレートは充実している。利用する場合は以下のようなクレートが候補になる。

*Enet4/faiss-rs
*lerouxrgd/ngt-rs
*rust-cv/hnsw
*hora-search/hora
*InstantDomain/instant-distance
*granne/granne
*qdrant/qdrant

Pythonでもしばしば利用されるfacebook researchのfaiss、Yahoo!のNGTのrust bindingsは強く候補に上がる。C++からGPUが触れる点から利用だけならfaissが活用しやすいだろう。

　
他にPure Rustで機能が充実しているクレートにhoraがある。
github.com

horaには、Python、Javascript、Javaのbindingsがあるだけでなく、Pure Rustである事でWebAssembly化などもサポートしている。
また、インデキシングアルゴリズムとして多く利用されているHNSWIndex以外にグラフベースのSatellite System Graph*1、直積量子化を行うProduct Quantization Inverted File*2が実装されており、開発が継続されている数少ないクレートである。
一部SIMDによる高速化が図られている（https://github.com/rust-lang/packed_simdによるもの）。

（horaの由来は「小さな恋の歌」とREADMEに書いてあるが、どういう経路で知られたのかよくわからない）

今回は、画像検索のwasm化を目指し、horaを利用する。
画像検索がwasm化する事で、API経由で行われていた画像検索の一部がエッジデバイス上で処理できる可能性などの幅が出る事を期待する。
例えば、ネット環境を扱えないや工場やサーバセンター、病院であったり、個人情報の観点でスマフォやカメラの外に出せない画像をその場で類似画像検索にかける事ができる可能性である。

　
画像特徴を抽出する部分でもwasm化を目指すため、wasmの利用実績が多いTensorFlowを利用する。

TensorFlowにはRust bindingsが存在する。
github.com

今回はこちらを利用してモデルを作成し、wasm化する。
他にもDNNのライブラリはいくつかあるが、開発が活発でないか、PyTorchのRust bindingsは現在中間層の出力を受け取る方法がないなど、機能的に難しい場合が多かった。

（実験時に作成したPyTorchのRust bindingsでpretrain modelのpredictを実行するdockerなども公開しているhttps://github.com/vaaaaanquish/tch-rs-pretrain-example-docker）

- pretrainモデルによる特徴量化 -

TensorFlow 2.xでのRustとPythonの相互運用に関する以下の記事を参考にした。

TensorFlow 2.xでのRustとPython

import tensorflowas tffrom keras.modelsimport Modelfrom tensorflow.python.framework.convert_to_constantsimport\    convert_variables_to_constants_v2# pretrainモデルの読み込みmodel = tf.keras.applications.resnet50.ResNet50(weights='imagenet')# 中間層の出力を得るモデルにするembedding_model = Model(inputs=model.layers[0].input, outputs=model.layers[-2].output)# tf.functionに変換しpbファイルとしてgraphを保存できる状態にするresnet = tf.TensorSpec(embedding_model.input_shape, tf.float32, name="resnet")concrete_function = tf.function(lambda x: embedding_model(x)).get_concrete_function(resnet)frozen_model = convert_variables_to_constants_v2(concrete_function)# fileをdumpするtf.io.write_graph(frozen_model.graph,'/app/model',"model.pb", as_text=False)

Rustのbindingsから読み込み、画像ファイルを特徴量に変換する。

// モデルファイルを読み込み、セッションを作るletmut graph=Graph::new();letmut proto=Vec::new();File::open("model/model.pb")?.read_to_end(&mut proto)?;graph.import_graph_def(&proto,&ImportGraphDefOptions::new())?;let session=Session::new(&SessionOptions::new(),&graph)?;// 入力画像を読み込み、リサイズしてTensorに変換するlet img=ImageReader::open("./img/example.jpeg")?.decode()?;let resized_img= img.resize_exact(224asu32,224asu32,FilterType::Lanczos3);let img_vec:Vec<f32>= resized_img.to_rgb8().to_vec().iter().map(|x|*xasf32).collect();let x=Tensor::new(&[1,224,224,3]).with_values(&img_vec)?;// DNNに入力するletmut args=SessionRunArgs::new();args.add_feed(&graph.operation_by_name_required("resnet")?,0,&x);let output= args.request_fetch(&graph.operation_by_name_required("Identity")?,0);session.run(&mut args)?;// check resultlet output_tensor: Tensor<f32>= args.fetch(output)?;let output_array:Vec<f32>= output_tensor.iter().map(|x| x.clone()).collect();println!("{:?}", output_array);

出力として、特徴量vectorが得られる。

- 画像特徴のインデックスと検索 -

horaを利用して画像検索を行う。

// init indexletmut index=hora::index::hnsw_idx::HNSWIndex::<f32,usize>::new(2048,&hora::index::hnsw_params::HNSWParams::<f32>::default(),);// 特定ディレクトリの画像ファイルをインデックスlet paths=fs::read_dir("img")?;letmut file_map=HashMap::new();for (i, path)in paths.into_iter().enumerate() {let file_path= path?.path();let path_str= file_path.to_str();if path_str.is_some() {        file_map.insert(i, path_str.unwrap().to_string().clone());// ファイル一覧を作成let emb_vec= emb.convert_from_img(path_str.unwrap())?;// 画像特徴を得るメソッド        index.add(emb_vec.as_slice(), i)?;// インデックス    }}index.build(hora::core::metrics::Metric::Euclidean).unwrap();// 画像をqueryとして検索let query_image=&file_map[&100]let emb_vec_target= emb.convert_from_img(&query_image.to_string())?;// 画像特徴を得るメソッドlet result= index.search(emb_vec_target.as_slice(),10);// 特徴量をqueryとし検索println!("neighbor images by query: {:?}", query_image);for rin result {println!("{:?}",&file_map[&r]);}