加密流量分类torch实践2：CNN+LSTM模型训练与测试

烟玉蓝田

已于 2023-10-31 15:27:26 修改

阅读量4.1k

点赞数 8

分类专栏：加密流量分类实验文章标签： cnn 深度学习分类

于 2023-03-09 16:28:41 首次发布

本文链接：https://blog.csdn.net/qq_45125356/article/details/129425682

版权

加密流量分类同时被 2 个专栏收录

18 篇文章 85 订阅

订阅专栏

实验

9 篇文章 2 订阅

订阅专栏

加密流量分类torch实践2：CNN+LSTM模型训练与测试

代码模板参考：CENTIME:A Direct Comprehensive Traffic Features Extraction for Encrypted Traffic Classification

1、原理

一维卷积处理负载数据，处理流前n个包的前m个字节
Bi-LSTM处理包长序列，取流前seq_length的长度序列
模型结构类似于APP-Net
模型代码：

"""
cnn处理负载
lstm处理包长序列
"""

import torch
import torch.nn as nn


class Cnn_Lstm(nn.Module):
    def __init__(self,input_size, hidden_size, num_layers,bidirectional,num_classes=12):
        super(Cnn_Lstm, self).__init__()
        # rnn配置
        self.bidirectional = bidirectional
        self.hidden_size = hidden_size
        self.num_layers = num_layers
        self.lstm = nn.LSTM(input_size, hidden_size, num_layers,bidirectional=bidirectional,batch_first=True)
        self.fc0 = nn.Linear(hidden_size, num_classes)
        self.fc1= nn.Linear(hidden_size*2,num_classes)

        self.cnn_feature = nn.Sequential(
            # 卷积层1
            nn.Conv1d(kernel_size=25, in_channels=1, out_channels=32, stride=1, padding=12),  # (1,1024)->(32,1024)
            nn.BatchNorm1d(32),  # 加上BN的结果
            nn.ReLU(),
            nn.MaxPool1d(kernel_size=3, stride=3, padding=1),  # (32,1024)->(32,342)

            # 卷积层2
            nn.Conv1d(kernel_size=25, in_channels=32, out_channels=64, stride=1, padding=12),  # (32,342)->(64,342)
            nn.BatchNorm1d(64),
            nn.ReLU(),
            nn.MaxPool1d(kernel_size=3, stride=3, padding=1),  # (64,342)->(64,114)
        )
        # 全连接层
        self.cnn_classifier = nn.Sequential(
            # 64*114
            nn.Flatten(),
            nn.Linear(in_features=64*114, out_features=1024), # 784:88*64, 1024:114*64, 4096:456*64
        )

        self.cnn=nn.Sequential(
            self.cnn_feature,
            self.cnn_classifier,
        )


        self.rnn = nn.Sequential(
            nn.LSTM(input_size, hidden_size, num_layers, bidirectional=bidirectional, batch_first=True),
        )
        self.classifier=nn.Sequential(
            nn.Linear(in_features=2048,out_features=num_classes),
            # nn.Dropout(p=0.7),
            # nn.Linear(in_features=1024,out_features=num_classes)
        )


    def forward(self, x_payload,x_sequence):
        x_payload=self.cnn(x_payload)
        x_sequence=self.rnn(x_sequence)
        x_sequence=x_sequence[0][:, -1, :]
        x=torch.cat((x_payload,x_sequence),1)
        x=self.classifier(x)
        return x


def cnn_rnn(model_path, pretrained=False, **kwargs):
    model = Cnn_Lstm(**kwargs)
    if pretrained:
        checkpoint = torch.load(model_path)
        model.load_state_dict(checkpoint['state_dict'])
    return model


# 仅仅是CNN
class Cnn(nn.Module):
    def __init__(self,input_size, hidden_size, num_layers,bidirectional,num_classes=12):
        super(Cnn, self).__init__()
        # rnn配置
        self.bidirectional = bidirectional
        self.hidden_size = hidden_size
        self.num_layers = num_layers
        self.lstm = nn.LSTM(input_size, hidden_size, num_layers,bidirectional=bidirectional,batch_first=True)
        self.fc0 = nn.Linear(hidden_size, num_classes)
        self.fc1= nn.Linear(hidden_size*2,num_classes)

        self.cnn_feature = nn.Sequential(
            # 卷积层1
            nn.Conv1d(kernel_size=25, in_channels=1, out_channels=32, stride=1, padding=12),  # (1,1024)->(32,1024)
            nn.BatchNorm1d(32),  # 加上BN的结果
            nn.ReLU(),
            nn.MaxPool1d(kernel_size=3, stride=3, padding=1),  # (32,1024)->(32,342)

            # 卷积层2
            nn.Conv1d(kernel_size=25, in_channels=32, out_channels=64, stride=1, padding=12),  # (32,342)->(64,342)
            nn.BatchNorm1d(64),
            nn.ReLU(),
            nn.MaxPool1d(kernel_size=3, stride=3, padding=1),  # (64,342)->(64,114)
        )
        # 全连接层
        self.cnn_classifier = nn.Sequential(
            # 64*114
            nn.Flatten(),
            nn.Linear(in_features=64*114, out_features=1024), # 784:88*64, 1024:114*64, 4096:456*64
        )

        self.cnn=nn.Sequential(
            self.cnn_feature,
            self.cnn_classifier,
        )
        self.classifier=nn.Sequential(
            nn.Linear(in_features=1024,out_features=num_classes),
            # nn.Dropout(p=0.7),
            # nn.Linear(in_features=1024,out_features=num_classes)
        )


    def forward(self, x_payload,x_sequence):
        x_payload=self.cnn(x_payload)
        x=self.classifier(x_payload)
        return x_payload


def cnn(model_path, pretrained=False, **kwargs):
    model = Cnn(**kwargs)
    if pretrained:
        checkpoint = torch.load(model_path)
        model.load_state_dict(checkpoint['state_dict'])
    return model


class Lstm(nn.Module):
    def __init__(self,input_size, hidden_size, num_layers,bidirectional,num_classes=12):
        super(Lstm, self).__init__()
        # rnn配置
        self.bidirectional = bidirectional
        self.hidden_size = hidden_size
        self.num_layers = num_layers
        self.lstm = nn.LSTM(input_size, hidden_size, num_layers,bidirectional=bidirectional,batch_first=True)
        self.fc0 = nn.Linear(hidden_size, num_classes)
        self.fc1= nn.Linear(hidden_size*2,num_classes)

        self.rnn = nn.Sequential(
            nn.LSTM(input_size, hidden_size, num_layers, bidirectional=bidirectional, batch_first=True),
        )
        self.classifier=nn.Sequential(
            nn.Linear(in_features=1024,out_features=num_classes),
            # nn.Dropout(p=0.7),
            # nn.Linear(in_features=1024,out_features=num_classes)
        )


    def forward(self, x_payload,x_sequence):
        x_sequence=self.rnn(x_sequence)
        x_sequence=x_sequence[0][:, -1, :]
        x=self.classifier(x_sequence)
        return x


def rnn(model_path, pretrained=False, **kwargs):
    model = Lstm(**kwargs)
    if pretrained:
        checkpoint = torch.load(model_path)
        model.load_state_dict(checkpoint['state_dict'])
    return model

2、运行

在自己的环境下修改路径，包括删除from sequence_payload.xx import xx下面的sequence_payload.
修改配置文件entry下面的traffic_classification.yaml的路径，与模型参数，名字
- 训练流程代码

from utils.helper import AverageMeter, accuracy
from TrafficLog.setLog import logger


def train_process(train_loader, model, criterion, optimizer, epoch, device, print_freq):
    """训练一个 epoch 的流程

    Args:
        train_loader (dataloader): [description]
        model ([type]): [description]
        criterion ([type]): [description]
        optimizer ([type]): [description]
        epoch (int): 当前所在的 epoch
        device (torch.device): 是否使用 gpu
        print_freq ([type]): [description]
    """
    losses = AverageMeter()  # 在一个 train loader 中的 loss 变化
    top1 = AverageMeter()  # 记录在一个 train loader 中的 accuracy 变化

    model.train()  # 切换为训练模型

    for i, (pcap, seq,target) in enumerate(train_loader):
        pcap = pcap.reshape(-1,1,1024)
        seq = seq.reshape(-1,64,1)
        pcap = pcap.to(device)
        seq = seq.to(device)
        target = target.to(device)

        output = model(pcap,seq)  # 得到模型预测结果
        loss = criterion(output, target)  # 计算 loss

        # 计算准确率, 记录 loss 和 accuracy
        # print(pcap.size(0))
        prec1 = accuracy(output.data, target)
        losses.update(loss.item(), pcap.size(0))
        top1.update(prec1[0].item(), pcap.size(0))

        # 反向传播
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        
        if (i+1) % print_freq == 0:
            logger.info('Epoch: [{0}][{1}/{2}], Loss {loss.val:.4f} ({loss.avg:.4f}), Prec@1 {top1.val:.3f} ({top1.avg:.3f})'.format(
                epoch, i, len(train_loader), loss=losses, top1=top1))
    return losses.val,top1.val

验证流程代码：

2.1 数据预处理

原理：

使用flowcontainer包提取含有tcp或者udp负载的包，提取负载与ip数据包序列长度

使用：

将原始pcap文件放在traffic_data下

格式:

  |---traffic_data
    |---bilibili
      |--- xx.pcap
      |--- xxx.pcap
    |---qq
      |--- xx.pcap
      |--- xxx.pcap
    |--- 今日头条
      |--- xx.pcap
      |--- xxx.pcap

运行entry/preprocess.py,完成后复制控制台输出的__label2index__,粘贴到traffic_classification.yaml/test/traffic_classification.yaml
得到处理好的 npy_data

2.2 训练

打开entry/train.py，注释或者取消注释40、41、42行，选择cnn、lstm、cnn+lstm进行训练，记得改配置文件的model_name
可以打开tensorboard查看loss与acc曲线
- loss:
- acc:
  
  两图为lstm处理序列数据的tensorboard示例