簡介一套基于聯(lián)邦學(xué)習(xí)與NSL-KDD數(shù)據(jù)集的網(wǎng)絡(luò)入侵檢測Python項目源碼及運行指南屬于經(jīng)導(dǎo)師指導(dǎo)并認(rèn)可的高分項目評審98分適合計算機相關(guān)專業(yè)學(xué)生用于課程設(shè)計、期末大作業(yè)以及想要進(jìn)行項目實戰(zhàn)的機器學(xué)習(xí)/網(wǎng)絡(luò)安全學(xué)習(xí)者。資源共63個文件壓縮包大小約26.19MB包含12個Python源代碼、26個Pyc編譯文件、10個Txt說明文檔、模型權(quán)重、CSV數(shù)據(jù)、結(jié)果對比PNG以及帶GUI界面的數(shù)據(jù)集等覆蓋數(shù)據(jù)預(yù)處理、模型構(gòu)建、聯(lián)邦訓(xùn)練與測試等模塊。已有88人學(xué)習(xí)下載。項目將聯(lián)邦學(xué)習(xí)與NSL-KDD數(shù)據(jù)集結(jié)合既演示了如何在隱私保護(hù)前提下進(jìn)行分布式訓(xùn)練也提供了帶圖形界面的數(shù)據(jù)操作方式借助附帶的運行說明、模型文件和對比圖學(xué)習(xí)者可以快速復(fù)現(xiàn)入侵檢測實驗深入理解從數(shù)據(jù)準(zhǔn)備到模型部署的完整流程為網(wǎng)絡(luò)安全與聯(lián)邦學(xué)習(xí)方向的實踐提供扎實參考。1. 聯(lián)邦學(xué)習(xí)與NSL-KDD做網(wǎng)絡(luò)入侵檢測先別急著跑代碼想清楚這三件事用聯(lián)邦學(xué)習(xí)與NSL-KDD數(shù)據(jù)集做網(wǎng)絡(luò)入侵檢測本質(zhì)上是把兩個成熟技術(shù)接到一起聯(lián)邦學(xué)習(xí)解決流量日志“數(shù)據(jù)不出域”的合規(guī)訴求NSL-KDD給出一道能反復(fù)驗證的基準(zhǔn)題Python則是把它們黏在一起的膠水。這幾年相關(guān)高分項目幾乎都從集中式往聯(lián)邦上靠因為評分重點已經(jīng)從“檢測準(zhǔn)不準(zhǔn)”變成了“數(shù)據(jù)隔離條件下還能不能準(zhǔn)”。拿到壓縮包先別急著解壓跑訓(xùn)練先想清楚三件事數(shù)據(jù)是天然按節(jié)點分片還是要人工模擬Non-IID模型做二分類還是五分類指南里說的聯(lián)邦是本機模擬還是真多機通信。想清楚后再動代碼每個參數(shù)都能說出改它的理由。這篇筆記寫給不滿足于“能運行”的Python從業(yè)者和安全方向同學(xué)重點放在可復(fù)現(xiàn)步驟和踩坑記錄上。2. 為什么要聯(lián)邦化FedAvg機制、數(shù)據(jù)不出域的價值與NSL-KDD的基準(zhǔn)定位2.1 聯(lián)邦學(xué)習(xí)在入侵檢測里的角色客戶端訓(xùn)練什么、服務(wù)器聚合什么入侵檢測的常規(guī)做法是把流量特征匯集到一個訓(xùn)練中心在全部數(shù)據(jù)上訓(xùn)一個全局模型。問題在于安全日志是最敏感的數(shù)據(jù)之一跨部門、跨機房、跨公司的流量特征往往不能直接匯總別說原始報文就連統(tǒng)計特征也要走審批。聯(lián)邦學(xué)習(xí)在這里的角色很直接各個節(jié)點用自己的日志訓(xùn)練本地模型訓(xùn)練完只上傳模型權(quán)重原始特征永遠(yuǎn)留在本地服務(wù)器把權(quán)重按某種策略加權(quán)平均后下發(fā)反復(fù)迭代。這里最常用的聚合策略就是FedAvg聯(lián)邦平均。每輪通信服務(wù)器把當(dāng)前全局權(quán)重廣播給參與節(jié)點每個節(jié)點在本地數(shù)據(jù)上做幾個epoch的梯度下降再把更新后的權(quán)重和本地樣本量一起傳回服務(wù)器按樣本量占比加權(quán)求平均得到下一輪的全局權(quán)重。整套機制里沒有原始數(shù)據(jù)流動傳輸?shù)闹皇歉Ｐ痛笮〉攘康囊欢褟埩窟@也是聯(lián)邦學(xué)習(xí)能過合規(guī)審查的核心原因。在入侵檢測場景里有個容易被忽略的點這種訓(xùn)練方式天然適配“分公司/節(jié)點/機房”的組織結(jié)構(gòu)。每個節(jié)點的流量分布不同有的節(jié)點DDoS流量多有的節(jié)點主要是掃描探測聯(lián)邦框架并不要求各節(jié)點數(shù)據(jù)同分布它只要求各節(jié)點把梯度朝各自數(shù)據(jù)的方向推一步整體模型再在“每個方向的平均值”上邁一步。實際項目里聯(lián)邦學(xué)習(xí)還能疊一層安全聚合Secure Aggregation讓服務(wù)器無法從收到的梯度反推某個節(jié)點的樣本信息。NSL-KDD規(guī)模小沒必要上這個復(fù)雜度但如果以后把這套代碼接到真實流量上安全聚合就是必須考慮的下一步。做課程項目時在報告里提一句“聯(lián)邦只解決了數(shù)據(jù)不出域梯度本身仍可能泄露分布信息”導(dǎo)師會認(rèn)為你理解到了這一層的邊界比空寫“隱私保護(hù)”要有說服力。2.2 從KDD99到NSL-KDD去冗余之后測評分?jǐn)?shù)才真實標(biāo)題里的NSL-KDD是KDDCUP99數(shù)據(jù)集的修正版官方提供Train、Test和Test-21三份文件每行是一條網(wǎng)絡(luò)連接記錄共41個特征外加一個標(biāo)簽。它與舊版KDD99最大的差別在于去掉了大量重復(fù)記錄舊版訓(xùn)練集里同一類記錄重復(fù)幾十萬條模型記性好一點的都能靠背答案拿高分而NSL-KDD的訓(xùn)練集規(guī)模和重復(fù)度都被控制測試集還額外按難度分了級。因此拿NSL-KDD報告的數(shù)字更接近真實泛化能力不再是一種“背題考高分”。對比項KDD99NSL-KDD訓(xùn)練集冗余記錄大量重復(fù)易被模型記憶已去重訓(xùn)練集規(guī)模約12萬條測試集難度分級無分低/中/高三檔實驗可復(fù)現(xiàn)性切分混亂官方切分明確訓(xùn)練/測試文件固定聯(lián)邦實驗適配度不推薦數(shù)據(jù)分布失真常用于Non-IID與聯(lián)邦消融實驗這個數(shù)據(jù)集之所以到今天還在被各類項目采用是因為它足夠小、有現(xiàn)成的二分/多分標(biāo)簽、還能和舊版KDD99貫通做論文做畢設(shè)都有參照系。指望它替代真實流量不現(xiàn)實但把聯(lián)邦方案先在NSL-KDD上驗證一輪再遷移到自己的流量特征上是一條性價比很高的技術(shù)路線。注意官方三份文件經(jīng)常被重新打包發(fā)布你拿到的csv可能是帶表頭、不帶表頭、多一列難度系數(shù)三種版本之一屬于正?，F(xiàn)象本文第3章會給出對應(yīng)處理。2.3 二分類還是五分類檢測率、誤報率與每類召回NSL-KDD的標(biāo)簽可以歸成正常加上四類攻擊DoS拒絕服務(wù)、Probe掃描探測、R2L遠(yuǎn)程到本地、U2R提權(quán)攻擊。二分類只管正常和異常最簡單訓(xùn)練快適合快速驗證聯(lián)邦鏈路通不通五分類要求模型分辨攻擊類型難度明顯上升尤其是R2L和U2R的樣本量極少經(jīng)常只占總量的百分之幾模型天然偏向多數(shù)類。實戰(zhàn)里我不會只看總體準(zhǔn)確率。入侵檢測更真實的指標(biāo)是檢測率也就是召回率攻擊樣本里被揪出來的比例和誤報率正常樣本里被冤枉的比例。在聯(lián)邦場景里這兩個指標(biāo)還要拆到每個客戶端上去看因為全局準(zhǔn)確率高完全有可能掩蓋某個節(jié)點上檢測率為零的翻車情況。我一般會固定評估五分類至少要在報告里給每類召回率不然項目答辯時很難回答“你的模型到底能不能防住U2R”這種問題。多分類在聯(lián)邦訓(xùn)練中有一個額外負(fù)擔(dān)各節(jié)點的類別分布不一致會導(dǎo)致所謂的客戶端漂移。一個節(jié)點全是R2L樣本它的本地模型更新方向就偏向R2L那條梯度另一個節(jié)點只有normal和probe方向就完全不同。這兩種方向平均在一起全局模型可能兩頭都學(xué)不好。所以第5章里的Non-IID模擬和第3章的數(shù)據(jù)分片方式是決定這個項目成敗的關(guān)鍵步驟不能跳過。最后談一下Python在這個方向幾乎是唯一選項的原因PyTorch處理神經(jīng)網(wǎng)絡(luò)訓(xùn)練、pandas做特征表、sklearn出混淆矩陣和分類報告三個庫一條鏈路能打通整個實驗。對Python剛?cè)腴T的人這個項目反而比純工程項目友好因為跑通最小例子的代碼量不到兩百行要做的雜活不過是先裝好Python環(huán)境再把numpy、pandas、scikit-learn、torch四個庫用pip裝到最新穩(wěn)定版而已。3. 數(shù)據(jù)預(yù)處理落地用Python把NSL-KDD的41維特征變成能喂PyTorch的張量3.1 讀入CSV與字段梳理先分清數(shù)值特征和符號特征讀文件這一步用pandas一行就能搞定但列名要跟官方順序?qū)R否則后續(xù)編碼全亂。NSL-KDD的csv沒有表頭需要手動指定列名。讀入之前先確認(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個特征名read_csv時用headerNone跳過默認(rèn)表頭names參數(shù)把列名掛上去第42列在代碼里命名為label。這里有一個本地容易踩的坑有的公開渠道放的NSL-KDD文件多了一列難度系數(shù)直接讀會報“列數(shù)不匹配”這時候給names多加一個level或者讀進(jìn)來后用drop列處理掉。參數(shù)說明path是訓(xùn)練集或測試集文件路徑函數(shù)返回的DataFrame要保持行順序不變因為后面做Non-IID分片時會用它來回放標(biāo)簽索引。csv分隔符是英文逗號文本列里出現(xiàn)引號也沒關(guān)系pandas會自動處理。訓(xùn)練集用load_nsl_kdd(KDDTrain.csv)測試集用load_nsl_kdd(KDDTest-21.csv)文件名以你實際解壓出來的為準(zhǔn)。3.2 類別特征one-hot與數(shù)值歸一化先編碼還是先縮放41維里有三個符號特征protocol_type協(xié)議類型tcp/udp/icmp三種、service服務(wù)類型約70種、flag連接狀態(tài)標(biāo)志約11種。剩下的38維都是數(shù)值或比率字段。符號特征不能直接喂給線性層常見做法是轉(zhuǎn)成one-hotservice種類太多全量one-hot會把維度頂?shù)?10以上我一般先按出現(xiàn)頻率篩出前20個其余并成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) # 三個符號特征統(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)計高頻 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時必須傳入訓(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歸一化、最后按列方向拼接成一個大矩陣。fit參數(shù)決定這次調(diào)用是“擬合scaler并統(tǒng)計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章避坑中會細(xì)講。TOP_SERVICE不是固定值看重區(qū)分度可以提到30看重維度稀疏可以降到10改了之后注意模型的in_dim要跟著變。數(shù)值特征這一列最容易出問題的是num_outbound_cmds在KDD99里全是0在NSL-KDD里接近全0StandardScaler對它做z-score得到一堆接近0的小數(shù)不會影響訓(xùn)練但別好奇地去刪列刪了維度就亂了。3.3 標(biāo)簽映射把攻擊名歸并成五分類NSL-KDD的原始標(biāo)簽是具體的攻擊名需要歸并才能做五分類。歸并邏輯很簡單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四個攻擊族映射為1到4。后面encode_and_normalize里用的就是這個字典如果只想做二分類把字典改成{normal: 0}后其他所有攻擊歸為1即可但那樣會丟掉攻擊類型信息聯(lián)邦輪次里各個類別的差異化學(xué)習(xí)就看不清了。參數(shù)說明這個映射表要覆蓋官方所有攻擊名。如果你拿到的NSL-KDD文件里有字典之外的標(biāo)簽encode_and_normalize的map會返回NaN訓(xùn)練時交叉熵直接報錯。保險做法是在map調(diào)用后加一行assert label_int.notna().all()。測試集里出現(xiàn)的攻擊名不一定在訓(xùn)練集里出現(xiàn)過這屬于NSL-KDD的刻意設(shè)計聯(lián)邦測試時模型對“沒見過”的攻擊名沒有先驗跨類別泛化會直接體現(xiàn)在召回率上別把這個當(dāng)成bug。3.4 Non-IID分片用Dirichlet分布模擬各節(jié)點數(shù)據(jù)不均真實聯(lián)邦場景里各節(jié)點的數(shù)據(jù)分布從來不是均勻的有的節(jié)點全是web服務(wù)有的節(jié)點只跑數(shù)據(jù)庫協(xié)議。要在本地驗證聯(lián)邦算法對“分布不均”的容忍度就用Dirichlet分布來控制每個類別在各客戶端上的占比。def non_iid_split(labels, n_clients5, alpha0.5, seed42): 按Dirichlet(alpha)把樣本分給n_clients個客戶端。 alpha越小各客戶端類別分布越傾斜alpha100時接近均勻。 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邏輯說明逐類別處理每個類別下的樣本按Dirichlet比例隨機分給各客戶端。rng.multinomial保證每個類別的樣本全部被分到某個客戶端不會因為浮點取整丟樣本。返回的client_ids和原始數(shù)據(jù)行一一對應(yīng)后面訓(xùn)練循環(huán)里用client_ids cid做布爾索引切片就行用不著復(fù)雜的數(shù)據(jù)集類。參數(shù)說明alpha是出鏡率最高的聯(lián)邦實驗參數(shù)。alpha0.5代表比較極端的Non-IID每個客戶端可能只擁有兩三個類別的樣本alpha100則幾乎均勻。課程項目建議兩端都跑一遍在非均勻分布下模型準(zhǔn)確率掉幾個點屬于正?，F(xiàn)象報告里把這個趨勢寫清楚反而是加分項。這里有個數(shù)組索引的小技巧如果想檢查每個客戶端分到了什么用np.bincount(client_ids)配合labels[client_ids cid]看一眼類別直方圖比打印幾十行日志直觀得多。4. Python實現(xiàn)FedAvg模型定義、客戶端訓(xùn)練與服務(wù)端加權(quán)聚合4.1 網(wǎng)絡(luò)結(jié)構(gòu)選擇MLP足夠把復(fù)雜度留給聯(lián)邦機制網(wǎng)絡(luò)結(jié)構(gòu)不需要多深。入侵檢測的特征雖然維度不低但大部分是歸一化后的統(tǒng)計量三到四層全連接加Dropout就能擬合得很好。不建議在這個數(shù)據(jù)集上上CNN或Transformer不是用不了而是模型變大之后聯(lián)邦通信開銷線性上漲訓(xùn)練時間拉長項目收益卻不明顯。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)邦聚合的時候BN層里的running_mean和running_var不能像普通權(quán)重那樣直接加權(quán)平均一旦引入就得在聚合時做特殊處理很多新手在這里栽跟頭。用Dropout代替BN聚合代碼就只剩weight和bias的加權(quán)求和復(fù)雜度低一截。參數(shù)說明in_dim在第3章預(yù)處理后得到一般九十多維num_classes二分類傳2、五分類傳5。Dropout取0.2是為了在幾十輪聯(lián)邦訓(xùn)練中既提供一點正則又不至于讓每輪本地更新太弱。如果你發(fā)現(xiàn)訓(xùn)練集上loss收斂慢先把Dropout調(diào)到0.1試試比動網(wǎng)絡(luò)層數(shù)見效快。4.2 客戶端本地訓(xùn)練一個類封裝一輪本地更新把每個參與方封裝成一個Client對象內(nèi)部只做一件事接收全局權(quán)重在本地數(shù)據(jù)上訓(xùn)幾個epoch返回新權(quán)重。這個設(shè)計讓后續(xù)從單機模擬切到真聯(lián)邦時改動量最小。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對象保存一個客戶端的數(shù)據(jù)和標(biāo)簽local_train接收服務(wù)器下發(fā)的全局權(quán)重在本地數(shù)據(jù)上做幾個epoch的梯度下降最后返回更新后的state_dict和樣本量。關(guān)鍵點是每次訓(xùn)練都從global_state開始而不是從上一輪本地狀態(tài)繼續(xù)這保證聯(lián)邦學(xué)習(xí)收斂的語義是“圍繞全局模型做局部修正”而不是各節(jié)點在自己的模型上一路跑到黑。定義函數(shù)和類的時候把變量名寫清楚后面跑消融實驗時就不用每次重讀代碼。參數(shù)說明epochs對應(yīng)本地訓(xùn)練輪數(shù)是聯(lián)邦里最敏感的超參一般取1到3。取20以上的話每個客戶端都過擬合到自己的局部數(shù)據(jù)上聚合出來的全局模型反而變差。batch_size用64lr用1e-3這兩個參數(shù)和集中訓(xùn)練差別不大但如果Non-IID程度高lr降到5e-4更穩(wěn)。device參數(shù)默認(rèn)cpu數(shù)據(jù)量大或要跑多輪實驗時改成cuda即可注意把x和y都放到同一個設(shè)備否則torch會報device mismatch。4.3 服務(wù)端聚合FedAvg的加權(quán)平均實現(xiàn)服務(wù)端聚合是整個聯(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初始化全零張量然后每個客戶端的權(quán)重按樣本數(shù)占比累加進(jìn)去。torch.zeros_like保證跟原模型張量形狀一致load_state_dict不會報錯。這個實現(xiàn)只支持weight和bias這類普通張量所以4.1節(jié)特意避開了BN層否則這里還要寫B(tài)N統(tǒng)計量的合并邏輯。參數(shù)說明client_sizes來自每個Client返回的len(dataset)。要做的是加權(quán)平均而不是簡單平均如果某個客戶端樣本量是另一個的100倍前者的梯度方向會在聚合結(jié)果里占絕對主導(dǎo)。課程項目里的典型錯誤是只想“公平”用簡單平均結(jié)果小客戶端的數(shù)據(jù)直接被淹沒聚合出來的模型對大節(jié)點過擬合。如果你希望加入權(quán)重衰減或模型剪枝在這個函數(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ìn)程內(nèi)順序執(zhí)行所有客戶端模擬的是服務(wù)端逐個接收客戶端上傳的過程真實分布式環(huán)境里這些調(diào)用會被網(wǎng)絡(luò)通信替代但聚合的數(shù)學(xué)完全一致。evaluate函數(shù)在測試集上做前向推理argmax取預(yù)測類別后和真實標(biāo)簽比均值。參數(shù)說明n_rounds是聯(lián)邦通信輪數(shù)這個項目里50到100輪基本收斂再多收益很小。epochs_per_client2是計算量與精度的折中每輪每個客戶端在本地只跑兩遍數(shù)據(jù)。n_clients要和非IID分片時的數(shù)字保持一致不一致時布爾索引就會漏掉一部分樣本。跑完看一下最后10輪的準(zhǔn)確率波動波動超過±0.5%說明還沒收斂把n_rounds往上加到100即可。5. 訓(xùn)練避坑與排查災(zāi)難性遺忘、模型漂移與測試集選擇的四個翻車現(xiàn)場5.1 坑一訓(xùn)練/測試特征不一致全局準(zhǔn)確率虛高現(xiàn)象訓(xùn)練完在KDDTest上一測準(zhǔn)確率接近99%換成KDDTest-21直接跌到91%同一個模型兩份測試集差出一大截。原因一個是測試集選擇的差異另一個是更隱蔽的預(yù)處理泄漏——有的教程讓scaler在訓(xùn)練集和測試集整份數(shù)據(jù)上一起fit測試集的信息已經(jīng)被偷走了。NSL-KDD的KDDTest本身包含大量與訓(xùn)練集重復(fù)的連接模型見過類似樣本分?jǐn)?shù)天然偏高Test-21去掉重復(fù)后才是真實泛化水平。解決報告數(shù)字一律用KDDTest-21并且scaler嚴(yán)格只在訓(xùn)練集上fit。第3.2節(jié)里fit參數(shù)就是為這個設(shè)計的測試集上必須傳fitFalse和訓(xùn)練集返回的scaler。做對比實驗時固定這套流程防止不同實驗間的預(yù)處理不一致影響結(jié)論。有一個小技巧如果某輪實驗得到的結(jié)果異常高先檢查是不是把測試集混進(jìn)fit了這是這個項目里最容易犯也是最不容易發(fā)現(xiàn)的錯誤。5.2 坑二本地epochs設(shè)太大聚合后模型漂移現(xiàn)象本地訓(xùn)練20個epoch每個客戶端本地準(zhǔn)確率都98%以上聚合后的全局模型在測試集上只有70%還不如只訓(xùn)2個epoch的結(jié)果。原因每個客戶端在本地數(shù)據(jù)上反復(fù)迭代模型被拉向本節(jié)點的局部最優(yōu)方向多個方向的平均變成了一個四不像。聯(lián)邦學(xué)習(xí)論文里管這叫client drift本地訓(xùn)練步數(shù)越多漂移越嚴(yán)重。解決把epochs降到1到3。這是一種“后悔藥”式的參數(shù)修正如果發(fā)現(xiàn)已經(jīng)跑了一輪epochs20的實驗不用重寫代碼把run_federated(epochs_per_client2)重跑一遍即可。如果降到1仍然不穩(wěn)可以在客戶端loss里加一個近端項懲罰本地權(quán)重偏離全局權(quán)重這個思路對應(yīng)FedProx在loss上再加一項mu/2 * ||w - w_global||^2mu取0.01到0.1之間。我一般先試epochs2不穩(wěn)再加近端項不急著動lr。5.3 坑三Non-IID分片后訓(xùn)練發(fā)散loss持續(xù)增大現(xiàn)象alpha0.5的Non-IID分片下全局和本地loss都不下降有時直接變成NaNalpha100時一切正常。原因一個客戶端只有一種攻擊類型時它的梯度方向跟全局方向幾乎垂直幾輪平均下來形成震蕩如果再用Adam的默認(rèn)lr個別客戶端梯度炸了就NaN了。解決兩個改動一起做。一是把學(xué)習(xí)率從1e-3降到3e-4二是在每個客戶端的交叉熵上做類別加權(quán)讓樣本量少的R2L、U2R類別即使在一個客戶端上出現(xiàn)次數(shù)很少也不至于被遺忘。也可以在分片時先用alpha1跑通鏈路再逐步調(diào)低alpha觀察模型退化曲線這個曲線本身就是項目報告里很好的素材。要記住在大規(guī)模聯(lián)邦里Non-IID下的發(fā)散可能來自單個客戶端的病態(tài)梯度打印每個客戶端各自的loss而不是只看聚合后的全局loss能更快定位到是誰拖垮了全局。5.4 坑四災(zāi)難性遺忘——輪次推進(jìn)后舊攻擊類型“突然不會了”現(xiàn)象第20輪時五分類的每類召回都正常第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ù)不平衡讓遺忘來得更快更隱蔽。解決三個手段按成本從低到高排列。第一每輪聚合后對每個類別單獨算召回一旦某類比上一輪跌超過5個點就回滾到上一輪權(quán)重再小步重訓(xùn)給項目加一個簡單的“剎車機制”。第二在服務(wù)端保存一份驗證集每輪用驗證集挑出最優(yōu)round的權(quán)重相當(dāng)于給模型存了后悔藥訓(xùn)練結(jié)束后加載最優(yōu)權(quán)重而不是最后一輪權(quán)重。第三如果類別不平衡是常態(tài)考慮在客戶端本地做少數(shù)類過采樣或者干脆改用二分類加一個專門檢測R2L/U2R的小模型別讓一個模型背所有鍋。我自己做實驗時最常用的是第二條簡單可靠而且在答辯時可以直接展示“最優(yōu)權(quán)重出現(xiàn)在第幾輪、為什么”。關(guān)于評估指標(biāo)還有一條硬規(guī)矩五分類下只報accuracy是不夠的至少要帶上每類recall和誤報率否則災(zāi)難性遺忘根本不會被你發(fā)現(xiàn)。用sklearn的classification_report一行就能輸出全部類別精確率、召回率、F1畫混淆矩陣時如果類別標(biāo)簽擠成一團(tuán)把figure size調(diào)大到(8, 6)以上就好。聯(lián)邦場景下我更建議額外打印每個客戶端在測試集上的單獨表現(xiàn)因為全局指標(biāo)好看不能掩蓋某個節(jié)點上的檢測失效這一點在答辯時也是容易被追問的角落。6. 從模擬到真聯(lián)邦消融實驗、通信開銷與三行改造清單模擬畢竟是模擬真聯(lián)邦在網(wǎng)絡(luò)不穩(wěn)定的真實環(huán)境里還要面對掉線、傳輸延遲和梯度過期問題。做驗證時有三個技巧值得用。第一個是消融實驗。把集中式訓(xùn)練、聯(lián)邦5客戶端、聯(lián)邦10客戶端三組跑在同一份預(yù)處理代碼下記錄測試準(zhǔn)確率、每類召回率和訓(xùn)練耗時。集中式可以看作聯(lián)邦的上限參照兩個聯(lián)邦實驗的差距體現(xiàn)數(shù)據(jù)分片帶來的信息損失。我習(xí)慣把結(jié)果畫成一張對比表集中式94.2%、聯(lián)邦5客戶端92.8%、聯(lián)邦10客戶端91.5%這種形式比文字描述“聯(lián)邦效果還行”有說服力得多。第二個是通信開銷的量化。單個模型參數(shù)量乘以2上傳加下載再乘以通信輪數(shù)就能估算傳輸字節(jié)量。以文中這個MLP為例in_dim約90、兩層64寬參數(shù)量大約一萬出頭float32一個參數(shù)4字節(jié)一輪通信約80KB60輪不到5MB。這個數(shù)字放進(jìn)報告里就能正面回應(yīng)“聯(lián)邦學(xué)習(xí)到底多耗帶寬”的質(zhì)疑。第三個是遷移真聯(lián)邦的改造清單把Client里的數(shù)據(jù)加載換成遠(yuǎn)程文件路徑、把local_train返回值序列化后走消息隊列傳輸、服務(wù)端fed_avg換成異步聚合以容納掉線節(jié)點。常見做法是直接用Flower這類聯(lián)邦框架改配置或者自己拉一套gRPC兩者都能保留第4章代碼本體的八到九成。我做這個方向時踩過最大的坑是默認(rèn)“本機模擬的結(jié)果等于分布式結(jié)果”真搬到多機后才發(fā)現(xiàn)網(wǎng)絡(luò)延時會放大每個round的耗時原來用60輪訓(xùn)練的代價在真環(huán)境下翻了不止三倍。后來我在模擬階段就把round數(shù)壓到40把epochs壓到1用更多客戶端替代更多輪次反而拿到更穩(wěn)定的結(jié)果。把“通信輪次”當(dāng)成跟學(xué)習(xí)率一樣的超參去調(diào)是聯(lián)邦項目跟普通深度學(xué)習(xí)項目最大的習(xí)慣差異。這個思路貫穿了我后來所有的聯(lián)邦實驗希望這些記錄對你有所幫助。本文還有配套的精品資源點擊獲取