# -*- coding: utf-8 -*- import torch import torch.nn as nn import matplotlib.pyplot as plt import numpy as np import pandas as pd import os, glob, time import pickle from timeit import default_timer as timer # matplot lib complains about librosa import warnings warnings.filterwarnings('ignore') # classes index traffic_dict ={ 0: 'Normal', 1: 'BFSSH', 2: 'Infilt', 3: 'HttpDoS', 4: 'DDoS' } # network traffic attributes traffic_attributes = { '01': 'normal', # 正常流量 '02': 'anomaly' # abnormal } """## Load Data """ # path to data for glob DATA_DIR = 'D:\\PyProject\\malware_traffic\\3_Packet\\' def load_data(): """ 加载数据 :return: """ t1 = timer() sessions = [] labels = [] num_pkls = len(glob.glob(DATA_DIR + 'ISCX2012_labels_*.pkl')) # 匹配路径 for i in range(num_pkls): # if i != 1: # continue session_pkl = DATA_DIR + 'ISCX2012_pcaps_' + str(i) + '.pkl' session_lists = pickle.load(open(session_pkl, 'rb')) # 反序列化对象 sessions.extend(session_lists.values.tolist()) # 追加元素 label_pkl = DATA_DIR + 'ISCX2012_labels_' + str(i) + '.pkl' label_lists = pickle.load(open(label_pkl, 'rb')) labels.extend(label_lists.values.tolist()) print(i) t2 = timer() print("load data tims: ", t2 - t1) labels = np.array(labels) normal_indices = np.where(labels == 0)[0] # 结果所在的 行, 是个array # 数据量太大, 不好训练。以下注释代码可以选择100000条正常流量进入训练（建议在数据预处理阶段选择一定数量的正常数据——节约内存开支） normal_indices = np.random.choice(normal_indices, 100000, replace=False) # 注释代码 attack_indices = [np.where(labels == i)[0] for i in range(1, 5)] # label 1~4 所在行, 是个 list # np.random.choice 会重复抽样, 若想不重复, 增加参数：replace=False test_normal_indices = np.random.choice(normal_indices, int(len(normal_indices) * 0.4), replace=False) test_attack_indices = np.concatenate( # 模态融合 [np.random.choice(attack_indices[i], int(len(attack_indices[i]) * 0.4), replace=False) for i in range(4)]) test_indices = np.concatenate([test_normal_indices, test_attack_indices]).astype(int) # train_indices = np.array(list(set(np.arange(len(labels))) - set(test_indices))) attack_indices = np.concatenate(attack_indices).astype(int) # 注释代码 indices = np.concatenate([normal_indices, attack_indices]).astype(int) # 注释代码 train_indices = np.array(list(set(indices) - set(test_indices))) # 注释代码 return sessions, labels, train_indices, test_indices """# Architecture Overview # CNN Motivation ** 构建两个并行卷积神经网络(CNN)来对流量数据进行空间特征表示 # Transformer-Encoder Motivation **使用了Transformer-Encoder层 **I maxpool 映射到Transformer, 以大大减少网络需要学习的参数数量 """ class ByteBlock(nn.Module): """ 1D FCN: 1维全卷积神经网络 in_channels：输入通道数, 在一维卷积中由于不存在通道数, 因此in_channels的数值为词向量的维度, 如果一个单词用128维向量表示, 那么in_channels = 128 out_channels：输出通道数, 表示经过卷积之后, 一个词向量嵌入维度应该为多少。如果out_channels = 64, 那么经过本次卷积之后的每个词的嵌入维度为64。 kernel_size：卷积核大小, 表示本次卷积核的维度, 一般是赋值为int类型。kernel_size=3, 表示每次卷积计算操作涉及到3个词, 也就是卷积核维度被设为(in_channels, kernel_size)。 - 在Pytorch中, 对于一条语句序列数据的每个词都是用一个列向量表示 stride：滑动步长, 表示在卷积方向上滑动的步长。stride=2, 表示在当前卷积的范围为123, 下一个卷积范围就是345。 padding：填补操作, 表示在对特征矩阵剩余部分不足卷积时的操作。 - str --> padding =“valid”:表示不填充, 剩余部分丢弃。 padding =“same”:表示在右侧填充之后要求输入输出序列长度一致 - int --> padding = k： 表示在右侧填充k列 """ def __init__(self, in_channels, nb_filter=(64, 100), filter_length=(3, 3), subsample=(2, 1), pool_length=(2, 2)): super(ByteBlock, self).__init__() layers = [] for i in range(len(nb_filter)): layers.append(nn.Conv1d(in_channels, nb_filter[i], kernel_size=filter_length[i], padding=0, stride=subsample[i])) layers.append(nn.Tanh()) if pool_length[i]: layers.append(nn.MaxPool1d(pool_length[i])) in_channels = nb_filter[i] self.block = nn.Sequential(*layers) self.global_pool = nn.AdaptiveMaxPool1d(1) def forward(self, x): x = self.block(x) x = self.global_pool(x).squeeze(dim=2) x = torch.nn.functional.leaky_relu(x) return x class FCN_Transformer(nn.Module): # Define all layers present in the network def __init__(self,num_emotions): super().__init__() ################ TRANSFORMER BLOCK ############################# # maxpool the input feature map/tensor to the transformer # a rectangular kernel worked better here for the rectangular input spectrogram feature map/tensor self.transformer_maxpool = nn.MaxPool1d(2) # define single transformer encoder layer # self-attention + feedforward network from "Attention is All You Need" paper # Input size: sequence length, batch size, feature size = 128 transformer_layer = nn.TransformerEncoderLayer( d_model=128, # input feature (frequency) dim after maxpooling 128*y -> 64*140 (freq*time) 输入特征维度 nhead=4, # 4 self-attention layers in each multi-head self-attention layer in each encoder block 注意力头数 dim_feedforward=512, # 2 linear layers in each encoder block's feedforward network: dim 64-->512--->64 前馈神经网络中隐藏层的维度。 dropout=0.4, activation='relu' # ReLU: avoid saturation/tame gradient/reduce compute time ) # Using 4 instead of the 6 identical stacked encoder layrs used in Attention is All You Need paper # Complete transformer block contains 4 full transformer encoder layers (each w/ multihead self-attention+feedforward) self.transformer_encoder = nn.TransformerEncoder(transformer_layer, num_layers=4) ############### 1ST PARALLEL 1D CONVOLUTION BLOCK ############ # 1 sequential conv1D layers: (1,128,282) --> (x, y, z) self.conv1Dblock1 = ByteBlock(128, (128, 256), (5, 5), (1, 1), (2, 2)) ############### 2ND PARALLEL 1D CONVOLUTION BLOCK ############ # 1 sequential conv1D layers: (1,128,282) --> (x, y, z) self.conv1Dblock2 = ByteBlock(128, (192, 320), (7, 5), (1, 1), (2, 2)) ################# FINAL LINEAR BLOCK #################### # Linear softmax layer to take final concatenated embedding tensor # from parallel 1D convolutional and transformer blocks, output classes logits self.fc1_linear = nn.Linear(512*2+40,num_emotions) ### Softmax layer for the 8 output logits from final FC linear layer self.softmax_out = nn.Softmax(dim=1) # dim==1 is the freq embedding # define one complete parallel fwd pass of input feature tensor thru 2*conv+1*transformer blocks def forward(self,x): ############ 1st parallel Conv1D block: 4 Convolutional layers ############################ # create final feature embedding from 1st convolutional layer # input features pased through 4 sequential 1D convolutional layers print("x: ",x.type()) conv1d_embedding1 = self.conv1Dblock1(x) # x == N/batch * channel * freq * time # flatten final 64*1*4 feature map from convolutional layers to length 256 1D array # skip the 1st (N/batch) dimension when flattening conv1d_embedding1 = torch.flatten(conv1d_embedding1, start_dim=1) ############ 2nd parallel Conv1D block: 4 Convolutional layers ############################# # create final feature embedding from 2nd convolutional layer # input features pased through 4 sequential 1D convolutional layers conv1d_embedding2 = self.conv1Dblock2(x) # x == N/batch * channel * freq * time # flatten final 64*1*4 feature map from convolutional layers to length 256 1D array # skip the 1st (N/batch) dimension when flattening conv1d_embedding2 = torch.flatten(conv1d_embedding2, start_dim=1) ########## 4-encoder-layer Transformer block w/ 64-->512-->64 feedfwd network ############## # maxpool input feature map: 1*40*282 w/ 1*4 kernel --> 1*40*70 x_maxpool = self.transformer_maxpool(x) # remove channel dim: 1*x*y --> x*y x_maxpool_reduced = torch.squeeze(x_maxpool,1) # convert maxpooled feature map format: batch * freq * time ---> time * batch * freq format # because transformer encoder layer requires tensor in format: time * batch * embedding (freq) x = x_maxpool_reduced.permute(2,0,1) print("x_maxpool_reduced: ",x_maxpool_reduced.shape) # finally, pass reduced input feature map x into transformer encoder layers transformer_output = self.transformer_encoder(x) # create final feature emedding from transformer layer by taking mean in the time dimension (now the 0th dim) # transformer outputs 64*140 (freq embedding*time) feature map, take mean of all columns i.e. take time average transformer_embedding = torch.mean(transformer_output, dim=0) # dim ############# concatenate freq embeddings from convolutional and transformer blocks ###### # concatenate embedding tensors output by parallel 2*conv and 1*transformer blocks complete_embedding = torch.cat([conv1d_embedding1, conv1d_embedding2,transformer_embedding], dim=1) ######### final FC linear layer, need logits for loss ######################### output_logits = self.fc1_linear(complete_embedding) ######### Final Softmax layer: use logits from FC linear, get softmax for prediction ###### output_softmax = self.softmax_out(output_logits) # need output logits to compute cross entropy loss, need softmax probabilities to predict class return output_logits, output_softmax """# 查看模型结构 """ from torchsummary import summary # need device to instantiate model device = torch.device("cuda" if torch.cuda.is_available() else "cpu") # instantiate model for 8 emotions and move to CPU for summary model = FCN_Transformer(len(traffic_dict)).to(device) # print("\nmodel: \n", model,"\n") data = torch.randint(255, size=(128, 128, 100)) # batch_size, flow_len, packet_len print("data: ", data.type()) model(data) # include input feature map dims in call to summary() # summary(model, input_size=(128,282))