本文是《使用PaddleNLP进行恶意网页识别》系列第五篇,该系列持续更新中……

系列背景介绍:《使用PaddleNLP进行恶意网页识别》系列项目,围绕当前网络安全管理中面临的最大挑战之一——如何防止用户有益/无意中访问恶意网页,进而给企业造成损失的问题展开,尝试使用人工智能技术,针对此问题提出相应的解决方案。

系列还有一个姊妹篇,《使用PaddleNLP识别垃圾邮件》,欢迎感兴趣的读者点进来交流评论https://ai-studio-static-online.cdn.bcebos.com/d500f483e55544a5b929abad59de208f180c068cc81648009fab60a0b6d9bda2

系列目录

  • 使用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网页链接进行识别分类的思路。

关于BERT模型的详细介绍与PaddleNLP对BERT模型的应用,在项目PaddleNLP2.0:BERT模型的应用中进行了详细总结,欢迎感兴趣的读者交流指导。

一、环境配置

本教程基于Paddle 2.0 编写,如果你的环境不是本版本,请先参考官网安装 Paddle 2.0 。

!pip install bs4 -i https://mirror.baidu.com/pypi/simple/
import os
import sys
import codecs
import chardet
import shutil
import re
import time
import numpy as np
import pandas as pd
import jieba
from tqdm import tqdm, trange
from bs4 import BeautifulSoup
from functools import partial
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

print(paddle.__version__)

二、数据加载

2.1 数据集加载

将使用 https://www.heywhale.com/ 提供的恶意网页分析样本作为数据集,来完成本任务。该数据集含有169个恶意网页。

当然,也可以直接从项目挂载的数据集中解压。

!unzip /home/aistudio/data/data69222/MaliciousWebpage.zip

2.2 读取文件信息

读取文件列表信息。

columns = ['id', 'flag', 'filename', 'url']
tempdf = pd.read_csv('MaliciousWebpage/file_list.txt', sep=',',skiprows=0, header=None, names=columns, skipfooter=0)

2.3 使用下采样处理数据集不均衡问题

在数据集中,正常网页样本有9700个,而恶意网页样本近169个,数据集严重不均衡,使用下采样的方法,随机筛选出500个正常网页的样本参与训练。

n_page = tempdf[tempdf['flag']=='n']
# 对正常页面进行随机采样
n_page = n_page.sample(n=500)
# 提取全部被黑页面样本
d_page = tempdf[tempdf['flag']=='d']
# 合并样本
train_page = pd.concat([n_page,d_page],axis=0)
# 合并样本
train_page = pd.concat([n_page,d_page],axis=0)
# 做一个乱序
train_page = train_page.sample(frac = 1) 

2.4 进行字符集编码处理

解析数据集中网页内容时,可能出现因字符集编码不一致导致的读取错误,因此,要先对进行批量字符集编码转换,对数据集进行清洗。

!mkdir TrainWebpage && mkdir TrainWebpage/file1
# 去掉非中文字符
def clean_str(string):
    string = re.sub(r"[^\u4e00-\u9fff]", " ", string)
    string = re.sub(r"\s{2,}", " ", string)
    return string.strip()
    
for filename in tqdm(train_page['filename']):
    # 这里要先做个判断,有file_list里面的文件不存在
    if os.path.exists('MaliciousWebpage/file1/'+filename):
        # 读取文件,获取字符集
        content = codecs.open('MaliciousWebpage/file1/'+filename,'rb').read()
        source_encoding = chardet.detect(content)['encoding']
        # 个别文件的source_encoding是None,这里要先进行筛选
        if source_encoding is None:
            pass
        # 只对字符集是gb2312格式的文件尝试转码
        elif source_encoding == 'gb2312':
            # 转码如果失败,就跳过该文件
            try:
                content = content.decode(source_encoding).encode('utf-8')
                codecs.open('TrainWebpage/file1/'+filename,'wb').write(content)
            except UnicodeDecodeError:
                print(filename + "读取失败")
                pass
        # 字符集是utf-8格式的文件直接保存
        elif source_encoding == 'utf-8':
            codecs.open('TrainWebpage/file1/'+filename,'wb').write(content)
        else:
            pass
    else:
        pass
100%|██████████| 669/669 [02:28<00:00,  4.52it/s]

2.5 提取网页内容,划分训练集、验证集、测试集

被黑网页的一个典型特征是恶意插入的内容大量集中在HTML页面底部,因此可以提取网页末尾的HTML内容作为输入LSTM的文本信息。

for i, filename in enumerate(train_page['filename']):
    # 这里要先做个判断,有file_list里面的文件不存在
    if os.path.exists('TrainWebpage/file1/'+filename):
        # 读取文件,解析HTML页面
        html = BeautifulSoup(open('TrainWebpage/file1/'+filename),'html.parser', from_encoding='utf-8')
        text = ''.join(list(html.stripped_strings)[-20:])
        # 去掉多余的换行符(部分数据最后解析结果为空格)
        text = text.replace("\n", "")
        text = text.replace(" ", ",")
         # 去掉非中文字符
        text = clean_str(text)
        # 这一步非常重要,否则会报错:清洗脏数据,如果文本为空,不参与训练
        if len(text) > 0:
            if i % 5 == 0:
                if train_page['flag'][train_page['filename']==filename].values[0] == 'n':
                    with open("webtest.txt","a+") as f:
                        f.write(text[-100:] + '\t' + '0' + '\n')
                elif train_page['flag'][train_page['filename']==filename].values[0] == 'd':
                    with open("webtest.txt","a+") as f:
                        f.write(text[-100:] + '\t' + '1' + '\n')
            elif i % 5 == 1:
                if train_page['flag'][train_page['filename']==filename].values[0] == 'n':
                    with open("webdev.txt","a+") as f:
                        f.write(text[-100:] + '\t' + '0' + '\n')
                elif train_page['flag'][train_page['filename']==filename].values[0] == 'd':
                    with open("webdev.txt","a+") as f:
                        f.write(text[-100:] + '\t' + '1' + '\n')
            else:
                if train_page['flag'][train_page['filename']==filename].values[0] == 'n':
                    with open("webtrain.txt","a+") as f:
                        f.write(text[-100:] + '\t' + '0' + '\n')
                elif train_page['flag'][train_page['filename']==filename].values[0] == 'd':
                    with open("webtrain.txt","a+") as f:
                        f.write(text[-100:] + '\t' + '1' + '\n')
    else:
        pass

2.6 自定义数据集

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('webtrain.txt')
devlst = txt_to_list('webdev.txt')
testlst = txt_to_list('webtest.txt')

train_ds, dev_ds, test_ds = SelfDefinedDataset.get_datasets([trainlst, devlst, testlst])
# 准备标签
label_list = train_ds.get_labels()
print(label_list)
#看看数据长什么样子,分别打印训练集、验证集、测试集的前3条数据。
print("训练集数据:{}\n".format(train_ds[0:3]))
print("验证集数据:{}\n".format(dev_ds[0:3]))
print("测试集数据:{}\n".format(test_ds[0:3]))

print("训练集样本个数:{}".format(len(train_ds)))
print("验证集样本个数:{}".format(len(dev_ds)))
print("测试集样本个数:{}".format(len(test_ds)))

2.7 数据预处理

#调用ppnlp.transformers.BertTokenizer进行数据处理,tokenizer可以把原始输入文本转化成模型model可接受的输入数据格式。
tokenizer = ppnlp.transformers.BertTokenizer.from_pretrained("bert-base-chinese")

#数据预处理
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-chinese", num_classes=2)

3.2 开始训练

#设置训练超参数

#学习率
learning_rate = 1e-4
#训练轮次
epochs = 20
#学习率预热比率
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
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 % 1 == 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)

可以发现,验证集上模型准确率可以轻松达到91.5%以上,最高达到95%,在BERT预训练模型上进行finetune,得到了目前在HTML恶意网页识别上的最好表现,远高于LSTM的81.5%,也高于基于Ernie预训练模型上finetune的87.9%。

四、预测效果

完成上面的模型训练之后,可以得到一个能够通过网页内容识别是否为恶意网页的模型。接下来查看模型在测试集上的泛化能力。

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
data = ['娱乐城 百家乐 澳门百家乐 金宝博 澳门赌场 汉庭网络 真人百家乐博彩网皇冠网全讯网 皇冠现金网赌博网站 官方网站皇冠足球比分', '登入 简 錯誤 錯誤 用戶帳號 密碼 記住我的戶口號碼登入登入重要告示 私隱政策桌面版本']
label_map = {0: '正常网页', 1: '恶意网页'}

predictions = predict(model, data, tokenizer, label_map, batch_size=32)
for idx, text in enumerate(data):
 錯誤 錯誤 用戶帳號 密碼 記住我的戶口號碼登入登入重要告示 私隱政策桌面版本']
label_map = {0: '正常网页', 1: '恶意网页'}

predictions = predict(model, data, tokenizer, label_map, batch_size=32)
for idx, text in enumerate(data):
    print('预测内容: {} \n网页标签: {}'.format(text, predictions[idx]))
预测内容: 娱乐城 百家乐 澳门百家乐 金宝博 澳门赌场 汉庭网络 真人百家乐博彩网皇冠网全讯网 皇冠现金网赌博网站 官方网站皇冠足球比分 
网页标签: 恶意网页
预测内容: 登入 简 錯誤 錯誤 用戶帳號 密碼 記住我的戶口號碼登入登入重要告示 私隱政策桌面版本 
网页标签: 正常网页

致谢

该系列项目能够顺利开发,要特别感谢社区@没入门的研究生大佬的指导,瞄一眼就能debug!https://ai-studio-static-online.cdn.bcebos.com/d500f483e55544a5b929abad59de208f180c068cc81648009fab60a0b6d9bda2

Logo

学大模型,用大模型上飞桨星河社区!每天8点V100G算力免费领!免费领取ERNIE 4.0 100w Token >>>

更多推荐