プライベートなwaifu2xバッチクラスタを構築する #jawsmeguro

ども、大瀧です。
ミクさんの画像が超美麗に2倍のサイズに！で話題の二次元画像向け超解像技術waifu2x、すごいですねぇ。人工知能による画像処理への応用という分野には全く明るくないのですが、このWebサービスはAmazon EC2で動いていて、しかも仮想マシンイメージ(AMI)が公開されていると！そこで、持ち前のAWSの知識を使ってプライベートかつスケーラブルなwaifu2xバッチクラスタを組んでみたので、ご紹介してみます。(なお、本構成は先週金曜日のJAWS-UG Meguro #0のLTで発表したものです)

waifu2x AMIについて

上記ブログ記事にAMIの概要が示されています。

AMI name: waifu2x server
AMI ID: ami-75f01931
Region: N. California
Instance: g2.2xlarge (require a GPU)
OS: Ubuntu 14.04
User: ubuntu

リージョンがN. Californiaなのが特徴的ですね。
一度N. Californiaリージョンでインスタンス→AMI作成→他のリージョンにコピーと手順を踏めば、東京リージョンに持ち込むことも可能と思われます。g2.2xlargeのスポット価格が安価で安定していることからN. Californiaを選択したのだろうと推測しました。

また、ブログ記事からリンクのあるGitHubページにはwaifu2xを実行するため要件やコマンドラインツールの紹介などがREADMEに記載されています。コードはLuaJITのCUDAライブラリであるTorchで実装されていて、WebアプリのほかCLIも用意されていることがわかります。CUDAのためにg2インスタンスファミリーが必要なわけですね。今回はバッチ処理をユーザーデータ(bash)で書きたかったので、CLIを利用することにしました。

構成

では、構築手順を追っていきましょう。まずは構成図を。

バッチジョブをAmazon SQSのキューに格納し、AutoScalingによるスポットインスタンスでwaifu2xを実行、キューにポーリングする形態です。こうすることで、スポットインスタンスがTerminateされてもAuto Scalingによる再作成が期待できますし、拡大ジョブを複数同時にこなすためにAuto Scalingでインスタンス数の増加(スケールアウト)を検討することもできるようになっています。

各サービスとも、N. California(us-west-1)リージョンに作成しました。

各サービスの準備(SQS、S3、IAM)

まずは、ジョブキューとなるSQSのキューを作成します。特に設定変更せず、キュー名waifu2x_inputで作成しました。

続いて、拡大した画像をアップロードするS3バケットを作成します。こちらも特に設定変更せず、waifu2x-sampleバケットを作成しました。

インスタンスからSQSキュー、S3バケットにアクセスできるようにIAMロールを作成します。今回は緩めに既存のマネージドポリシーでS3とSQSのフルアクセスを付与したwaifu2xロールを作成しました。必要に応じてポリシーを絞ってもよいでしょう。

インスタンス以外の準備はOKです。Auto Scalingの起動設定は、AMI IDにami-75f01931、IAMロールに作成したwaifu2xロール、セキュリティグループは任意でOKです。作成した起動設定sqs-waifu2xのプロパティは以下のようになります。

あと、起動設定には、今回のキモとなるユーザーデータを指定します。

#!/bin/bash

pip install awscli cat < /sqs-waifu2x.bash #!/bin/bash

INTERVAL=30 REGION=us-west-1 QUEUE_URL=https://sqs.us-west-1.amazonaws.com/XXXXXXXXXXXX/waifu2x_input BUCKET_NAME=waifu2x-sample export HOME=/root

cd /home/ubuntu/waifu2x/ if [ ! -d out/ ]; then mkdir out/ fi

while : do # SQSからメッセージを受け取る aws sqs receive-message --queue-url \$QUEUE_URL --region \$REGION --output text > result.txt if [ -s result.txt ]; then SOURCE_URL=\`cut -f 2 < result.txt\` RECEIPT_HANDLE=\`cut -f 5 < result.txt\` # 拡大前の画像データをダウンロード wget \$SOURCE_URL IMG_FILENAME=\`basename \$SOURCE_URL\` # waifu2xを実行、2倍拡大しノイズリダクションを適用 th waifu2x.lua -i \$IMG_FILENAME -o out/\${IMG_FILENAME} # 適用したファイルをS3にアップロード aws s3 cp --region \$REGION out/\${IMG_FILENAME} s3://\${BUCKET_NAME}/\${IMG_FILENAME} # SQSからメッセージを削除 aws sqs delete-message --receipt-handle \$RECEIPT_HANDLE --queue-url \$QUEUE_URL --region \$REGION else sleep \$INTERVAL fi done EOF nohup bash /sqs-waifu2x.bash [/bash]

Auto Scalingグループの設定は任意で構いません。作成した起動設定を紐付け、ひとまず希望/最小/最大1台で動かして見ましょう。これで準備OKです！

動作確認

本構成では、インスタンスからアクセスできる場所に画像ファイルを配置とSQSキューへのメッセージ送信を手作業で行います。このあたりはWebサーバーなどをフロントに立てて、Webアプリで実装しても良いですね。 SQSのメッセージボディのURLに拡大元の画像ファイルを指定、メッセージを送信して少し待つと、インスタンスがメッセージを取得、画像を拡大してS3バケットにアップロードします。

今回はクラスメソッドの非公式マスコットキャラクター、くらめそちゃんの画像を2倍に拡大してみました。左がOS Xのプレビューでそのまま拡大したもの、右がwaifu2xで拡大したものです。(クリックして拡大表示して見ましょう)

今回CLIではデフォルトのオプション指定にしているので、ノイズ除去のレベル(-noise_levelオプション)を上げることでより良い効果も期待できそうです。

まとめ

AWSのサービスを組み合わせて、waifu2xのバッチクラスタ構成を組んでみました。JAWS-UG Meguroの発表後に@suzuryoさんや@ijinさんからSQSへのキューイングをS3イベント通知からキックしたらどう？という指摘があったり、まだまだ改善・工夫できるところがあるんじゃないかと思います。思い思いの画像拡大ライフをエンジョイしましょう！

参考 : LTの発表資料