iOSで超解像 - その3：パフォーマンス改善 #Swift

超解像とは、画像の解像度を高める技術のことで、100x100ピクセルの画像から200x200ピクセルの画像を生成する、というようなものです。

そしてwaifu2xは深層畳み込みニューラルネットワークを用いた超解像システムで、従来の補間処理によるリサイズとはまったく別次元の効果を発揮します。

（左）従来法による拡大、（右）waifu2xによる拡大

（写真はitSnap様より特別に許可をいただき使用させていただきました¹）

「その1」では、waifu2xのMetal実装であるwaifu2x-metalを手元の画像で試し、実際に効果があることを確認しました。（それが上の画像）

waifu2x-metalはmacOSで動くように実装されていますが、MetalなのでiOSでも動作するはずです。

これをiOSでオンデバイスで利用できれば、アプリのUX改善にとってその効果は絶大です。たとえば、サーバーからは非Retinaサイズで取得しておいて、デバイス側でRetinaサイズに拡大して表示、ということが考えられます。画像のクオリティは維持しつつ通信量が4分の1になれば、通信速度というUXにおける大きなボトルネックが改善し、ネットのインフラが十分でない環境でもサクサク画像が表示されるようになる、というような効果が期待されます。

そんなわけで前回のその2では、waifu2x-metalをiOSに移植してみました。

iOSへの移植は無事うまくいき、オンデバイスで動くには動いたのですが、前回ネタバレをした通り、現時点での実装では

リアルタイム処理には程遠い

です。

この回では実際の処理時間の計測結果について書いたのちに、改善の展望を示したいと思います。

処理時間の計測

計測方法

計測は次のようにapply(for:)メソッドの処理時間を計測しており、

let startTime = CACurrentMediaTime()
texture = Waifu2x().apply(for: "image")
let endTime = CACurrentMediaTime()
print("Time to apply: \(endTime - startTime)")

当該メソッドは、

jsonファイルからのニューラルネットワークのパラメータ読み込み
入力画像のロード
waifu2x処理

を含みます²。処理結果テクスチャのMTKViewへの描画は含んでいません。

計測結果

iPhone X/ iOS 11で220 x 220ピクセルの画像を超解像処理させたところ、計測結果は

Time to apply: 11.0397451249883

となりました。220x220ピクセルの画像を最新のiPhoneで超解像処理するのに11秒です。

100倍高速化できたとしてやっと10fpsです。コアなアルゴリズムと関係ないところだけ見ても、jsonデータの読み込み等の明らかな改善ポイントがありますが、これはそういうレベルでの最適化ではリアルタイムには辿り着けそうにありません。

ちなみに処理時間の内訳としては、上記3の"waifu2x処理"が 10.57 secでした。残り 0.5 sec の処理もリアルタイム処理の観点からいうと最適化すべきですが、まずはこの「10秒」を抜本的に高速化できるような改善ポイントがないか、コードを見てみました。

問題点1: 処理を十分に並列化していない

コードを読んでいくうちに、この箇所に気付きました。

let outputs = zip(info.weight, info.bias).map { (a,b) -> MTLTexture in
    return waifu2x(inputs, weight: a, bias: b)
}

CNNの各層において、1チャネルごとに別々に処理しています。これではGPUの並列化の恩恵を十分に受けられていません。

このwaifu2xでは、各層のチャネル数は順番に、

32 → 32 → 64 → 64 → 128 → 128 → 3

となっています。ちゃんとベクタライズしてこれらをすべて並列に処理できれば、パフォーマンスは圧倒的に改善される可能性があると思います。

どうベクタライズするかですが、

Metal Performance Shaders (MPSCNN)を使って実装しなおす
- MPSImageは4チャネル以上のテクスチャを並列処理するためのフォーマットなので、やはりこういうCNNの処理を実装するには一番向いてそう
BNNSを使って実装しなおす
Core MLを使って実装しなおす
- 元のモデルデータをcoremltoolsで変換できるならこれが一番簡単かもしれない
このMetal実装を修正する
- 現時点ではwaifu2xのアルゴリズムを理解しないまま書いてるので、本当にMPSCNNやCore MLで置き換えられるのか確信がない。
- ので、結局これが一番近道かもしれない

というあたりが考えられます。

問題点2: GPUをフルに働かせてない

現実装では、GPUにコマンドをコミットする度に、その完了を待ってから、次の処理を行っています。

commandBuf.commit()
commandBuf.waitUntilCompleted()

トリプルバッファリングでもっとGPUを働かせるようにすれば、処理時間はもっと短くなるはずです。

Metal Best Practices Guide: Triple Buffering

トリプルバッファリングの実装自体は簡単なのですが、問題点1の改善をしないままにこれをやっても、並列化されてないコマンドを次々GPUにプッシュすることになりGPUにとっては徒労なので、まずは問題点1の解決が必要です。

他の改善ポイント

waifu2x-metalの作者の方の記事に、いくつかの改善指針が示されています。

Safx: Metalでwaifu2xを実装してみた

そのうち、

オリジナルwaifu2xや ueshita/waifu2x-converter-glsl では入力画像を128x128ピクセルごとに分割して処理しているようですが、Metal版ではそれを行っていません。

は実用的にしていく段階で必要そうだなと思いました。

その他、Metal 2の新機能である、Argument Buffersや、Non-uniform threadgroup sizeも、もしかしたら効くかもしれません。（前者はこのケースではあまり関係ないか）

次回

改善実装を行いたいと思います。リアルタイムとは行かないまでも、ある程度速くすることに成功したら何らかの形で公開したいと思っております。

Metalの参考書籍

技術書「iOS 11 Programming」の、「第13章 Metal」を執筆しました。

書籍のタイトルにはiOS 11とありますが、Metalについては新機能の紹介だけではなくて、基礎からじっくり解説しています。Metalの章だけで37ページもあり、日本語ではこれだけのまとまったMetalの解説はレアかと思います。

13.1 はじめに
13.2 Metalの概要
13.3 Metalの基礎
13.4 MetalKit
13.5 Metal入門その１ - 画像を描画する
13.6 Metal入門その２ - シェーダを利用する
13.7 Metal入門その３ - シェーダでテクスチャを描画する
13.8 ARKit+Metalその１ - マテリアルをMetalで描画する
13.9 ARKit+Metalその２ - MetalによるARKitのカスタムレンダリング
13.10 Metal 2
13.11 Metalを動作させるためのハードウェア要件

他の章も他著やネットではなかなか得られない濃い情報が詰まっているので、気になった方はぜひサンプルPDFもあるので見てみてください。

第1章　iOS 11 概要
第2章　ARKit
第3章　Core ML
第4章　Swift 4の新機能とアップデート
第5章　Xcode 9 の新機能
第6章　Drag and Drop
第7章　FilesとDocument Based Application
第8章　レイアウト関連の新機能及び変更点
第9章　Core NFC
第10章　PDF Kit
第11章　SiriKit
第12章　HomeKit入門とiOS 11のアップデート
第13章　Metal
第14章　Audio関連アップデート

PEAKSのサイトにて電子書籍・紙の書籍ともに販売されています。

iOS 11 Programming

著者：堤修一,吉田悠一,池田翔,坂田晃一,加藤尋樹,川邉雄介,岸川克己,所友太,永野哲久,加藤寛人,
発行日：2017年11月16日
対応フォーマット：製本版,PDF
PEAKSで購入する

執筆を担当したARKit、Metalの章の詳細、あるいは全体的なおすすめポイントは以下の記事にも書きました。

「iOS 11 Programming」を共著で執筆しました／本書のおすすめポイント - Over&Out その後

dealforestさんありがとうございます！ ↩
iOS移植版のソースコードはまだ公開していませんが、元となるmacOS版のアルゴリズム自体にはまだ手を入れていません。 ↩