本記事について
こんにちは。私はプリンストン大学で、宇宙論の研究をしています、山本と申します。私は現在研究の傍ら東京大学鈴村研究室で、データサイエンティストとして、AI（人工知能）を用いた気象予測・気候変動予測のリサーチをお手伝いしています。
前回まで、ClimaXを実際に使用して全球気象予測について紹介しましたが、本シリーズの最終回として、日本の気象予測をClimaXで行うのに必要なことについて考察していきたいと思います。本記事を読んで、実際にClimaXを日本の気象予測に使ってみたいと思う人がいた時に、このシリーズを読んで始めてもらえるきっかけになればと思います。
格子粒度について
ClimaXモデルは、地球の緯度180度・経度360度それぞれを1.40625度（1辺約150km）に分割したものと、5.625度（1辺約600km）に分割したものがあります。1.40625度の粒度では、地球は128 X 256の格子で表され、5.625度の粒度では、地球は32 X 64の格子で表されます。トレーニングにかかる想定時間と必要とするメモリから、ClimaXでは1.40625度以下の粒度の例えば0.25度（720 X 1440）ではトレーニングは行われませんでした。
日本を緯度20度〜45度、経度122度〜153度の範囲だと仮定すると、1.40625度では格子間隔がかなり粗く、有意義な予測ができないと考えます（以下図を参照）。ですので、pretrainingやfinetuningを日本用にやり直す際は、0.25度（1辺約28km）のデータセットを用いるのが良いと考えます。0.25度の粒度は100 X 124の格子で表されますから、PretrainingやFinetuningにかかる時間は1.40625度のケースより短時間で行うことができるはずです。それを考慮した上で、日本の気象をClimaXで予測するために以下のケースについて考察していきます。
ケース１: Pretrainされたモデルを使用し、地域気象予測を行うケース
本シリーズでは取り扱いませんでしたが、ClimaXのタスクの一つとして、地域限定的な気象予測をするというタスクがあります。これは、pretrainされたモデルは全球ですが、気象予測したい地域の気象データをfinetuning時に使うことで、地域限定的な気象予測ができるというものです。論文では、アメリカ大陸を対象として実験が行われましたが、同じアイデアを利用して、日本の気象予測も同様にできると考えます。しかし、ClimaXがパブリックにしているpretrainされたモデルは1.40625度が最低ですので、そちらを利用し、finetuningは0.25度で行い、気象予測するのに必要な手順を以下で考えます。
想定される手順
- 1.40625度のpretrainedモデルをダウンロードする。
- ERA5のデータをダウンロードする。(* ClimaXもERA5のデータを使用して、気象予測用のFinetuningをしている。 ** ERA5のデータは粒度が元々0.25度なので、粒度は変えなくて良い。)
- ERA5のデータは全球なので、日本用にデータを切り取る。
- Finetuningを行う。（* Pretraining時とfinetuning時の粒度は異なっても構わない。参考：https://github.com/microsoft/ClimaX/issues/20）
- 予測を行う。
Finetuningにかかる想定時間
全球の1.40625度での気象予測のfinetuningはNVIDIA V100マシンを8 GPUsで約1日かかったので、0.25度での日本のエリアのfinetuningには、同じマシンを使用すると、1日以下でfinetuneされることが予想されます。H100マシンを4 GPUs使用した場合は、V100がH100の性能の約1/9ということを考慮すると、約6時間以下でfinetuningできることが予想されます。
ケース２: 日本の気象・気候用にモデルをpretrainingし直し、全域気象予測を行うケース
ケース１では、時間的制限から、0.25度の粒度のデータを使用して、ClimaXを全球でpretrainingし直さないケースを扱いました。しかし、pretrainingするエリアを日本だけに絞れば、pretrainingし直すことも可能であると考えます。この際どのデータセットを使用するかは、考察の余地があると思います。
ClimaXのpretrainingに使用されたCMIP6データセットは、格子間隔が100kmのデータ（約0.9度）が最も解像度が高いものになるので、CMIP6を使うのは難しいでしょう。一方で、ERA5のデータセットは元々の格子粒度が0.25度なので、こちらをpretrainingに使用することはできると考えます。この際に、考慮すべきなのは気象モデルの多様性です。様々なモデルを使用したシミュレーションデータセットを使用することでモデルの柔軟性が上がり、精度は良くなるので、日本国内の研究グループで様々な気象モデルを用いたデータセットがあればそちらを利用することが考えられるでしょう。
一方で、finetuningに関しては、使用するデータセットは行いたいタスクによると考えて良いでしょう。日本全域の気象予測であれば、0.25度の解像度でpretrainingとは別のデータセットを使用して、finetuningができますし、地域的な気象予測・異常気象予測はより解像度の高いデータセットを用いて、finetuningすることができると考えます。
想定される手順
- Pretraining用のデータを日本用に切り取る。
- こちらに従い、ClimaXモデルをpretrainingする。
- 行うタスクに合ったデータセットを準備し、finetuningを行う。
- 予測を行う。
Pretraining & Finetuningにかかる想定時間
全球の1.40625度でのpretrainingはNVIDIA V100マシンを80 GPUsで約5~7日かかったので、0.25度での日本のエリアのpretrainingには、同じマシンを使用すると、それ以下でpretrainingされることが予想されます。H100マシンを4 GPUs使用した場合は、V100がH100の性能の約1/9ということを考慮すると、約11~15日以下でpretrainingできることが予想されます。Finetuningにかかる時間は、ケース１と同様と考えられます。
Discussion