:從GraphCast到FastAPI完整指南)
簡介面向氣象AI研究與開發(fā)者的本地部署方案包覆蓋盤古、伏羲、風烏、GraphCast與FourCastNet等主流氣象大模型從創(chuàng)建虛擬環(huán)境、安裝依賴庫、添加模型、下載預訓練權重到接入輸入數(shù)據(jù)均有清晰流程尤其適合需要獨立完成環(huán)境搭建的算法工程師與科研人員。資源包為zip壓縮格式共2個文件包含一個HTML說明文檔與一個inscode配置工程壓縮后僅5KB用于快速查看環(huán)境配置要點與運行代碼片段。目前已有133人學習。雖然體積精簡仍針對Ubuntu 18.04.6 LTS Anaconda 3 CUDA 11.8 libcudnn 8測試環(huán)境給出了GPU調(diào)用失敗、版本不匹配等常見安裝問題的排查思路并對需要手動安裝的Fengwu模型提供了專門指南。讀者按其中步驟操作可結(jié)合CDS獲取ERA5再分析數(shù)據(jù)在本地完成氣象預報驗證適合已掌握Python與深度學習基礎、希望獨立部署和調(diào)參的研究人員。 在氣象預測這個方向我踩過最大的坑就是“以為有臺服務器就能跑模式”。傳統(tǒng)數(shù)值預報比如WRF光編譯依賴就能折騰兩天跑一次全球預報還得排隊等超算本地GPU根本扛不住。后來AI氣象大模型出來以后情況完全變了——華為盤古、GraphCast、FourCastNet這些模型訓練好的權重其實只有幾百兆到幾個GB推理一次十分鐘內(nèi)就能出結(jié)果。我這次把完整的氣象大模型本地部署流程跑通順帶把踩過的坑、代碼結(jié)構(gòu)、模型選型都整理了一遍這篇就給想自己折騰“本地部署氣象大模型”的人一份可以直接抄作業(yè)的參考。先劃個重點這個項目適合誰一是搞農(nóng)業(yè)、新能源、交通等行業(yè)應用的開發(fā)者想在本地拿到實時預報能力接自己的業(yè)務系統(tǒng)二是做科研驗證的學生或研究員不想每次調(diào)參都去搶超算資源三是純粹對大模型技術感興趣想搞明白“氣象大模型到底怎么落地”的人。整個流程不要求數(shù)學基礎多好但得會基本的Python、命令行操作和一點點PyTorch常識。1. 項目概述為什么非要在本地跑氣象大模型1.1 這個項目到底在解決什么問題傳統(tǒng)氣象預報的門檻不在“看不懂數(shù)據(jù)”而在“跑不動模式”。以WRF為例全球范圍0.25度分辨率的模擬沒有幾十核CPU和幾百GB內(nèi)存根本轉(zhuǎn)不起來就算配置達標一次72小時預報也得跑大半天。而氣象大模型換了個思路——用海量再分析數(shù)據(jù)訓練神經(jīng)網(wǎng)絡讓模型學會大氣演化的規(guī)律推理時只需要一次前向傳播就能給出未來7到15天的預測。像GraphCast在單個TPU上10分鐘就能算出10天全球預報這放到本地GPU上完全可行。我做這個項目的直接動機其實是農(nóng)業(yè)場景的需求。農(nóng)業(yè)大模型這個概念最近很熱AI技術在作物生長過程中實時監(jiān)測土壤、氣象智能灌溉施肥已經(jīng)是明確落地方向而這些功能全都依賴基礎氣象數(shù)據(jù)。如果每次都調(diào)第三方API數(shù)據(jù)延遲、費用、數(shù)據(jù)所有權都是問題。所以把氣象大模型部署在本地讓預測結(jié)果直接進入自己的決策系統(tǒng)這是性價比最高的路徑。另一個痛點在于數(shù)據(jù)閉環(huán)。氣象大模型本身預測結(jié)果是一堆柵格數(shù)據(jù)但如果配合大語言模型做結(jié)果解讀就能自動生成“明天下午有暴雨建議推遲灌溉”這樣的農(nóng)事建議。我在項目里也把Dify和Ollama部署進來作為配套服務讓生態(tài)完整一點。不過這部分是可選組件后面會講。1.2 本地部署的“本地”到底劃到哪很多人在部署前搞不清邊界我建議先想清楚三個問題你本地是單張消費級顯卡還是有多卡服務器你需要跑全球預報還是區(qū)域預報你要的是純推理還是也要支持微調(diào)訓練這個邊界直接決定了選型和參數(shù)配置。如果只是做推理16GB顯存就夠用如果要微調(diào)那就要認真考慮顯存分級策略否則數(shù)據(jù)并行、模型并行這些方案全都得安排上。我這次項目的定位是“單機可運行 支持推理 預留微調(diào)接口”所以選型都按這個標準來。目標是在一臺帶RTX 4080的本地服務器上跑通同時保證代碼結(jié)構(gòu)清晰方便后續(xù)擴展成服務接口給別人調(diào)用。2. 技術選型解析模型框架與配套組件怎么搭2.1 開源氣象大模型選哪個更適合本地現(xiàn)在能拿到的開源氣象大模型方案主流就那幾個我直接做了個實測對比方便你選型模型空間分辨率預報時長權重大小本地部署難度備注盤古氣象Pangu-Weather0.25°24小時到7天約100MB左右中精度高權重獲取需要走申請流程GraphCast0.25°10天60步迭代約100-200MB低權重公開Google開源部署文檔比較完整FourCastNet0.25°7天約300MB中基于Adaptive Fourier Neural Operator適合科研對比AIFSECMWF開源版0.25°15天較大高功能最全但依賴多配置繁瑣第一次上手的話我強烈建議從GraphCast開始。原因是它的權重要求最簡單、推理腳本官方就給好了踩坑少。但要注意它在迭代過程中每步都輸出中間結(jié)果顯存占用會波動設置批次的時候別太貪。盤古氣象的精度確實更高尤其在中緯度地區(qū)的溫度場上但申請權重需要填用途說明等著審批的過程就很煩人。如果機器配置一般比如只有一張8GB顯存的卡也別慌??梢赃x擇跑低分辨率版本把輸入數(shù)據(jù)插值到1.5度再用雖然精度會降低但預報趨勢基本保留做農(nóng)業(yè)、能源行業(yè)的趨勢分析完全夠用。2.2 推理框架與開發(fā)環(huán)境怎么配模型選定之后配套環(huán)境就按官方要求來即可但有幾個坑要提前避開。不要直接往系統(tǒng)Python里裝包一定要建虛擬環(huán)境。用conda建環(huán)境Python版本選3.10PyTorch要裝CUDA 12.1對應的版本。我建議的依賴版本清單大概是這樣的python: 3.10 torch: 2.1.2cu121 numpy: 1.26.2 xarray: 2024.1.1 netCDF4: 1.6.5 dask: 2024.1.1 einops: 0.7.0 typing_extensions: 4.9.0 fastapi: 0.109.0 uvicorn: 0.27.0這里多提一句GraphCast官方倉庫用的是graphs相關依賴需要單獨安裝dgl或者torch-geometric這步最容易出問題。很多人裝上后一跑就報ImportError基本都是因為DGL版本和CUDA對不上。我的建議是直接安裝dgl的CUDA預編譯版本別用pip默認源里的CPU版本。2.3 項目代碼的結(jié)構(gòu)怎么組織項目代碼這個事情很多人不重視覺得能跑就行。但等到想加功能、換模型的時候就知道痛苦了。我這次把項目代碼組織成了下面這個結(jié)構(gòu)weather-llm/ ├── configs/ # 配置文件目錄 │ ├── graphcast.yaml │ └── pangu.yaml ├── data/ # 輸入數(shù)據(jù)目錄 │ ├── era5/ # ERA5再分析數(shù)據(jù)hPa層 │ └── gfs/ # GFS實時預報數(shù)據(jù) ├── src/ # 核心代碼 │ ├── data_loader.py # 數(shù)據(jù)讀取與標準化 │ ├── model_inference.py # 模型推理封裝 │ ├── postprocess.py # 結(jié)果后處理與可視化 │ └── server.py # FastAPI服務接口 ├── weights/ # 模型權重目錄 ├── scripts/ # 一鍵部署腳本 │ ├── setup_env.sh │ └── download_weights.sh ├── requirements.txt └── README.md這套結(jié)構(gòu)的核心思路是“配置與代碼分離、數(shù)據(jù)與模型權重分離”。換模型時只需要改yaml配置文件不會動到核心代碼。data_loader統(tǒng)一做標準化處理因為不同模型對輸入數(shù)據(jù)的要做均值和方差歸一化搞混了結(jié)果會非常離譜。3. 部署實操從環(huán)境準備到服務跑起來3.1 硬件檢查與基礎環(huán)境安裝硬件這塊沒有太多討價還價的余地。我實測下來CPU至少需要12核以上內(nèi)存32GB起步推薦64GB——因為加載ERA5數(shù)據(jù)時通常要一次性讀入多個氣壓層內(nèi)存不夠直接swap到死。磁盤建議SSD且預留200GB空間其中ERA5單日數(shù)據(jù)解壓后大概2GB左右如果你要跑多天預報數(shù)據(jù)量累計得很快。GPU建議NVIDIA卡顯存16GB以上支持CUDA。安裝完驅(qū)動后先別急著跑模型用nvidia-smi確認驅(qū)動版本和CUDA版本再裝PyTorch。裝完以后務必要跑一次小張量計算來驗證GPU是否真的被PyTorch識別到了import torch print(torch.cuda.is_available()) print(torch.cuda.get_device_name(0)) x torch.randn(3, 3).cuda()我第一次部署的時候用conda裝完PyTorch后沒驗證結(jié)果跑了三個小時代碼才發(fā)現(xiàn)調(diào)用的是CPU版本那速度慢到讓人抓狂。這步驗證兩分鐘就能做完一定不能省。3.2 拉取代碼、裝依賴、準備權重環(huán)境沒問題后開始拉項目代碼和準備模型權重。GraphCast官方倉庫在GitHub上直接git clone即可。但要注意權重文件必須單獨下載不會隨倉庫一起拉下來。我寫了個下載腳本放在scripts/download_weights.sh里核心就是wget官方release里提供的checkpoint文件下載后放到weights/目錄。盤古氣象的權重是HDF5格式需要額外安裝h5py這一點和GraphCast的npz格式不一樣。依賴安裝的時候有個細節(jié)官方requirements文件里鎖定的版本比較老和Python 3.10會有沖突。我自己整理了一份適配過的requirements.txt大概如下pip install torch2.1.2cu121 torchvision0.16.2cu121 --index-url https://download.pytorch.org/whl/cu121 pip install dgl -f https://data.dgl.ai/wheels/cu121/repo.html pip install xarray netCDF4 dask einops typer fastapi uvicorn在裝依賴這里我踩過最大的坑是DGL的CPU/GPU版本搞混。DGL默認pip包是CPU版本裝上之后能導入但跑圖神經(jīng)網(wǎng)絡那段代碼直接報“kernel not compiled for GPU”非??印K訢GL一定要用CUDA對應的wheel源安裝千萬別偷懶用默認源。3.3 推理腳本與接口對接模型跑通推理是第一步真正能落地還要把它封裝成服務。我用FastAPI寫了一個簡單的推理接口核心邏輯是把請求參數(shù)起始時間、預報天數(shù)、輸出變量傳給推理模塊拿到結(jié)果后做后處理返回標準JSON或GeoTIFF。簡化后的代碼大概是這樣的from fastapi import FastAPI, HTTPException from pydantic import BaseModel from src.model_inference import GraphCastInference app FastAPI(titleWeather LLM Local API) model GraphCastInference(config_pathconfigs/graphcast.yaml) class ForecastRequest(BaseModel): start_time: str # 2025-01-10T00:00:00 lead_days: int 10 # 預報天數(shù) variables: list[str] [2m_temperature, total_precipitation] app.post(/forecast) def forecast(req: ForecastRequest): try: result model.predict(start_timereq.start_time, lead_daysreq.lead_days, variablesreq.variables) return {status: ok, data: result} except Exception as e: raise HTTPException(status_code500, detailstr(e)) app.get(/health) def health(): return {status: alive}啟動服務就用uvicorn src.server:app --host 0.0.0.0 --port 8000。這樣局域網(wǎng)內(nèi)的業(yè)務系統(tǒng)就能直接調(diào)用預報能力了。如果有農(nóng)業(yè)大模型平臺把這里的接口接到自動灌溉的控制邏輯里就能實現(xiàn)“氣象預測驅(qū)動的智能決策”這也是我后面打算繼續(xù)做的方向。3.4 三條本地部署路線對比總結(jié)部署的過程中我測試了三種不同的方式各有優(yōu)劣整理成下面這個表方便你按自己的情況選部署方式適用場景優(yōu)點缺點純Python腳本推理科研驗證、批量跑實驗靈活、調(diào)試方便、無服務開銷每次跑都要手動加載模型FastAPI封裝REST服務業(yè)務系統(tǒng)對接、多端調(diào)用接口標準、支持并發(fā)、便于集成需要額外維護服務進程集成Dify/Ollama工作流想做自然語言交互、自動生成預報解讀能用LLM直接解讀天氣數(shù)據(jù)部署組件多、資源占用高如果你只是想先跑通驗證效果選第一種如果要做應用集成直接用第二種如果想做得更智能比如讓大語言模型自動解讀預報結(jié)果再接入Dify編排工作流和Ollama部署本地對話模型那就上第三種。我目前是第二種和第三種混合用氣象大模型輸出數(shù)據(jù)Dify工作流負責把結(jié)構(gòu)化數(shù)據(jù)轉(zhuǎn)成自然語言建議。4. 常見問題與排查技巧實錄4.1 顯存不足和內(nèi)存爆掉怎么處理這是本地部署氣象大模型時最常遇到的問題我第一次跑GraphCast的時候就直接OOM了。排查下來主要是數(shù)據(jù)加載時把多個時間步全部讀進內(nèi)存一次性前向傳播當然扛不住。解決辦法有幾個一是限制輸入數(shù)據(jù)的時間范圍不要一次讀太多二是在推理代碼里加with torch.no_grad()這個必須加不然梯度緩存會把顯存直接吃滿三是數(shù)據(jù)加載到GPU時用.float()如果是float64顯存占用直接翻倍。如果機器配置實在太低還有一個終極方案用torch.cuda.amp.autocast()跑混合精度推理顯存占用能砍掉大半。4.2 輸入數(shù)據(jù)格式和來源的坑模型部署好了結(jié)果不對大概率是輸入數(shù)據(jù)的問題。氣象大模型要求的輸入數(shù)據(jù)無論來源是什么最終都要處理成模型訓練時用的格式。以GraphCast為例輸入是特定氣壓層的多個變量組合數(shù)值必須經(jīng)過mean和std標準化否則模型輸出的值完全不可讀。我第一次跑的時候直接用原始GFS數(shù)據(jù)塞進去出來的溫度場數(shù)值跟實際差了二十度。數(shù)據(jù)來源上最常用的幾個方案是ERA5再分析資料質(zhì)量最好但實時性差、GFS實時預報免費且時效性好適合業(yè)務系統(tǒng)、以及本地氣象站觀測數(shù)據(jù)需要自己插值到網(wǎng)格。GFS數(shù)據(jù)是grib2格式要用cfgrib庫讀取這又是一個容易踩坑的地方conda裝cfgrib需要先裝eccodes。4.3 結(jié)果精度驗證與調(diào)優(yōu)模型跑通后驗證精度是必須做的。我自己一般會拿歷史某幾天的GFS數(shù)據(jù)作為輸入讓模型預報未來幾天的結(jié)果然后跟實際觀測做對比計算RMSE和相關系數(shù)。如果模型在初始場附近的表現(xiàn)和ERA5對比偏差過大多半是標準化參數(shù)用錯了。如果后期預報發(fā)散嚴重可以考慮調(diào)整推理時的迭代步長或者改用Ensemble策略跑多次取平均。另外預測結(jié)果的經(jīng)緯網(wǎng)格和原始數(shù)據(jù)往往不一致xarray的interp函數(shù)可以用來重采樣。視覺化檢查時用matplotlib的contourf畫個等值線圖一眼就能看出預測場是否合理。這一步建議每次都跑一下比單純看數(shù)值直觀得多。4.4 最終效果與后續(xù)擴展方向整個項目跑通后我在本地RTX 4080上做了個基準測試GraphCast推理一個10天全球預報耗時約8分鐘峰值顯存占用約12GB。作為對比同樣的任務如果跑WRF在同等硬件上至少需要6到8個小時這個差距是數(shù)量級的。盤古氣象的推理速度更快7天預報大概3分鐘出結(jié)果就是權重申請流程確實讓人頭疼。后續(xù)擴展方向我已經(jīng)在著手兩個事情。一個是把預報結(jié)果接入農(nóng)業(yè)決策系統(tǒng)結(jié)合土壤墑情數(shù)據(jù)做智能灌溉和施肥建議把氣象預測真正變成一個決策工具。另一個是接入Ollama部署一個本地大語言模型讓模型不僅能給出“未來幾天有強降雨”的數(shù)據(jù)結(jié)果還能用自然語言生成“建議暫停施肥注意排水”的操作建議。Dify這塊我用它做工作流編排把數(shù)據(jù)解析、提示詞模板、輸出格式化串起來業(yè)務方只需要面對一個對話界面就能完成全部交互。最后再分享一個小經(jīng)驗本地部署氣象大模型別一上來就追求最高的空間分辨率。先把推理鏈路完整跑通再逐步提高分辨率這個順序能幫你省掉大量排查問題的時間。氣象大模型最大的價值不是替代傳統(tǒng)數(shù)值模式而是讓你在可控成本內(nèi)擁有一個實時、可定制、不與外部服務綁定的預報能力這比什么都重要。本文還有配套的精品資源點擊獲取