🛠

ぼくがかんがえたさいきょうのWasmビルド環境

2021/08/09に公開

TensorFlow

TL;DR

VS Code 前提になってしまいますが、VS Code Remote Containers 拡張最強です、という話。実際にどういう風に作れば良いかは参考リポジトリ作ったので見て下さい。

はじめに

背景ぼかしのような推論結果を使ってカメラ画像にフィルタを入れる処理を書くとき、今だと選択肢は大きく分けて 2 つあります。

tfjs を使う
Wasm 経由で tflite やその他ライブラリを使う

この記事では後者の Wasm を使う方に注目して、その開発環境について記載します。

先行例

w-okadaさんのリポジトリには複数の Wasm を使った事例が紹介されています。
例を挙げるとTFLite Wasm for Google Meet SegmentationやTFLite Wasm for ESPCNなどが Wasm です。

これらのビルド環境はどうなっているかというと、Dockerfileが置いてあります。package.jsonの scripts で Docker イメージのビルド、スタート、Docker 内に exec して Wasm のビルド、という手順になっています。
OpenCV を Wasm ライブラリとしてビルドした libopencv*.a ファイルなどを Docker イメージに押し込めて管理を楽にすることが出来ています。

辛いところ

しかし、同じような構成で開発をしてみると色々と辛いところが出てきます。

プロジェクト始めるごとにそれぞれ Dockerfile を置いて毎度 OpenCV の Wasm ビルドからしないとならない
docker を落とすと bazel のビルドキャッシュが消えて再ビルド（最近組み立てたつよつよマシンで 10 分）
使ってるライブラリのヘッダファイルが見つけられなくて補完が効かない
ライブラリバージョンアップする度にそれぞれのプロジェクトを更新するのは大変

これらの辛いところを解決して、快適な開発体験を得られるのが、VS Code の Remote Containers 拡張です。

VS Code Remote Containers

詳細は公式の記事などを見てもらうとして、簡単に言うと以下のような特徴があります。

Docker 内で VS Code Server を動かす(Remote SSH や WSL と一緒)
ソースコードはボリュームマウントするから Docker 内からも見えるし VS Code で編集できる
拡張機能も含め、VS Code が Docker 内で動くので中のインクルードヘッダとかビルドツールが使える

実際にやってみる

準備段階として Remote - Containers 拡張を入れるとかの作業は必要…かな？　最近のデフォルトがどうなってるのかわからないです。
実際にやってみた結果は参考リポジトリを作ったので見て下さい。

.devcontainerディレクトリの準備

各リポジトリ（というか VS Code で開くフォルダ）の直下に.devcontainerというディレクトリを作り、その下にdevcontainer.jsonを置きます。
中身はこんな感じになっています。

{
	"name": "C++",
	"build": {
		"dockerfile": "Dockerfile",
	},
	"runArgs": [ "--cap-add=SYS_PTRACE", "--security-opt", "seccomp=unconfined"],

	// Set *default* container specific settings.json values on container create.
	"settings": {},

	// Add the IDs of extensions you want installed when the container is created.
	"extensions": [
		"ms-vscode.cpptools",
		"bazelbuild.vscode-bazel",
		"twxs.cmake",
        "ms-vscode.cmake-tools"
	],

	// Use 'forwardPorts' to make a list of ports inside the container available locally.
	// "forwardPorts": [3000],

	// Use 'postCreateCommand' to run commands after the container is created.
	// persistence $HOME/.cache (bazel cache directory)
	"postCreateCommand": "ln -sf $(pwd)/build_cache ${HOME}/.cache",

	// Comment out connect as root instead. More info: https://aka.ms/vscode-remote/containers/non-root.
	"remoteUser": "vscode"
}

.build.dockerfileでDockerfileを指定しているので、同じディレクトリにDockerfileを置きます。
この Dockerfile で開発に必要なライブラリなどが入った Docker イメージを作ります。

ただし、これだと辛いところの 1.と 4.が解決できないので、もう一ひねりして GitHub Container Registry に共通のイメージを置くことにします。

FROM ghcr.io/kounoike/imgproc-wasm:latest

ghcr にイメージが置いてあるのでそれをベースにしてるだけです。必要ならプロジェクトごとにライブラリ追加などをすると良いでしょう。

ghcr のイメージはこんな感じのDockerfileで作っています。

.devcontainer ディレクトリが準備出来たらコマンドパレットから「Remote Containers: Reopen in Container」を選択します。その設定に従って自動的にコンテナがビルドされ、中に VS Code をインストールして接続してくれます。

VS Code の左下に Remote の接続先が表示されているはずです。私は普段 Windows+WSL で使っているのでこのように WSL と表示されています。

コンテナ内の VS Code に接続するとこのように Dev Container となります。C++の部分は多分 devcontainer.json のnameで指定したのが使われているのでしょう。

ソースコード（開いたディレクトリの中身）

Docker 内の/workspaces/<folder name> にマウントされていて、そのディレクトリが Docker 内の VS Code で開かれています。

VS Code C/C++拡張の設定

このままではまだ辛いところの 3.のヘッダの解決が出来ていません。例えば下図は Emscripten のヘッダを見つけられず波線が出ています。

Quick Fix から includePath の編集を選びます。「パスを含める」（訳がちょっと…）に/emsdk/upstream/emscripten/system/includeを追加します。

波線が消えて emscripten.h 由来のマクロなども指摘されなくなりました。

他のライブラリも標準インクルードパスにない場合は追加していきます。

ちゃんと追加出来ると C/C++拡張の補完などが効くようになります。例えば OpenCV の resize 関数にマウスオーバーするとこんな感じです。

bazelのキャッシュ再利用

tensorflow-lite をビルドしようとすると、CMake と bazel の 2 つの手段があります（make はさすがに更新されてなさすぎらしい）。最終的なパフォーマンスやビルドまでの楽さで bazel を使うことが多いです。

bazel はビルド結果を$HOME/.cache/bazel以下に持っていて、その下の適切なディレクトリへのシンボリックリンクをソースのあるフォルダに生成します。
$HOME/.cache ということは docker イメージの中にしかないので、停止すると消えてしまいます。消えてしまうと再ビルド（10 分以上）なので出来るだけ残しておきたいです。また、.cache ディレクトリは bazel の他にも色々なツールがキャッシュディレクトリとして使うのでホストのディレクトリにしておいた方が色々と便利です。

やり方はとてもシンプルで、上で紹介した devcontainer.json の中で半分設定出来ています。postCreateCommandに ln -sf $(pwd)/build_cache ${HOME}/.cacheと書きます。この設定で Docker イメージを作る度に$HOME/.cache をシンボリックリンクとして、カレントディレクトリの build_cache ディレクトリを指すようにしています。カレントディレクトリは VS Code が開いているディレクトリそのものですので、その直下に build_cache ディレクトリを作り、中身が*だけの.gitignoreを置いておくと良いです。

これからやりたいこと

しっかりと環境が出来てきたので、もっと改善していきたくなってしまいます。今までは Wasm でビルドしてはブラウザで開いて確認、という流れで開発していました。いっそネイティブでビルドして main 関数から実行したり、google test などでテストコードをしっかり書いて行くのも良さそうです。

あと、コンテナイメージの中には入れてあるのですがHalideとの連携もやってみたいです。

また、まだパフォーマンスが劣化していて悩ましいものの tensorflow-lite の CMake でのスタティックライブラリのビルドにも成功しました。これを活用する方法とかも考えていきたいですし、パフォーマンスを bazel 版と同等にする方法も調べていきたいです。

今のところ tflite ファイルを tensorflow-lite で処理しているだけです。しかし、onnnxruntime や pytorch など別系統のライブラリも使えるようになると幅が広がりそうです。