はじめに
M:CPPで機械学習環境の調査と整備を担当している坂本です。
M:CPPとは、AMDのGPUを搭載したサーバをクラウドとして提供するプラットフォームです。これまで基本的な使い方をブログで発信してきましたが、現実の機械学習タスクにどれくらい使用できるのか、という観点ではご紹介できていませんでした。
そこで、Solafuneというプラットフォームで行われた「市街地画像の超解像化」コンペティションに出場してみました。
本記事では、この「市街地画像の超解像化」コンペティションで5位に入賞した解法を紹介します。
コンペの概要
今回私が参加したのは衛星画像を超解像化するタスクの精度を競うコンペです。超解像とはその名の通り、画像の解像度を高くするように変換する技術です。
このコンペでは、375×300ピクセルの画像を4倍の1500×1200ピクセルに拡大したときにどれくらい元の画像を復元できているかを争いました。
イメージを持っていただくために、以下の画像をご覧ください。この画像はコンペで使用されたものではなく、裾野市のオープンデータを使用しています。
一番左の画像が正解画像です。真ん中の画像が、公開されているモデルを用いて超解像を行った画像です。一番右の画像は、単純にリサイズした画像です。このように、現在の技術を用いると単純にリサイズするより遥かに高画質の画像を復元できます。
画像引用元：令和元年度 航空写真オルソ画像データ（令和2年1⽉1⽇時点） by 裾野市 under CC-BY 4.0 / Adapted
しかし、公開されているモデルは衛星画像⽤に特化していないため、例えば家の屋根、街路の模様、⽩い⾞が2台横並びになっているところなどをうまく復元できていません。このような箇所をどれだけうまく復元できるようなモデルを作成するか、が問われるコンペでした。
解法
概要
SwinIR（https://arxiv.org/pdf/2108.10257.pdf）という、Swin Transformerを用いた超解像モデルをFine Tuningしました。
SwinIRは現在の超解像タスクのState of the Artであり、なおかつ実装が公開されている（https://github.com/cszn/KAIR）ため取り組むのに最適でした。
すでに公開されている001_classicalSR_DF2K_s64w8_SwinIR-M_x4.pthを初期値として、3モデルを学習させた出力を平均したものを提出しました。
公式の学習データ60枚中54枚（90%）と、外部データから作成した学習データ477枚、合計533枚を学習に使用しています。
低解像度の画像を64×64にCropし、それを256×256に超解像するような学習を実施しました。
2300Epochまでは上記の533枚で学習、そこから3000epochまでは公式の学習データのみで学習させています。
工夫した点
モデルを作成するにあたって工夫した点を紹介します。上から順に効果があったと感じています。
外部データの使用
今回は公式に提供された画像が60枚と少なかったため、外部データの活用が鍵になる戦いでした。
外部データを使用しなかった場合、5位の私のモデルは10位相当まで下がってしまいます。
今回は練馬区が公開している航空写真のデータセット（https://www.city.nerima.tokyo.jp/kusei/tokei/opendata/opendatasite/tokei_kusei/koku-shashin.html）を用いました。
衛星画像で使用可能なライセンスのちょうどよいデータセットが見つけられなかったこと、見た目が衛星写真と航空写真であまり変わらなかったことからこのデータを使用しました。
外部データの追加により、横断歩道のような、衛星写真には頻繁に含まれるが模様が細かく復元の難易度が高いものを上手に超解像できるようになったようです。
Data Augmentation
以下のData Augmentationを実施しました。
- Filp4通り x 縦横軸入れ替え2通り = 8通りの画像変換
- MixUp (alpha = 0.2)
MixUpにたどり着いたのはCutblurの解説記事（https://ai-scholar.tech/articles/super-resolution/Data-Augmentation）を読み、実装コストに対して精度への寄与が大きそうだなと判断したためです。本当はCutblurのほうが高精度だったかもしれません。
Learning Rate Schedulingなどの調整
Learning Rate、MixUpの有無、外部データの使用の有無を学習が進むにつれて変化させました。
具体的には、学習の終盤に行くにつれてLeraning Rateを下げ、分布がテストデータに近くなるようにMixUpを外したり外部データを外したりしています。
- ~2000epoch 1e-4 with MixUp
- ~2100epoch 1e-5 with MixUp
- ~2200epoch 1e-5 without MixUp
- ~2300epoch 1e-6 without MixUp
- ~3000epoch 1e-5 without MixUp, 公式の学習データのみ
SSIM Lossのカスタム
今回はSSIMが評価指標でした。skimageのデフォルト値がコンペの評価指標でしたが、一般的に超解像タスクで使用されるSSIM Lossはskimageのデフォルト値とは少し計算が異なるようでした。具体的にはSSIMを計算する前にBlurをかける処理に違いがありそうです。
そこで、skimageのSSIMを直接最適化できるようにSSIM Lossをカスタムしました。
うまく行かなかったこと
- SwinIRの**001_classicalSR_DF2K_s64w8_SwinIR-M_x4.pth**以外のモデルも試しましたが、これより精度が高いものを発見できませんでした。もう少しパラメータが多いモデルもあったのですが、現実的な時間内で良いモデルが作成できませんでした。これは事前学習のデータセットの問題なのかもしれなません。
- 色調を変化させるようなData Augmentationはうまく働きませんでした。
- 外部データに手を出す前はCross Pseudo Labelingを試していましたが、外部データを増やしたほうが勝てる見込みが高そうだったので断念してしまいました。もう少しやり込めた部分かと思います。
結果
3ヶ月間のコンペの結果、参加者200名中5位に入賞することができました。
人生初入賞なのでとても嬉しいです！
使用したインフラについて
今回のコンペにはもちろんM:CPPを使用しました。繰り返しになってしまいますが、M:CPPとはAMDのGPUを搭載したサーバをクラウドとして提供するプラットフォームです。
以下の構成のサーバを通常1台、終盤でもう1台追加して2台使用しました。
- OS: Ubuntu 18.04
- GPU: AMD Vega56 x 1 (Memory 8GB)
- Driver: ROCm4.2
NVIDIAのGPUで学習をさせるときとAMDのGPUで学習をさせるときでPythonコードを書き換える必要がないため、公開されている実装をそのまま流用できるのが非常にありがたかったです。
初期の環境構築をDockerで済ませてしまえば、あとは動かしているGPUを気にせずにモデル開発ができるのは良いと感じました。
特別なコーディングや煩雑な環境構築を必要とせずに2021年のState of the Artなモデルを長時間学習させることができたので、M:CPPが現実の機械学習タスクに適用できることを今回のコンペを通じて確認できました。
まとめ
本記事では、M:CPPを用いて「市街地画像の超解像化」コンペで5位に入賞した解法をご紹介しました。
以下のような学びが共有できておりましたら幸いです。
- SwinIRは精度が高いモデルである
- 衛星画像を扱うタスクでは、航空写真を外部データとして使えることがある
- 超解像でもMixUpは効くことがある
- 評価指標を直接最適化することは大事なこともある
機械学習ビジネスに携わっている方でM:CPPにご興味がある方は、是非以下のページをご覧ください。
サポートしているライブラリについての記述や、AMDのGPUで機械学習をする際お力になれるかもしれないブログへのリンクが掲載されています。
また、コストや利用方法についてご興味がある方は、こちらの「FAQ」をご覧ください。
コストパフォーマンスに優れたサーバをご提供できると考えておりますので、なにかご不明な点がございましたら、お気軽にご連絡いただけますと幸いです。