習(xí)的智能合約漏洞檢測(cè)：操作碼序列分類(lèi)實(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)行。隨附源碼與部署教程可讓用戶(hù)直接啟動(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è)改寫(xiě)成「序列分類(lèi)任務(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ī)制三種模型做分類(lèi)數(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]) # 輸出類(lèi)似 [PUSH1, 0x80, PUSH1, 0x40, MSTORE, CALLVALUE, DUP1, ISZERO]這段代碼只做一件事把十六進(jìn)制字節(jié)碼變成人類(lèi)能讀的操作碼助記符。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é)果遇到新的編譯器版本、新的合約寫(xiě)法操作碼對(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è)性?xún)r(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。滑窗的步長(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ù)類(lèi)別而應(yīng)該給一個(gè)與漏洞類(lèi)型數(shù)量等長(zhǎng)的 0/1 向量。7 種常見(jiàn)漏洞就是 7 維向量某一位是 1 表示這個(gè)窗口命中了對(duì)應(yīng)類(lèi)別。這種多標(biāo)簽設(shè)定下?lián)p失函數(shù)必須用BCEWithLogitsLoss不能隨手拿CrossEntropyLoss頂上去否則模型只會(huì)學(xué)到“每個(gè)樣本恰好一個(gè)類(lèi)”的強(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 基線門(mén)控單元怎么讀操作碼上下文LSTM 的核心單元有三道門(mén)遺忘門(mén)決定丟掉多少舊信息輸入門(mén)決定寫(xiě)入多少新信息輸出門(mén)決定放出多少信息到下一步。在操作碼序列里一個(gè)CALL出現(xiàn)之后是否跟著SSTORE以及中間隔了多少步正是這種門(mén)控設(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)簽分類(lèi)的最后一層不用加 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之前的檢查條件寫(xiě)在更早的位置而“這個(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é)果”。我通常的做法是寫(xiě)一個(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)解包。如果只寫(xiě)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能避免大部分翻車(chē)但不建議設(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à)值原因是類(lèi)別不平衡很常見(jiàn)模型只需要把多數(shù)類(lèi)全預(yù)測(cè)為 0準(zhǔn)確率就能到九成以上。我更推薦看 macro-F1它給每個(gè)類(lèi)別相同的權(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)的提升主要提升在那些依賴(lài)后續(xù)上下文的漏洞類(lèi)別上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)致部署階段寫(xiě)了一大堆臨時(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é)果翻車(chē)的真實(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 類(lèi)別不平衡準(zhǔn)確率虛高關(guān)鍵漏洞幾乎全漏現(xiàn)象總準(zhǔn)確率顯示 0.9但打開(kāi)混淆矩陣發(fā)現(xiàn)某一個(gè)冷門(mén)漏洞類(lèi)別的 recall 是 0模型把所有樣本都預(yù)測(cè)成了“無(wú)漏洞”。原因數(shù)據(jù)集中無(wú)漏洞樣本或常見(jiàn)漏洞占了壓倒比例模型學(xué)會(huì)無(wú)腦預(yù)測(cè)多數(shù)類(lèi)就能拿到高準(zhǔn)確率。交叉熵或普通 BCE 對(duì)少數(shù)類(lèi)沒(méi)有額外關(guān)注梯度完全被多數(shù)類(lèi)主導(dǎo)。解決換成 macro-F1 做評(píng)估指標(biāo)同時(shí)在損失函數(shù)里給少數(shù)類(lèi)更高的權(quán)重。PyTorch 的BCEWithLogitsLoss有pos_weight參數(shù)可以直接按類(lèi)別頻率反比設(shè)置。# 假設(shè) class_freq 是各漏洞類(lèi)別在訓(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è)類(lèi)別出現(xiàn)得越少權(quán)重越大。比如一個(gè)類(lèi)別只占 5%那它的權(quán)重就是 20。這會(huì)逼迫模型在梯度更新時(shí)把注意力分給少數(shù)類(lèi)。但注意不要把這個(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í)直接取用不寫(xiě)任何硬編碼的維度數(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è)類(lèi)別的概率是一個(gè) 0 到 1 之間的數(shù)字。默認(rèn)取 0.5 作為判定閾值但冷門(mén)類(lèi)別可能因?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è)明顯高于周?chē)臋?quán)重那這個(gè)判斷就有了可解釋的依據(jù)。這一點(diǎn)在畢業(yè)設(shè)計(jì)答辯時(shí)非常加分也是我推薦做注意力模型而不是只做 LSTM 的最重要原因。6.3 這個(gè)方案值不值得投入以及往深走的三個(gè)方向作為畢業(yè)設(shè)計(jì)這個(gè)方向性?xún)r(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ǔ)義遇到全新類(lèi)型的漏洞時(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èi)似文本領(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)擊獲取 基于深度學(xué)習(xí)的智能合約漏洞檢測(cè)：操作碼序列分類(lèi)實(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)行。隨附源碼與部署教程可讓用戶(hù)直接啟動(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è)改寫(xiě)成「序列分類(lèi)任務(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ī)制三種模型做分類(lèi)數(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]) # 輸出類(lèi)似 [PUSH1, 0x80, PUSH1, 0x40, MSTORE, CALLVALUE, DUP1, ISZERO]這段代碼只做一件事把十六進(jìn)制字節(jié)碼變成人類(lèi)能讀的操作碼助記符。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é)果遇到新的編譯器版本、新的合約寫(xiě)法操作碼對(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è)性?xún)r(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。滑窗的步長(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ù)類(lèi)別而應(yīng)該給一個(gè)與漏洞類(lèi)型數(shù)量等長(zhǎng)的 0/1 向量。7 種常見(jiàn)漏洞就是 7 維向量某一位是 1 表示這個(gè)窗口命中了對(duì)應(yīng)類(lèi)別。這種多標(biāo)簽設(shè)定下?lián)p失函數(shù)必須用BCEWithLogitsLoss不能隨手拿CrossEntropyLoss頂上去否則模型只會(huì)學(xué)到“每個(gè)樣本恰好一個(gè)類(lèi)”的強(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 基線門(mén)控單元怎么讀操作碼上下文LSTM 的核心單元有三道門(mén)遺忘門(mén)決定丟掉多少舊信息輸入門(mén)決定寫(xiě)入多少新信息輸出門(mén)決定放出多少信息到下一步。在操作碼序列里一個(gè)CALL出現(xiàn)之后是否跟著SSTORE以及中間隔了多少步正是這種門(mén)控設(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)簽分類(lèi)的最后一層不用加 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之前的檢查條件寫(xiě)在更早的位置而“這個(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é)果”。我通常的做法是寫(xiě)一個(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)解包。如果只寫(xiě)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能避免大部分翻車(chē)但不建議設(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à)值原因是類(lèi)別不平衡很常見(jiàn)模型只需要把多數(shù)類(lèi)全預(yù)測(cè)為 0準(zhǔn)確率就能到九成以上。我更推薦看 macro-F1它給每個(gè)類(lèi)別相同的權(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)的提升主要提升在那些依賴(lài)后續(xù)上下文的漏洞類(lèi)別上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)致部署階段寫(xiě)了一大堆臨時(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é)果翻車(chē)的真實(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 類(lèi)別不平衡準(zhǔn)確率虛高關(guān)鍵漏洞幾乎全漏現(xiàn)象總準(zhǔn)確率顯示 0.9但打開(kāi)混淆矩陣發(fā)現(xiàn)某一個(gè)冷門(mén)漏洞類(lèi)別的 recall 是 0模型把所有樣本都預(yù)測(cè)成了“無(wú)漏洞”。原因數(shù)據(jù)集中無(wú)漏洞樣本或常見(jiàn)漏洞占了壓倒比例模型學(xué)會(huì)無(wú)腦預(yù)測(cè)多數(shù)類(lèi)就能拿到高準(zhǔn)確率。交叉熵或普通 BCE 對(duì)少數(shù)類(lèi)沒(méi)有額外關(guān)注梯度完全被多數(shù)類(lèi)主導(dǎo)。解決換成 macro-F1 做評(píng)估指標(biāo)同時(shí)在損失函數(shù)里給少數(shù)類(lèi)更高的權(quán)重。PyTorch 的BCEWithLogitsLoss有pos_weight參數(shù)可以直接按類(lèi)別頻率反比設(shè)置。# 假設(shè) class_freq 是各漏洞類(lèi)別在訓(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è)類(lèi)別出現(xiàn)得越少權(quán)重越大。比如一個(gè)類(lèi)別只占 5%那它的權(quán)重就是 20。這會(huì)逼迫模型在梯度更新時(shí)把注意力分給少數(shù)類(lèi)。但注意不要把這個(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í)直接取用不寫(xiě)任何硬編碼的維度數(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è)類(lèi)別的概率是一個(gè) 0 到 1 之間的數(shù)字。默認(rèn)取 0.5 作為判定閾值但冷門(mén)類(lèi)別可能因?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è)明顯高于周?chē)臋?quán)重那這個(gè)判斷就有了可解釋的依據(jù)。這一點(diǎn)在畢業(yè)設(shè)計(jì)答辯時(shí)非常加分也是我推薦做注意力模型而不是只做 LSTM 的最重要原因。6.3 這個(gè)方案值不值得投入以及往深走的三個(gè)方向作為畢業(yè)設(shè)計(jì)這個(gè)方向性?xún)r(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ǔ)義遇到全新類(lèi)型的漏洞時(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èi)似文本領(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)擊獲取