【AI】恶意文件静态检测模型检验及小结

sidiot

发布于 2023-08-30 14:54:30

28400

代码可运行

文章被收录于专栏：技术大杂烩技术大杂烩

运行总次数：0

代码可运行

前言

在之前的博文【AI】浅析恶意文件静态检测及部分问题解决思路中，博主提及过恶意文件静态检测的一种方法，并因此训练了模型，由于样本量巨大以及资源有限，训练一个 epoch 就需要一周多的时间，因此就先拿训练过一个 epoch 的模型来进行测试；

拉取测试集

既然是要用来测试的样本，那么我们要尽可能的与训练集以及验证集中的样本不一样，因此，最好在一开始就做好分类，或者可以借用集合 set 的特性来整合；

我们先用训练集和验证集进行测试，伪代码如下：

with open('...pkl', 'rb') as f:
    train_data = pickle.load(f)

with open('...pkl', 'rb') as f:
    val_data = pickle.load(f)

train_data_ = [x[0] for x in train_data]
val_data_ = [x[0] for x in val_data]
zz = set(train_data_) - set(val_data_)

可以看到数量是完全相同的，因此训练集和验证集没有交集，即两者之间没有重复的样本；

接下来我们就开始拉去测试集，先从 Metadata_PE 表中去获取到 path 和 sha256 字段，然后在根据 sha256 去 Event_PE_lab_22_11_24 表中进行查询 lab 标签；

这里的话，可以根据联合索引，直接从数据库中将全部数据导入，借用 pymsql 和 pandas 的包，在 python 中处理的速度比原生 SQL 要快不少，不过因为数据量较大，导入也消耗的一定的时间：

导入完成之后就是对数据进行处理：

new_sample_df = sample_df[sample_df['date'] >= pd.Timestamp('2022-12-01')]
pd.merge(new_sample_df, label_df, on="sha256")

这里的话，根据入库时间进行拉取，选取 2022-12-01 之后入库的样本：

因为这里只需要 exe 类型的文件，所以还需要再进行一次判断，样本量过大可采取多线程 ThreadPoolExecutor：

if pefile.PE(path).is_exe():
    ...

全部完成之后，就是我们需要的测试集了；

检验模型

我们先将刚刚准备好的测试集进行导入：

test_loader = DataLoader(...)

然后传入到模型中，获取预测值，并计算损失：

# 特征值，模型的输入
test_x = batch_data[0].to(torch.float32).to(device)
# 预测值，模型的输出，两个值分别为黑白样本概率，如 [0.4052, -0.3841]
out = model(test_x)
# 标签值，用于计算损失
label = batch_data[1].to(device)
# 预测值与真实值之间的损失
loss = criterion(out, label.long())
# 一个 batch 的大小
val_size += label.size(0)
# 一个 batch 的损失，loss.item() 每个样本的平均损失
running_loss += loss.item() * label.size(0)

因为是检验模型，我们需要去评价模型的好坏，判断是否为恶意文件其实就是个二分类问题，这里的话使用混淆矩阵：

	预测值0	预测值1
真实值0	TN	FP
真实值1	FN	TP

TN：真实值是0，预测值是0，即我们预测是 negative，预测正确了。
FP：真实值是0，预测值是1，即我们预测是 positive，预测错误了。
FN：真实值是1，预测值是0，即我们预测是 negative，预测错误了。
TP：真实值是1，预测值是1，即我们预测是 positive，预测正确了。

accuracy_score = (TP+TN) / (TP+TN+FP+FN)：函数计算分类准确率，返回被正确分类的样本比例（default）或者是数量（normalize=False）。

精准率（查准率）和召回率（查全率）等指标对衡量机器学习的模型性能在某些场合下要比 accuracy 更好。

精准率：**precision = TP / (TP+FP)。所谓的精准率是：分母为所有预测为1的个数，分子是其中预测对了的个数，即预测为正的样本中，实际为正的比例。**

召回率：**recall = TP / (TP+FN)。所谓的召回率是：所有真实值为1的数据中，预测对了的个数，也就是我们关注的那个事件真实的发生情况下，我们成功预测的比例是多少。**

接下来，我们就根据预测值和标签值来进行计算：

preds_n = preds_sg
label_n = label_sg
# zes: [0. 0. 0. 0. 0. 0. 0. 0. 0. 0.] 
zes = np.zeros(label.size(0))
# ons: [1. 1. 1. 1. 1. 1. 1. 1. 1. 1.]
ons = np.ones(label.size(0))
preds_np = preds_n
label_np = label_n.reshape(-1)
train_correct01 = int(((preds_np == zes) & (label_np == ons)).sum())
train_correct10 = int(((preds_np == ons) & (label_np == zes)).sum())
train_correct11 = int(((preds_np == ons) & (label_np == ons)).sum())
train_correct00 = int(((preds_np == zes) & (label_np == zes)).sum())
FN += train_correct01
FP += train_correct10
TP += train_correct11
TN += train_correct00
accuracy_score = (TP+TN) / (TP+TN+FP+FN)
precision = TP / (TP+FP)
recall = TP / (TP+FN)

这里的话就用几个 batch 来略作检验：