加密流量分类torch实践5：TrafficClassificationPandemonium项目更新3

更新日志

代码已经推送开源至露露云的github，如果能帮助你，就给鼠鼠点一个star吧！！！

我的CSDN博客

我的Github Page博客

3/23日更新（此部分上节已经介绍）

当前更新对运行项目是无影响的，也就是说如果你是仅仅使用项目而不进行扩展的话，此处更新是透明的，对当前仓库版本的代码可以不进行同步。
代码已经推送开源至露露云的github，如果能帮助你，就给鼠鼠点一个star吧！！！

模型结构更新

简要：由原先各个模型独立抽象出了一个base_model模型基类，由该基类继承nn.Module类，定义抽象方法forward与data_trans，方便不同模型进行各自的数据变换

为什么要改？

dataloader给模型输入的数据格式是固定死的，给每一个模型设定不同的dataloader违背了项目多个模型统一代码原则，而不同模型对于数据的输入样式是不同的，为了适用于之后会加入项目的模型，抽象出一个基类，设定一个data_trans抽象方法，每一个模型都根据模型的输入去实现该方法即可，这样做到了不更改dataloader的目的，做到代码复用

dataloader给定的数据样式?

分析日志可以给出以下各个维度下dataloader给定的数据shape
[2024-03-23 17:19:38,802 INFO] 是否使用 GPU 进行训练, cuda
[2024-03-23 17:19:44,781 INFO] 成功初始化模型.
[2024-03-23 17:19:44,814 INFO] pcap 文件大小, torch.Size([404, 1, 1024]); seq文件大小:torch.Size([404, 128, 1]); sta文件大小: torch.Size([404, 1024]); label 文件大小: torch.Size([404])
[2024-03-23 17:19:44,851 INFO] pcap 文件大小, torch.Size([404, 1, 1024]); seq文件大小:torch.Size([404, 128, 1]); sta文件大小: torch.Size([404, 1024]); label 文件大小: torch.Size([404])
[2024-03-23 17:19:44,851 INFO] 成功加载数据集.
负载pay: [batch_size,1,m*n]

包长seq: [batch_size,seq_len,1]

统计sta: [batch_size,sta_len]

m*n是预处理的前m个包的前n个字节，这里目前写的是4*256也就是1024
seq_len是预处理的前ip_length个包长，这里目前是128
sta_len是预处理的统计维度，在10号更新的数据下是26

3/28日更新

增加模型二维卷积神经网络CNN2d

由于前期中的使用继承改善了模型结构，这里只需要写一个py文件就可以了

"""
@Description: 二维卷积神经网络
"""
from math import sqrtimport torch
import torch.nn as nn
from models.base_model import BaseModelclass Cnn2d(BaseModel):def __init__(self, num_classes=12):super(Cnn2d, self).__init__()# 卷积层+池化层self.features = nn.Sequential(nn.Conv2d(kernel_size=5,in_channels=1,out_channels=32,stride=1,padding=2), # b,32,32,32nn.MaxPool2d(kernel_size=2), # b,32,16,16nn.Conv2d(kernel_size=5,in_channels=32,out_channels=64,stride=1,padding=2), # b,64,16,16nn.MaxPool2d(kernel_size=2), # b,64,8,8)# 全连接层self.classifier = nn.Sequential(# 29*64nn.Flatten(),nn.Linear(in_features=64 * 64, out_features=1024),  # 1024:64*64nn.Dropout(0.5),nn.Linear(in_features=1024, out_features=num_classes))def forward(self, pay, seq, sta):pay, seq, sta = self.data_trans(pay, seq, sta)pay = self.features(pay)  # 卷积层, 提取特征pay = self.classifier(pay)  # 分类层, 用来分类return pay, Nonedef data_trans(self, x_payload, x_sequence, x_sta):# 转换x_0,x_1,x_2 = x_payload.shape[0],x_payload.shape[1],x_payload.shape[2]x_payload = x_payload.reshape(x_0,x_1,int(sqrt(x_2)),int(sqrt(x_2)))return x_payload, x_sequence, x_stadef cnn2d(model_path, pretrained=False, **kwargs):"""CNN 1D model architectureArgs:pretrained (bool): if True, returns a model pre-trained model"""model = Cnn2d(**kwargs)if pretrained:checkpoint = torch.load(model_path)model.load_state_dict(checkpoint['state_dict'])return modeldef main():a = sqrt(1024)x_pay = torch.rand(8,1,1024)cnn = Cnn2d()x = cnn(x_pay,x_pay,x_pay)if __name__=="__main__":main()

模型结构：

两个卷积+池化的组合，卷积核大小都是5X5，池化层的核大小都是2X2

在train_test_model.py中，改动

from utils.set_config import setup_config
# from models.cnn1d import cnn1d as train_model
# from models.app_net import app_net as train_model
from models.cnn2d import cnn2d as train_model

即可！

开始训练！

[2024-03-28 21:20:53,317 INFO] Epoch: [47][2/4], Loss 0.0001 (0.0005), Prec@1 100.000 (100.000)
[2024-03-28 21:20:53,345 INFO] Epoch: [47][3/4], Loss 0.0000 (0.0005), Prec@1 100.000 (100.000)
[2024-03-28 21:20:53,544 INFO]  * Prec@1 100.000
[2024-03-28 21:20:53,716 INFO] Epoch: [48][1/4], Loss 0.0001 (0.0002), Prec@1 100.000 (100.000)
[2024-03-28 21:20:53,723 INFO] Epoch: [48][3/4], Loss 0.0000 (0.0003), Prec@1 100.000 (100.000)
[2024-03-28 21:20:54,066 INFO] Epoch: [48][0/4], Loss 0.0000 (0.0000), Prec@1 100.000 (100.000)
[2024-03-28 21:20:54,105 INFO] Epoch: [48][1/4], Loss 0.0014 (0.0007), Prec@1 100.000 (100.000)
[2024-03-28 21:20:54,146 INFO] Epoch: [48][2/4], Loss 0.0001 (0.0005), Prec@1 100.000 (100.000)
[2024-03-28 21:20:54,153 INFO] Epoch: [48][3/4], Loss 0.0000 (0.0005), Prec@1 100.000 (100.000)
[2024-03-28 21:20:54,331 INFO]  * Prec@1 100.000
[2024-03-28 21:20:54,537 INFO] Epoch: [49][1/4], Loss 0.0080 (0.0055), Prec@1 99.219 (99.609)
[2024-03-28 21:20:54,558 INFO] Epoch: [49][3/4], Loss 0.0000 (0.0058), Prec@1 100.000 (99.505)
[2024-03-28 21:20:54,880 INFO] Epoch: [49][0/4], Loss 0.0000 (0.0000), Prec@1 100.000 (100.000)
[2024-03-28 21:20:54,929 INFO] Epoch: [49][1/4], Loss 0.0001 (0.0001), Prec@1 100.000 (100.000)
[2024-03-28 21:20:54,970 INFO] Epoch: [49][2/4], Loss 0.0013 (0.0005), Prec@1 100.000 (100.000)
[2024-03-28 21:20:54,982 INFO] Epoch: [49][3/4], Loss 0.0000 (0.0005), Prec@1 100.000 (100.000)
[2024-03-28 21:20:55,147 INFO]  * Prec@1 100.000

修改测试文件相应的配置项，并且切换为测试模式

Model Classification report:
[2024-03-28 21:26:19,166 INFO] ------------------------------
[2024-03-28 21:26:19,172 INFO]               precision    recall  f1-score   supportqq       1.00      1.00      1.00        90微信       1.00      1.00      1.00       206淘宝       1.00      1.00      1.00       108accuracy                           1.00       404macro avg       1.00      1.00      1.00       404
weighted avg       1.00      1.00      1.00       404[2024-03-28 21:26:19,175 INFO] 
Prediction Confusion Matrix:
[2024-03-28 21:26:19,175 INFO] ------------------------------
[2024-03-28 21:26:19,845 INFO]            Predicted:          qq   微信   淘宝
Actual: qq         90    0    0微信          0  206    0淘宝          0    0  108