TrafficLLM的儲存庫是一個通用的 LLM 適應框架,用於學習現實場景中所有開源 LLM 的穩健流量表示,並增強跨不同流量分析任務的泛化能力。
注意:此程式碼基於 ChatGLM2 和 Llama2。非常感謝作者。
[2024.10.28]我們更新了使用GLM4建構TrafficLLM的適配程式碼,它比ChatGLM2具有更快的調優和推理速度。請造訪 Adapt2GLM4 以了解更多詳細資訊。
TrafficLLM建立在使用自然語言和流量資料的複雜微調框架之上,它提出了以下技術來增強大型語言模型在網路流量分析中的實用性。
流量域標記化。為了克服自然語言和異質流量資料之間的模態差距,TrafficLLM 引入了流量域標記化來處理流量偵測和生成任務的不同輸入,以實現 LLM 適應。該機制透過在大規模流量域語料庫上訓練專門的標記化模型,有效地擴展了 LLM 的原生標記器。
雙級調整管道。 TrafficLLM 採用雙階段調整管道來實現 LLM 跨不同流量域任務的穩健表示學習。此管道訓練 LLM 理解指令並學習不同階段與任務相關的流量模式,這建立在 TrafficLLM 任務理解和流量推理能力的基礎上,用於各種流量檢測和生成任務。
具有參數有效微調的可擴展適應(EA-PEFT)。為了使 LLM 泛化到新的交通環境中,TrafficLLM 提出了一種可擴展的參數有效微調自適應方法 (EA-PEFT),以低開銷更新模型參數。該技術將模型功能拆分為不同的 PEFT 模型,這有助於最大限度地降低因流量模式變化而產生的動態場景的成本。
我們發布了 TrafficLLM 的訓練資料集,其中包含超過 0.4M 的流量資料和 9K 的人工指令,用於 LLM 適應不同的流量分析任務。
Instruction Datasets
:指令資料集用於幫助LLM學習流量偵測或產生任務的領域知識,並了解在不同場景下應該執行哪些任務。
Traffic Datasets
:流量資料集包含我們從公共流量資料集中提取的流量調整數據,這有助於LLM了解不同下游任務中的流量模式。
為了建構自然語言語料庫作為 TrafficLLM 中的人類指令,我們收集了由專家和人工智慧助理監督的 9,209 條特定任務指令。統計結果如下:
主流任務 | 下游任務 | 簡稱。 | #樣本 |
---|---|---|---|
流量偵測 | 惡意軟體流量偵測 | MTD | 1.0K |
殭屍網路偵測 | 聯邦國防軍 | 1.1K | |
惡意 DoH 檢測 | MDD | 0.6K | |
網路攻擊偵測 | 華德 | 0.6K | |
APT攻擊偵測 | 亞德 | 0.6K | |
加密 VPN 偵測 | 伊波拉病毒病 | 1.2K | |
Tor 行為偵測 | 待定 | 0.6K | |
加密應用分類 | 歐洲聯盟 | 0.6K | |
網站指紋識別 | WF | 0.6K | |
概念漂移 | 光碟 | 0.6K | |
流量生成 | 惡意軟體流量生成 | MTG | 0.6K |
殭屍網路流量生成 | BTG | 0.1K | |
加密 VPN 生成 | EVG | 0.4K | |
加密應用程式生成 | EAG | 0.6K |
為了評估 TrafficLLM 在各種網路情境下的效能,我們從公開的流量資料集中提取了超過 0.4M 的調優數據,以衡量 TrafficLLM 偵測或產生惡意和良性流量的能力。統計結果如下:
數據集 | 任務 | 簡稱。 | #樣本 |
---|---|---|---|
中科大TFC 2016 | 惡意軟體流量偵測 | MTD | 50.7K |
ISCX 殭屍網路 2014 | 殭屍網路偵測 | 聯邦國防軍 | 25.0K |
2020 年 DHBrw | 惡意 DoH 檢測 | MDD | 47.8K |
中船重工2010 | 網路攻擊偵測 | 華德 | 34.5K |
2020年DAPT | APT攻擊偵測 | 亞德 | 10.0K |
ISCX VPN 2016 | 加密 VPN 偵測 | 伊波拉病毒病 | 64.8K |
ISCX 托爾 2016 | Tor 行為偵測 | 待定 | 40.0K |
科技網2023 | 加密應用分類 | 歐洲聯盟 | 97.6K |
CW-100 2018 | 網站指紋識別 | WF | 7.4K |
APP-53 2023 | 概念漂移 | 光碟 | 109.8K |
1. 環境準備 2. 訓練 TrafficLLM 2.1.準備預訓練檢查點 2.2。 2.3. 預處理資料集 訓練流量域標記器(選購) 2.4。神經語言指令調整 2.5。特定任務的流量調整 2.6。 PEFT 的可擴展適應 (EA-PEFT) 3. 評估 TrafficLLM 3.1。 3.2. 準備檢查點和數據 運行評估目錄:
一、環境準備【回到頂部】
請透過執行以下命令克隆儲存庫並安裝所需的環境。
conda 創建-n Trafficllm python=3.9 conda activate Trafficllm# 複製我們的 TrafficLLMgit 克隆 https://github.com/ZGC-LLM-Safety/TrafficLLM.gitcd TrafficLLM# 安裝所需的庫pip install -r requests.
TrafficLLM 採用三種核心技術:處理指令和流量資料的流量域標記化、用於理解文字語義和學習跨不同任務的流量模式的雙階段調整管道以及用於更新的EA -PEFT新場景適應的模型參數。
TrafficLLM 是基於現有開源 LLM 進行訓練的。請依照說明準備檢查點。
ChatGLM2
:準備基本模型ChatGLM,這是一個具有輕型部署需求的開源LLM。請在此下載其權重。我們通常使用具有 6B 參數的 v2 模型。
Other LLMs
:要讓其他LLM適應流量分析任務,您可以重複使用儲存庫中的訓練資料並根據官方說明修改其訓練腳本。例如,Llama2 需要在配置中註冊新資料集。
為了從原始流量資料集中提取適合 LLM 學習的訓練數據,我們設計了專門的提取器來預處理不同任務的流量資料集。預處理程式碼包含以下參數進行配置。
input
:原始流量資料集路徑(包含標記子目錄的主目錄路徑。每個標記子目錄包含要預處理的原始 .pcap 檔案)。
dataset_name
:原始流量資料集名稱(有助於判斷該名稱是否已在 TrafficLLM 的程式碼中註冊)。
traffic_task
:偵測任務或產生任務。
granularity
:資料包級或流級粒度。
output_path
:輸出訓練資料集路徑。
output_name
:輸出訓練資料集名稱。
這是一個為資料包級流量偵測任務預處理原始流量資料集的實例。
光碟預處理 python preprocess_dataset.py --輸入/您的/原始/資料集/路徑--dataset_name /您的/原始/資料集/名稱--traffic_task 偵測--粒度包--output_path /您的/輸出/資料集/路徑--output_name /您的/Output/Dataset/Name
TrafficLLM 引入了流量域分詞器來處理神經語言和流量資料。如果您想使用自己的資料集訓練自訂分詞器,請修改程式碼中的model_name
和data_path
。
model_name
:包含本機分詞器的基本模型路徑。
data_path
:從預處理過程中提取的訓練資料集。
請依照指令使用代碼。
CD標記化 python Traffic_tokenizer.py
準備資料:神經語言指令調優資料是我們收集的用於流量分析任務理解的指令資料集。
開始調優:完成上述步驟後,就可以使用trafficllm_stage1.sh開始第一階段調優了。有一個例子如下:
PRE_SEQ_LEN=128 LR=2e-2 NUM_GPUS=1匯出 CUDA_VISIBLE_DEVICES=1 torchrun --standalone --nnodes=1 --nproc-per-node=$NUM_GPUS main.py --do_train --train_file ../datasets/instructions/instructions.json --validation_file ../datasets/instructions/instructions.json --preprocessing_num_workers 10 --prompt_column 指令 --response_column 輸出 --overwrite_cache --cache_dir ../快取 --模型名稱或路徑../models/chatglm2/chatglm2-6b --output_dir ../models/chatglm2/peft/指令 --overwrite_output_dir --最大源長度 1024 --最大目標長度 32 --per_device_train_batch_size 1 --per_device_eval_batch_size 1 --gradient_accumulation_steps 16 --predict_with_generate --max_steps 20000 --logging_steps 10 --save_steps 4000 --學習率 $LR --pre_seq_len $PRE_SEQ_LEN
準備資料:特定於任務的流量調整資料集是從不同下游任務的預處理步驟中提取的訓練資料集。
開始調優:完成上述步驟後,就可以使用trafficllm_stage2.sh開始第二階段調優了。有一個例子如下:
PRE_SEQ_LEN=128 LR=2e-2 NUM_GPUS=1匯出 CUDA_VISIBLE_DEVICES=1 torchrun --standalone --nnodes=1 --nproc-per-node=$NUM_GPUS main.py --do_train --train_file ../datasets/ustc-tfc-2016/ustc-tfc-2016_detection_packet_train.json --validation_file ../datasets/ustc-tfc-2016/ustc-tfc-2016_detection_packet_train.json --preprocessing_num_workers 10 --prompt_column 指令 --response_column 輸出 --overwrite_cache --cache_dir ../快取 --模型名稱或路徑../models/chatglm2/chatglm2-6b --output_dir ../models/chatglm2/peft/ustc-tfc-2016-檢測封包 --overwrite_output_dir --最大源長度 1024 --最大目標長度 32 --per_device_train_batch_size 1 --per_device_eval_batch_size 1 --gradient_accumulation_steps 16 --predict_with_generate --max_steps 20000 --logging_steps 10 --save_steps 4000 --學習率 $LR --pre_seq_len $PRE_SEQ_LEN
TrafficLLM 採用 EA-PEFT 來組織具有可擴展適配的參數有效微調 (PEFT) 模型,這可以幫助 TrafficLLM 輕鬆適應新環境適配的參數有效微調 (PEFT) 模型,這可以幫助 TrafficLLM 輕鬆適應新環境適配的參數有效微調 (PEFT) 模型,這可以幫助 TrafficLLM 輕鬆適應新環境適配的參數有效微調 (PEFT) 模型,這可以幫助 TrafficLLM 輕鬆適應新環境適配的參數有效微調 (PEFT) 模型,這可以幫助 TrafficLLM 輕鬆適應新環境適配的參數有效微調 (PEFT) 模型,這可以幫助 TrafficLLM 輕鬆適應新環境適應。 TrafficLLM適配器允許靈活的操作來更新舊模型或註冊新任務。
model_name
:基礎模型的路徑。
tuning_data
:新的環境資料集。
adaptation_task
:更新或註冊(更新舊模型或註冊新任務)。
task_name
:要更新或插入的下游任務名稱。
有一個使用惡意軟體流量偵測 (MTD) 任務更新 TrafficLLM 的範例。
電弧光等離子體增強傅立葉變換 python ea-peft.py --model_name /Your/Base/Model/Path --tuning_data /Your/New/Dataset/Path --adaptation_task update --task_name MTD
檢查點:您可以嘗試使用您自己的模型或我們發布的檢查點來評估 TrafficLLM。
資料:在預處理步驟中,我們分割測試資料集並為不同資料集建立標籤檔案以進行評估。請參考預處理代碼。
要測量 TrafficLLM 對不同下游任務的有效性,請執行評估代碼。
model_name
:基礎模型的路徑。
traffic_task
:偵測任務或產生任務。
test_file
:在預處理步驟中擷取的測試資料集。
label_file
:在預處理步驟中提取的標籤檔案。
ptuning_path
:用於特定任務評估的 PEFT 模型路徑。
有一個對 MTD 任務進行評估的範例。
python評價.py --model_name /您的/基礎/模型/路徑--traffic_task檢測--test_file資料集/ustc-tfc-2016/ustc-tfc-2016_detection_packet_test.json --label_file資料集/ustc-tfc-2016 /ustc- tfc-2016_label.json --ptuning_path models/chatglm2/peft/ustc-tfc-2016-detection-packet/checkpoints-20000/
您可以在本機裝置上部署 TrafficLLM。首先,在 config.json 中配置模型路徑以註冊從訓練步驟收集的 PEFT 模型。有一個在 TrafficLLM 中註冊 6 個任務的範例:
{ "model_path": "models/chatglm2/chatglm2-6b/", "peft_path": "models/chatglm2/peft/", "peft_set": { "NLP": "models/chatglm2/peft/", "peft_set": { "NLP": "models/chatglm2/peft/", "peft_set": { "NLP": "models/chatglm2/peft/", "peft_set": { "NLP": "models/chatglm2/peft/", "peft_set": { "NLP": "指令/checkpoint-8000/”,“MTD”:“ustc-tfc-2016-檢測套件/checkpoint-10000/”,“BND”:“iscx-botnet-2014-偵測套件/checkpoint-5000/”, “ WAD”:“csic-2010-檢測-數據包/checkpoint-6000/”,“AAD”:“dapt-2020-檢測-數據包/checkpoint-20000/”,“EVD”:“iscx-vpn-2016-檢測-packet/checkpoint-4000/”,“TBD”:“iscx-tor-2016-檢測-packet/checkpoint-10000/” }, "tasks": { "惡意軟體流量偵測": "MTD", "殭屍網路偵測": "BND", "Web 攻擊偵測": "WAD", "APT 攻擊偵測": "AAD", "加密VPN 偵測”:“EVD”、“Tor 行為偵測”:“TBD” } 然後
你應該在inference.py和trafficllm_server.py的prepromt
函式中加入preprompt。前置提示是在特定於任務的流量調整期間訓練資料中使用的前綴文字。
要在終端模式下與 TrafficLLM 聊天,您可以執行以下命令:
python inference.py --config=config.json --prompt="Your instructions Text +: + Traffic Data"
您可以使用下列指令啟動TrafficLLM 的網站示範:
streamlit run Trafficllm_server.py
此示範執行 TrafficLLM 的 Web 伺服器。造訪http://Your-Server-IP:Port
在聊天框中聊天。
非常感謝 ChatGLM2 和 Llama2 的相關工作,它們為我們的框架和程式碼奠定了基礎。 TrafficLLM的建造設計受到ET-BERT和GraphGPT的啟發。感謝他們的精彩作品。