習(xí)的智能合約漏洞檢測(cè):操作碼序列分類實(shí)戰(zhàn))
簡(jiǎn)介面向區(qū)塊鏈安全與智能合約方向的畢業(yè)設(shè)計(jì)場(chǎng)景這份資源提供了一套基于深度學(xué)習(xí)的漏洞檢測(cè)完整方案。項(xiàng)目采用長(zhǎng)短時(shí)記憶網(wǎng)絡(luò)LSTM與雙向長(zhǎng)短時(shí)記憶網(wǎng)絡(luò)BLSTM作為主干并引入注意力機(jī)制聚焦漏洞關(guān)鍵特征可對(duì)Solidity合約代碼進(jìn)行自動(dòng)預(yù)測(cè)適合承擔(dān)相關(guān)課題或課程設(shè)計(jì)的高校學(xué)生、開(kāi)發(fā)者參考。壓縮包共1428個(gè)文件包含1388個(gè)sol合約樣本、11個(gè)Python腳本及日志、圖表、配置文件等整體大小僅3MB便于本地部署與快速運(yùn)行。隨附源碼與部署教程可讓用戶直接啟動(dòng)系統(tǒng)界面簡(jiǎn)潔美觀、操作易上手幫助理清數(shù)據(jù)預(yù)處理、特征提取、模型訓(xùn)練與漏洞預(yù)測(cè)的完整鏈路。當(dāng)前已有133人學(xué)習(xí)瀏覽該資源整套資料在智能合約安全分析與深度學(xué)習(xí)序列建模方面具有較高實(shí)踐價(jià)值目錄結(jié)構(gòu)清晰適合作為畢業(yè)設(shè)計(jì)或課程設(shè)計(jì)的起點(diǎn)也便于后續(xù)在此基礎(chǔ)上進(jìn)行功能擴(kuò)展與改進(jìn)。1. 為什么把合約漏洞檢測(cè)改寫成「序列分類任務(wù)」別一上來(lái)就盯著合約源碼逐行讀很多漏洞并不會(huì)在單行代碼里露出破綻。真正上線跑的是編譯后的字節(jié)碼反匯編出的每一條操作碼就像指令流漏洞的特征藏在 CALL、SSTORE、BALANCE 這些關(guān)鍵操作碼的組合順序里。比如一個(gè)典型的重入漏洞往往表現(xiàn)為“余額更新操作碼出現(xiàn)在外部調(diào)用之后”這種上下文關(guān)系用人眼找容易漏用序列模型反而能穩(wěn)定總結(jié)出來(lái)。這篇筆記講的正是畢業(yè)設(shè)計(jì)里最常見(jiàn)的一條閉環(huán)基于深度學(xué)習(xí)的智能合約漏洞檢測(cè)把操作碼序列當(dāng)文本輸入用 LSTM、BLSTM、BLSTM注意力機(jī)制三種模型做分類數(shù)據(jù)準(zhǔn)備、訓(xùn)練、本地部署一條線走完。適合正在找題目、希望方案能快速出效果的同學(xué)也適合已經(jīng)有一個(gè)粗糙數(shù)據(jù)集、想換成序列模型做對(duì)比的人。2. 操作碼序列數(shù)據(jù)集的構(gòu)建從源碼到可訓(xùn)練的樣本這個(gè)方向能不能跑通七成取決于數(shù)據(jù)怎么造。模型吃的是序列原文是合約的二進(jìn)制字節(jié)碼這兩者之間還隔著一道反匯編的工序。先把這一步走穩(wěn)后面的模型才不會(huì)被垃圾輸入帶偏。2.1 從編譯產(chǎn)物到操作碼序列反匯編與詞表映射我的習(xí)慣是先積累一批合約源碼用編譯器統(tǒng)一版本編譯拿到字節(jié)碼之后反匯編成操作碼序列。為什么不用源碼直接喂模型因?yàn)樵创a經(jīng)過(guò)編譯器優(yōu)化最終執(zhí)行語(yǔ)義都在字節(jié)碼里而且不同合約的源碼風(fēng)格差異很大直接建模源碼等于讓模型先學(xué)會(huì)“讀不同人的代碼風(fēng)格”對(duì)一個(gè)小數(shù)據(jù)集來(lái)說(shuō)太奢侈。# 反匯編把十六進(jìn)制字節(jié)碼拆成一條條可讀的操作碼 from pyevmasm import disassemble_all BIN_CODE 60806040526004361061004b576000357f... # 編譯器輸出的合約字節(jié)碼 ops [str(op) for op in disassemble_all(BIN_CODE)] print(ops[:8]) # 輸出類似 [PUSH1, 0x80, PUSH1, 0x40, MSTORE, CALLVALUE, DUP1, ISZERO]這段代碼只做一件事把十六進(jìn)制字節(jié)碼變成人類能讀的操作碼助記符。disassemble_all返回的是指令對(duì)象str()轉(zhuǎn)出來(lái)的文本正好可以作為序列元素。注意我沒(méi)有保留操作數(shù)比如PUSH1 0x80只取PUSH1這個(gè)助記符原因在后面詞表小節(jié)里說(shuō)明。下一步是把操作碼文本映射成數(shù)字 ID。常見(jiàn)做法是維護(hù)一張固定詞表0留給 padding1留給未知操作碼。這樣訓(xùn)練和推理時(shí)用的是同一套映射不會(huì)出現(xiàn)訓(xùn)練集有一套詞表、現(xiàn)場(chǎng)部署又另起一套的尷尬。# 固定詞表示意0 是 padding1 是 UNK其余按 EVM 操作碼規(guī)范排 OPCODE2ID { PAD: 0, UNK: 1, STOP: 2, ADD: 3, MSTORE: 4, SSTORE: 5, CALL: 6, BALANCE: 7, JUMPI: 8, # ... 按需要覆蓋完整操作碼集合 } def ops_to_ids(op_list): return [OPCODE2ID.get(op, OPCODE2ID[UNK]) for op in op_list]這里把詞表固定下來(lái)很重要。很多初學(xué)者喜歡從訓(xùn)練集里動(dòng)態(tài)統(tǒng)計(jì)操作碼再建詞表結(jié)果遇到新的編譯器版本、新的合約寫法操作碼對(duì)不上號(hào)就崩了。固定詞表配合UNK兜底雖然單條未知指令丟失了信息但模型至少不會(huì)因?yàn)檩斎刖S度不一致直接報(bào)錯(cuò)。為什么不把PUSH的不同長(zhǎng)度當(dāng)成不同詞PUSH1、PUSH2在語(yǔ)義上確實(shí)不同但也有不少方案只保留PUSH助記符。我自己的經(jīng)驗(yàn)是PUSH 后面的立即數(shù)承載了跳轉(zhuǎn)目標(biāo)、地址等重要信息但直接把操作數(shù)當(dāng)詞表成員會(huì)讓詞表膨脹一堆無(wú)意義的數(shù)字所以更穩(wěn)妥的做法是把PUSH1、PUSH2這種帶長(zhǎng)度的助記符保留下來(lái)丟棄后面的十六進(jìn)制數(shù)值。這樣詞表仍然可控信息也損失得最少。2.2 滑窗切分與定長(zhǎng)填充窗口長(zhǎng)度怎么選一整份合約反匯編出來(lái)可能有幾百上千條操作碼不能整條喂給模型。常見(jiàn)做法是滑窗切分每個(gè)窗口就是一個(gè)訓(xùn)練樣本。窗口長(zhǎng)度太短漏洞上下文被截?cái)嗵L(zhǎng)訓(xùn)練時(shí)間和顯存翻倍。我用下來(lái)的經(jīng)驗(yàn)是128 是一個(gè)性價(jià)比比較高的起點(diǎn)50 到 200 之間都可以試但不要一上來(lái)就追 256。import torch from torch.utils.data import Dataset SEQ_LEN 128 def seq_to_ids(op_list, seq_lenSEQ_LEN): ids ops_to_ids(op_list) if len(ids) seq_len: return ids[:seq_len] return ids [OPCODE2ID[PAD]] * (seq_len - len(ids)) class OpcodeDataset(Dataset): def __init__(self, samples, labels, seq_lenSEQ_LEN): self.samples [seq_to_ids(seq, seq_len) for seq in samples] self.labels labels def __len__(self): return len(self.samples) def __getitem__(self, idx): x self.samples[idx] y self.labels[idx] return torch.tensor(x, dtypetorch.long), torch.tensor(y, dtypetorch.float)這個(gè)數(shù)據(jù)集的邏輯很簡(jiǎn)單切窗、截?cái)?、補(bǔ)齊最后返回(token_ids, labels)。__getitem__里把標(biāo)簽轉(zhuǎn)成float是為了后面用BCEWithLogitsLoss多標(biāo)簽場(chǎng)景下標(biāo)簽就是 0/1 向量不需要轉(zhuǎn)成 long?;暗牟介L(zhǎng)也是一個(gè)能影響結(jié)果的小參數(shù)。步長(zhǎng)等于窗口長(zhǎng)度時(shí)樣本之間沒(méi)有重疊數(shù)據(jù)量最小步長(zhǎng)短一些會(huì)有重疊相當(dāng)于做了數(shù)據(jù)增強(qiáng)。我一般用步長(zhǎng)等于窗口長(zhǎng)度的一半128 的窗口配 64 的步長(zhǎng)。再小就很不好看了相鄰樣本幾乎一樣訓(xùn)練集和測(cè)試集的真實(shí)差異會(huì)變小評(píng)估結(jié)果虛高。2.3 多標(biāo)簽標(biāo)注與按合約劃分防同源數(shù)據(jù)泄漏一個(gè)合約片段可能同時(shí)存在多個(gè)漏洞比如重入之外還帶未受保護(hù)的selfdestruct。標(biāo)注就不能只給一個(gè)整數(shù)類別而應(yīng)該給一個(gè)與漏洞類型數(shù)量等長(zhǎng)的 0/1 向量。7 種常見(jiàn)漏洞就是 7 維向量某一位是 1 表示這個(gè)窗口命中了對(duì)應(yīng)類別。這種多標(biāo)簽設(shè)定下?lián)p失函數(shù)必須用BCEWithLogitsLoss不能隨手拿CrossEntropyLoss頂上去否則模型只會(huì)學(xué)到“每個(gè)樣本恰好一個(gè)類”的強(qiáng)假設(shè)。數(shù)據(jù)劃分上有一條血淚經(jīng)驗(yàn)先按合約劃分再切窗順序不能反。如果先切窗再把所有窗口打亂劃分訓(xùn)練集和測(cè)試集同一個(gè)合約的相鄰窗口會(huì)同時(shí)出現(xiàn)在兩邊。這相當(dāng)于把題目和答案一起放進(jìn)考場(chǎng)序列模型對(duì)這種“鄰居片段”的識(shí)別能力又特別強(qiáng)結(jié)果就是驗(yàn)證集 F1 高得離譜到真實(shí)場(chǎng)景一測(cè)就崩。from sklearn.model_selection import train_test_split contracts list(contract_data.keys()) train_cids, test_cids train_test_split(contracts, test_size0.2, random_state42) train_samples [] train_labels [] test_samples [] test_labels [] for cid in train_cids: for seq, label in contract_data[cid][windows]: train_samples.append(seq) train_labels.append(label) for cid in test_cids: for seq, label in contract_data[cid][windows]: test_samples.append(seq) test_labels.append(label)這段代碼的關(guān)鍵點(diǎn)在于先按合約 ID 拆再取窗口。這樣測(cè)試集里出現(xiàn)的任何窗口片段都不會(huì)在訓(xùn)練集里有一個(gè)“幾乎一模一樣的鄰居”。尤其要注意的是窗口重疊生成的樣本也屬于同一個(gè)合約必須在拆分之前做好分組。3. 三種序列模型的結(jié)構(gòu)拆解LSTM、BLSTM、BLSTM注意力機(jī)制代碼層面三個(gè)模型不是三個(gè)互不相干的文件而是同一個(gè)骨架上的三種配置。我更喜歡把它們做成可以切換的實(shí)現(xiàn)這樣訓(xùn)練腳本只需要傳一個(gè)名字就能跑出三組結(jié)果。本章先把三種結(jié)構(gòu)的差異講清楚下一章直接落到可復(fù)現(xiàn)的訓(xùn)練命令。3.1 LSTM 基線門控單元怎么讀操作碼上下文LSTM 的核心單元有三道門遺忘門決定丟掉多少舊信息輸入門決定寫入多少新信息輸出門決定放出多少信息到下一步。在操作碼序列里一個(gè)CALL出現(xiàn)之后是否跟著SSTORE以及中間隔了多少步正是這種門控設(shè)計(jì)擅長(zhǎng)處理的上下文關(guān)系。LSTM 不需要像 N-gram 那樣硬掃固定窗口它能把更早之前的位置信息保留在隱狀態(tài)里所以對(duì)“前面埋了個(gè)可疑操作幾步之后才暴露”的模式更友好。import torch.nn as nn class LSTMDetector(nn.Module): def __init__(self, vocab_size, embed_dim128, hidden_dim128, num_layers2, num_labels7): super().__init__() self.embedding nn.Embedding(vocab_size, embed_dim, padding_idx0) self.lstm nn.LSTM( input_sizeembed_dim, hidden_sizehidden_dim, num_layersnum_layers, batch_firstTrue, dropout0.3 if num_layers 1 else 0, ) self.classifier nn.Sequential( nn.Dropout(0.5), nn.Linear(hidden_dim, num_labels), ) def forward(self, x): emb self.embedding(x) # (batch, seq_len, embed_dim) out, (h_n, c_n) self.lstm(emb) last_h h_n[-1] # 取最后一層的最后一個(gè)時(shí)間步 return self.classifier(last_h)h_n[-1]取的是 LSTM 最后一層最后的隱狀態(tài)。有一個(gè)容易被忽視的細(xì)節(jié)h_n的形狀是(num_layers, batch, hidden_dim)如果num_layers2h_n[0]是第一層的信息h_n[1]才是第二層的最終輸出。直接取h_n[-1]是穩(wěn)妥的。dropout只在層數(shù)大于 1 時(shí)生效這是 PyTorch 的規(guī)則單層 LSTM 傳dropout會(huì)被靜默忽略。多標(biāo)簽分類的最后一層不用加 sigmoid因?yàn)橛?xùn)練時(shí)用BCEWithLogitsLoss它內(nèi)部已經(jīng)包含 sigmoid 和數(shù)值穩(wěn)定處理推理時(shí)再單獨(dú)sigmoid。3.2 BLSTM雙向上下文為什么更適合漏洞定位LSTM 只能看到過(guò)去的信息但對(duì)操作碼序列來(lái)說(shuō)一個(gè)位置是否危險(xiǎn)往往右文比左文更關(guān)鍵。比如SSTORE之前的檢查條件寫在更早的位置而“這個(gè) CALL 是不是外部合約調(diào)用”要從后面的指令行為推斷。BLSTM 的常見(jiàn)做法是正序和反序各跑一個(gè) LSTM輸出在特征維度上拼接讓每個(gè)時(shí)間步同時(shí)擁有前后兩個(gè)方向的上下文。class BLSTMDetector(nn.Module): def __init__(self, vocab_size, embed_dim128, hidden_dim128, num_layers2, num_labels7): super().__init__() self.embedding nn.Embedding(vocab_size, embed_dim, padding_idx0) self.lstm nn.LSTM( input_sizeembed_dim, hidden_sizehidden_dim, num_layersnum_layers, batch_firstTrue, bidirectionalTrue, dropout0.3 if num_layers 1 else 0, ) self.classifier nn.Sequential( nn.Dropout(0.5), nn.Linear(hidden_dim * 2, num_labels), ) def forward(self, x): emb self.embedding(x) out, _ self.lstm(emb) # (batch, seq_len, hidden_dim * 2) batch_size, seq_len, hidden_x2 out.shape hidden_dim hidden_x2 // 2 # 正序取序列最后一步反序取序列第一步拼出一個(gè)整體 last_forward out[:, -1, :hidden_dim] last_backward out[:, 0, hidden_dim:] feat torch.cat([last_forward, last_backward], dim-1) return self.classifier(feat)這里沒(méi)有直接取h_n[-1]而是把輸出張量按位置切開(kāi)拼。原因在于雙向模型里正向和反向的“最后一步”不是同一個(gè)時(shí)間點(diǎn)正向的最后一步是序列末尾反向的最后一步其實(shí)是序列開(kāi)頭。如果直接拿h_n[-1]乍一看維度是對(duì)的但內(nèi)部等于是把兩個(gè)方向各自最后一個(gè)時(shí)間點(diǎn)的信息拼起來(lái)語(yǔ)義上容易混淆。out[:, -1, :hidden_dim]取正序在序列末尾的隱狀態(tài)out[:, 0, hidden_dim:]取反序在序列開(kāi)頭的隱狀態(tài)這樣拼出來(lái)的特征才各自覆蓋了兩個(gè)方向的終點(diǎn)信息。雙向帶來(lái)的代價(jià)是顯存幾乎翻倍、訓(xùn)練時(shí)間接近兩倍。如果你的實(shí)驗(yàn) GPU 只有 4G 顯存建議先跑 LSTM 基線把流程打通再切 BLSTM 調(diào)參。3.3 注意力機(jī)制讓模型自己指出關(guān)鍵操作碼BLSTM 最后只用兩個(gè)隱狀態(tài)代表整條序列信息量還是太“擠”。注意力機(jī)制可以看作是對(duì)全部時(shí)間步做加權(quán)求和權(quán)重由模型自己學(xué)哪幾個(gè)操作碼對(duì)判定漏洞貢獻(xiàn)大權(quán)重就高。這一步不只提升 F1更大的價(jià)值是可解釋性——拿到注意力權(quán)重你能告訴別人模型是因?yàn)槟臈lCALL或SSTORE判了重入漏洞。class AttentionPooling(nn.Module): def __init__(self, hidden_dim): super().__init__() self.W nn.Linear(hidden_dim, hidden_dim, biasFalse) self.v nn.Linear(hidden_dim, 1, biasFalse) def forward(self, lstm_out, maskNone): # lstm_out: (batch, seq_len, hidden_dim) energy torch.tanh(self.W(lstm_out)) scores self.v(energy).squeeze(-1) # (batch, seq_len) if mask is not None: scores scores.masked_fill(mask 0, float(-inf)) attn_weights torch.softmax(scores, dim-1) # (batch, seq_len) context (attn_weights.unsqueeze(-1) * lstm_out).sum(dim1) return context, attn_weights注意masked_fill(mask 0, float(-inf))這一行。padding 位置的注意力分?jǐn)?shù)不能參與 softmax否則模型會(huì)把“什么都沒(méi)有”的位置當(dāng)成重要特征。這個(gè) mask 可以直接從輸入 token 里生成mask (x ! PAD_IDX)。有了AttentionPoolingBLSTM注意力模型的結(jié)構(gòu)就很清晰了class BLSTMAttnDetector(nn.Module): def __init__(self, vocab_size, embed_dim128, hidden_dim128, num_layers2, num_labels7): super().__init__() self.embedding nn.Embedding(vocab_size, embed_dim, padding_idx0) self.lstm nn.LSTM(...) # 配置與 BLSTMDetector 相同bidirectionalTrue self.attn AttentionPooling(hidden_dim * 2) self.classifier nn.Sequential( nn.Dropout(0.5), nn.Linear(hidden_dim * 2, num_labels), ) def forward(self, x): mask (x ! PAD_IDX).unsqueeze(-1) emb self.embedding(x) out, _ self.lstm(emb) # (batch, seq_len, hidden_dim * 2) context, attn_weights self.attn(out, mask.squeeze(-1)) return self.classifier(context), attn_weightsAttentionPooling的hidden_dim傳入的是hidden_dim * 2因?yàn)?BLSTM 的輸出特征翻倍了。推理時(shí)除了 logits 還會(huì)拿到attn_weights這個(gè)輸出保留下來(lái)后面做漏洞定位解釋時(shí)有用。4. 本地訓(xùn)練與結(jié)果對(duì)比穩(wěn)定復(fù)現(xiàn)的一套參數(shù)模型結(jié)構(gòu)定下來(lái)之后訓(xùn)練環(huán)節(jié)要解決的是“怎么省事地跑出三組結(jié)果”。我通常的做法是寫一個(gè)帶模型名字參數(shù)的單文件訓(xùn)練腳本三組實(shí)驗(yàn)的隨機(jī)種子、數(shù)據(jù)加載、優(yōu)化器設(shè)置完全一致然后只改模型名。4.1 數(shù)據(jù)加載器與訓(xùn)練循環(huán)的最小骨架三個(gè)模型只有一個(gè)構(gòu)造函數(shù)的區(qū)別可以用一個(gè)工廠函數(shù)統(tǒng)一創(chuàng)建避免在三份腳本里復(fù)制粘貼模型參數(shù)然后改漏一處。def build_model(model_name, vocab_size, num_labels): if model_name lstm: return LSTMDetector(vocab_size, 128, 128, 2, num_labels) elif model_name blstm: return BLSTMDetector(vocab_size, 128, 128, 2, num_labels) elif model_name blstm_attn: return BLSTMAttnDetector(vocab_size, 128, 128, 2, num_labels) raise ValueError(funknown model: {model_name})訓(xùn)練循環(huán)里的默認(rèn)參數(shù)我建議統(tǒng)一為batch_size64、lr1e-3、epochs30、weight_decay1e-4。不是最優(yōu)但足夠穩(wěn)先跑通再調(diào)。from torch.utils.data import DataLoader from torch.optim import AdamW import torch.nn as nn model build_model(blstm_attn, len(OPCODE2ID), 7) optimizer AdamW(model.parameters(), lr1e-3, weight_decay1e-4) criterion nn.BCEWithLogitsLoss() loader DataLoader(train_ds, batch_size64, shuffleTrue) for epoch in range(30): model.train() total_loss 0.0 for x, y in loader: optimizer.zero_grad() logits, _ model(x) loss criterion(logits, y.float()) loss.backward() # 梯度裁剪雙向模型很容易在訓(xùn)練早期梯度爆炸 nn.utils.clip_grad_norm_(model.parameters(), max_norm5.0) optimizer.step() total_loss loss.item() print(fepoch {epoch:02d} loss {total_loss / len(loader):.4f})這里有個(gè)細(xì)節(jié)BLSTM注意力模型返回的是(logits, attn_weights)所以訓(xùn)練時(shí)要用logits, _ model(x)解包。如果只寫logits model(x)而不解包PyTorch 在反向傳播時(shí)會(huì)把a(bǔ)ttn_weights也當(dāng)作輸出張量雖然不一定直接報(bào)錯(cuò)但損失計(jì)算和梯度傳播的邊界會(huì)變得很混亂。梯度裁剪clip_grad_norm_是雙向模型里很實(shí)用的技巧。BLSTM 反向傳播路徑變長(zhǎng)早期梯度容易過(guò)大導(dǎo)致 loss 直接變成nan。設(shè)置max_norm5.0能避免大部分翻車但不建議設(shè)太小太小會(huì)讓模型收斂變慢。4.2 三個(gè)必調(diào)參數(shù)與經(jīng)驗(yàn)范圍參數(shù)參考值我的判斷seq_len128從 64 加到 128 效果提升明顯從 128 加到 256 收益遞減訓(xùn)練時(shí)間接近翻倍hidden_dim128數(shù)據(jù)集在幾千個(gè)樣本這個(gè)規(guī)模時(shí)256 的 hidden_dim 容易過(guò)擬合提升有限lr1e-33e-4 到 2e-3 都是可以接受的范圍建議先用 1e-3 跑 10 個(gè) epoch 看 loss 是否持續(xù)下降batch_size不算最敏感的參數(shù)但要注意它和顯存的關(guān)系。BLSTM 在 128 序列長(zhǎng)度下batch_size64大概需要 6G 左右的顯存如果把序列長(zhǎng)度調(diào)到 256batch_size不減半的話很容易 OOM。顯存不足時(shí)優(yōu)先降batch_size不要急著降序列長(zhǎng)度因?yàn)樾蛄虚L(zhǎng)度對(duì)效果的貢獻(xiàn)通常更大。4.3 評(píng)估指標(biāo)與三組實(shí)驗(yàn)怎么對(duì)比多標(biāo)簽場(chǎng)景里準(zhǔn)確率基本沒(méi)有參考價(jià)值原因是類別不平衡很常見(jiàn)模型只需要把多數(shù)類全預(yù)測(cè)為 0準(zhǔn)確率就能到九成以上。我更推薦看 macro-F1它給每個(gè)類別相同的權(quán)重。precision 和 recall 也要一起看這兩個(gè)指標(biāo)一個(gè)偏高一個(gè)偏低時(shí)說(shuō)明模型在“寧可錯(cuò)殺”和“盡量不漏”之間需要調(diào)閾值。from sklearn.metrics import f1_score, precision_score, recall_score model.eval() all_preds, all_labels [], [] with torch.no_grad(): for x, y in loader_test: logits, _ model(x) preds (torch.sigmoid(logits) 0.5).int() all_preds.append(preds.cpu()) all_labels.append(y.cpu()) all_preds torch.cat(all_preds).numpy() all_labels torch.cat(all_labels).numpy() print(macro-F1, f1_score(all_labels, all_preds, averagemacro)) print(precision, precision_score(all_labels, all_preds, averagemacro)) print(recall, recall_score(all_labels, all_preds, averagemacro))三組實(shí)驗(yàn)跑出來(lái)的趨勢(shì)通常是LSTM 作為基線能到七八成的 macro-F1BLSTM 在 LSTM 基礎(chǔ)上有幾個(gè)點(diǎn)的提升主要提升在那些依賴后續(xù)上下文的漏洞類別上BLSTM注意力機(jī)制再往上走一點(diǎn)并且給出可解釋的注意力分?jǐn)?shù)。如果 BLSTM 沒(méi)有明顯提升 LSTM先檢查數(shù)據(jù)劃分和窗口長(zhǎng)度不用急著懷疑模型結(jié)構(gòu)。4.4 保存最優(yōu)模型模型參數(shù)和元信息一起存訓(xùn)練結(jié)束后要保存的不只是模型權(quán)重還有詞表、標(biāo)簽名、模型名。單條推理時(shí)這三樣缺一個(gè)都會(huì)卡殼。我是把一個(gè) dict 整體存到torch.save里。torch.save({ model_name: blstm_attn, model_state_dict: model.state_dict(), opcode2id: OPCODE2ID, label_names: label_names, }, best_model.pt)model.state_dict()只存參數(shù)不存結(jié)構(gòu)所以必須帶上model_name加載時(shí)才能重建出相同的模型結(jié)構(gòu)。opcode2id和label_names這兩樣經(jīng)常有人漏導(dǎo)致部署階段寫了一大堆臨時(shí)兼容代碼。加載時(shí)代碼是反過(guò)來(lái)的checkpoint torch.load(best_model.pt, map_locationcpu) model build_model(checkpoint[model_name], len(checkpoint[opcode2id]), len(checkpoint[label_names])) model.load_state_dict(checkpoint[model_state_dict]) model.eval()為什么要加map_locationcpu如果 checkpoint 是在 GPU 上存的直接torch.load到無(wú) GPU 的機(jī)器會(huì)報(bào)錯(cuò)或者強(qiáng)行占用顯存。加了map_location之后同一份訓(xùn)練產(chǎn)物既能在帶顯卡的機(jī)器上繼續(xù)訓(xùn)練也能在驗(yàn)證機(jī)上做純推理。5. 避坑與排查五個(gè)讓結(jié)果翻車的真實(shí)問(wèn)題這個(gè)方向看起來(lái)流程短實(shí)際上每個(gè)環(huán)節(jié)都有暗坑。下面的五條是我自己跑這個(gè)方案時(shí)踩過(guò)的按影響程度排序每一條我都給出“現(xiàn)象、原因、解決”三段式記錄方便你對(duì)照自己的情況排查。5.1 數(shù)據(jù)泄漏窗口切完才按窗口分割現(xiàn)象驗(yàn)證集 macro-F1 高到 0.9 以上自己都覺(jué)得離譜換一批新合約一測(cè)直接跌回 0.6。原因我最早是先把所有合約的窗口全部切出來(lái)打亂之后再劃分訓(xùn)練集和測(cè)試集。由于滑窗有重疊同一個(gè)合約的相鄰窗口內(nèi)容高度相似模型等于在訓(xùn)練時(shí)見(jiàn)過(guò)測(cè)試樣本的“鄰居片段”自然學(xué)到的不是泛化特征而是對(duì)特定字節(jié)碼片段的記憶。解決改成先按合約 ID 劃分再切窗。這一步改動(dòng)之后驗(yàn)證集數(shù)字立刻變得“難看”許多但才是真實(shí)水平。5.2 詞表前后不一致新測(cè)試樣本直接 KeyError現(xiàn)象訓(xùn)練一切正常換一批測(cè)試合約后推理代碼在OPCODE2ID.get(op)處崩潰報(bào)出某個(gè) PyTorch 不認(rèn)識(shí)的字符串。原因我最初是直接對(duì)訓(xùn)練集做set()構(gòu)建詞表訓(xùn)練集里見(jiàn)過(guò)的操作碼有限測(cè)試集里出現(xiàn)了一個(gè)從未見(jiàn)過(guò)的操作碼詞表查不到就炸了。更隱蔽的是同一個(gè)操作碼在不同 solc 版本里的助記符可能略有差異。解決固定一張按操作碼規(guī)范生成的完整詞表并給所有映射不到的詞一個(gè)UNK兜底。不要信任動(dòng)態(tài)詞表所有代碼里唯一允許出現(xiàn)dict.get默認(rèn)值的地方就在這里。5.3 類別不平衡準(zhǔn)確率虛高關(guān)鍵漏洞幾乎全漏現(xiàn)象總準(zhǔn)確率顯示 0.9但打開(kāi)混淆矩陣發(fā)現(xiàn)某一個(gè)冷門漏洞類別的 recall 是 0模型把所有樣本都預(yù)測(cè)成了“無(wú)漏洞”。原因數(shù)據(jù)集中無(wú)漏洞樣本或常見(jiàn)漏洞占了壓倒比例模型學(xué)會(huì)無(wú)腦預(yù)測(cè)多數(shù)類就能拿到高準(zhǔn)確率。交叉熵或普通 BCE 對(duì)少數(shù)類沒(méi)有額外關(guān)注梯度完全被多數(shù)類主導(dǎo)。解決換成 macro-F1 做評(píng)估指標(biāo)同時(shí)在損失函數(shù)里給少數(shù)類更高的權(quán)重。PyTorch 的BCEWithLogitsLoss有pos_weight參數(shù)可以直接按類別頻率反比設(shè)置。# 假設(shè) class_freq 是各漏洞類別在訓(xùn)練集里的占比列表 pos_weight torch.tensor([1.0 / f for f in class_freq]) criterion nn.BCEWithLogitsLoss(pos_weightpos_weight)pos_weight的取值邏輯很簡(jiǎn)單某個(gè)類別出現(xiàn)得越少權(quán)重越大。比如一個(gè)類別只占 5%那它的權(quán)重就是 20。這會(huì)逼迫模型在梯度更新時(shí)把注意力分給少數(shù)類。但注意不要把這個(gè)值設(shè)得過(guò)大否則模型會(huì)走向另一個(gè)極端把大量正常合約誤報(bào)成漏洞precision 反而崩掉。5.4 序列長(zhǎng)度和批量大小互相打架顯存 OOM 頻繁現(xiàn)象把seq_len從 128 調(diào)到 256 之后原來(lái)穩(wěn)定的batch_size64在訓(xùn)練中途報(bào) CUDA out of memory。原因雙向 LSTM 需要保存前向后向兩條傳播路徑的中間結(jié)果序列長(zhǎng)度翻倍顯存占用不止翻倍注意力機(jī)制還要保存所有時(shí)間步的打分矩陣顯存壓力進(jìn)一步增大。解決先降batch_size到 32 或 16確認(rèn)能跑通再觀察顯存利用率。如果仍需 256 的序列長(zhǎng)度可以先用 128 調(diào)好其他參數(shù)最后再拉長(zhǎng)序列做終版實(shí)驗(yàn)避免每次調(diào)參都在等 OOM 后重啟。5.5 固定隨機(jī)種子后結(jié)果仍然不穩(wěn)定現(xiàn)象設(shè)置了random.seed、np.random.seed連續(xù)兩次跑同一組實(shí)驗(yàn)loss 曲線不完全一致。原因PyTorch 的 GPU 算子本身有非確定性DataLoader 的 shuffle 也會(huì)引入隨機(jī)性只固定 Python 和 NumPy 的種子覆蓋不到這些環(huán)節(jié)。解決使用一個(gè)完整的種子設(shè)置函數(shù)并把 DataLoader 的隨機(jī)源也固定下來(lái)。import random import numpy as np import torch def set_seed(seed42): random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) torch.cuda.manual_seed_all(seed) torch.backends.cudnn.deterministic True torch.backends.cudnn.benchmark False如果你的模型對(duì)數(shù)字極其敏感還要把 DataLoader 的generator一起固定。torch.Generator()傳入一個(gè)固定種子的示例g torch.Generator() g.manual_seed(42) loader DataLoader(train_ds, batch_size64, shuffleTrue, generatorg)注意一點(diǎn)cudnn.deterministic True會(huì)犧牲一部分運(yùn)行速度只建議在復(fù)現(xiàn)最終結(jié)果時(shí)開(kāi)啟日常調(diào)參可以關(guān)掉讓訓(xùn)練跑得更快。6. 把模型部署成本地推理從 checkpoint 到單條字節(jié)碼預(yù)測(cè)訓(xùn)練只是這個(gè)方向的一半另一半是把模型真正用起來(lái)。一個(gè)模型如果不部署永遠(yuǎn)只是實(shí)驗(yàn)報(bào)告里的數(shù)字。本章給出一個(gè)能在本地直接運(yùn)行的推理鏈路從加載模型到輸出漏洞概率再到解釋推理依據(jù)。6.1 用 checkpoint 恢復(fù)模型并做單條預(yù)測(cè)推理腳本的模型構(gòu)建必須和訓(xùn)練時(shí)完全一致。我已經(jīng)把opcode2id、label_names保存在 checkpoint 里加載時(shí)直接取用不寫任何硬編碼的維度數(shù)字。checkpoint torch.load(best_model.pt, map_locationcpu) model build_model( checkpoint[model_name], len(checkpoint[opcode2id]), len(checkpoint[label_names]), ) model.load_state_dict(checkpoint[model_state_dict]) model.eval() def predict_bin(bin_code, seq_len128, threshold0.5): ops [str(op) for op in disassemble_all(bin_code)] ids [OPCODE2ID.get(op, OPCODE2ID[UNK]) for op in ops] ids ids[:seq_len] [OPCODE2ID[PAD]] * max(0, seq_len - len(ids)) x torch.tensor([ids], dtypetorch.long) with torch.no_grad(): logits, attn_weights model(x) probs torch.sigmoid(logits).numpy()[0] result {} for i, label in enumerate(checkpoint[label_names]): result[label] float(probs[i]) return result, attn_weights單條推理的關(guān)鍵是全程不計(jì)算梯度用torch.no_grad()包住前向傳播顯存占用會(huì)小一個(gè)量級(jí)。推理速度一般在毫秒到十幾毫秒取決于序列長(zhǎng)度和模型維度本地部署完全夠用。6.2 閾值不必固定 0.5注意力權(quán)重是解釋入口多標(biāo)簽輸出經(jīng)過(guò) sigmoid 之后每個(gè)類別的概率是一個(gè) 0 到 1 之間的數(shù)字。默認(rèn)取 0.5 作為判定閾值但冷門類別可能因?yàn)闃颖旧俣w概率偏低此時(shí)可以考慮把閾值降到 0.35 或 0.4 來(lái)提升召回反過(guò)來(lái)如果誤報(bào)太多就調(diào)高閾值。這個(gè)閾值應(yīng)該放在測(cè)試集上試幾組值選 macro-F1 最高的那個(gè)而不是拍腦袋定。注意力權(quán)重是這個(gè)模型比普通 BLSTM 多出來(lái)的紅利。推理時(shí)我已經(jīng)把a(bǔ)ttn_weights帶出來(lái)了可以找到權(quán)重最大的那幾條操作碼告訴使用者模型依據(jù)什么做的判斷。top_idx attn_weights[0].argmax().item() print(f模型最關(guān)注的位置: seq {top_idx}, opcode {ops[top_idx] if top_idx len(ops) else PAD})如果模型在CALL附近給了一個(gè)明顯高于周圍的權(quán)重那這個(gè)判斷就有了可解釋的依據(jù)。這一點(diǎn)在畢業(yè)設(shè)計(jì)答辯時(shí)非常加分也是我推薦做注意力模型而不是只做 LSTM 的最重要原因。6.3 這個(gè)方案值不值得投入以及往深走的三個(gè)方向作為畢業(yè)設(shè)計(jì)這個(gè)方向性價(jià)比很高數(shù)據(jù)構(gòu)造方法清晰、三種模型對(duì)比有梯度、指標(biāo)設(shè)計(jì)有講究、部署落地不復(fù)雜足夠撐起一個(gè)完整的技術(shù)報(bào)告。但也要誠(chéng)實(shí)說(shuō)清楚它的邊界基于操作碼序列的模型更像一個(gè)“模式識(shí)別器”能發(fā)現(xiàn)和訓(xùn)練集中已知漏洞相似的模式但它并不理解操作碼背后的語(yǔ)義遇到全新類型的漏洞時(shí)它的召回能力有限。想把這個(gè)方案做得更扎實(shí)有三個(gè)方向值得投入第一個(gè)是把字節(jié)碼層面的純序列建模升級(jí)為融合語(yǔ)義特征的多輸入模型比如把靜態(tài)分析器輸出的“調(diào)用圖結(jié)構(gòu)”和操作碼序列并行輸入讓模型不只看到指令順序還看到函數(shù)之間的調(diào)用關(guān)系。第二個(gè)是利用注意力權(quán)重做漏洞定位的可視化把模型認(rèn)為危險(xiǎn)的操作碼高亮出來(lái)再結(jié)合人工審計(jì)確認(rèn)這樣模型的輸出就不只是一堆概率而是可以直接進(jìn)入審計(jì)工作流的提示。第三個(gè)是預(yù)訓(xùn)練方向用大量未標(biāo)注的合約操作碼序列先預(yù)訓(xùn)練一個(gè)序列編碼器再在小規(guī)模標(biāo)注數(shù)據(jù)上微調(diào)原理類似文本領(lǐng)域的做法。這條路能明顯緩解標(biāo)注樣本不足的問(wèn)題但訓(xùn)練成本和數(shù)據(jù)收集成本都比較高適合作為后續(xù)延伸研究。我自己把這個(gè)方案從數(shù)據(jù)處理到部署完整跑過(guò)一遍之后最大的體會(huì)是數(shù)據(jù)劃分和詞表一致性這兩件事比調(diào)模型結(jié)構(gòu)參數(shù)更影響最終結(jié)果。很多同學(xué)一上來(lái)就折騰模型層數(shù)結(jié)果數(shù)據(jù)泄漏導(dǎo)致指標(biāo)虛高而不自知白白消耗了時(shí)間。走這個(gè)方向先把數(shù)據(jù)鏈路做干凈模型部分反而是一個(gè)按部就班就能出結(jié)果的過(guò)程。希望這篇筆記能幫你少走幾步彎路。本文還有配套的精品資源點(diǎn)擊獲取