はじめに

TensorFlowのRustバインディングは公式に存在しているものの、ドキュメントは乏しく、情報も少ない状態です。
そこで、本記事では動作環境をDocker上に構築して公式のexamplesを試してみます。

また、比較のために「PyTorchのRustバインディングを使ってみる」という記事も書いています。Rustでディープラーニングをしてみたいと考えている方は是非そちらもご参照ください。

環境構築手順

プロジェクトの作成

任意のディレクトリにcargoでプロジェクトを作成します。

cargo new tf-example
cd tf-example

依存関係を追加する

Cargo.tomlのdependenciesに下記を追加します。

Cargo.toml

[dependencies]
tensorflow = "0.16.0"

Dockerfileの作成

tf-example下にDockerfileを作成し、下記をコピーします。

Dockerfile

FROM tensorflow/tensorflow

WORKDIR /home

# 必要なパッケージとRustをインストール
RUN apt update && apt install -y pkg-config libssl-dev && \
    curl -sSf https://sh.rustup.rs | sh -s -- -y
ENV PATH /root/.cargo/bin:$PATH

# ビルド
COPY Cargo.toml Cargo.toml
COPY ./src ./src
RUN cargo build --release

Dockerイメージのビルド

前の手順で作成したDockerfileをビルドします。

docker build -t tf-example .

サンプル１：xor.rs

このサンプルではモデル構築、学習、推論の一通りの動作を確認することが出来ます。
examplesに含まれるxor.rsの内容をコピーしてtf-example/src/main.rsに上書き、下記のコマンドを実行します。

docker run --rm -v "$PWD"/src:/home/src tf-example cargo run --release

以下のような結果が出力されれば正しく動いているかと思います。

   Compiling tf-example v0.1.0 (/home)
    Finished release [optimized] target(s) in 3.99s
     Running `target/release/tf-example`
2021-02-23 05:50:13.585369: I tensorflow/core/platform/cpu_feature_guard.cc:142] Your CPU supports instructions that this TensorFlow binary was not compiled to use: AVX2 FMA
2021-02-23 05:50:13.589896: I tensorflow/core/platform/profile_utils/cpu_utils.cc:94] CPU Frequency: 1996250000 Hz
2021-02-23 05:50:13.590687: I tensorflow/compiler/xla/service/service.cc:168] XLA service 0x55cd952bdb60 initialized for platform Host (this does not guarantee that XLA will be used). Devices:
2021-02-23 05:50:13.590734: I tensorflow/compiler/xla/service/service.cc:176]   StreamExecutor device (0): Host, Default Version
Error: 0.000009141944
Error: 0.00000010482859
Error: 0.0000047976396
Error: 0.0000007003914
2021-02-23 05:50:16.004730: I tensorflow/cc/saved_model/reader.cc:31] Reading SavedModel from: /tmp/tf-rust-example-xor-saved-model2
2021-02-23 05:50:16.005479: I tensorflow/cc/saved_model/reader.cc:54] Reading meta graph with tags { train serve }
2021-02-23 05:50:16.006944: I tensorflow/cc/saved_model/loader.cc:202] Restoring SavedModel bundle.
2021-02-23 05:50:16.026940: I tensorflow/cc/saved_model/loader.cc:311] SavedModel load for tags { train serve }; Status: success. Took 22207 microseconds.
Error: 0.000009141944
Error: 0.000007655446
Error: 0.00000042667978
Error: 0.000005234325

サンプル２：addition.rs

このサンプルでは、Pythonで作成したモデルをRustで読み込んで使用します。
本来であれば同じDockerイメージでexamples/addition/addition.pyを使用してモデルを生成したいのですが、このスクリプトではTF1系で書かれており動きません。
（ベースモデルを変えればよいのですが、今更TF1系を使うつもりもないので…）

そこで、下記のコマンドでリポジトリからモデルをダウンロードします。

curl -OL https://raw.githubusercontent.com/tensorflow/rust/master/examples/addition/model.pb

モデルのダウンロード以降は先の手順と同様に、addition.rsをtf-example/src/main.rsに上書きし、下記のコマンドを実行します。

docker run --rm \
  -v "$PWD"/src:/home/src \
  -v "$PWD"/model.pb:/home/examples/addition/model.pb \
  tf-example cargo run --release

以下のような結果が出力されれば正しく動いているかと思います。

   Compiling tf-example v0.1.0 (/home)
    Finished release [optimized] target(s) in 1.02s
     Running `target/release/tf-example`
2021-02-23 05:51:13.670859: I tensorflow/core/platform/cpu_feature_guard.cc:142] Your CPU supports instructions that this TensorFlow binary was not compiled to use: AVX2 FMA
2021-02-23 05:51:13.675355: I tensorflow/core/platform/profile_utils/cpu_utils.cc:94] CPU Frequency: 1996250000 Hz
2021-02-23 05:51:13.676066: I tensorflow/compiler/xla/service/service.cc:168] XLA service 0x55568a445b60 initialized for platform Host (this does not guarantee that XLA will be used). Devices:
2021-02-23 05:51:13.676130: I tensorflow/compiler/xla/service/service.cc:176]   StreamExecutor device (0): Host, Default Version
42

まとめ

examplesを動かす環境は作成することが出来ました。
examplesにはモデル構築、学習、推論の一通りの例が含まれており、Rustのみでディープラーニングをすることは可能ではあると言えそうでした。

一方で、examplesにはPythonのTF2.0で作成したモデルを使用する例が存在せず、Pythonで学習してRustで推論するという用途に使えるかはもう少し検証が必要かと思います。

また、本記事ではGPUを使用した際の動作は確認していません。
ベースモデルにGPUに対応したイメージを使用することで検証可能だと思うので、そういった情報が増えると良いかと思います。