77 lines
4.1 KiB
Python
77 lines
4.1 KiB
Python
import torch
|
||
import torch.nn as nn
|
||
import torch.nn.functional as F
|
||
|
||
|
||
class PLA_Attention_Model(nn.Module):
|
||
def __init__(self, byte_hidden_size, packet_hidden_size, packet_output_size, max_packet=1500, max_flow=128):
|
||
super().__init__()
|
||
self.byte_emb_size = 50
|
||
self.byte_hidden_size = byte_hidden_size # 100
|
||
self.packet_emb_size = 50
|
||
self.packet_hidden_size = packet_hidden_size # 100
|
||
self.packet_output_size = packet_output_size # 50
|
||
self.max_packet = max_packet
|
||
self.max_flow = max_flow
|
||
|
||
self.embedding = nn.Embedding(256 * 256, self.byte_emb_size)
|
||
self.byte_GRU = nn.GRU(self.byte_emb_size, self.byte_hidden_size, bias=False, bidirectional=True)
|
||
self.byte_attn = nn.Linear(self.byte_hidden_size * 2, self.packet_emb_size)
|
||
self.packet_GRU = nn.GRU(self.packet_emb_size, self.packet_hidden_size, bias=False, bidirectional=True)
|
||
self.packet_attn = nn.Linear(self.packet_hidden_size * 2, self.packet_output_size)
|
||
self.classify = nn.Linear(self.packet_output_size, 11)
|
||
|
||
# lstm_output : [batch_size, n_step, self.byte_hidden_size * num_directions(=2)], F matrix
|
||
def attention_net(self, lstm_output, final_state):
|
||
# print("lstm_output: ", lstm_output.shape)
|
||
# print("final_state: ", final_state.shape)
|
||
# hidden : [batch_size, self.byte_hidden_size * num_directions(=2), 1(=n_layer)]
|
||
hidden = final_state.view(-1, self.byte_hidden_size * 2, 1) # Tensor维度的重构，-1表示该维度取决于其他维度
|
||
# attn_weights : [batch_size, n_step]
|
||
attn_weights = torch.bmm(lstm_output, hidden).squeeze(2) # 加权求和，第三维降维
|
||
soft_attn_weights = F.softmax(attn_weights, 1)
|
||
# [batch_size, self.byte_hidden_size * num_directions(=2), n_step] * [batch_size, n_step, 1]
|
||
# = [batch_size, self.byte_hidden_size * num_directions(=2), 1]
|
||
context = torch.bmm(lstm_output.transpose(1, 2), soft_attn_weights.unsqueeze(2)).squeeze(2)
|
||
return context, soft_attn_weights # context : [batch_size, self.byte_hidden_size * num_directions(=2)]
|
||
|
||
def forward(self, flow):
|
||
num_packet = flow.shape[0]
|
||
batch_size = flow.shape[1]
|
||
embedded_bytes_list = self.embedding(flow)
|
||
encoding_bytes_list = torch.zeros((num_packet, batch_size, self.packet_emb_size))#.cuda() # [4, 3, 120, 50]
|
||
for i, embedded_bytes in enumerate(embedded_bytes_list): # [3, 120, 50]
|
||
h0_bytes = torch.randn(2, batch_size, self.byte_hidden_size)#.cuda() # [2,3,100]
|
||
embedded_bytes = embedded_bytes.transpose(0, 1) # 交换 dim_0 和 dim_1 : [120, 3, 50]
|
||
# print(embedded_bytes.shape, h0_bytes.shape)
|
||
output, final_hidden_state = self.byte_GRU(embedded_bytes, h0_bytes) # [120, 3, 200], [2, 3, 100]
|
||
output = output.permute(1, 0, 2) # [3, 120, 200]
|
||
# print(output.shape)
|
||
# print("final: ", final_hidden_state.shape)
|
||
attn_output, attention = self.attention_net(output, final_hidden_state) # [3, 200]
|
||
# print("attn_output", attn_output.shape)
|
||
encoding_bytes_list[i] = self.byte_attn(attn_output) # [4, 3, 50] 最后一维维持不变
|
||
|
||
# print("2: ", encoding_bytes_list.shape) # [4, 3, 50]
|
||
h0_packet = torch.randn(2, batch_size, self.packet_hidden_size)#.cuda() # [2,3,100]
|
||
# print(h0_packet.shape)
|
||
output, final_hidden_state = self.packet_GRU(encoding_bytes_list, h0_packet)
|
||
# print("final_hidden_state: ", final_hidden_state.shape) # [2, 3, 100]
|
||
output = output.permute(1, 0, 2) # 维度排列转置 [3, 4, 200]
|
||
attn_output, attention = self.attention_net(output, final_hidden_state) # [3, 200]
|
||
print("attn_output2: ", attn_output.shape)
|
||
output = self.packet_attn(attn_output)
|
||
classify = self.classify(output)
|
||
|
||
return classify
|
||
|
||
|
||
if __name__ == '__main__':
|
||
batch_size = 3
|
||
model = PLA_Attention_Model(100, 100, 50)#.cuda()
|
||
print(model)
|
||
data = torch.randint(255, size=(4, batch_size, 120))#.cuda() # flow_len, batch_size, packet_len
|
||
print(data.shape)
|
||
res = model(data)
|
||
print(res.shape)
|