OlmoEarth 模型是我們一系列的地球觀測基礎模型,已在約 10 TB 的多模態衛星資料上進行預訓練。各國政府、非政府組織及其他以使命為導向的機構,正將 OlmoEarth 應用於森林砍伐監測、糧食安全和野火風險等領域。
在 Ai2,我們深知如何訓練並發布強大的開源模型,對於擁有強大工程團隊的組織而言,一個開源模型就足以讓他們開始運作。然而,大多數環境領域的組織——儘管他們最適合應用這些模型——卻缺乏管理完整生命週期所需的基礎設施或工程團隊,這包括資料標註、模型微調及執行大規模推論。
我們在營運 Skylight 和 EarthRanger 等平台方面投入了十多年,這些軟體每天都受到全球用戶的信賴,因此必須確保其日常運作。這段經驗讓我們了解到,要產生實際影響,需要什麼:在正確的時間和地點以具成本效益的方式運行模型、監控效能、將原始輸出轉化為可操作的洞察,並驗證這些輸出能達成合作夥伴期望的成果。
這就是我們建立 OlmoEarth 平台的原因:它是一個基礎設施,能將地理空間模型從微調和評估階段,推進到大規模推論。
如此規模的推論帶來了一系列獨特的挑戰。衛星影像必須從多個供應商處尋找和存取,並在投影和解析度上進行對齊,然後高效處理。接著,結果必須拼接成地理上一致的地圖,同時基礎設施還要從分散式運算中常見的例行故障中恢復。
如今,該平台能在約一天內對洲際規模的區域執行推論,處理數十 TB 的影像,每平方公里的成本僅為幾分錢。開發它意味著要面對一系列工程挑戰,這些挑戰很可能也會被其他從事大規模地理空間系統的人遇到。本文將深入探討這些挑戰以及我們所採用的解決方案。
大多數機器學習模型接收幾 MB 的資料,並在一秒內產生結果——例如大型語言模型處理一段文字,或電腦視覺模型分析智慧型手機照片。然而,地球觀測推論的運作規模截然不同,單一任務若要微調基礎模型以達到最佳效能,可能需要處理數 TB 的資料並運行數小時。
輸入資料可能涵蓋廣大地理區域內的多個光譜波段、感測器類型和時間步長。它們可能來自多個供應商,每個供應商使用不同的投影和解析度,並且可能包含遺失或被雲層遮蔽的觀測結果。輸出本身就是一張地圖,因此每個預測都必須與周圍區域保持相同的投影和座標網格,精確對齊。
即使是資料獲取本身也是一大挑戰。預測任務通常花費更多時間在下載和準備影像上,而非執行模型本身,因此高效的資料管線至關重要。這些管線必須處理大容量的輸入/輸出,同時提供重投影和重採樣影像所需的運算能力。
由於資料獲取和準備通常佔據推論任務的大部分運行時間,將這項工作分配給 GPU 會讓系統中最昂貴的硬體執行更適合 CPU 的任務。因此,我們將每個任務分為三個階段,每個階段都匹配不同的硬體配置:
資料獲取和預處理(CPU,高 I/O):獲取、重投影、對齊和正規化影像,然後以優化格式寫入,以便在推論期間快速載入。
推論(GPU):執行模型的前向傳播,並將最少處理的輸出直接寫入儲存。
後處理(CPU):將每個視窗的輸出拼接在一起,應用遮罩或重新縮放,並以使用者友善的格式(例如 Zarr、GeoTIFF 或 GeoJSON)匯出。
OlmoEarth 平台將這些階段分佈在多台機器上,同時確保 GPU 得到充分利用。多進程資料載入器持續為每個 GPU 供料,而完成的輸出則直接串流到 blob 儲存。
OlmoEarth Run 是該平台用於大規模推論任務的執行層,它將每個任務所涵蓋的地理區域劃分為適合單個運算實例(工作者)大小的分區,然後再將這些分區細分為 OlmoEarth 模型處理的更小視窗。由於每個視窗都可以在獨立的前向傳播中單獨處理,因此地圖某一部分的工作無需等待另一部分。
實際上,一個州大小的區域可能變成約一百個分區,而洲際規模的運行則可能達到數千個分區。相鄰分區會略微重疊,我們在組裝輸出時會協調這些重疊,以確保最終的柵格影像沒有接縫。
由於分區是獨立的,同一階段可以同時在數千個運算實例上運行。我們最近利用這種方法生成了一張涵蓋整個北美洲的野火風險地圖。在高峰期,這次運行同時使用了約 19,600 個 CPU 和 994 個 GPU,網路吞吐量超過 168 GB/s。這種平行處理水準將估計的 4,737 小時串行運算時間縮短到約 30.5 小時的實際運行時間——實現了 155 倍的加速。
然而,扇出並非無限。更多的工件會觸及雲端配額限制,因此平行處理是每個運行任務的一個可調參數,也是我們可以在個別任務上調整的幾個參數之一。輸出解析度以資料量和運算換取細節;模型大小以 GPU 時間換取準確性;原始影像快取則以儲存空間換取在同一區域重複運行時的速度。正確的設定取決於手邊的任務和預算。
給定一個地理區域和時間範圍,平台首先必須確定哪些衛星場景應該輸入模型。這意味著要識別不同供應商在不同目錄、格式和發布延遲下,何時何地捕獲了什麼。選擇標準也取決於來源:對於 Sentinel-2 等光學影像,我們通常希望選擇雲量最少的可用場景;而對於合成孔徑雷達,可用的極化通道可能更為重要。
我們盡可能依賴公共 STAC 目錄和開放標準。然而,一個大型推論任務可能同時產生數千個中繼資料查詢——這遠遠超過了 ESA 或 Microsoft Planetary Computer 的 STAC API 等外部服務所設計的並發處理能力。
為了避免這些服務過載,OlmoEarth 平台維護自己的中繼資料索引,並在發布新影像時進行更新。對於透過 AWS Open Data 託管的資料集,我們會在每個新場景發布時收到 SNS 通知。當供應商不提供變更串流時,我們會每隔幾分鐘輪詢其上游索引。因此,我們對外部服務的請求遵循新發布的穩定節奏,而非大型推論任務所產生的突然爆發。
每個索引條目都儲存場景中繼資料,以及指向底層像素可用位置的指標。在運行時,平台會選擇最佳來源,並對 COG 或 Zarr 等雲端優化格式執行視窗讀取,僅檢索給定分區所需的位元組,而非下載整個場景。
該索引也支援我們的註釋工具。由於它維護著指向 Sentinel-1、Sentinel-2、Landsat 和 NISAR 影像的雲端優化格式指標,我們可以透過相同的視窗讀取系統,從任何索引場景提供圖塊,而無需建立單獨的攝取管線。
讓此工作流程最簡便的供應商,都具備我們推薦作為地球觀測資料發布最佳實踐的三個特點:當新影像可用時提供基於佇列的通知、在主要雲端平台上儲存且沒有客製化速率限制或可用性瓶頸,以及支援範圍讀取的雲端優化格式。
OlmoEarth 平台旨在自動從故障中恢復。對於每個階段和地理分區內的每個任務,它會動態配置一台運行我們 runner Docker 容器的虛擬機器。該 runner 會檢索任務參數、執行工作、返回結果,然後關閉。由於每個任務都是可重入且冪等的,因此間歇性故障可以透過重新運行來安全處理。
在這種規模下,故障是預料之中的:供應商可能緩慢或暫時不可用;中繼資料可能顯示影像存在,即使所需的波段或視窗遺失;雲層覆蓋可能導致可用觀測資料過少;或者任務可能直接崩潰。平台透過任務追蹤、自動重試、在可用時回退到替代供應商,以及明確區分可重試錯誤和致命錯誤來應對。一個獨立的監控程序會檢測已停滯或停止的 runner 並重新啟動其任務。
我們的發展藍圖是由合作夥伴發現的差距以及他們最需要的功能所塑造。我們正在努力的領域包括:
自動化模型運行:提前安排推論任務,或在影像索引註冊感興趣區域的新場景時觸發它們。
變更偵測和警報:當用戶監測的景觀發生變化時通知他們,這樣森林砍伐或洪水等事件就會以警報的形式呈現,而不是需要手動查找和檢查的柵格影像。
代理工具和介面:代理可以降低使用地理空間模型的門檻,從資料策劃和特徵工程,到識別改進微調模型的方法。我們希望任何技術水平的用戶都能完成以前需要經驗豐富的機器學習研究人員才能做的工作。
更快的模型:與我們的研究團隊合作,開發更高效的架構,以減少每個視窗的 GPU 時間。
更多模態:為 OlmoEarth 模型和支援它們的影像索引添加新的衛星感測器和資料來源。我們目前專注於整合天氣資料(ERA-5)和提供更多環境因素細節的衛星。
嵌入:開發專用的嵌入模型,並在全球範圍內預先計算嵌入。對於許多任務,對這些嵌入執行推論可以取代對原始影像的完整前向傳播,從而顯著加快工作負載並降低成本。微調和直接推論對於具挑戰性任務的最大效能仍然很重要,但嵌入可以開啟更廣泛的高效應用。
隨處運行:OlmoEarth Run 只需要能夠運行 Docker 映像檔的虛擬機器和對 blob 儲存的存取權限。我們目前在 Google Cloud 上運行它,但其架構設計旨在支援多個雲端以及在合作夥伴自己的帳戶和運算環境中部署。
這僅僅是個開始。地理空間基礎模型,特別是將其投入實際運作,仍然是一項新興技術。許多最適合使用這些模型的組織——例如從事保育、糧食安全、災害應變和氣候相關工作的單位——從未有機會接觸到這樣的基礎設施。
