423 lines
16 KiB
Python
423 lines
16 KiB
Python
|
|
|||
|
|
# 主网络
|
|||
|
|
import torch
|
|||
|
|
import torch.nn as nn
|
|||
|
|
import torch.optim as optim
|
|||
|
|
import torch.utils.data as data
|
|||
|
|
from torch.nn.functional import normalize
|
|||
|
|
import numpy as np
|
|||
|
|
import pickle
|
|||
|
|
import glob
|
|||
|
|
from timeit import default_timer as timer
|
|||
|
|
import time
|
|||
|
|
import warnings
|
|||
|
|
|
|||
|
|
# 用 BiGRU 网络试试
|
|||
|
|
from model import *
|
|||
|
|
|
|||
|
|
# 不输出“Warning”信息
|
|||
|
|
# warnings.filterwarnings("ignore")
|
|||
|
|
"""
|
|||
|
|
出现过错误的地方，反映了基础不牢
|
|||
|
|
交叉熵：https://blog.csdn.net/dss_dssssd/article/details/84036913
|
|||
|
|
label 转 one-hot：https://blog.csdn.net/qq_34914551/article/details/88700334
|
|||
|
|
|
|||
|
|
实验自适应据类
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
LSTM_UNITS = 92
|
|||
|
|
MINI_BATCH = 10
|
|||
|
|
TRAIN_STEPS_PER_EPOCH = 12000
|
|||
|
|
VALIDATION_STEPS_PER_EPOCH = 800
|
|||
|
|
DATA_DIR = 'D:\\PyProject\\malware_traffic\\3_Packet\\'
|
|||
|
|
CHECKPOINTS_DIR = './checkpoints/'
|
|||
|
|
|
|||
|
|
dict_5class = {0: 'Normal', 1: 'BFSSH', 2: 'Infilt', 3: 'HttpDoS', 4: 'DDoS'}
|
|||
|
|
class_num = 5
|
|||
|
|
|
|||
|
|
PACKET_NUM_PER_SESSION = 14
|
|||
|
|
PACKET_LEN = 100
|
|||
|
|
|
|||
|
|
|
|||
|
|
class ByteBlock(nn.Module):
|
|||
|
|
"""
|
|||
|
|
卷积神经网络
|
|||
|
|
"""
|
|||
|
|
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)
|
|||
|
|
self.fc = nn.Linear(nb_filter[-1], 128)
|
|||
|
|
|
|||
|
|
def forward(self, x):
|
|||
|
|
x = self.block(x)
|
|||
|
|
x = self.global_pool(x).squeeze(dim=2)
|
|||
|
|
x = torch.nn.functional.relu(self.fc(x))
|
|||
|
|
return x
|
|||
|
|
|
|||
|
|
|
|||
|
|
def update_confusion_matrix(confusion_matrix, actual_lb, predict_lb):
|
|||
|
|
for idx, value in enumerate(actual_lb):
|
|||
|
|
p_value = predict_lb[idx]
|
|||
|
|
confusion_matrix[value, p_value] += 1
|
|||
|
|
return confusion_matrix
|
|||
|
|
|
|||
|
|
|
|||
|
|
def truncate(f, n):
|
|||
|
|
trunc_f = np.math.floor(f * 10 ** n) / 10 ** n
|
|||
|
|
return '{:.2f}'.format(trunc_f)
|
|||
|
|
|
|||
|
|
|
|||
|
|
def binarize(x, sz=256):
|
|||
|
|
return torch.nn.functional.one_hot(x, sz).float()
|
|||
|
|
|
|||
|
|
|
|||
|
|
class OneHotEncodingLayer(nn.Module):
|
|||
|
|
def __init__(self, sz=256):
|
|||
|
|
super(OneHotEncodingLayer, self).__init__()
|
|||
|
|
self.size = sz
|
|||
|
|
|
|||
|
|
def forward(self, x):
|
|||
|
|
return torch.nn.functional.one_hot(x, num_classes=self.size).float()
|
|||
|
|
|
|||
|
|
|
|||
|
|
class MyDataset(data.Dataset):
|
|||
|
|
"""
|
|||
|
|
数据转换编码（独热编码）
|
|||
|
|
"""
|
|||
|
|
def __init__(self, sessions, labels, indices):
|
|||
|
|
self.sessions = sessions
|
|||
|
|
self.labels = labels
|
|||
|
|
self.indices = indices
|
|||
|
|
self.packet_num_per_session = PACKET_NUM_PER_SESSION
|
|||
|
|
self.packet_len = PACKET_LEN
|
|||
|
|
|
|||
|
|
def __getitem__(self, index):
|
|||
|
|
idx = self.indices[index]
|
|||
|
|
X = torch.ones((self.packet_num_per_session, self.packet_len), dtype=torch.int64) * -1
|
|||
|
|
|
|||
|
|
for i, packet in enumerate(self.sessions[idx]):
|
|||
|
|
if i < self.packet_num_per_session: # 每个数据会话留只取前 num 个
|
|||
|
|
for j, byte in enumerate(packet[:self.packet_len]): # 每个数据包取前 PACKET_LEN 字节
|
|||
|
|
X[i, (self.packet_len - 1 - j)] = byte
|
|||
|
|
|
|||
|
|
# label = self.labels[idx].astype(np.int64) # CrossEntropyLoss 会自动把标签转换成onehot形式
|
|||
|
|
# y = torch.nn.functional.one_hot(torch.from_numpy(label), num_classes=5)[0]
|
|||
|
|
y = self.labels[idx][0].astype(np.int64)
|
|||
|
|
return X, y
|
|||
|
|
|
|||
|
|
def __len__(self):
|
|||
|
|
return len(self.indices)
|
|||
|
|
|
|||
|
|
|
|||
|
|
class FEMnet(nn.Module):
|
|||
|
|
"""
|
|||
|
|
Feature extraction module 定义
|
|||
|
|
"""
|
|||
|
|
def __init__(self, flow_len, packet_len, gru_units):
|
|||
|
|
super(FEMnet, self).__init__()
|
|||
|
|
self.packet_len = packet_len
|
|||
|
|
self.flow_len = flow_len
|
|||
|
|
self.batch_size = 10
|
|||
|
|
self.gru_hidden_size = gru_units
|
|||
|
|
self.rep_dim = gru_units
|
|||
|
|
|
|||
|
|
self.embedding = OneHotEncodingLayer(sz=256) # 独热编码
|
|||
|
|
self.block2 = ByteBlock(self.packet_len, (128, 256), (5, 5), (1, 1), (2, 2))
|
|||
|
|
self.block3 = ByteBlock(self.packet_len, (192, 320), (7, 5), (1, 1), (2, 2))
|
|||
|
|
|
|||
|
|
self.lstm_layer = nn.GRU(256, self.gru_hidden_size, dropout=0.1, bidirectional=True)
|
|||
|
|
# self.dense_layer = nn.Linear(self.gru_hidden_size * 2, 5)
|
|||
|
|
# self.output = nn.Softmax(dim=1) # Linear 自带Softmax 分类
|
|||
|
|
|
|||
|
|
def forward(self, x): # x: [batch_size, flow_len, packet_len] = [10, 14, 100]
|
|||
|
|
embeddings_list = self.embedding(x) # [10, 14, 100, 256]
|
|||
|
|
encoder_list = torch.zeros((self.batch_size, self.flow_len, 256))
|
|||
|
|
for ix, embeddings in enumerate(embeddings_list): # [14, 100, 256]
|
|||
|
|
# print("embeddings 1: ", embeddings.shape)
|
|||
|
|
# embeddings = embeddings.permute(0, 2, 1)
|
|||
|
|
# print("embeddings 2: ", embeddings.shape)
|
|||
|
|
encoder1 = self.block2(embeddings) # [14, 128]
|
|||
|
|
# print("encoder 1: ", encoder1.shape)
|
|||
|
|
encoder2 = self.block3(embeddings) # [14, 128]
|
|||
|
|
# print("encoder 2: ", encoder2.shape)
|
|||
|
|
encoder = torch.cat([encoder1, encoder2], 1) # [14, 256]
|
|||
|
|
# print("encoder : ", encoder.shape)
|
|||
|
|
encoder_list[ix] = encoder
|
|||
|
|
|
|||
|
|
# rnn 网络输入：[seq_len, batch_size, hidden_size]
|
|||
|
|
encoder_list = encoder_list.permute(1, 0, 2) # [10, 14, 256] -> [14, 10 ,256]
|
|||
|
|
biLSTM, final_hidden_state = self.lstm_layer(encoder_list) # [14, 10, 184], [2, 10, 92]
|
|||
|
|
biLSTM = biLSTM.permute(1, 0, 2) # [10, 14, 184]
|
|||
|
|
# print("biLSTM: ", biLSTM.shape)
|
|||
|
|
attn_output, attention = self.attention_net(biLSTM, final_hidden_state)
|
|||
|
|
dense = self.dense_layer(attn_output)
|
|||
|
|
# out = self.output(dense)
|
|||
|
|
# out = dense[:, -1, :]
|
|||
|
|
return dense
|
|||
|
|
|
|||
|
|
# 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.gru_hidden_size * 2, 1) # Tensor维度的重构，-1表示该维度取决于其他维度
|
|||
|
|
# attn_weights : [batch_size, n_step]
|
|||
|
|
attn_weights = torch.bmm(lstm_output, hidden).squeeze(2) # 加权求和，第三维降维
|
|||
|
|
soft_attn_weights = torch.nn.functional.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)]
|
|||
|
|
|
|||
|
|
|
|||
|
|
class MyModel(nn.Module):
|
|||
|
|
"""
|
|||
|
|
Feature extraction module定义
|
|||
|
|
"""
|
|||
|
|
def __init__(self, binet, feature_dim, class_n):
|
|||
|
|
super(MyModel, self).__init__()
|
|||
|
|
self.binet = binet
|
|||
|
|
self.feature_dim = feature_dim
|
|||
|
|
self.cluster_num = class_n
|
|||
|
|
""" 对比聚类 """
|
|||
|
|
self.instance_projector = nn.Sequential( # 实例-单例
|
|||
|
|
nn.Linear(self.binet.rep_dim, self.binet.rep_dim),
|
|||
|
|
nn.ReLU(),
|
|||
|
|
nn.Linear(self.binet.rep_dim, self.feature_dim),
|
|||
|
|
)
|
|||
|
|
self.cluster_projector = nn.Sequential( # 聚类-集群
|
|||
|
|
nn.Linear(self.binet.rep_dim, self.binet.rep_dim),
|
|||
|
|
nn.ReLU(),
|
|||
|
|
nn.Linear(self.binet.rep_dim, self.cluster_num),
|
|||
|
|
nn.Softmax(dim=1)
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
# 训练时执行
|
|||
|
|
def forward(self, x_i, x_j): # x_i, x_j 来自同一数据的两种不同预处理方式得到的嵌入矩阵
|
|||
|
|
h_i = self.resnet(x_i)
|
|||
|
|
h_j = self.resnet(x_j)
|
|||
|
|
|
|||
|
|
z_i = normalize(self.instance_projector(h_i), dim=1)
|
|||
|
|
z_j = normalize(self.instance_projector(h_j), dim=1)
|
|||
|
|
|
|||
|
|
c_i = self.cluster_projector(h_i)
|
|||
|
|
c_j = self.cluster_projector(h_j)
|
|||
|
|
|
|||
|
|
return z_i, z_j, c_i, c_j
|
|||
|
|
|
|||
|
|
# 测试时执行
|
|||
|
|
def forward_cluster(self, x):
|
|||
|
|
h = self.resnet(x)
|
|||
|
|
c = self.cluster_projector(h) # 聚类分布
|
|||
|
|
c = torch.argmax(c, dim=1) # 得到每个样本的聚类标签。
|
|||
|
|
return c
|
|||
|
|
|
|||
|
|
|
|||
|
|
def save_result(cf_ma, len_test, epochs=8):
|
|||
|
|
"""
|
|||
|
|
计算 metrics_list，保存测试/训练结果
|
|||
|
|
:param cf_ma:
|
|||
|
|
:param len_test:
|
|||
|
|
:param epochs:
|
|||
|
|
:return:
|
|||
|
|
"""
|
|||
|
|
metrics_list = []
|
|||
|
|
for i in range(5):
|
|||
|
|
if i == 0:
|
|||
|
|
metrics_list.append(
|
|||
|
|
[dict_5class[i], str(i), str(cf_ma[i, 0]), str(cf_ma[i, 1]), str(cf_ma[i, 2]), str(cf_ma[i, 3]),
|
|||
|
|
str(cf_ma[i, 4]), '--', '--', '--'])
|
|||
|
|
else:
|
|||
|
|
acc = truncate((float(len_test - cf_ma[:, i].sum() - cf_ma[i, :].sum() + cf_ma[i, i] * 2) / len_test) * 100,
|
|||
|
|
2)
|
|||
|
|
tpr = truncate((float(cf_ma[i, i]) / cf_ma[i].sum()) * 100, 2)
|
|||
|
|
fpr = truncate((float(cf_ma[0, i]) / cf_ma[0].sum()) * 100, 2)
|
|||
|
|
metrics_list.append(
|
|||
|
|
[dict_5class[i], str(i), str(cf_ma[i, 0]), str(cf_ma[i, 1]), str(cf_ma[i, 2]), str(cf_ma[i, 3]),
|
|||
|
|
str(cf_ma[i, 4]), str(acc), str(tpr), str(fpr)])
|
|||
|
|
overall_acc = truncate(
|
|||
|
|
(float(cf_ma[0, 0] + cf_ma[1, 1] + cf_ma[2, 2] + cf_ma[3, 3] + cf_ma[4, 4]) / len_test) * 100, 2)
|
|||
|
|
overall_tpr = truncate((float(cf_ma[1, 1] + cf_ma[2, 2] + cf_ma[3, 3] + cf_ma[4, 4]) / cf_ma[1:].sum()) * 100, 2)
|
|||
|
|
overall_fpr = truncate((float(cf_ma[0, 1:].sum()) / cf_ma[0, :].sum()) * 100, 2)
|
|||
|
|
|
|||
|
|
# 获取当前日期
|
|||
|
|
current_date = time.strftime("%Y%m%d", time.localtime())
|
|||
|
|
# 构造文件名
|
|||
|
|
file_name = f"evaluate_{current_date}.txt"
|
|||
|
|
|
|||
|
|
with open(file_name, 'a') as f:
|
|||
|
|
f.write("\n")
|
|||
|
|
t = time.strftime('%Y-%m-%d %X', time.localtime())
|
|||
|
|
f.write(t + "\n")
|
|||
|
|
f.write('CLASS_NUM: 5\n')
|
|||
|
|
f.write('PACKET_LEN: ' + str(PACKET_LEN) + "\n")
|
|||
|
|
f.write('PACKET_NUM_PER_SESSION: ' + str(PACKET_NUM_PER_SESSION) + "\n")
|
|||
|
|
f.write('MINI_BATCH: ' + str(MINI_BATCH) + "\n")
|
|||
|
|
f.write('TRAIN_EPOCHS: ' + str(epochs) + "\n")
|
|||
|
|
f.write('DATA_DIR: ' + DATA_DIR + "\n")
|
|||
|
|
f.write("label\tindex\t0\t1\t2\t3\t4\tACC\tTPR\tFPR\n")
|
|||
|
|
for metrics in metrics_list:
|
|||
|
|
f.write('\t'.join(metrics) + "\n")
|
|||
|
|
f.write('Overall accuracy: ' + str(overall_acc) + "\n")
|
|||
|
|
f.write('Overall TPR: ' + str(overall_tpr) + "\n")
|
|||
|
|
f.write('Overall FPR: ' + str(overall_fpr) + "\n")
|
|||
|
|
# f.write('Train time(second): ' + str(int(train_time)) + "\n")
|
|||
|
|
# f.write('Test time(second): ' + str(int(test_time)) + "\n\n")
|
|||
|
|
f.write("\n\n")
|
|||
|
|
|
|||
|
|
|
|||
|
|
def train(model, dataloader, criterion, optimizer, device):
|
|||
|
|
model.train()
|
|||
|
|
running_loss = 0.0
|
|||
|
|
num = 0
|
|||
|
|
for inputs, labels in dataloader:
|
|||
|
|
inputs = inputs.to(device)
|
|||
|
|
labels = labels.to(device)
|
|||
|
|
|
|||
|
|
optimizer.zero_grad() # 梯度清零
|
|||
|
|
|
|||
|
|
outputs = model(inputs)
|
|||
|
|
num += 1
|
|||
|
|
# print(“ {}/{} outputs & label shape: ".format(num,len(dataloader.data_loader)),
|
|||
|
|
# outputs.shape, labels.shape)
|
|||
|
|
loss = criterion(outputs, labels)
|
|||
|
|
|
|||
|
|
loss.backward()
|
|||
|
|
optimizer.step()
|
|||
|
|
|
|||
|
|
running_loss += loss.item() * inputs.size(0)
|
|||
|
|
|
|||
|
|
return running_loss / len(dataloader.dataset)
|
|||
|
|
|
|||
|
|
|
|||
|
|
def evaluate(model, dataloader, criterion, device):
|
|||
|
|
model.eval()
|
|||
|
|
running_loss = 0.0
|
|||
|
|
cf_ma = np.zeros((5, 5), dtype=int)
|
|||
|
|
num = 0
|
|||
|
|
with torch.no_grad():
|
|||
|
|
for inputs, labels in dataloader:
|
|||
|
|
inputs = inputs.to(device)
|
|||
|
|
labels = labels.to(device)
|
|||
|
|
|
|||
|
|
outputs = model(inputs)
|
|||
|
|
loss = criterion(outputs, labels)
|
|||
|
|
|
|||
|
|
running_loss += loss.item() * inputs.size(0)
|
|||
|
|
predicted_labels = torch.argmax(outputs, dim=1).cpu().numpy()
|
|||
|
|
true_labels = labels.cpu().numpy()
|
|||
|
|
|
|||
|
|
cf_ma = update_confusion_matrix(cf_ma, true_labels, predicted_labels)
|
|||
|
|
num += predicted_labels.shape[0]
|
|||
|
|
|
|||
|
|
print("num len: ", num)
|
|||
|
|
save_result(cf_ma, num)
|
|||
|
|
return running_loss / len(dataloader.dataset), cf_ma
|
|||
|
|
|
|||
|
|
|
|||
|
|
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
|
|||
|
|
|
|||
|
|
|
|||
|
|
# 定义超参数
|
|||
|
|
NUM_EPOCHS = 8
|
|||
|
|
LEARNING_RATE = 0.001
|
|||
|
|
|
|||
|
|
# 创建模型实例
|
|||
|
|
fe_model = FEMnet(PACKET_NUM_PER_SESSION, PACKET_LEN, LSTM_UNITS)
|
|||
|
|
model = MyModel(fe_model, 128, class_num)
|
|||
|
|
# model = PLA_Attention_Model(100, 100, 50, 100, 14)
|
|||
|
|
|
|||
|
|
# 将模型移到GPU上（如果可用）
|
|||
|
|
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
|||
|
|
model.to(device)
|
|||
|
|
|
|||
|
|
# 定义损失函数和优化器
|
|||
|
|
criterion = nn.CrossEntropyLoss()
|
|||
|
|
optimizer = optim.Adam(model.parameters(), lr=LEARNING_RATE)
|
|||
|
|
|
|||
|
|
# 创建训练集和验证集的数据加载器
|
|||
|
|
sessions, labels, train_indices, test_indices = load_data()
|
|||
|
|
|
|||
|
|
train_dataset = MyDataset(sessions, labels, train_indices) # # num_workers=2，Windows不允许多个进程加载数据
|
|||
|
|
train_dataloader = data.DataLoader(train_dataset, batch_size=MINI_BATCH, shuffle=True, drop_last=True)
|
|||
|
|
|
|||
|
|
val_dataset = MyDataset(sessions, labels, test_indices) # 丢弃不成 batch 数据，否则会报错
|
|||
|
|
val_dataloader = data.DataLoader(val_dataset, batch_size=MINI_BATCH, shuffle=False, drop_last=True)
|
|||
|
|
|
|||
|
|
# 开始训练
|
|||
|
|
start_time = timer()
|
|||
|
|
|
|||
|
|
for epoch in range(NUM_EPOCHS):
|
|||
|
|
# 训练模型
|
|||
|
|
train_loss = train(model, train_dataloader, criterion, optimizer, device)
|
|||
|
|
# print(f"Train Loss: {truncate(train_loss, 4)}")
|
|||
|
|
# 在验证集上评估模型
|
|||
|
|
val_loss, confusion_matrix = evaluate(model, val_dataloader, criterion, device)
|
|||
|
|
|
|||
|
|
# 输出训练和验证的损失以及混淆矩阵
|
|||
|
|
print(f"Epoch {epoch + 1}/{NUM_EPOCHS}")
|
|||
|
|
print(f"Train Loss: {truncate(train_loss, 4)}")
|
|||
|
|
print(f"Validation Loss: {truncate(val_loss, 4)}")
|
|||
|
|
print("Confusion Matrix:")
|
|||
|
|
for i in range(5):
|
|||
|
|
for j in range(5):
|
|||
|
|
print(confusion_matrix[i][j], end="\t")
|
|||
|
|
print()
|
|||
|
|
print("---------------------------------")
|
|||
|
|
|
|||
|
|
end_time = timer()
|
|||
|
|
training_time = end_time - start_time
|
|||
|
|
|
|||
|
|
print(f"Training Time: {truncate(training_time, 2)} seconds")
|
|||
|
|
|