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