加密流量分類torch實踐1：1D-CNN模型訓練與測試

模型：model.py

論文參數：

import torch.nn as nn
import torch
import torch.nn.functional as F

class OneCNNC(nn.Module):
    def __init__(self,label_num):
        super(OneCNNC,self).__init__()
        self.layer_1 = nn.Sequential(
            # 輸入784*1
            nn.Conv2d(1,32,(1,25),1,padding='same'),
            nn.ReLU(),
            # 輸出262*32
            nn.MaxPool2d((1, 3), 3, padding=(0,1)),
        )
        self.layer_2 = nn.Sequential(
            # 輸入262*32
            nn.Conv2d(32,64,(1,25),1,padding='same'),
            nn.ReLU(),
            # 輸入262*64
            nn.MaxPool2d((1, 3), 3, padding=(0,1))
        )
        self.fc1=nn.Sequential(
            # 輸入88*64
            nn.Flatten(),
            nn.Linear(88*64,1024),
            # 自主加了兩個dropout層
            nn.Dropout(p=0.5),
            nn.Linear(1024,label_num),
            nn.Dropout(p=0.3)
        )
    def forward(self,x):
        # print("x.shape:",x.shape)
        x=self.layer_1(x)
        # print("x.shape:",x.shape)
        x=self.layer_2(x)
        # print("x.shape:",x.shape)
        x=self.fc1(x)
        # print("x.shape:",x.shape)
        return x



# x=torch.tensor([[1, 1,  0,  1,  2,  3],
#                 [1, 1,  4,  5,  6,  7],
#                 [1, 10, 8,  9, 10, 11]],dtype=torch.float32)
# x=x.reshape(1,3,-1)


# out_tensor=F.max_pool2d(x,(3,1),stride=3,padding=0)

# print(out_tensor)
           

資料部分：

資料讀取：data.py

import os
from torch.utils.data import Dataset
import gzip
import numpy as np
class DealDataset(Dataset):
    """
        讀取資料、初始化資料
    """
 
    def __init__(self, folder, data_name, label_name, transform=None):
        (train_set, train_labels) = load_data(folder, data_name,label_name)  
        self.train_set = train_set
        self.train_labels = train_labels
        self.transform = transform
 
    def __getitem__(self, index):
        img, target = self.train_set[index], int(self.train_labels[index])
        # 這裡要copy一下不然會報錯
        img=img.copy()
        # 28*28 -> 764
        img=img.reshape(1,1,-1)
        # target=target.copy()
        if self.transform is not None:
            img = self.transform(img)
        return img, target
 
    def __len__(self):
        return len(self.train_set)
 
 
def load_data(data_folder, data_name, label_name):
    with gzip.open(os.path.join(data_folder, label_name), 'rb') as lbpath:  
        y_train = np.frombuffer(lbpath.read(), np.uint8, offset=8)
 
    with gzip.open(os.path.join(data_folder, data_name), 'rb') as imgpath:
        x_train = np.frombuffer(
            imgpath.read(), np.uint8, offset=16).reshape(len(y_train), 28, 28)
    return (x_train, y_train)

           

模型訓練主子產品：

from random import shuffle
import time
import sys
import torch.nn as nn
import numpy as np
import os

import torchvision

from model import OneCNN,CNNImage,OneCNNC
from torchvision import datasets,transforms
import gzip
import torch
from data import DealDataset


def main():
    # Device configuration
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

    # 設定超參數
    batch_size = 50
    lr = 1.0e-4
    num_epochs = 40
    label_num = 12



    # 導入資料
    folder_path_list=[
        r"2.encrypted_traffic_classification/3.PerprocessResults/12class/FlowAllLayers",
        r"2.encrypted_traffic_classification/3.PerprocessResults/12class/FlowL7",
        r"2.encrypted_traffic_classification/3.PerprocessResults/12class/SessionAllLayers",
        r"2.encrypted_traffic_classification/3.PerprocessResults/12class/SessionL7"
                      ]

	# 選擇哪個資料集
    task_index = 0

    folder_path = folder_path_list[task_index]
    train_data_path = "train-images-idx3-ubyte.gz"
    train_label_path = "train-labels-idx1-ubyte.gz"
    test_data_path = "t10k-images-idx3-ubyte.gz"
    test_label_path = "t10k-labels-idx1-ubyte.gz"
    trainDataset = DealDataset(folder_path,train_data_path,train_label_path)
    testDataset = DealDataset(folder_path,test_data_path,test_label_path
   
    train_loader = torch.utils.data.DataLoader(
        dataset=trainDataset,
        batch_size=batch_size,
        shuffle=True
    )
    
    test_loader = torch.utils.data.DataLoader(
        dataset=testDataset,
        batch_size=batch_size,
        shuffle=False
    )


    # 定義模型
    model = OneCNNC(label_num)
    model = model.to(device)
    # model = CNNImage()

    # Loss and optimizer
    criterion = nn.CrossEntropyLoss()
    optimizer = torch.optim.SGD(model.parameters(), lr=lr)

    # Train the model
    total_step = len(train_loader)
    for epoch in range(num_epochs):
        for i, (images, labels) in enumerate(train_loader):
            # images=images.reshape(-1,1,28,28)
            images = images.to(device)
            labels = labels.to(device)
            
            # Forward pass
            outputs = model(images.to(torch.float32))
            loss = criterion(outputs, labels)
            
            # Backward and optimize
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            
            if (i+1) % 100 == 0:
                print ('Epoch [{}/{}], Step [{}/{}], Loss: {:.4f}' 
                    .format(epoch+1, num_epochs, i+1, total_step, loss.item()))
    # Test the model
    model.eval()
    with torch.no_grad():
        correct = 0
        total = 0
        test_length = len(testDataset)
        for images, labels in test_loader:
            images = images.to(device)
            labels = labels.to(device)
            outputs = model(images.to(torch.float32))
            _, predicted = torch.max(outputs.data, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()

        print('Test Accuracy of the model on the {} test images: {} %'.format(test_length,100 * correct / total)) 

    # Save the model checkpoint
    torch.save(model.state_dict(), 'model.ckpt')

if __name__ == '__main__':
    main()


           

運作結果：

項目位址：https://github.com/lulu-cloud/Pytorch-Encrypted-Traffic-Classification-with-1D_CNN