Solafune の衛星画像の超解像コンペで 1 位を獲得しました

2023/4/21 (Fri)

コンペの概要
参加のきっかけ
Team N のベースライン
Team N の解法
素早く実験を回すためのインフラ
終わりに

日経イノベーション・ラボの青田、山田です。

先日、弊社のデータサイエンティスト 2 名と記者 1 名で「Team N」として Solafune で開催された衛星画像の超解像コンペに参加し、 1 位を獲得しました。コンペへの参加者数は 94 人でした。 Solafune は主に衛星データ解析コンテストを開催しているプラットフォームで、過去にも夜間光データから土地価格を予測するコンペや、市街地の高解像度画像を活用した超解像画像の生成をするコンペの開催実績があります。

「Team N」のメンバー

青田雅輝 (日経イノベーション・ラボ)
山田健太 (日経イノベーション・ラボ)
並木亮 (編集データビジュアルセンター)

本記事ではコンペに取り組んだきっかけから 1 位解法まで簡単に記載します。なお、2023 年 4 月 21 日時点では運営による最終審査中です。具体的なコンテンツは以下のとおりです。

コンペの概要
コンペに取り組んだきっかけ
問題の基本的な解き方 (Team N のベースライン)
更に工夫した点 (Team N の解法)
素早く実験を回すためのインフラ

対象とする読者は

超解像の技術に興味のある読者
コンペでのアプローチの一例を知りたい読者

です。ただ、詳細なコードの解説はしないので、理解を深めるには以下の知識があるとより理解しやすいと思います。

Python に関する基本的な知識
PyTorch や Transformers に関する実装の知識
深層学習や画像処理でよく使われる語彙

コンペの概要

今回私達が挑戦したのは Solafune の主催する「衛星画像の5倍超解像度化 (for OSS)」というコンペです。

コンペページから引用しますが、超解像とは以下の例に示すように低解像度な画像から高解像度な画像に変換するタスクのことです。

本コンペでは参加者に人工衛星から撮影された画像をできるだけ高精度に 5 倍の大きさに超解像を行うモデルを作成するタスクが課されました。高解像度な衛星画像は入手が難しく、価格が高い傾向という課題点から本コンペは開催されました。比較的安価に利用可能な低解像度な衛星画像を高解像度に変換できれば、経済的にもとても嬉しいわけです。

また評価指標は SSIM という、画像の構造がどれだけ似ているかを表す指標です。 SSIM は 0 から 1 までの値を取り、1 に近いほど似ているということを表します。参加者のモデルの出力と正解の画像の SSIM を計算し、その平均値によって順位が決まります。

参加のきっかけ

本コンペに参加した理由は業務への応用可能性です。衛星画像は機械学習との親和性が高いとされ、機械学習分野における研究対象の一つとされています。satellite-image-deep-learning で例示されているように、著名な情報科学分野の国際学会でも衛星画像をテーマにした発表がなされています。加えて、衛星画像は報道上でも高い価値を持ちます。昨今の調査報道やデータジャーナリズムへの強い関心が背景にあります。日経で衛星画像を活用した記事の例はこちらです。

衛星画像について扱ってみたいという気持ちがチームメンバーにはありました。そう思っていたときに本コンペの存在に気づき、興味のある人同士でチームを組みました。参加して衛星画像の扱い方に慣れることで超解像のみならず他のタスクへの造詣を深めることも目的の一つでした。

Team N のベースライン

超解像のモデルを学習する手順を下記に示します。

我々はバックボーンとして事前学習済みの Swin2SR のモデルを用いました。このモデルは画像を 4 倍に超解像するモデルです。そのため、モデルの出力部分で Bicubic 補間で 5 倍まで画像を拡大しました。

def forward(self, x: Tensor):
    f = self.backbone(x)["reconstruction"]
    f = f[:, :, 0:520, 0:520]
    f = TF.resize(f, size=650, interpolation=T.InterpolationMode.BICUBIC) # ここで画像を拡大
    f = torch.clamp(f, min=0.0, max=1.0)
    f = f * 255
    return f

その後正解データである高解像度画像と 5 倍推論画像の SSIM を計算し、1-SSIM を損失関数として追加学習を行いました。なお SSIM の計算には torchmetrics.StructuralSimilarityIndexMeasure を採用しました。損失関数の定義は以下のようになります。

from torchmetrics import StructuralSimilarityIndexMeasure as SSIM
class SSIMLoss(nn.Module):
    def __init__(self, kernel_size: int = 7, sigma: float = 1.5):
        super().__init__()
        self.kernel_size = kernel_size
        self.sigma = sigma
        self.ssim = SSIM(kernel_size=self.kernel_size, sigma=self.sigma)

    def forward(self, x: Tensor, y: Tensor) -> Tensor:
        return 1 - self.ssim(x, y)

また学習時、Epoch 数は固定で最後に学習率が 0 になるように Cosine Annealing で学習率をスケジュールしました。 Early Stopping 等は特に行いませんでした。

以上が基本的なアプローチです。

Team N の解法

上記の基本的なアプローチ (ベースラインと呼称) に加えて、我々のチームが行ったことを解法概要図にまとめました。ベースラインに更なる工夫を加えることで CV を伸ばし、様々なモデルをアンサンブルすることで最終提出としました。以下ではそれぞれのテクニックについて、解説していきます。

モデル	CV スコア	Public スコア	Private スコア
シングル	0.77996	0.77761	-
アンサンブル	0.78120	0.77790	0.78345

Lion Optimizer

Optimizer として、コンペ期間中にちょうど発表された Lion Optimizer を用いました。論文中に

Based on our experience, a suitable learning rate for Lion is typically 3-10x smaller than that for AdamW. Since the effective weight decay is lr * λ, the value of decoupled weight decay λ used for Lion is 3-10x larger than that for AdamW in order to maintain a similar strength.

とあるので AdamW で用いていた学習率の 1/3 に設定しています。結果 AdamW を用いていたときと比較し CV +0.0002 の効果がありました。

正直 AdamW でもパラメータ次第で同様の効果はあったのかもしれませんが、新しい Optimizer を試してみようということで Lion で実装を進めました。

×4推論×2学習

5 倍解像度にするために 2 つの超解像モデルを組み合わせたものです。スコアが伸びただけでなく、アンサンブル時の多様性にも貢献しました。なお、データ拡張などで学習した強い 4 倍シングルモデルには効果がない傾向でした。

処理の概要を図にすると以下になります。

データ拡張

データ拡張では、超解像分野のデータ拡張手法として有効であるとされている cutblur の論文に記載された手法を、比較手法も含め概ね試しました。結果データ拡張としてシンプルな幾何変換 (flip+transpose+rotate90) と cutmix の組み合わせが有効でした。 cutmix については、使用していたデータ拡張ライブラリに実装がなかったため、次のように実装しました。

def cutmix(
    LR: torch.Tensor,
    HR: torch.Tensor,
    refLR: torch.Tensor,
    refHR: torch.Tensor,
    p: float = 1.0,
    alpha: float = 0.7,
) -> tuple[torch.Tensor, torch.Tensor]:

    if torch.rand(1) >= p:
        return LR, HR

    scale = HR.size(1) // LR.size(1)
    v = np.random.normal(alpha, 0.01)
    h, w = LR.size()[1:]
    ch, cw = np.int(h * v), np.int(w * v)

    fcy, fcx = np.random.randint(0, h - ch + 1), np.random.randint(0, w - cw + 1)
    tcy, tcx = np.random.randint(0, h - ch + 1), np.random.randint(0, w - cw + 1)

    LR[:, tcy : tcy + ch, tcx : tcx + cw] = refLR[:, fcy : fcy + ch, fcx : fcx + cw]
    HR[:, tcy * scale : (tcy + ch) * scale, tcx * scale : (tcx + cw) * scale] = refHR[
        :, fcy * scale : (fcy + ch) * scale, fcx * scale : (fcx + cw) * scale
    ]

    return LR, HR

AWP

Adversarial Weight Perturbation と言われる学習手法です。 Kaggle では自然言語処理コンペでよく使用されるのですが、画像コンペでもスコアを伸ばすことができました。

詳しくは青田の解説スライドを御覧ください。

外部データ

過去コンペでも言及されていた練馬区の航空写真データセットを利用しました。

航空写真：練馬区公式ホームページ

また、外部データセット作成にあたり、運営より与えられたデータセットの低解像度画像は高解像度画像を縮小したものであるという仮説のもと、縮小方法を検討しました。結果外部データセットの低解像度画像作成に当たって以下の処理を施しました。

low_chunk = cv2.resize(high_chunk, (130,130), fx=0.5, fy=0.5, interpolation=cv2.INTER_AREA)

モデル学習 (exp103)

我々のチームでシングルモデルで最高の CV である exp103 の学習概要です。

学習は 40 epoch 行いました。学習データについては 4 fold で切り、画像に対して様々な変換を施すことによるデータの水増しや、外部データを用いデータセットを拡張しました。強力な水増し手法や外部データなどを学習が進むにつれ取り除き、最後には本来の学習データのみで学習するようにしています。このようにすることでデータ拡張を行わなかった場合や終始データ拡張を行った場合よりもスコアが改善しました。以下に epoch ごとの学習戦略をまとめます。

~14 epoch: 学習データ+外部データ+flip+transpose+rotate90+cutmix
~20 epoch: 学習データ+flip+transpose+rotate90+cutmix
~26 epoch: 学習データ+flip+transpose+rotate90
~33 epoch: 学習データ
~40 epoch: 学習データ+AWP

素早く実験を回すためのインフラ

実験の入力ディレクトリ構成、fold 分割、モデル/Out-of-fold (oof) の保存、学習の進捗管理、等々をチーム内で共有しておくと、実験を効率的に進めることが出来ます。他人の実験の実装を自分の実験に組み込みやすくなったり、過去の実験を選択するだけでアンサンブルが出来たりと楽に色々試すことが出来ます。そのようなインフラを GCP 上で作成したので、一例としてご紹介します。

expXXX というのは実験番号です。 1 実験 1 script という戦略でコンペに取り組んでいます。

大きな図ですが、左上の学習パートを起点に情報が処理されていきます。学習パートでは、ちゃんとログを吐き出します。これにより実装のミスなどで学習の挙動が変になっているときにすぐに気づくことが出来ます。また学習時に出力される生成物 (モデル、oof の推論結果、test の推論結果、学習ログ) 等をすぐさま GCS (モデル保存パート 1) に保存しています。これにより、GCS を oof を読み込むだけでアンサンブル処理につなぐことが出来たり、GCS からダウンロードするだけでコンペへサブミットしたりすることができます。

チームメンバーがこの様なインフラの上で作業したおかげで、アイデアのマージやアンサンブルの取り組みが非常にスムーズに進んだと感じました。

終わりに

今回のコンペでは超解像分野の知見を獲得できたほか、1 位を取ることができて嬉しかったです。これに驕ることなく引き続き画像分野についても学び続けていきたいです。また、今回は超解像に特化した画像処理タスクを解きましたが、別ドメインでも応用可能な知見を蓄えることができたので、学んだ内容を今後業務で活かすのが楽しみです。

日本経済新聞社には、これらの知見を存分に活かせるだけの整ったデータ基盤が存在します。そして私たちは、記事データなどの大規模なテキスト情報を処理し、ユーザーに新たな価値を提供するため、データサイエンティストの仲間を募集しております。少しでもご興味のある方は、ぜひ採用情報ページやインターン情報ページをチェックしてみてください。

本ブログに登場した Solafune に関する情報について、当コンペの参加以外を目的とした利用及び商用利用は禁止されています。商用利用・その他当コンペ以外で利用したい場合は Solafune までお問い合わせください。

DATA SCIENTIST青田雅輝

DATA SCIENTIST山田健太

Blog