使用PaddleNLP进行恶意网页识别(六):用BERT识别网页标签序列
本文是《使用PaddleNLP进行恶意网页识别》系列第六篇,该系列持续更新中……系列背景介绍:《使用PaddleNLP进行恶意网页识别》系列项目,围绕当前网络安全管理中面临的最大挑战之一——如何防止用户有益/无意中访问恶意网页,进而给企业造成损失的问题展开,尝试使用人工智能技术,针对此问题提出相应的解决方案。系列还有一个姊妹篇,《使用PaddleNLP识别垃圾邮件》,欢迎感兴趣的读者点进来交流评论
本文是《使用PaddleNLP进行恶意网页识别》系列第六篇,该系列持续更新中……
系列背景介绍:《使用PaddleNLP进行恶意网页识别》系列项目,围绕当前网络安全管理中面临的最大挑战之一——如何防止用户有益/无意中访问恶意网页,进而给企业造成损失的问题展开,尝试使用人工智能技术,针对此问题提出相应的解决方案。
系列还有一个姊妹篇,《使用PaddleNLP识别垃圾邮件》,欢迎感兴趣的读者点进来交流评论
系列目录
- 使用PaddleNLP进行恶意网页识别(一)
- 使用PaddleNLP的文本分类模型,做正常网页与被黑网页的简单二分类,根据HTML网页内容处理结果判断网页是否正常。
- 使用PaddleNLP进行恶意网页识别(二)
- 使用PaddleNLP的Ernie预训练模型Fine-tune,大幅提高根据HTML网页内容处理结果判断网页准确率。
- 使用PaddleNLP进行恶意网页识别(三)
- 使用PaddleNLP的文本分类模型,做正常网页与恶意网页的简单二分类,提取HTML标签信息判断网页是否正常。
- 使用PaddleNLP进行恶意网页识别(四)
- 尝试使用人工判断条件,设计提取HTML标签信息识别恶意网页的流程。
- 使用PaddleNLP的预训练模型Fine-tune,尝试提高提取HTML标签信息判断网页是否正常的效果。
- 将动态图训练的网页分类模型导出并使用Python部署。
- 使用PaddleNLP进行恶意网页识别(五)
- 该项目直接对标系列第二篇,对比BERT中文预训练模型和Ernie预训练模型在HTML网页内容分类流程和效果上的差异。
- 项目进一步完善和优化了HTML网页内容提取和数据清洗流程。
- 验证集上模型准确率可以轻松达到91.5%以上,最高达到95%,在BERT预训练模型上进行finetune,得到了目前在该HTML网页内容分类任务上的最好表现。
- 使用PaddleNLP进行恶意网页识别(六)
- 该项目直接对标系列第四篇,对比BERT中文预训练模型和Ernie预训练模型在HTML网页标签分类效果上的差异。
- 项目进一步完善和优化了HTML的tag内容提取和数据清洗流程。
- 验证集上模型准确率可以轻松达到96.5%以上,测试集上准确率接近97%,在BERT预训练模型上进行finetune,得到了目前在该HTML网页标签序列分类任务上的最好表现。
- 使用PaddleNLP进行恶意网页识别(七)
- 介绍了使用自动化测试工具selenium进行网页快照抓取的方法。
- 介绍了使用zxing开源库进行网页快照二维码定位和解析的方法。
- 介绍使用系列第六篇训练的模型,对二维码中包含的url网页链接进行识别分类的思路。
项目介绍
本项目使用的数据集来自《基于数据科学的恶意软件分析》约书亚·萨克斯(Joshua Saxe),希拉里•桑德斯(Hillary Sanders)著
在该书的第11章中,有一个数量大约为10万的正常和恶意HTML文件集。
数据集包括两个文件夹,分别是:
- 正常HTML文件
ch11/data/html/benign_files/
- 恶意HTML文件
ch11/data/html/malicious_files/
记住不要在浏览器中打开这些文件!
原始数据集经过了精心准备,数据分布如下:
- 训练集
- 正常HTML网页文件45000个
- 恶意HTML网页文件45000个
- 验证集
- 正常HTML网页文件5000个
- 恶意HTML网页文件5000个
这些HTML文件的内容,涵盖了英文、中文、俄文等各国语言,如果尝试从网页内容的角度进行恶意网页分类,将面临多语言处理的巨大挑战。
因此,在项目使用PaddleNLP进行恶意网页识别(三)中,通过尝试提取多语言HTML页面的“共性”——HTML标签序列进而进行文本分类,实现了较好的效果(demo准确率大于84%)。
在项目使用PaddleNLP进行恶意网页识别(四)中,基于Ernie预训练模型,HTML标签序列分类的识别准确率进一步提升,达到94%以上。
那么,如果换成BERT,效果又会如何?本项目将进行探讨。
# 解压并提取数据集
!unzip data/data69222/64652.zip
!mv malware_data_science/ch11/chapter_11_UNDER_40/data/html ./
一、环境配置
!pip install bs4
!pip install html5lib
!pip install HTMLParser
import os
import sys
import codecs
import chardet
import shutil
import time
from tqdm import tqdm, trange
from bs4 import BeautifulSoup
import re
from html.parser import HTMLParser
from functools import partial
import numpy as np
import paddle
import paddlenlp as ppnlp
from paddlenlp.data import Stack, Pad, Tuple
import paddle.nn.functional as F
import paddle.nn as nn
from visualdl import LogWriter
import random
二、准备数据集
2.1 优化HTML标签提取结果
在运行上一个项目使用PaddleNLP进行恶意网页识别(三)时发现,其实有个别样本的标签提取结果不太理想,而在项目使用PaddleNLP进行恶意网页识别(四)中。处理方式又过于简单粗暴。
在本项目中,则找到了一种更简单准确的标签提取方法:
text = []
# 找一个网页解析看看
html = BeautifulSoup(open('html/malicious_files/validation/99f756bfa3f3c3be65550b2d27abccb90496089b24252130ce16edc7b83c1ba5'),'html.parser', from_encoding='utf-8')
for tag in html.find_all(True):
text.append(tag.name)
print(','.join(text))
html,head,meta,title,body,script,script
2.2 字符编码处理
def change_code(input_path, output_path):
for filename in tqdm(os.listdir(input_path)):
# 这里要先做个判断,有file_list里面的文件不存在
if os.path.exists(input_path+filename):
# 读取文件,获取字符集
content = codecs.open(input_path+filename,'rb').read()
source_encoding = chardet.detect(content)['encoding']
# 个别文件的source_encoding是None,这里要先进行筛选
if source_encoding is None:
pass
# 对字符集不是utf-8格式的文件尝试转码
elif source_encoding != 'utf-8':
# 转码如果失败,就跳过该文件
try:
content = content.decode(source_encoding).encode('utf-8')
codecs.open(output_path+filename,'wb').write(content)
except UnicodeDecodeError:
print(filename + "读取失败")
pass
# 字符集是utf-8格式的文件直接保存
else:
codecs.open(output_path+filename,'wb').write(content)
else:
pass
!mkdir -p change_code/benign_files/training/
!mkdir -p change_code/benign_files/validation/
!mkdir -p change_code/malicious_files/training/
!mkdir -p change_code/malicious_files/validation/
change_code('html/benign_files/validation/', 'change_code/benign_files/validation/')
change_code('html/benign_files/training/', 'change_code/benign_files/training/')
change_code('html/malicious_files/training/', 'change_code/malicious_files/training/')
change_code('html/malicious_files/validation/', 'change_code/malicious_files/validation/')
2.3 批量提取网页的HTML标签序列
# 创建提取的训练集
def creat_trainset(path, label):
for i, filename in enumerate(tqdm(os.listdir(path))):
tag_list = []
text = ''
html = BeautifulSoup(open(path + filename),'html.parser', from_encoding='utf-8')
for tag in html.find_all(True):
tag_list.append(tag.name)
text = ','.join(tag_list)
with open("webtrain.txt","a+") as f:
f.write(text + '\t' + label + '\n')
creat_trainset('change_code/benign_files/training/', '1')
creat_trainset('change_code/malicious_files/training/', '0')
# 创建提取的训练集
def creat_valset(path, label):
for i, filename in enumerate(tqdm(os.listdir(path))):
tag_list = []
text = ''
html = BeautifulSoup(open(path + filename),'html.parser', from_encoding='utf-8')
for tag in html.find_all(True):
tag_list.append(tag.name)
text = ','.join(tag_list)
with open("webval.txt","a+") as f:
f.write(text + '\t' + label + '\n')
creat_valset('change_code/benign_files/validation/', '1')
creat_valset('change_code/malicious_files/validation/', '0')
2.4 乱序处理,重新划分数据集
#分别统计原训练集与测试集的总数
train_num = 0
non_train_num = 0
all_train_list = []
all_non_train_list = []
f = open("webval.txt","r")
lines = f.readlines() #读取全部内容 ,并以列表方式返回
for line in lines:
all_non_train_list.append(line)
non_train_num +=1
f = open("webtrain.txt","r")
lines = f.readlines() #读取全部内容 ,并以列表方式返回
for line in lines:
all_train_list.append(line)
train_num +=1
data_list_path="./"
all_data_path=data_list_path + "all_data.txt"
all_data_list = all_train_list + all_non_train_list
random.shuffle(all_data_list)
#在生成all_data.txt之前,首先将其清空
with open(all_data_path, 'w') as f:
f.seek(0)
f.truncate()
with open(all_data_path, 'a') as f:
for data in all_data_list:
f.write(data)
with open(os.path.join(data_list_path, 'eval_list.txt'), 'w', encoding='utf-8') as f_eval:
f_eval.seek(0)
f_eval.truncate()
with open(os.path.join(data_list_path, 'train_list.txt'), 'w', encoding='utf-8') as f_train:
f_train.seek(0)
f_train.truncate()
with open(os.path.join(data_list_path, 'test_list.txt'), 'w', encoding='utf-8') as f_test:
f_test.seek(0)
f_test.truncate()
with open(os.path.join(data_list_path, 'all_data.txt'), 'r', encoding='utf-8') as f_data:
lines = f_data.readlines()
i = 0
with open(os.path.join(data_list_path, 'eval_list.txt'), 'a', encoding='utf-8') as f_eval, open(os.path.join(data_list_path, 'test_list.txt'), 'w', encoding='utf-8') as f_test,open(os.path.join(data_list_path, 'train_list.txt'), 'a', encoding='utf-8') as f_train:
for line in lines:
label= line.split('\t')[-1].replace('\n', '')
words= line.split('\t')[0]
labs = ""
if len(words) > 0:
if i % 10 == 0:
labs = words + '\t' + label + '\n'
f_eval.write(labs)
elif i % 10 ==1:
labs = words + '\t' + label + '\n'
f_test.write(labs)
else:
labs = words + '\t' + label + '\n'
f_train.write(labs)
i += 1
print("数据列表生成完成!")
2.5 自定义数据集
class SelfDefinedDataset(paddle.io.Dataset):
def __init__(self, data):
super(SelfDefinedDataset, self).__init__()
self.data = data
def __getitem__(self, idx):
return self.data[idx]
def __len__(self):
return len(self.data)
def get_labels(self):
return ["0", "1"]
def txt_to_list(file_name):
res_list = []
for line in open(file_name):
res_list.append(line.strip().split('\t'))
return res_list
trainlst = txt_to_list('train_list.txt')
devlst = txt_to_list('eval_list.txt')
testlst = txt_to_list('test_list.txt')
train_ds, dev_ds, test_ds= SelfDefinedDataset.get_datasets([trainlst, devlst, testlst])
# 准备标签
label_list = train_ds.get_labels()
print(label_list)
#看看数据长什么样子,分别打印训练集、验证集、测试集的第1条数据。
print("训练集数据:{}\n".format(train_ds[0:1]))
print("验证集数据:{}\n".format(dev_ds[0:1]))
print("测试集数据:{}\n".format(test_ds[0:1]))
print("训练集样本个数:{}".format(len(train_ds)))
print("验证集样本个数:{}".format(len(dev_ds)))
print("测试集样本个数:{}".format(len(test_ds)))
['0', '1']
训练集数据:[['html,head,script,script', '0']]
验证集数据:[['html,head,meta,title,script,title,base,link,body,table,tr,td,a,img,td,a,img,a,img,a,img,a,img,table,tr,td,a,a,a,td,a,a,a,table,tr,td,table,tr,td,table,tr,td,img,td,td,img,table,tr,td,table,tr,td,img,tr,td,a,br,a,br,a,br,a,br,a,br,a,br,a,br,a,br,a,br,a,br,a,br,a,br,a,br,tr,td,img,tr,td,table,tr,td,img,td,td,img,table,tr,td,table,tr,td,img,tr,td,form,select,option,option,option,option,option,option,option,option,option,option,option,option,option,option,option,option,option,option,option,option,option,option,option,option,option,option,option,option,option,option,option,option,option,option,option,input,tr,td,img,tr,td,table,tr,td,img,td,td,a,img,img,table,tr,td,table,tr,td,img,tr,td,a,img,br,a,br,tr,td,img,tr,td,table,tr,td,img,td,td,img,table,tr,td,table,tr,td,img,tr,td,form,input,input,input,br,br,a,b,tr,td,img,tr,td,table,tr,td,img,td,td,img,table,tr,td,table,tr,td,img,tr,td,a,br,a,tr,td,img,td,table,tr,td,table,tr,td,td,img,tr,td,img,tr,td,table,tr,td,tr,td,img,tr,td,table,tr,td,table,tr,td,img,td,a,img,td,img,td,table,tr,td,table,tr,td,img,td,td,a,img,img,table,tr,td,table,tr,td,img,tr,td,tr,td,img,tr,td,table,tr,td,img,td,td,img,table,tr,td,table,tr,td,img,tr,td,table,tr,td,td,a,tr,td,td,a,tr,td,td,a,tr,td,td,a,tr,td,td,a,tr,td,td,a,tr,td,td,a,tr,td,td,a,tr,td,td,a,tr,td,td,a,tr,td,img,tr,td,table,tr,td,img,td,td,a,img,img,table,tr,td,table,tr,td,img,tr,td,tr,td,img,tr,td,table,tr,td,img,td,td,img,table,tr,td,table,tr,td,img,tr,td,a,img,a,img,tr,td,img,tr,td,table,tr,td,img,td,td,img,table,tr,td,table,tr,td,img,tr,td,form,select,option,option,option,input,tr,td,img,table,tr,td,td,br,table,tr,td,br,table,tr,td,br', '0']]
测试集数据:[['html,head,meta,title,script,link,link,link,link,link,body,style,div,a,div,div,div,div,div,div,div,div,div,div,script,br,br,br,br,a,br,br,br,br,script,script,script,script,script,script,script,script,script,script,script,script,script', '1']]
训练集样本个数:78075
验证集样本个数:9767
测试集样本个数:9760
2.6 数据预处理
#调用ppnlp.transformers.BertTokenizer进行数据处理,tokenizer可以把原始输入文本转化成模型model可接受的输入数据格式。
tokenizer = ppnlp.transformers.BertTokenizer.from_pretrained("bert-base-cased")
#数据预处理
def convert_example(example,tokenizer,label_list,max_seq_length=256,is_test=False):
if is_test:
text = example
else:
text, label = example
#tokenizer.encode方法能够完成切分token,映射token ID以及拼接特殊token
encoded_inputs = tokenizer.encode(text=text, max_seq_len=max_seq_length)
input_ids = encoded_inputs["input_ids"]
#注意,在早前的PaddleNLP版本中,token_type_ids叫做segment_ids
segment_ids = encoded_inputs["token_type_ids"]
if not is_test:
label_map = {}
for (i, l) in enumerate(label_list):
label_map[l] = i
label = label_map[label]
label = np.array([label], dtype="int64")
return input_ids, segment_ids, label
else:
return input_ids, segment_ids
#数据迭代器构造方法
def create_dataloader(dataset, trans_fn=None, mode='train', batch_size=1, use_gpu=False, pad_token_id=0, batchify_fn=None):
if trans_fn:
dataset = dataset.apply(trans_fn, lazy=True)
if mode == 'train' and use_gpu:
sampler = paddle.io.DistributedBatchSampler(dataset=dataset, batch_size=batch_size, shuffle=True)
else:
shuffle = True if mode == 'train' else False #如果不是训练集,则不打乱顺序
sampler = paddle.io.BatchSampler(dataset=dataset, batch_size=batch_size, shuffle=shuffle) #生成一个取样器
dataloader = paddle.io.DataLoader(dataset, batch_sampler=sampler, return_list=True, collate_fn=batchify_fn)
return dataloader
#使用partial()来固定convert_example函数的tokenizer, label_list, max_seq_length, is_test等参数值
trans_fn = partial(convert_example, tokenizer=tokenizer, label_list=label_list, max_seq_length=128, is_test=False)
batchify_fn = lambda samples, fn=Tuple(Pad(axis=0,pad_val=tokenizer.pad_token_id), Pad(axis=0, pad_val=tokenizer.pad_token_id), Stack(dtype="int64")):[data for data in fn(samples)]
#训练集迭代器
train_loader = create_dataloader(train_ds, mode='train', batch_size=64, batchify_fn=batchify_fn, trans_fn=trans_fn)
#验证集迭代器
dev_loader = create_dataloader(dev_ds, mode='dev', batch_size=64, batchify_fn=batchify_fn, trans_fn=trans_fn)
#测试集迭代器
test_loader = create_dataloader(test_ds, mode='test', batch_size=64, batchify_fn=batchify_fn, trans_fn=trans_fn)
三、模型训练
3.1 加载BERT预训练模型
#加载预训练模型Bert用于文本分类任务的Fine-tune网络BertForSequenceClassification, 它在BERT模型后接了一个全连接层进行分类。
#由于本任务中的恶意网页识别是二分类问题,设定num_classes为2
model = ppnlp.transformers.BertForSequenceClassification.from_pretrained("bert-base-cased", num_classes=2)
3.2 开始训练
#设置训练超参数
#学习率
learning_rate = 5e-5
#训练轮次
epochs = 10
#学习率预热比率
warmup_proption = 0.1
#权重衰减系数
weight_decay = 0.01
num_training_steps = len(train_loader) * epochs
num_warmup_steps = int(warmup_proption * num_training_steps)
def get_lr_factor(current_step):
if current_step < num_warmup_steps:
return float(current_step) / float(max(1, num_warmup_steps))
else:
return max(0.0,
float(num_training_steps - current_step) /
float(max(1, num_training_steps - num_warmup_steps)))
#学习率调度器
lr_scheduler = paddle.optimizer.lr.LambdaDecay(learning_rate, lr_lambda=lambda current_step: get_lr_factor(current_step))
#优化器
optimizer = paddle.optimizer.AdamW(
learning_rate=lr_scheduler,
parameters=model.parameters(),
weight_decay=weight_decay,
apply_decay_param_fun=lambda x: x in [
p.name for n, p in model.named_parameters()
if not any(nd in n for nd in ["bias", "norm"])
])
#损失函数
criterion = paddle.nn.loss.CrossEntropyLoss()
#评估函数
metric = paddle.metric.Accuracy()
#评估函数,设置返回值,便于VisualDL记录
def evaluate(model, criterion, metric, data_loader):
model.eval()
metric.reset()
losses = []
for batch in data_loader:
input_ids, segment_ids, labels = batch
logits = model(input_ids, segment_ids)
loss = criterion(logits, labels)
losses.append(loss.numpy())
correct = metric.compute(logits, labels)
metric.update(correct)
accu = metric.accumulate()
print("eval loss: %.5f, accu: %.5f" % (np.mean(losses), accu))
model.train()
metric.reset()
return np.mean(losses), accu
#开始训练
global_step = 0
max_acc = 0
with LogWriter(logdir="./log") as writer:
for epoch in range(1, epochs + 1):
for step, batch in enumerate(train_loader, start=1): #从训练数据迭代器中取数据
input_ids, segment_ids, labels = batch
logits = model(input_ids, segment_ids)
loss = criterion(logits, labels) #计算损失
probs = F.softmax(logits, axis=1)
correct = metric.compute(probs, labels)
metric.update(correct)
acc = metric.accumulate()
global_step += 1
if global_step % 100 == 0 :
print("global step %d, epoch: %d, batch: %d, loss: %.5f, acc: %.5f" % (global_step, epoch, step, loss, acc))
#记录训练过程
writer.add_scalar(tag="train/loss", step=global_step, value=loss)
writer.add_scalar(tag="train/acc", step=global_step, value=acc)
loss.backward()
optimizer.step()
lr_scheduler.step()
optimizer.clear_gradients()
eval_loss, eval_acc = evaluate(model, criterion, metric, dev_loader)
#记录评估过程
writer.add_scalar(tag="eval/loss", step=epoch, value=eval_loss)
writer.add_scalar(tag="eval/acc", step=epoch, value=eval_acc)
# 保存最佳模型
if eval_loss>max_acc:
max_acc = eval_loss
print('saving the best_model...')
paddle.save(model.state_dict(), 'best_model')
# 保存最终模型
paddle.save(model.state_dict(),'final_model')
3.3 保存模型和网络结构
# Convert to static graph with specific input description
model = paddle.jit.to_static(
model,
input_spec=[
paddle.static.InputSpec(
shape=[None, None], dtype="int64"), # input_ids
paddle.static.InputSpec(
shape=[None, None], dtype="int64") # segment_ids
])
# Save in static graph model.
paddle.jit.save(model, './static_graph_params')
可以看到,使用BERT预训练模型进行finetune,10个epoch内验证集准确率已经可以达到96.7%以上。
在VisualDL中,也可以查看BERT模型的网络结构,方法如下:
四、预测效果
完成上面的模型训练之后,可以得到一个能够通过HTML标签序列识别是否存在恶意网页的模型。接下来查看模型在测试集上的表现,非常接近97%。
# 评估模型在测试集上的表现
evaluate(model, criterion, metric, test_loader)
eval loss: 0.10088, accu: 0.96998
(0.10087659, 0.9699795081967213)
def predict(model, data, tokenizer, label_map, batch_size=1):
examples = []
for text in data:
input_ids, segment_ids = convert_example(text, tokenizer, label_list=label_map.values(), max_seq_length=128, is_test=True)
examples.append((input_ids, segment_ids))
batchify_fn = lambda samples, fn=Tuple(Pad(axis=0, pad_val=tokenizer.pad_token_id), Pad(axis=0, pad_val=tokenizer.pad_token_id)): fn(samples)
batches = []
one_batch = []
for example in examples:
one_batch.append(example)
if len(one_batch) == batch_size:
batches.append(one_batch)
one_batch = []
if one_batch:
batches.append(one_batch)
results = []
model.eval()
for batch in batches:
input_ids, segment_ids = batchify_fn(batch)
input_ids = paddle.to_tensor(input_ids)
segment_ids = paddle.to_tensor(segment_ids)
logits = model(input_ids, segment_ids)
probs = F.softmax(logits, axis=1)
idx = paddle.argmax(probs, axis=1).numpy()
idx = idx.tolist()
labels = [label_map[i] for i in idx]
results.extend(labels)
return results
import requests
# 获取某知名网站首页链接
r = requests.get("https://www.csdn.net")
demo = r.text
soup=BeautifulSoup(demo,"html.parser")
tags = []
for tag in soup.find_all(True):
tags.append(tag.name)
data = []
data.append(','.join(tags))
label_map = {0: '恶意网页', 1: '正常网页'}
predictions = predict(model, data, tokenizer, label_map, batch_size=64)
for idx, text in enumerate(data):
ttps://www.csdn.net")
demo = r.text
soup=BeautifulSoup(demo,"html.parser")
tags = []
for tag in soup.find_all(True):
tags.append(tag.name)
data = []
data.append(','.join(tags))
label_map = {0: '恶意网页', 1: '正常网页'}
predictions = predict(model, data, tokenizer, label_map, batch_size=64)
for idx, text in enumerate(data):
print('预测网页: {} \n网页标签: {}'.format("https://www.csdn.net", predictions[idx]))
预测网页: https://www.csdn.net
网页标签: 正常网页
小结
- 在该项目中,进一步完善了HTML标签序列的提取流程,项目也优化了日志文件写入方式,在ViusalDL中训练过程将连续显示。
- 使用BERT预训练模型Finetune后,两种分类模型预测准确率已经接近97%。
- 在项目最后,已经出现了网页识别工程化的雏形(获取网页链接——提取标签序列——判断网页类型),接下来将进一步探索。
更多推荐
所有评论(0)