pix2pixのために高速な線画変換環境を構築してみた

pix2pixのためにHED(Holistically-Nested Edge Detection)による高速な線画変換環境を構築してみました。

pix2pixは、比較的少ない学習データでも線画→画像変換モデルが構築できるようなのですが、このためには、画像を線画に変換するツールが必要になります。

オリジナルのTorch版pix2pix(https://github.com/phillipi/pix2pix)では、線画変換のために、有料のMATLABが必要なので諦めていたのですが、TensorFlow版pix2pix(https://github.com/affinelayer/pix2pix-tensorflow)では、MATLABの代わりにOctaveを使用することで、無料ツールだけで線画変換を行うことができます。

ただし、線画変換にかなり時間がかかっています。

ちなみに、線画変換前(左側)と変換後(右側)は、以下のようになります。

スポンサーリンク

HED

HED(https://github.com/s9xie/hed)は、改造Caffe上で動作する線画変換のニューラルネットワークです。ただし、本体部分は、2年くらい前から更新されていないため、CUDA 8.0 + cuDNN 5の環境では、ビルドに失敗します。

Docker版pix2pix-tensorflowでは、HED(改造Caffe)がCPUモードでビルドされているため、線画変換に時間がかかるようです。

Docker版pix2pix-tensorflowで、256×256にリサイズした画像10枚の線画化を試してみると、以下のように約3分(約0.05画像/秒、または、約20秒/画像)掛かっています。

# python tools/process.py \
 --input_dir hed/resized \
 --output_dir hed/edged-1 \
 --operation edges
...
1/10 complete  0.04 images/sec  0m27s elapsed  4m8s remaining
2/10 complete  0.04 images/sec  0m45s elapsed  3m0s remaining
3/10 complete  0.05 images/sec  1m2s elapsed  2m26s remaining
...
8/10 complete  0.05 images/sec  2m30s elapsed  0m37s remaining
9/10 complete  0.05 images/sec  2m48s elapsed  0m18s remaining
10/10 complete  0.05 images/sec  3m5s elapsed  0m0s remaining

なお、線画化は以下のように--workersオプションによるマルチスレッド化に対応していますが、その高速化(--workers 4で約0.20画像/秒、--workers 8で約0.40画像/秒)には限界があります。

# python tools/process.py \
 --input_dir hed/resized \
 --output_dir hed/edged-2 \
 --operation edges \
 --workers 4
...
1/10 complete  0.05 images/sec  0m22s elapsed  3m19s remaining
2/10 complete  0.09 images/sec  0m22s elapsed  1m29s remaining
3/10 complete  0.13 images/sec  0m22s elapsed  0m52s remaining
4/10 complete  0.18 images/sec  0m22s elapsed  0m33s remaining
5/10 complete  0.13 images/sec  0m39s elapsed  0m39s remaining
6/10 complete  0.15 images/sec  0m40s elapsed  0m26s remaining
7/10 complete  0.17 images/sec  0m40s elapsed  0m17s remaining
8/10 complete  0.20 images/sec  0m40s elapsed  0m10s remaining
9/10 complete  0.16 images/sec  0m57s elapsed  0m6s remaining
10/10 complete  0.17 images/sec  0m58s elapsed  0m0s remaining

高速版HED

下記の方法で構築した高速版HED(改造Dockerと改造Docker用の線画化のコードの組み合わせ)では、256×256にリサイズした画像10枚の線画化が、以下のように約6秒(約1.50画像/秒)なります。

# python tools/process_edges.py \
 --input_dir hed/resized \
 --output_dir hed/edged-3
skipping 0 files that already exist
processing 10 files
1/10 complete  0.55 images/sec  0m1s elapsed  0m16s remaining
2/10 complete  0.85 images/sec  0m2s elapsed  0m9s remaining
3/10 complete  1.03 images/sec  0m2s elapsed  0m6s remaining
4/10 complete  1.16 images/sec  0m3s elapsed  0m5s remaining
5/10 complete  1.25 images/sec  0m4s elapsed  0m4s remaining
6/10 complete  1.32 images/sec  0m4s elapsed  0m3s remaining
7/10 complete  1.37 images/sec  0m5s elapsed  0m2s remaining
8/10 complete  1.41 images/sec  0m5s elapsed  0m1s remaining
9/10 complete  1.45 images/sec  0m6s elapsed  0m0s remaining
10/10 complete  1.48 images/sec  0m6s elapsed  0m0s remaining
WARNING: Logging before InitGoogleLogging() is written to STDERR
F0725 12:36:26.956246   171 syncedmem.cpp:19] Check failed: error == cudaSuccess (29 vs. 0)  driver shutting down
*** Check failure stack trace: ***
Aborted (core dumped)

なお、処理が終わった後で、Check failed: ... driver shutting downと出力されていますが、Torch版pix2pix(のbatch_hed.py)でも、時々発生しているようなので、ここでは気にしないことにします。

高速版HED環境の構築方法

HEDのIssuesを調べてみると、CUDA 7.0 + cuDNN 3であれば、GPUモードでビルドできるようなのですが、Docker版pix2pix-tensorflowのCUDA 8.0 + cuDNN 5環境での動作を目標としました。

「CUDAあり・cuDNNなしならビルドできるかも」という直感を信じました。

以下のコマンドで、pix2pix-tensorflowのDockerイメージをpullし、pix2pix-tensorflowをcloneします。

$ docker pull affinelayer/pix2pix-tensorflow
$ cd ~/github
$ git clone https://github.com/affinelayer/pix2pix-tensorflow
$ cd pix2pix-tensorflow

次に、適当なディレクトリ(ここでは、modified_docker)を作成し、このディレクトリにGithub(https://github.com/SoraLab/pix2pix-tensorflow/tree/develop/modified_docker)からダウンロードした改造Docker用のDockerfileをコピーします。

下に示すように、改造Docker用のDockerfileでは、pix2pix-tensorflowのDockerイメージの作成時に行われたloss_layers.hppの変更を元に戻し、CPUモードの代わりにGPUモード(CUDAあり・cuDNNなし)でHED(改造Caffe)を再ビルドしています。

FROM affinelayer/pix2pix-tensorflow

RUN cd $CAFFE_ROOT && \
    sed -i "647s|vir|//vir|" include/caffe/loss_layers.hpp && \
    sed -i "648s|    con|//    con|" include/caffe/loss_layers.hpp && \
    cd build && \
    cmake -DCPU_ONLY=0 -DUSE_CUDNN=0 .. && \
    make -j"$(nproc)"

以下のコマンドで、改造Dockerイメージ(pix2pix-tensorflow-mod)を作成します。

$ cd modified_docker
$ docker build -t pix2pix-tensorflow-mod .

また、toolsディレクトリに、Github(https://raw.githubusercontent.com/SoraLab/pix2pix-tensorflow/develop/tools/process_edges.py)からダウンロードした改造Docker用の線画化のコードをコピーします。

以下のコマンドで、pix2pix-tensorflowディレクトリに戻って、改造Dockerイメージ(pix2pix-tensorflow-mod)からDockerコンテナを起動します。

$ cd ~/github/pix2pix-tensorflow
$ nvidia-docker run -it \
 -p 8888:8888 -p 6006:6006 \
 --volume $PWD:/prj \
 --workdir /prj \
 --env PYTHONUNBUFFERED=x \
 pix2pix-tensorflow-mod

ここで起動したDockerコンテナは、上記の高速版HEDを利用することができます。また、従来のtools/process.pyは変更していないので、従来通りの利用も可能です。

まとめ

pix2pix-tensorflowのDockerイメージを改造してHED(Holistically-Nested Edge Detection)による高速な線画変換環境を構築してみました。