【科研】ET-BERT代码分析

0. 数据集

论文使用的TLS1.3数据集是从 2021 年 3 月到 2021 年 7 月在中国科技网 (CSTNET) 上收集的。

如果要使用自己的数据集的话，需要检查数据格式是否与datasets/cstnet-tls1.3/目录下一致，并且在data_process/目录下指定数据集路径。

用于微调的数据集，分为train、test、valid。微调数据集的链接—>fine-tuning cstnet-tls 1.3

1. 数据预处理【最难的】

数据预处理的md说明文件

1.1 预训练阶段

主程序：dataset_generation.py

函数：pretrain_dataset_generation，get_burst_feature

初始化。将变量pcap_path（line:616）设置为要处理的 PCAP数据的目录。设置变量word_dir（line:23）和word_name（line:24）为预训练数据集的存放目录。
预处理PCAP。设置变量output_split_path（line:583）和pcap_output_path（line:584）。表示PCAP数据的pcapng格式转换为pcap格式的pcap_output_path存放目录。output_split_path代表PCAP数据切片成session格式的存储目录。
生成预训练数据集。PCAP数据处理完成后，程序生成由BURST组成的预训练数据集。

def pretrain_dataset_generation(pcap_path):
    # pcap--->session的存储路径
    output_split_path = ".\\data\\dataset\\"
    # pcapng--->pcap的存储路径
    pcap_output_path = ".\\data\\dataset\\"    
    if not os.listdir(pcap_output_path):
        print("Begin to convert pcapng to pcap.")
        for _parent,_dirs,files in os.walk(pcap_path):
            for file in files:
                if 'pcapng' in file:
                    #print(_parent + file)
                    convert_pcapng_2_pcap(_parent, file, pcap_output_path)
                else:
                    shutil.copy(_parent+"\\"+file, pcap_output_path+file)
    
    if not os.path.exists(output_split_path + "splitcap"):
        print("Begin to split pcap as session flows.")
        
        for _p,_d,files in os.walk(pcap_output_path):
            for file in files:
                split_cap(output_split_path,_p+file,file)
    print("Begin to generate burst dataset.")
    # burst sample
    for _p,_d,files in os.walk(output_split_path + "splitcap"):
        for file in files:
            get_burst_feature(_p+"\\"+file, payload_len=64)
    return 0
# 预训练阶段dataset_generation.py  main函数主要的子函数
if __name__ == '__main__':
    # pretrain
    pcap_path = "I:\\pcaps\\"
    # pretrain data
    pretrain_dataset_generation(pcap_path)

input：原始的pcap流量
output：word_dir+word_name的burst数据--->encrypted_burst.txt文件
word_dir = "I:/corpora/"
word_name = "encrypted_burst.txt"

1.2 微调阶段

主程序：data_process/main.py

函数：data_preprocess.py、dataset_generation.py、open_dataset_deal.py、dataset_cleanning.py

处理公共PCAP数据集微调阶段的核心思想是先对数据集中不同标签的数据进行文件夹区分，然后对数据进行session切分，最后根据样本需要生成包级或流级数据集.

注意：由于可能存在原始PCAP数据的复杂性，建议在报错时执行以下步骤检查代码执行情况。

初始化。 pcap_path, dataset_save_path, samples, features,dataset_level(line:28)为基础变量，分别代表原始数据目录、存储的生成数据目录、样本数量、特征类型和数据级别。open_dataset_not_pcap(line:215)表示将不是PCAP数据转换为pcap格式的处理，例如pcapng到pcap。file2dir(line:226) 表示当一个pcap文件为类别时，生成类别目录存放PCAP数据。
预处理。数据预处理主要是将目录中的PCAP数据拆分成session数据。请将该splitcap_finish参数设置为0，初始化样本数数组，sample此时设置的值不要超过最小样本数。然后您可以设置splitcap=True（line:54）并运行拆分 PCAP 数据的代码。拆分的会话将保存在pcap_path\splitcap.
生成数据。数据预处理完成后，需要更改变量以生成微调的训练数据。pcap_path应该是拆分数据的路径并且修改splitcap=False。现在sample的大小可以不受最小样本量的限制。open_dataset_not_pcap和file2dir应该是False。然后生成微调的数据集并保存在dataset_save_path.

2. 微调

step 1. 下载预训练模型，这个预训练模型是作者提供的pretrained_model.bin
step 2. ET-BERT 可以通过使用标记的网络流量在数据包级别进行微调来应用于特定任务
corpora链接—>encrypted_traffic_burst.txt

【困惑点1】执行代码里使用的encryptd_vocab.txt不造从哪里生成的，，待我看看代码

python3 fine-tuning/run_classifier.py --pretrained_model_path models/pre-trained_model.bin \
                                   --vocab_path models/encryptd_vocab.txt \
                                   --train_path datasets/cstnet-tls1.3/packet/train_dataset.tsv \
                                   --dev_path datasets/cstnet-tls1.3/packet/valid_dataset.tsv \
                                   --test_path datasets/cstnet-tls1.3/packet/test_dataset.tsv \
                                   --epochs_num 10 --batch_size 32 --embedding word_pos_seg \
                                   --encoder transformer --mask fully_visible \
                                   --seq_length 128 --learning_rate 2e-5

step 3. 微调分类器模型的默认路径为models/finetuned_model.bin. 然后您可以使用微调模型进行预测：

【困惑点2】finetuned_model.bin不知道哪来的，，待我看看代码

python3 inference/run_classifier_infer.py --load_model_path models/finetuned_model.bin \
                                          --vocab_path models/encryptd_vocab.txt \
                                          --test_path datasets/cstnet-tls1.3/packet/nolabel_test_dataset.tsv \
                                          --prediction_path datasets/cstnet-tls1.3/packet/prediction.tsv \
                                          --labels_num 120 \
                                          --embedding word_pos_seg --encoder transformer --mask fully_visible

总结，要仔细阅读run_classifier.py和run_classifier_infer.py两份代码！！！

3. 预训练

3.1 预处理

要复现在网络流量数据上预训练 ET-BERT，执行以下步骤：

运行vocab_process/main.py生成加密后的流量语料或者直接使用生成的语料在corpora/.
请注意，您需要更改文件路径和文件代码顶部的一些配置。
运行main/preprocess.py对加密流量burst语料库进行预处理。

   python3 preprocess.py --corpus_path corpora/encrypted_traffic_burst.txt \
                         --vocab_path models/encryptd_vocab.txt \
                         --dataset_path dataset.pt --processes_num 8 --target bert

data_process/main.py如果有pcap格式的数据集需要处理，运行生成下游任务的数据。这个过程包括两个步骤：
第一步是分割。通过设置datasets/main.py:54中的splitcap=True来拆分流量文件并另存为npy数据集；
第二步是生成微调数据。如果您使用共享数据集，则需要在dataset_save_path下创建一个文件夹dataset并将数据集复制到此处。

3.2 预训练

要复现在标记数据上微调 ET-BERT 所需的步骤，运行pretrain.py以进行预训练。

   python3 pre-training/pretrain.py --dataset_path dataset.pt --vocab_path models/encryptd_vocab.txt \
                       --output_model_path models/pre-trained_model.bin \
                       --world_size 8 --gpu_ranks 0 1 2 3 4 5 6 7 \
                       --total_steps 500000 --save_checkpoint_steps 10000 --batch_size 32 \
                       --embedding word_pos_seg --encoder transformer --mask fully_visible --target bert

3.3 微调

要查看如何使用 ET-BERT 执行加密流量分类任务的示例，运行fine-tuning下的run_classifier.py脚本。

具体看上文的1. 微调章节

注意：您需要更改程序中的路径。

文章出处登录后可见！

已经登录？立即刷新