簡介一套基于聯(lián)邦學(xué)習(xí)與NSL-KDD數(shù)據(jù)集的網(wǎng)絡(luò)入侵檢測Python項(xiàng)目源碼及運(yùn)行指南屬于經(jīng)導(dǎo)師指導(dǎo)并認(rèn)可的高分項(xiàng)目評審98分適合計(jì)算機(jī)相關(guān)專業(yè)學(xué)生用于課程設(shè)計(jì)、期末大作業(yè)以及想要進(jìn)行項(xiàng)目實(shí)戰(zhàn)的機(jī)器學(xué)習(xí)/網(wǎng)絡(luò)安全學(xué)習(xí)者。資源共63個(gè)文件壓縮包大小約26.19MB包含12個(gè)Python源代碼、26個(gè)Pyc編譯文件、10個(gè)Txt說明文檔、模型權(quán)重、CSV數(shù)據(jù)、結(jié)果對比PNG以及帶GUI界面的數(shù)據(jù)集等覆蓋數(shù)據(jù)預(yù)處理、模型構(gòu)建、聯(lián)邦訓(xùn)練與測試等模塊。已有88人學(xué)習(xí)下載。項(xiàng)目將聯(lián)邦學(xué)習(xí)與NSL-KDD數(shù)據(jù)集結(jié)合既演示了如何在隱私保護(hù)前提下進(jìn)行分布式訓(xùn)練也提供了帶圖形界面的數(shù)據(jù)操作方式借助附帶的運(yùn)行說明、模型文件和對比圖學(xué)習(xí)者可以快速復(fù)現(xiàn)入侵檢測實(shí)驗(yàn)深入理解從數(shù)據(jù)準(zhǔn)備到模型部署的完整流程為網(wǎng)絡(luò)安全與聯(lián)邦學(xué)習(xí)方向的實(shí)踐提供扎實(shí)參考。1. 聯(lián)邦學(xué)習(xí)與NSL-KDD做網(wǎng)絡(luò)入侵檢測先別急著跑代碼想清楚這三件事用聯(lián)邦學(xué)習(xí)與NSL-KDD數(shù)據(jù)集做網(wǎng)絡(luò)入侵檢測本質(zhì)上是把兩個(gè)成熟技術(shù)接到一起聯(lián)邦學(xué)習(xí)解決流量日志“數(shù)據(jù)不出域”的合規(guī)訴求NSL-KDD給出一道能反復(fù)驗(yàn)證的基準(zhǔn)題Python則是把它們黏在一起的膠水。這幾年相關(guān)高分項(xiàng)目幾乎都從集中式往聯(lián)邦上靠因?yàn)樵u分重點(diǎn)已經(jīng)從“檢測準(zhǔn)不準(zhǔn)”變成了“數(shù)據(jù)隔離條件下還能不能準(zhǔn)”。拿到壓縮包先別急著解壓跑訓(xùn)練先想清楚三件事數(shù)據(jù)是天然按節(jié)點(diǎn)分片還是要人工模擬Non-IID模型做二分類還是五分類指南里說的聯(lián)邦是本機(jī)模擬還是真多機(jī)通信。想清楚后再動(dòng)代碼每個(gè)參數(shù)都能說出改它的理由。這篇筆記寫給不滿足于“能運(yùn)行”的Python從業(yè)者和安全方向同學(xué)重點(diǎn)放在可復(fù)現(xiàn)步驟和踩坑記錄上。2. 為什么要聯(lián)邦化FedAvg機(jī)制、數(shù)據(jù)不出域的價(jià)值與NSL-KDD的基準(zhǔn)定位2.1 聯(lián)邦學(xué)習(xí)在入侵檢測里的角色客戶端訓(xùn)練什么、服務(wù)器聚合什么入侵檢測的常規(guī)做法是把流量特征匯集到一個(gè)訓(xùn)練中心在全部數(shù)據(jù)上訓(xùn)一個(gè)全局模型。問題在于安全日志是最敏感的數(shù)據(jù)之一跨部門、跨機(jī)房、跨公司的流量特征往往不能直接匯總別說原始報(bào)文就連統(tǒng)計(jì)特征也要走審批。聯(lián)邦學(xué)習(xí)在這里的角色很直接各個(gè)節(jié)點(diǎn)用自己的日志訓(xùn)練本地模型訓(xùn)練完只上傳模型權(quán)重原始特征永遠(yuǎn)留在本地服務(wù)器把權(quán)重按某種策略加權(quán)平均后下發(fā)反復(fù)迭代。這里最常用的聚合策略就是FedAvg聯(lián)邦平均。每輪通信服務(wù)器把當(dāng)前全局權(quán)重廣播給參與節(jié)點(diǎn)每個(gè)節(jié)點(diǎn)在本地?cái)?shù)據(jù)上做幾個(gè)epoch的梯度下降再把更新后的權(quán)重和本地樣本量一起傳回服務(wù)器按樣本量占比加權(quán)求平均得到下一輪的全局權(quán)重。整套機(jī)制里沒有原始數(shù)據(jù)流動(dòng)傳輸?shù)闹皇歉Ｐ痛笮〉攘康囊欢褟埩窟@也是聯(lián)邦學(xué)習(xí)能過合規(guī)審查的核心原因。在入侵檢測場景里有個(gè)容易被忽略的點(diǎn)這種訓(xùn)練方式天然適配“分公司/節(jié)點(diǎn)/機(jī)房”的組織結(jié)構(gòu)。每個(gè)節(jié)點(diǎn)的流量分布不同有的節(jié)點(diǎn)DDoS流量多有的節(jié)點(diǎn)主要是掃描探測聯(lián)邦框架并不要求各節(jié)點(diǎn)數(shù)據(jù)同分布它只要求各節(jié)點(diǎn)把梯度朝各自數(shù)據(jù)的方向推一步整體模型再在“每個(gè)方向的平均值”上邁一步。實(shí)際項(xiàng)目里聯(lián)邦學(xué)習(xí)還能疊一層安全聚合Secure Aggregation讓服務(wù)器無法從收到的梯度反推某個(gè)節(jié)點(diǎn)的樣本信息。NSL-KDD規(guī)模小沒必要上這個(gè)復(fù)雜度但如果以后把這套代碼接到真實(shí)流量上安全聚合就是必須考慮的下一步。做課程項(xiàng)目時(shí)在報(bào)告里提一句“聯(lián)邦只解決了數(shù)據(jù)不出域梯度本身仍可能泄露分布信息”導(dǎo)師會(huì)認(rèn)為你理解到了這一層的邊界比空寫“隱私保護(hù)”要有說服力。2.2 從KDD99到NSL-KDD去冗余之后測評分?jǐn)?shù)才真實(shí)標(biāo)題里的NSL-KDD是KDDCUP99數(shù)據(jù)集的修正版官方提供Train、Test和Test-21三份文件每行是一條網(wǎng)絡(luò)連接記錄共41個(gè)特征外加一個(gè)標(biāo)簽。它與舊版KDD99最大的差別在于去掉了大量重復(fù)記錄舊版訓(xùn)練集里同一類記錄重復(fù)幾十萬條模型記性好一點(diǎn)的都能靠背答案拿高分而NSL-KDD的訓(xùn)練集規(guī)模和重復(fù)度都被控制測試集還額外按難度分了級。因此拿NSL-KDD報(bào)告的數(shù)字更接近真實(shí)泛化能力不再是一種“背題考高分”。對比項(xiàng)KDD99NSL-KDD訓(xùn)練集冗余記錄大量重復(fù)易被模型記憶已去重訓(xùn)練集規(guī)模約12萬條測試集難度分級無分低/中/高三檔實(shí)驗(yàn)可復(fù)現(xiàn)性切分混亂官方切分明確訓(xùn)練/測試文件固定聯(lián)邦實(shí)驗(yàn)適配度不推薦數(shù)據(jù)分布失真常用于Non-IID與聯(lián)邦消融實(shí)驗(yàn)這個(gè)數(shù)據(jù)集之所以到今天還在被各類項(xiàng)目采用是因?yàn)樗銐蛐?、有現(xiàn)成的二分/多分標(biāo)簽、還能和舊版KDD99貫通做論文做畢設(shè)都有參照系。指望它替代真實(shí)流量不現(xiàn)實(shí)但把聯(lián)邦方案先在NSL-KDD上驗(yàn)證一輪再遷移到自己的流量特征上是一條性價(jià)比很高的技術(shù)路線。注意官方三份文件經(jīng)常被重新打包發(fā)布你拿到的csv可能是帶表頭、不帶表頭、多一列難度系數(shù)三種版本之一屬于正常現(xiàn)象本文第3章會(huì)給出對應(yīng)處理。2.3 二分類還是五分類檢測率、誤報(bào)率與每類召回NSL-KDD的標(biāo)簽可以歸成正常加上四類攻擊DoS拒絕服務(wù)、Probe掃描探測、R2L遠(yuǎn)程到本地、U2R提權(quán)攻擊。二分類只管正常和異常最簡單訓(xùn)練快適合快速驗(yàn)證聯(lián)邦鏈路通不通五分類要求模型分辨攻擊類型難度明顯上升尤其是R2L和U2R的樣本量極少經(jīng)常只占總量的百分之幾模型天然偏向多數(shù)類。實(shí)戰(zhàn)里我不會(huì)只看總體準(zhǔn)確率。入侵檢測更真實(shí)的指標(biāo)是檢測率也就是召回率攻擊樣本里被揪出來的比例和誤報(bào)率正常樣本里被冤枉的比例。在聯(lián)邦場景里這兩個(gè)指標(biāo)還要拆到每個(gè)客戶端上去看因?yàn)槿譁?zhǔn)確率高完全有可能掩蓋某個(gè)節(jié)點(diǎn)上檢測率為零的翻車情況。我一般會(huì)固定評估五分類至少要在報(bào)告里給每類召回率不然項(xiàng)目答辯時(shí)很難回答“你的模型到底能不能防住U2R”這種問題。多分類在聯(lián)邦訓(xùn)練中有一個(gè)額外負(fù)擔(dān)各節(jié)點(diǎn)的類別分布不一致會(huì)導(dǎo)致所謂的客戶端漂移。一個(gè)節(jié)點(diǎn)全是R2L樣本它的本地模型更新方向就偏向R2L那條梯度另一個(gè)節(jié)點(diǎn)只有normal和probe方向就完全不同。這兩種方向平均在一起全局模型可能兩頭都學(xué)不好。所以第5章里的Non-IID模擬和第3章的數(shù)據(jù)分片方式是決定這個(gè)項(xiàng)目成敗的關(guān)鍵步驟不能跳過。最后談一下Python在這個(gè)方向幾乎是唯一選項(xiàng)的原因PyTorch處理神經(jīng)網(wǎng)絡(luò)訓(xùn)練、pandas做特征表、sklearn出混淆矩陣和分類報(bào)告三個(gè)庫一條鏈路能打通整個(gè)實(shí)驗(yàn)。對Python剛?cè)腴T的人這個(gè)項(xiàng)目反而比純工程項(xiàng)目友好因?yàn)榕芡ㄗ钚±拥拇a量不到兩百行要做的雜活不過是先裝好Python環(huán)境再把numpy、pandas、scikit-learn、torch四個(gè)庫用pip裝到最新穩(wěn)定版而已。3. 數(shù)據(jù)預(yù)處理落地用Python把NSL-KDD的41維特征變成能喂PyTorch的張量3.1 讀入CSV與字段梳理先分清數(shù)值特征和符號特征讀文件這一步用pandas一行就能搞定但列名要跟官方順序?qū)R否則后續(xù)編碼全亂。NSL-KDD的csv沒有表頭需要手動(dòng)指定列名。讀入之前先確認(rèn)Python環(huán)境就緒缺庫就pip install numpy pandas scikit-learn torch一條命令裝齊Windows和Linux下沒有區(qū)別。import pandas as pd import numpy as np FEATURES [ duration, protocol_type, service, flag, src_bytes, dst_bytes, land, wrong_fragment, urgent, hot, num_failed_logins, logged_in, num_compromised, root_shell, su_attempted, num_root, num_file_creations, num_shells, num_access_files, num_outbound_cmds, is_host_login, is_guest_login, count, srv_count, serror_rate, srv_serror_rate, rerror_rate, srv_rerror_rate, same_srv_rate, diff_srv_rate, srv_diff_host_rate, dst_host_count, dst_host_srv_count, dst_host_same_srv_rate, dst_host_diff_srv_rate, dst_host_same_src_port_rate, dst_host_srv_diff_host_rate, dst_host_serror_rate, dst_host_srv_serror_rate, dst_host_rerror_rate, dst_host_srv_rerror_rate ] def load_nsl_kdd(path): df pd.read_csv(path, headerNone, namesFEATURES [label]) return df邏輯說明FEATURES按官方文檔順序列出41個(gè)特征名read_csv時(shí)用headerNone跳過默認(rèn)表頭names參數(shù)把列名掛上去第42列在代碼里命名為label。這里有一個(gè)本地容易踩的坑有的公開渠道放的NSL-KDD文件多了一列難度系數(shù)直接讀會(huì)報(bào)“列數(shù)不匹配”這時(shí)候給names多加一個(gè)level或者讀進(jìn)來后用drop列處理掉。參數(shù)說明path是訓(xùn)練集或測試集文件路徑函數(shù)返回的DataFrame要保持行順序不變因?yàn)楹竺孀鯪on-IID分片時(shí)會(huì)用它來回放標(biāo)簽索引。csv分隔符是英文逗號文本列里出現(xiàn)引號也沒關(guān)系pandas會(huì)自動(dòng)處理。訓(xùn)練集用load_nsl_kdd(KDDTrain.csv)測試集用load_nsl_kdd(KDDTest-21.csv)文件名以你實(shí)際解壓出來的為準(zhǔn)。3.2 類別特征one-hot與數(shù)值歸一化先編碼還是先縮放41維里有三個(gè)符號特征protocol_type協(xié)議類型tcp/udp/icmp三種、service服務(wù)類型約70種、flag連接狀態(tài)標(biāo)志約11種。剩下的38維都是數(shù)值或比率字段。符號特征不能直接喂給線性層常見做法是轉(zhuǎn)成one-hotservice種類太多全量one-hot會(huì)把維度頂?shù)?10以上我一般先按出現(xiàn)頻率篩出前20個(gè)其余并成other這一類把維度壓在可控范圍。from sklearn.preprocessing import StandardScaler TOP_SERVICE 20 def encode_and_normalize(df, scalerNone, top_serviceNone, fitFalse): label_int df[label].map(build_label_map()) label_int label_int.values.astype(np.int64) # 三個(gè)符號特征統(tǒng)一轉(zhuǎn) one-hot前綴區(qū)分來源 proto pd.get_dummies(df[protocol_type], prefixproto) flag pd.get_dummies(df[flag], prefixflag) if fit: # 只在訓(xùn)練集上統(tǒng)計(jì)高頻 service避免測試集信息泄漏 top_service df[service].value_counts().index[:TOP_SERVICE].tolist() df df.copy() df[service] df[service].apply( lambda s: s if s in top_service else other) service pd.get_dummies(df[service], prefixsvc) # 剩余38列都是數(shù)值型注意先轉(zhuǎn)換類型再縮放 numeric_cols [c for c in FEATURES if c not in (protocol_type, service, flag)] numeric df[numeric_cols].astype(np.float32).values if fit: scaler StandardScaler().fit(numeric) elif scaler is None: raise ValueError(fitFalse時(shí)必須傳入訓(xùn)練集fit好的scaler) numeric_scaled scaler.transform(numeric) x np.concatenate([ numeric_scaled, proto.values.astype(np.float32), flag.values.astype(np.float32), service.values.astype(np.float32) ], axis1) return x, label_int, scaler, top_service邏輯說明encode_and_normalize做的事是符號特征轉(zhuǎn)one-hot、數(shù)值特征用StandardScaler做z-score歸一化、最后按列方向拼接成一個(gè)大矩陣。fit參數(shù)決定這次調(diào)用是“擬合scaler并統(tǒng)計(jì)service”還是“復(fù)用訓(xùn)練集返回的scaler和top_service”。用astype(np.float32)做數(shù)據(jù)類型轉(zhuǎn)換是為了跟PyTorch默認(rèn)的float32對齊順便把內(nèi)存減半。參數(shù)說明scaler和top_service的傳遞是這套代碼的命門。訓(xùn)練集上傳入fitTrue拿到scaler和top_service測試集上必須傳fitFalse并原樣代入否則訓(xùn)練和測試的特征分布不一致測試集準(zhǔn)確率虛高這在第5章避坑中會(huì)細(xì)講。TOP_SERVICE不是固定值看重區(qū)分度可以提到30看重維度稀疏可以降到10改了之后注意模型的in_dim要跟著變。數(shù)值特征這一列最容易出問題的是num_outbound_cmds在KDD99里全是0在NSL-KDD里接近全0StandardScaler對它做z-score得到一堆接近0的小數(shù)不會(huì)影響訓(xùn)練但別好奇地去刪列刪了維度就亂了。3.3 標(biāo)簽映射把攻擊名歸并成五分類NSL-KDD的原始標(biāo)簽是具體的攻擊名需要?dú)w并才能做五分類。歸并邏輯很簡單normal是一類其余按DoS、Probe、R2L、U2R四大族歸類。DOS {back, land, neptune, pod, smurf, teardrop} PROBE {ipsweep, nmap, portsweep, satan} R2L {ftp_write, guess_passwd, imap, multihop, phf, spy, warezclient, warezmaster} U2R {buffer_overflow, loadmodule, perl, rootkit} def build_label_map(): label_map {normal: 0} for idx, group in enumerate([DOS, PROBE, R2L, U2R], start1): for attack in group: label_map[attack] idx return label_map邏輯說明build_label_map返回的字典把normal映射為0四個(gè)攻擊族映射為1到4。后面encode_and_normalize里用的就是這個(gè)字典如果只想做二分類把字典改成{normal: 0}后其他所有攻擊歸為1即可但那樣會(huì)丟掉攻擊類型信息聯(lián)邦輪次里各個(gè)類別的差異化學(xué)習(xí)就看不清了。參數(shù)說明這個(gè)映射表要覆蓋官方所有攻擊名。如果你拿到的NSL-KDD文件里有字典之外的標(biāo)簽encode_and_normalize的map會(huì)返回NaN訓(xùn)練時(shí)交叉熵直接報(bào)錯(cuò)。保險(xiǎn)做法是在map調(diào)用后加一行assert label_int.notna().all()。測試集里出現(xiàn)的攻擊名不一定在訓(xùn)練集里出現(xiàn)過這屬于NSL-KDD的刻意設(shè)計(jì)聯(lián)邦測試時(shí)模型對“沒見過”的攻擊名沒有先驗(yàn)跨類別泛化會(huì)直接體現(xiàn)在召回率上別把這個(gè)當(dāng)成bug。3.4 Non-IID分片用Dirichlet分布模擬各節(jié)點(diǎn)數(shù)據(jù)不均真實(shí)聯(lián)邦場景里各節(jié)點(diǎn)的數(shù)據(jù)分布從來不是均勻的有的節(jié)點(diǎn)全是web服務(wù)有的節(jié)點(diǎn)只跑數(shù)據(jù)庫協(xié)議。要在本地驗(yàn)證聯(lián)邦算法對“分布不均”的容忍度就用Dirichlet分布來控制每個(gè)類別在各客戶端上的占比。def non_iid_split(labels, n_clients5, alpha0.5, seed42): 按Dirichlet(alpha)把樣本分給n_clients個(gè)客戶端。 alpha越小各客戶端類別分布越傾斜alpha100時(shí)接近均勻。 rng np.random.default_rng(seed) n len(labels) client_ids np.zeros(n, dtypeint) for cat in np.unique(labels): idx np.where(labels cat)[0] if len(idx) 0: continue p rng.dirichlet([alpha] * n_clients) counts rng.multinomial(len(idx), p) counts[-1] len(idx) - counts[:-1].sum() order rng.permutation(idx) start 0 for cid, cnt in enumerate(counts): client_ids[order[start:start cnt]] cid start cnt return client_ids邏輯說明逐類別處理每個(gè)類別下的樣本按Dirichlet比例隨機(jī)分給各客戶端。rng.multinomial保證每個(gè)類別的樣本全部被分到某個(gè)客戶端不會(huì)因?yàn)楦↑c(diǎn)取整丟樣本。返回的client_ids和原始數(shù)據(jù)行一一對應(yīng)后面訓(xùn)練循環(huán)里用client_ids cid做布爾索引切片就行用不著復(fù)雜的數(shù)據(jù)集類。參數(shù)說明alpha是出鏡率最高的聯(lián)邦實(shí)驗(yàn)參數(shù)。alpha0.5代表比較極端的Non-IID每個(gè)客戶端可能只擁有兩三個(gè)類別的樣本alpha100則幾乎均勻。課程項(xiàng)目建議兩端都跑一遍在非均勻分布下模型準(zhǔn)確率掉幾個(gè)點(diǎn)屬于正常現(xiàn)象報(bào)告里把這個(gè)趨勢寫清楚反而是加分項(xiàng)。這里有個(gè)數(shù)組索引的小技巧如果想檢查每個(gè)客戶端分到了什么用np.bincount(client_ids)配合labels[client_ids cid]看一眼類別直方圖比打印幾十行日志直觀得多。4. Python實(shí)現(xiàn)FedAvg模型定義、客戶端訓(xùn)練與服務(wù)端加權(quán)聚合4.1 網(wǎng)絡(luò)結(jié)構(gòu)選擇MLP足夠把復(fù)雜度留給聯(lián)邦機(jī)制網(wǎng)絡(luò)結(jié)構(gòu)不需要多深。入侵檢測的特征雖然維度不低但大部分是歸一化后的統(tǒng)計(jì)量三到四層全連接加Dropout就能擬合得很好。不建議在這個(gè)數(shù)據(jù)集上上CNN或Transformer不是用不了而是模型變大之后聯(lián)邦通信開銷線性上漲訓(xùn)練時(shí)間拉長項(xiàng)目收益卻不明顯。import torch import torch.nn as nn class NIDSMLP(nn.Module): def __init__(self, in_dim, num_classes5): super().__init__() self.net nn.Sequential( nn.Linear(in_dim, 64), nn.ReLU(), nn.Dropout(0.2), nn.Linear(64, 64), nn.ReLU(), nn.Linear(64, num_classes) ) def forward(self, x): return self.net(x)邏輯說明用64寬度的兩層MLP刻意不用BatchNorm。這里是有意為之——聯(lián)邦聚合的時(shí)候BN層里的running_mean和running_var不能像普通權(quán)重那樣直接加權(quán)平均一旦引入就得在聚合時(shí)做特殊處理很多新手在這里栽跟頭。用Dropout代替BN聚合代碼就只剩weight和bias的加權(quán)求和復(fù)雜度低一截。參數(shù)說明in_dim在第3章預(yù)處理后得到一般九十多維num_classes二分類傳2、五分類傳5。Dropout取0.2是為了在幾十輪聯(lián)邦訓(xùn)練中既提供一點(diǎn)正則又不至于讓每輪本地更新太弱。如果你發(fā)現(xiàn)訓(xùn)練集上loss收斂慢先把Dropout調(diào)到0.1試試比動(dòng)網(wǎng)絡(luò)層數(shù)見效快。4.2 客戶端本地訓(xùn)練一個(gè)類封裝一輪本地更新把每個(gè)參與方封裝成一個(gè)Client對象內(nèi)部只做一件事接收全局權(quán)重在本地?cái)?shù)據(jù)上訓(xùn)幾個(gè)epoch返回新權(quán)重。這個(gè)設(shè)計(jì)讓后續(xù)從單機(jī)模擬切到真聯(lián)邦時(shí)改動(dòng)量最小。class Client: def __init__(self, cid, x, y, devicecpu, lr1e-3): self.cid cid self.x torch.tensor(x, dtypetorch.float32, devicedevice) self.y torch.tensor(y, dtypetorch.long, devicedevice) self.device device self.lr lr def local_train(self, global_state, epochs2, batch_size64): model NIDSMLP(self.x.shape[1]).to(self.device) model.load_state_dict(global_state) optimizer torch.optim.Adam(model.parameters(), lrself.lr) loss_fn nn.CrossEntropyLoss() dataset torch.utils.data.TensorDataset(self.x, self.y) loader torch.utils.data.DataLoader( dataset, batch_sizebatch_size, shuffleTrue) model.train() for _ in range(epochs): for xb, yb in loader: optimizer.zero_grad() out model(xb) loss loss_fn(out, yb) loss.backward() optimizer.step() return model.state_dict(), len(dataset)邏輯說明Client對象保存一個(gè)客戶端的數(shù)據(jù)和標(biāo)簽local_train接收服務(wù)器下發(fā)的全局權(quán)重在本地?cái)?shù)據(jù)上做幾個(gè)epoch的梯度下降最后返回更新后的state_dict和樣本量。關(guān)鍵點(diǎn)是每次訓(xùn)練都從global_state開始而不是從上一輪本地狀態(tài)繼續(xù)這保證聯(lián)邦學(xué)習(xí)收斂的語義是“圍繞全局模型做局部修正”而不是各節(jié)點(diǎn)在自己的模型上一路跑到黑。定義函數(shù)和類的時(shí)候把變量名寫清楚后面跑消融實(shí)驗(yàn)時(shí)就不用每次重讀代碼。參數(shù)說明epochs對應(yīng)本地訓(xùn)練輪數(shù)是聯(lián)邦里最敏感的超參一般取1到3。取20以上的話每個(gè)客戶端都過擬合到自己的局部數(shù)據(jù)上聚合出來的全局模型反而變差。batch_size用64lr用1e-3這兩個(gè)參數(shù)和集中訓(xùn)練差別不大但如果Non-IID程度高lr降到5e-4更穩(wěn)。device參數(shù)默認(rèn)cpu數(shù)據(jù)量大或要跑多輪實(shí)驗(yàn)時(shí)改成cuda即可注意把x和y都放到同一個(gè)設(shè)備否則torch會(huì)報(bào)device mismatch。4.3 服務(wù)端聚合FedAvg的加權(quán)平均實(shí)現(xiàn)服務(wù)端聚合是整個(gè)聯(lián)邦學(xué)習(xí)的數(shù)學(xué)核心FedAvg的原理一句話就能說清按各客戶端本地樣本量占比對模型參數(shù)做加權(quán)平均。def fed_avg(global_model, client_states, client_sizes): new_state {} for k, v in global_model.state_dict().items(): new_state[k] torch.zeros_like(v) total sum(client_sizes) for state, size in zip(client_states, client_sizes): weight size / total for k, v in state.items(): new_state[k] v * weight global_model.load_state_dict(new_state) return new_state邏輯說明先按key初始化全零張量然后每個(gè)客戶端的權(quán)重按樣本數(shù)占比累加進(jìn)去。torch.zeros_like保證跟原模型張量形狀一致load_state_dict不會(huì)報(bào)錯(cuò)。這個(gè)實(shí)現(xiàn)只支持weight和bias這類普通張量所以4.1節(jié)特意避開了BN層否則這里還要寫B(tài)N統(tǒng)計(jì)量的合并邏輯。參數(shù)說明client_sizes來自每個(gè)Client返回的len(dataset)。要做的是加權(quán)平均而不是簡單平均如果某個(gè)客戶端樣本量是另一個(gè)的100倍前者的梯度方向會(huì)在聚合結(jié)果里占絕對主導(dǎo)。課程項(xiàng)目里的典型錯(cuò)誤是只想“公平”用簡單平均結(jié)果小客戶端的數(shù)據(jù)直接被淹沒聚合出來的模型對大節(jié)點(diǎn)過擬合。如果你希望加入權(quán)重衰減或模型剪枝在這個(gè)函數(shù)里對new_state統(tǒng)一操作即可位置就在load_state_dict之前。4.4 主訓(xùn)練循環(huán)與端到端跑通把前面的塊拼起來主循環(huán)只需要做三件事按client_ids切片、調(diào)用local_train、調(diào)用fed_avg。為了能看到收斂過程每10輪打印一次測試集準(zhǔn)確率。def evaluate(model, x, y): model.eval() x_t torch.tensor(x, dtypetorch.float32) with torch.no_grad(): pred model(x_t).argmax(dim1).numpy() return np.mean(pred y) def run_federated(x_train, y_train, client_ids, x_test, y_test, n_rounds60, epochs_per_client2, n_clients5): in_dim x_train.shape[1] global_model NIDSMLP(in_dim) global_state global_model.state_dict() for rnd in range(1, n_rounds 1): states, sizes [], [] for cid in range(n_clients): mask client_ids cid client Client(cid, x_train[mask], y_train[mask]) state, size client.local_train(global_state, epochsepochs_per_client) states.append(state) sizes.append(size) global_state fed_avg(global_model, states, sizes) if rnd % 10 0: acc evaluate(global_model, x_test, y_test) print(fround {rnd}: test acc {acc:.4f}) return global_model邏輯說明run_federated在循環(huán)里依次執(zhí)行切片、訓(xùn)練、聚合三步。這里是在單機(jī)進(jìn)程內(nèi)順序執(zhí)行所有客戶端模擬的是服務(wù)端逐個(gè)接收客戶端上傳的過程真實(shí)分布式環(huán)境里這些調(diào)用會(huì)被網(wǎng)絡(luò)通信替代但聚合的數(shù)學(xué)完全一致。evaluate函數(shù)在測試集上做前向推理argmax取預(yù)測類別后和真實(shí)標(biāo)簽比均值。參數(shù)說明n_rounds是聯(lián)邦通信輪數(shù)這個(gè)項(xiàng)目里50到100輪基本收斂再多收益很小。epochs_per_client2是計(jì)算量與精度的折中每輪每個(gè)客戶端在本地只跑兩遍數(shù)據(jù)。n_clients要和非IID分片時(shí)的數(shù)字保持一致不一致時(shí)布爾索引就會(huì)漏掉一部分樣本。跑完看一下最后10輪的準(zhǔn)確率波動(dòng)波動(dòng)超過±0.5%說明還沒收斂把n_rounds往上加到100即可。5. 訓(xùn)練避坑與排查災(zāi)難性遺忘、模型漂移與測試集選擇的四個(gè)翻車現(xiàn)場5.1 坑一訓(xùn)練/測試特征不一致全局準(zhǔn)確率虛高現(xiàn)象訓(xùn)練完在KDDTest上一測準(zhǔn)確率接近99%換成KDDTest-21直接跌到91%同一個(gè)模型兩份測試集差出一大截。原因一個(gè)是測試集選擇的差異另一個(gè)是更隱蔽的預(yù)處理泄漏——有的教程讓scaler在訓(xùn)練集和測試集整份數(shù)據(jù)上一起fit測試集的信息已經(jīng)被偷走了。NSL-KDD的KDDTest本身包含大量與訓(xùn)練集重復(fù)的連接模型見過類似樣本分?jǐn)?shù)天然偏高Test-21去掉重復(fù)后才是真實(shí)泛化水平。解決報(bào)告數(shù)字一律用KDDTest-21并且scaler嚴(yán)格只在訓(xùn)練集上fit。第3.2節(jié)里fit參數(shù)就是為這個(gè)設(shè)計(jì)的測試集上必須傳fitFalse和訓(xùn)練集返回的scaler。做對比實(shí)驗(yàn)時(shí)固定這套流程防止不同實(shí)驗(yàn)間的預(yù)處理不一致影響結(jié)論。有一個(gè)小技巧如果某輪實(shí)驗(yàn)得到的結(jié)果異常高先檢查是不是把測試集混進(jìn)fit了這是這個(gè)項(xiàng)目里最容易犯也是最不容易發(fā)現(xiàn)的錯(cuò)誤。5.2 坑二本地epochs設(shè)太大聚合后模型漂移現(xiàn)象本地訓(xùn)練20個(gè)epoch每個(gè)客戶端本地準(zhǔn)確率都98%以上聚合后的全局模型在測試集上只有70%還不如只訓(xùn)2個(gè)epoch的結(jié)果。原因每個(gè)客戶端在本地?cái)?shù)據(jù)上反復(fù)迭代模型被拉向本節(jié)點(diǎn)的局部最優(yōu)方向多個(gè)方向的平均變成了一個(gè)四不像。聯(lián)邦學(xué)習(xí)論文里管這叫client drift本地訓(xùn)練步數(shù)越多漂移越嚴(yán)重。解決把epochs降到1到3。這是一種“后悔藥”式的參數(shù)修正如果發(fā)現(xiàn)已經(jīng)跑了一輪epochs20的實(shí)驗(yàn)不用重寫代碼把run_federated(epochs_per_client2)重跑一遍即可。如果降到1仍然不穩(wěn)可以在客戶端loss里加一個(gè)近端項(xiàng)懲罰本地權(quán)重偏離全局權(quán)重這個(gè)思路對應(yīng)FedProx在loss上再加一項(xiàng)mu/2 * ||w - w_global||^2mu取0.01到0.1之間。我一般先試epochs2不穩(wěn)再加近端項(xiàng)不急著動(dòng)lr。5.3 坑三Non-IID分片后訓(xùn)練發(fā)散loss持續(xù)增大現(xiàn)象alpha0.5的Non-IID分片下全局和本地loss都不下降有時(shí)直接變成NaNalpha100時(shí)一切正常。原因一個(gè)客戶端只有一種攻擊類型時(shí)它的梯度方向跟全局方向幾乎垂直幾輪平均下來形成震蕩如果再用Adam的默認(rèn)lr個(gè)別客戶端梯度炸了就NaN了。解決兩個(gè)改動(dòng)一起做。一是把學(xué)習(xí)率從1e-3降到3e-4二是在每個(gè)客戶端的交叉熵上做類別加權(quán)讓樣本量少的R2L、U2R類別即使在一個(gè)客戶端上出現(xiàn)次數(shù)很少也不至于被遺忘。也可以在分片時(shí)先用alpha1跑通鏈路再逐步調(diào)低alpha觀察模型退化曲線這個(gè)曲線本身就是項(xiàng)目報(bào)告里很好的素材。要記住在大規(guī)模聯(lián)邦里Non-IID下的發(fā)散可能來自單個(gè)客戶端的病態(tài)梯度打印每個(gè)客戶端各自的loss而不是只看聚合后的全局loss能更快定位到是誰拖垮了全局。5.4 坑四災(zāi)難性遺忘——輪次推進(jìn)后舊攻擊類型“突然不會(huì)了”現(xiàn)象第20輪時(shí)五分類的每類召回都正常第40輪開始U2R的召回率從80%掉到30%后續(xù)輪次再也沒有恢復(fù)。原因這是聯(lián)邦學(xué)習(xí)里災(zāi)難性遺忘的典型表現(xiàn)。全局模型在新輪次里被大多數(shù)客戶端的大類樣本主導(dǎo)少數(shù)類別尤其樣本量極小的U2R和R2L的梯度信號被淹沒模型把參數(shù)空間里原先識別U2R的區(qū)域逐漸覆蓋掉。集中式訓(xùn)練同樣有遺忘問題但聯(lián)邦環(huán)境里各客戶端數(shù)據(jù)不平衡讓遺忘來得更快更隱蔽。解決三個(gè)手段按成本從低到高排列。第一每輪聚合后對每個(gè)類別單獨(dú)算召回一旦某類比上一輪跌超過5個(gè)點(diǎn)就回滾到上一輪權(quán)重再小步重訓(xùn)給項(xiàng)目加一個(gè)簡單的“剎車機(jī)制”。第二在服務(wù)端保存一份驗(yàn)證集每輪用驗(yàn)證集挑出最優(yōu)round的權(quán)重相當(dāng)于給模型存了后悔藥訓(xùn)練結(jié)束后加載最優(yōu)權(quán)重而不是最后一輪權(quán)重。第三如果類別不平衡是常態(tài)考慮在客戶端本地做少數(shù)類過采樣或者干脆改用二分類加一個(gè)專門檢測R2L/U2R的小模型別讓一個(gè)模型背所有鍋。我自己做實(shí)驗(yàn)時(shí)最常用的是第二條簡單可靠而且在答辯時(shí)可以直接展示“最優(yōu)權(quán)重出現(xiàn)在第幾輪、為什么”。關(guān)于評估指標(biāo)還有一條硬規(guī)矩五分類下只報(bào)accuracy是不夠的至少要帶上每類recall和誤報(bào)率否則災(zāi)難性遺忘根本不會(huì)被你發(fā)現(xiàn)。用sklearn的classification_report一行就能輸出全部類別精確率、召回率、F1畫混淆矩陣時(shí)如果類別標(biāo)簽擠成一團(tuán)把figure size調(diào)大到(8, 6)以上就好。聯(lián)邦場景下我更建議額外打印每個(gè)客戶端在測試集上的單獨(dú)表現(xiàn)因?yàn)槿种笜?biāo)好看不能掩蓋某個(gè)節(jié)點(diǎn)上的檢測失效這一點(diǎn)在答辯時(shí)也是容易被追問的角落。6. 從模擬到真聯(lián)邦消融實(shí)驗(yàn)、通信開銷與三行改造清單模擬畢竟是模擬真聯(lián)邦在網(wǎng)絡(luò)不穩(wěn)定的真實(shí)環(huán)境里還要面對掉線、傳輸延遲和梯度過期問題。做驗(yàn)證時(shí)有三個(gè)技巧值得用。第一個(gè)是消融實(shí)驗(yàn)。把集中式訓(xùn)練、聯(lián)邦5客戶端、聯(lián)邦10客戶端三組跑在同一份預(yù)處理代碼下記錄測試準(zhǔn)確率、每類召回率和訓(xùn)練耗時(shí)。集中式可以看作聯(lián)邦的上限參照兩個(gè)聯(lián)邦實(shí)驗(yàn)的差距體現(xiàn)數(shù)據(jù)分片帶來的信息損失。我習(xí)慣把結(jié)果畫成一張對比表集中式94.2%、聯(lián)邦5客戶端92.8%、聯(lián)邦10客戶端91.5%這種形式比文字描述“聯(lián)邦效果還行”有說服力得多。第二個(gè)是通信開銷的量化。單個(gè)模型參數(shù)量乘以2上傳加下載再乘以通信輪數(shù)就能估算傳輸字節(jié)量。以文中這個(gè)MLP為例in_dim約90、兩層64寬參數(shù)量大約一萬出頭float32一個(gè)參數(shù)4字節(jié)一輪通信約80KB60輪不到5MB。這個(gè)數(shù)字放進(jìn)報(bào)告里就能正面回應(yīng)“聯(lián)邦學(xué)習(xí)到底多耗帶寬”的質(zhì)疑。第三個(gè)是遷移真聯(lián)邦的改造清單把Client里的數(shù)據(jù)加載換成遠(yuǎn)程文件路徑、把local_train返回值序列化后走消息隊(duì)列傳輸、服務(wù)端fed_avg換成異步聚合以容納掉線節(jié)點(diǎn)。常見做法是直接用Flower這類聯(lián)邦框架改配置或者自己拉一套gRPC兩者都能保留第4章代碼本體的八到九成。我做這個(gè)方向時(shí)踩過最大的坑是默認(rèn)“本機(jī)模擬的結(jié)果等于分布式結(jié)果”真搬到多機(jī)后才發(fā)現(xiàn)網(wǎng)絡(luò)延時(shí)會(huì)放大每個(gè)round的耗時(shí)原來用60輪訓(xùn)練的代價(jià)在真環(huán)境下翻了不止三倍。后來我在模擬階段就把round數(shù)壓到40把epochs壓到1用更多客戶端替代更多輪次反而拿到更穩(wěn)定的結(jié)果。把“通信輪次”當(dāng)成跟學(xué)習(xí)率一樣的超參去調(diào)是聯(lián)邦項(xiàng)目跟普通深度學(xué)習(xí)項(xiàng)目最大的習(xí)慣差異。這個(gè)思路貫穿了我后來所有的聯(lián)邦實(shí)驗(yàn)希望這些記錄對你有所幫助。本文還有配套的精品資源點(diǎn)擊獲取