[网络安全篇] 引入Vision Transformer

近年来，恶意软件的出现尤为猖獗。所以解决恶意软件的问题也成为了刻不容缓的事情，当然，现在运用机器学习去完成恶意软件分类的手段已经比较的成熟了，但是，机器学习去特征的处理要求是比较高的，这里我采用深度学习的思维去解决。与此同时，Canadian Institute for Cybersecurity这个研究所提供了近几年的安卓恶意软件的提取特征文件，我在此基础上进行改进，初步完成一个较为入门级别的恶意软件分类项目。

0 背景

本项目仍然是以处理图像的手段完成恶意软件的分类问题，但是加入了目前比较流行的VIT模型结构。

ViT算法综述

论文地址：An Image is Worth 16x16 Words:Transformers for Image Recognition at Scale

之前的算法大都是保持CNN整体结构不变，在CNN中增加attention模块或者使用attention模块替换CNN中的某些部分。ViT算法中，作者提出没有必要总是依赖于CNN，仅仅使用Transformer结构也能够在图像分类任务中表现很好。

受到NLP领域中Transformer成功应用的启发，ViT算法中尝试将标准的Transformer结构直接应用于图像，并对整个图像分类流程进行最少的修改。具体来讲，ViT算法中，会将整幅图像拆分成小图像块，然后把这些小图像块的线性嵌入序列作为Transformer的输入送入网络，然后使用监督学习的方式进行图像分类的训练。ViT算法的整体结构如 图1 所示。

项目的前半部分，仍然是和之前相同的数据处理方式，这次为了更好的契合VIT的切片，我调整了图像了生成的大小为20x20，这样的话，我就可以设置vit中的patch大小为4x4的大小。

当然这里还有个思考，就是图像太小了，后续会考虑自行去反汇编提取特征，或者使用作者提供的大的那个数据集进行图像构建，这个项目仍然使用471的csv文件数据。

当然，项目中对VIT部分的代码不做过多的解释，大家可以自行去查看VIT的paddle复现

PaddleVit:

https://github.com/BR-IDL/PaddleViT

1 数据集

详情参照CIC研究所的具体发布内容。

https://www.unb.ca/cic/datasets/maldroid-2020.html

这里我采用了其中470维度的一个数据进行处理，数据集已经挂载在项目数据集当中。

# 对数据进行一个copy
!cp data/data157200/syscallsbinders_frequency_5_Cat.csv work/

2 数据处理

2.1 数据展示

因为是csv文件，我这里采用pandas的方式进行读取，更为方便快速。

可以通过表头看到，是各种权限、调用函数、绑定器的使用频次统计。

每一行的最后一列是该行数据所属于的分类类别。

import pandas as pd
# 按文件名读取整个文件
data_471 = pd.read_csv("work/syscallsbinders_frequency_5_Cat.csv") 
data_471.head(2)     # 展示前几行数据，通常为展示前5行

	ACCESS_PERSONAL_INFO___	ALTER_PHONE_STATE___	ANTI_DEBUG_____	CREATE_FOLDER_____	CREATE_PROCESS`_____	CREATE_THREAD_____	DEVICE_ACCESS_____	EXECUTE_____	FS_ACCESS____	FS_ACCESS()____	...	utimes	vfork	vibrate	vibratePattern	wait4	watchRotation	windowGainedFocus	write	writev	Class
0	1	0	0	3	0	14	2	0	3	0	...	0	0	0	0	0	0	0	37	10	1
1	3	0	0	6	0	42	91	0	32	0	...	0	0	0	0	0	0	2	2838	46	1

2 rows × 471 columns

2.2 无效数据过滤

我这里的思维就是，竖向的每一列，如果求和他的值为0，那么他的这一列的数值存在将是毫无意义的。

换句话来讲，其实这里也是一种特征降维。

因为我渴望通过图像的方式去解决问题，那么我就可以将每一行的数据保留到$NxN$的一个维度。

那么，这里是不是可以放大这个0值，每一列小于某个阈值的话，我们都可以视为这一列起不到作用。

这里我采用3这个阈值，这样可以把数据降到403维，最后一位为分类类别，相当于就是

20 * 20 = 400 + 1 = 401

那么其实我们还需要再删除两个维度的数据，才能到401，这里我自行看了新生成的403的csv文件

表头为uname、umask的这两列也是没有什么太大影响的列，可以考虑去除

（这里不用太过于在意去除的影响是不是负面的）

import pandas as pd
# 按文件名读取整个文件
data_471 = pd.read_csv("work/syscallsbinders_frequency_5_Cat.csv") 
df = pd.DataFrame(data_471)
len(df.columns)
# 记录每一列求和小于阈值的列
threshold = 3
de = []
de_max = 0
de_va = []
for i in range(0,470):
    name = df.columns.values[i]
    if(df[name].sum() <= threshold):
        de.append(name)
    if(df[name].max() >= de_max):
        de_max = df[name].max()

# 删除小于阈值的列
for j in de:
    df.drop(j,axis=1,inplace=True)

df.drop('umask',axis=1,inplace=True)
df.drop('uname',axis=1,inplace=True)

len_new = len(df.columns)

len_new

401

PS： 不合理点

这里有些频次的调用很大，这里统计了最大的数值de_max，为3697410.

那么，我们将其转换为图像时，范围是0-255才是合理的。

所以需要使用归一化操作，但是这里，这个值过于庞大，肯定是会将其他数值缩放之后几乎为0的状态。

我这里采取，大值统一设置为10000，当然大家可以尝试其他值看看结果的差异有多大。

de_max

3697410

# 新数据写入csv
df.to_csv("new_date_"+str(len_new)+".csv",index=False,sep=',')

import numpy as np
import pandas as pd

mydata = pd.read_csv("new_date_"+str(len_new)+".csv")
mydata_array = np.array(mydata)
mydata_array

array([[   1,    0,    0, ...,   37,   10,    1],
       [   3,    0,    0, ..., 2838,   46,    1],
       [   2,    0,    0, ...,  111,   20,    1],
       ...,
       [   0,    0,    0, ...,  241,   67,    5],
       [   1,    0,    0, ..., 1703,  774,    5],
       [   0,    0,    0, ..., 3102,  186,    5]])

mydata_array.shape

(11598, 401)

# 获取最后一维的数据，也就是分类类别 1-5
mydata_array[0][-1]

1

# 删除最后一列
del_arr = np.delete(mydata_array, -1, axis=1)

del_arr

array([[   1,    0,    0, ...,    0,   37,   10],
       [   3,    0,    0, ...,    2, 2838,   46],
       [   2,    0,    0, ...,    1,  111,   20],
       ...,
       [   0,    0,    0, ...,    5,  241,   67],
       [   1,    0,    0, ...,    3, 1703,  774],
       [   0,    0,    0, ...,   13, 3102,  186]])

# 将其转换为11598条17*17的数据
del_arr = del_arr.reshape(-1,20,20)

del_arr.shape

(11598, 20, 20)

import os
def init_mkdir():
    for i in range(1,6):
        data_path = os.path.join('mal_data',str(i)) # 文件夹路径'data\1'
        if not os.path.exists(data_path): # 判断文件夹是否存在
            os.makedirs(data_path) # 不存在则新建文件夹

2.3 生成灰度图

这里使用了阈值去将大值设置为阈值的操作。

同时进行了数据的归一化

import cv2
import numpy as np
import matplotlib.pylab as pylab
from sklearn.utils import shuffle

# 忽略（垃圾）警告信息
import warnings
warnings.filterwarnings("ignore")
# 在 notebook 画图展示
%matplotlib inline

# 大于10000的数据会被变成10000
img_yuzhi = 10000

# 生成灰度图图像
def mk_img(): 
    # 初始化文件夹
    init_mkdir()
    print("灰度图生成ing...")
    with_label_txt = []
    for i in range(11598):
        # 使用归一化
        x=  np.clip(del_arr[i],0,img_yuzhi)
        x = (x-np.min(x))/(np.max(x)-np.min(x)) # 缩放到0-1
        img = x * 255 # 放大到0-255
        label = mydata_array[i][-1]
        img_path = "mal_data/" + str(label) + "/" +str(i+1) + ".jpg"
        # 生成带有标签的txt 例： data/1/1.jpg 1
        with_label_txt.append(img_path+" " +str(label-1) + "\n")
        cv2.imwrite(img_path,img)
    
    shuf_with_label_txt = shuffle(with_label_txt)
    all_str = ''.join(shuf_with_label_txt)
    f = open('shuf_with_label_txt.txt','w',encoding='utf-8')
    f.write(all_str)
    
    print("标签文件TXT已生成，灰度图生成完成！")

    return shuf_with_label_txt,len(shuf_with_label_txt)

shuf_txt, length = mk_img()

print(length)

灰度图生成ing...
标签文件TXT已生成，灰度图生成完成！
11598

到这里位置，其实大家就可以采用任何的图像分类的方式去实现软件的分类了。

灰度图可视化：

# 按照比例划分数据集 总数据有11598张图片，我这里采用8:2来划分数据
train_size = int(length * 0.8)
train_list = shuf_txt[:train_size]
val_list = shuf_txt[train_size:]

print(len(train_list))
print(len(val_list))

9278
2320

# 运行cell，生成训练集txt 
train_txt = ''.join(train_list)
f_train = open('train_list.txt','w',encoding='utf-8')
f_train.write(train_txt)
f_train.close()
print("train_list.txt 生成成功！")

# 运行cell，生成验证集txt
val_txt = ''.join(val_list)
f_val = open('val_list.txt','w',encoding='utf-8')
f_val.write(val_txt)
f_val.close()
print("val_list.txt 生成成功！")

train_list.txt 生成成功！
val_list.txt 生成成功！

# ['maldata/4/7659.jpg 4\n',...]
train_img = []
train_label = []
for img_label in train_list:
    img,label = img_label.split("\n")[0].split(" ")
    train_img.append(img)
    train_label.append(label)

val_img = []
val_label = []
for img_label in val_list:
    img,label = img_label.split("\n")[0].split(" ")
    val_img.append(img)
    val_label.append(label)

print("-----train sample demo------")
print(train_list[0])
print(train_img[0])
print(train_label[0])
print("-------val sample demo------")
print(val_list[0])
print(val_img[0])
print(val_label[0])

-----train sample demo------
mal_data/5/10229.jpg 4

mal_data/5/10229.jpg
4
-------val sample demo------
mal_data/1/618.jpg 0

mal_data/1/618.jpg
0

3 dataset类编写

from paddle.io import Dataset
from PIL import Image

class MyDataset(Dataset):
    def __init__(self, mode = 'train'):
        # 训练样本数量
        self.training_data, self.training_label, self.test_data, self.test_label = train_img, train_label, val_img, val_label
        if mode  == 'train':
            self.num_samples = len(train_img)
        else:
            self.num_samples = len(val_img)
        self.mode = mode

    def __getitem__(self, idx):
        if self.mode == 'train':
            image = self.training_data[idx]
            label = self.training_label[idx]
        else:
            image = self.test_data[idx]
            label = self.test_label[idx]
        
        im = Image.open(image).convert('L')
        im = im.resize((20, 20), Image.ANTIALIAS)
        img = np.array(im).astype('float32')
        
        return img, np.array(label, dtype='int64')

    def __len__(self):
        # 返回样本总数量
        return self.num_samples

# 训练的数据提供器
train_dataset = MyDataset(mode='train')
# 测试的数据提供器
eval_dataset = MyDataset(mode='val')

# 查看训练和测试数据的大小
print('train大小：', train_dataset.__len__())
print('eval大小：', eval_dataset.__len__())

# 查看图片数据、大小及标签
for data, label in train_dataset:
    print(np.array(data).shape)
    break

train大小： 9278
eval大小： 2320
(20, 20)

4 VIT 网络搭建

4.1 Multihead self-Attention

将图像转化为序列后，就可以将其输入到 Tranformer 结构中进行特征提取了。Tranformer 结构中最重要的结构就是 Multi-head Attention，即多头注意力结构.

# 实现Multihead self-Attention模块
import paddle
import paddle.nn as nn

class Attention(nn.Layer):
    def __init__(self, embed_dim, num_heads, qkv_bias=False, qk_scale=None, droupout=0.):
        super().__init__()
        self.embed_dim = embed_dim  # 每个token的维度，对应于Q的宽度
        self.num_heads = num_heads  # 头
        self.head_dim = int(self.embed_dim / num_heads)  # 每个头需要将token投射到什么长度
        self.all_head_dim = self.head_dim * num_heads  # 这个参数的目的只是为了方便编写待会的fc层
        self.qkv = nn.Linear(embed_dim, 
                            self.all_head_dim * 3, 
                            bias_attr=False if qkv_bias is False else None,  # bias_attr默认为None，此时框架提供初始全零的偏置
                            )
        self.scale = self.embed_dim ** -0.5 if qk_scale is None else qk_scale  # 为了对计算的注意力进行缩放的根号dk，这里提供了可选比例
        self.softmax = nn.Softmax(-1)  # 沿最后一维做softmax，也就是将每个q与所有k的内积进行概率化
        self.proj = nn.Linear(self.all_head_dim, embed_dim)  # 投射层
    
    def transpose_multi_head(self, x):
        # 这个函数就是对输入的所有头的x分离一下，输出每个头的 k/q/v
        # x:[batch, num_patches, all_head_dim]
        new_shape = x.shape[:-1] + [self.num_heads, self.head_dim]
        x = x.reshape(new_shape)  # x:[batch, num_patches, num_heads, head_dim]
        x = x.transpose([0, 2, 1, 3])  # x:[batch, num_heads, num_patches, head_dim]
        # 这个转置吧head提前，目的是保证每个头单独做自注意力计算
        return x

    def forward(self, x):
        batch, num_patches, _ = x.shape  # 输入x的形状为[batch, num_patches, embed_dim]
        qkv = self.qkv(x).chunk(3, -1)  # 线性层默认只对最后一个维度进行全连接计算，经过qkv后的最后一维维度是all_head_dim * 3, 切分返回列表含三个元素
        q, k, v = map(self.transpose_multi_head, qkv)  # 列表每个元素分别经过前面函数句柄的处理返回三个值
        # q, k, v : [batch, num_heads, num_patches, head_dim]
        attn = paddle.matmul(q, k, transpose_y=True)  # 计算Q * K'
        attn = self.softmax(attn * self.scale)  # 计算softmax
        # attn : [batch, num_heads, num_patches, num_patches]
        out = paddle.matmul(attn, v)  # attn为n*n, v为n*dk, 直接做矩阵乘法
        # out:[batch, num_heads, num_patches, embed_dim]
        out = out.transpose([0, 2, 1, 3])  
        # out:[batch, num_patches, num_heads, embed_dim]
        # 多头结果合并
        out = out.reshape([batch, num_patches, -1])  # 输出x的形状为[batch, num_patches, embed_dim]
        # 投射到原始维度，接下来准备进行MLP操作
        out = self.proj(out)
        return out

# 测试注意力模块:输入和输出形状一样
MSA = Attention(embed_dim=20, num_heads=6)
x = paddle.randn([8, 16, 20])
out = MSA(x)
print(out.shape)

W0720 10:51:16.782967   103 gpu_context.cc:278] Please NOTE: device: 0, GPU Compute Capability: 7.0, Driver API Version: 11.2, Runtime API Version: 10.1
W0720 10:51:16.786649   103 gpu_context.cc:306] device: 0, cuDNN Version: 7.6.


[8, 16, 20]

4.2 EncoderLayer

这里实现的是上图的transformer的encoderLayer

# 实现EncoderLayer
import paddle
import paddle.nn as nn
class EncoderLayer(nn.Layer):
    def __init__(self, embed_dim=400, num_heads=6, mlp_ratio=2.0):
        super().__init__()
        self.attn_norm = nn.LayerNorm(embed_dim)
        self.attn = Attention(embed_dim, num_heads)
        self.mlp_norm = nn.LayerNorm(embed_dim)
        self.mlp = Mlp(embed_dim, mlp_ratio)

    def forward(self, x):
        # 采用prenorm，也就是说layernorm层在MSA和MLP之前实现
        # 下面就是一个encoder的逻辑
        h = x
        x = self.attn_norm(x)
        x = self.attn(x)
        x = x + h

        h = x
        x = self.mlp_norm(x)
        x = self.mlp(x)
        x = x + h
        return x

# 实现Encoder
class Encoder(nn.Layer):
    def __init__(self, embed_dim, depth):
        super().__init__()
        layer_list = []
        # depth就是堆叠多少个encoder
        for i in range(depth):
            encoder_layer = EncoderLayer()
            layer_list.append(encoder_layer)
        self.layers = nn.LayerList(layer_list)
        self.norm = nn.LayerNorm(embed_dim)

    def forward(self, x):
        for layer in self.layers:
            x = layer(x)
        x = self.norm(x)
        return x

4.3 PatchEmbedding

# 实现Positionembedding
class PatchEmbedding(nn.Layer):
    def __init__(self, image_size=20, patch_size=4, in_channels=1, embed_dim=400, dropout=0.):
        super().__init__()
        # 计算patches的数量，默认是5x5
        num_patches = (image_size // patch_size) * (image_size // patch_size)
        self.patch_embedding = nn.Conv2D(in_channels=in_channels, out_channels=embed_dim, kernel_size=patch_size, stride=patch_size)
        self.dropout = nn.Dropout(dropout)
        self.embed_dim = embed_dim
        # 添加cls_token
        self.class_token = paddle.create_parameter(
                    shape=[1, 1, embed_dim],
                    dtype='float32',
                    default_initializer=nn.initializer.Constant(0.))
        
        # 添加位置编码
        self.position_embedding = paddle.create_parameter(
            shape=[1, num_patches + 1, embed_dim],
            dtype='float32',
            default_initializer=nn.initializer.TruncatedNormal(std=(0.2)))
        

    def forward(self, x):
        # x:[N,C,H,W]
        cls_token = self.class_token.expand([x.shape[0], 1, self.embed_dim])
        x = self.patch_embedding(x)  # x:[N, embed_dim, h', w']
        x = x.flatten(2)
        x = x.transpose([0, 2, 1])  # x:[N, h'*w', embed_dim]
        x = paddle.concat([cls_token, x], axis=1)
        x = x + self.position_embedding
        return x

4.4 MLP模块

多层感知机由输入层、输出层和至少一层的隐藏层构成。

网络中各个隐藏层中神经元可接收相邻前序隐藏层中所有神经元传递而来的信息，经过加工处理后将信息输出给相邻后续隐藏层中所有神经元。

在多层感知机中，相邻层所包含的神经元之间通常使用“全连接”方式进行连接。

多层感知机可以模拟复杂非线性函数功能，所模拟函数的复杂性取决于网络隐藏层数目和各层中神经元数目。

# 实现MLP模块
class Mlp(nn.Layer):
    def __init__(self, embed_dim, mlp_ratio, dropout=0.):
        super().__init__()
        self.fc1 = nn.Linear(embed_dim, int(embed_dim*mlp_ratio))
        self.fc2 = nn.Linear(int(embed_dim*mlp_ratio), embed_dim)
        self.act = nn.GELU()
        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        x = self.fc1(x)
        x = self.act(x)
        x = self.dropout(x)
        x = self.fc2(x)
        x = self.dropout(x)
        return x

4.5 组装VIT模块

# 实现ViT
import paddle
import paddle.nn as nn
class ViT(nn.Layer):
    def __init__(self,
                image_size=20,  # 图片大小
                patch_size=4,  # 图片的patches数量
                in_channels=1,  # 原始图片通道
                num_classes=5,  # 目标类别
                embed_dim=400,  # 编码长度
                depth=3,  # encoder堆叠的数量
                num_heads=6,  # 头
                mlp_ratio=2,  # mlp维度缩放比例
                qkv_bias=True,  # qkv投射层的偏置
                dropout=0.,  # dropout
                attention_dropout=0.,  # attn的dropout，其实就是qkv投射层之后是否加入dropout层
                droppath=0.):
        super().__init__()
        self.patch_embedding = PatchEmbedding(image_size, patch_size, in_channels, embed_dim)  # 将图像使用一个卷积编码为patches，然后加入clstoken和位置编码
        self.encoder = Encoder(embed_dim, depth)  # 这个堆叠了depth个encoder的编码器，输入等于输出，最后加一个norm
        self.classfier = nn.Linear(embed_dim, num_classes)  # 分类头
        
    def forward(self, x):
        x = paddle.reshape(x,[-1,1,20,20])
        x = self.patch_embedding(x)  # x:[N, num_patches, embed_dim]
        x = self.encoder(x)  # x:[N, num_patches, embed_dim]
        x = self.classfier(x[:, 0])  # 相当于只对所有的cls_token做映射
        return x

# 测试
vit = ViT()

# 模型结构展示
paddle.summary(vit, (-1, 1, 20, 20))

----------------------------------------------------------------------------
  Layer (type)       Input Shape          Output Shape         Param #    
============================================================================
    Conv2D-2       [[1, 1, 20, 20]]      [1, 400, 5, 5]         6,800     
PatchEmbedding-2   [[1, 1, 20, 20]]       [1, 26, 400]         10,800     
  LayerNorm-8       [[1, 26, 400]]        [1, 26, 400]           800      
   Linear-16        [[1, 26, 400]]       [1, 26, 1188]         475,200    
   Softmax-5       [[1, 6, 26, 26]]      [1, 6, 26, 26]           0       
   Linear-17        [[1, 26, 396]]        [1, 26, 400]         158,800    
  Attention-5       [[1, 26, 400]]        [1, 26, 400]            0       
  LayerNorm-9       [[1, 26, 400]]        [1, 26, 400]           800      
   Linear-18        [[1, 26, 400]]        [1, 26, 800]         320,800    
     GELU-4         [[1, 26, 800]]        [1, 26, 800]            0       
   Dropout-6        [[1, 26, 400]]        [1, 26, 400]            0       
   Linear-19        [[1, 26, 800]]        [1, 26, 400]         320,400    
     Mlp-4          [[1, 26, 400]]        [1, 26, 400]            0       
 EncoderLayer-4     [[1, 26, 400]]        [1, 26, 400]            0       
  LayerNorm-10      [[1, 26, 400]]        [1, 26, 400]           800      
   Linear-20        [[1, 26, 400]]       [1, 26, 1188]         475,200    
   Softmax-6       [[1, 6, 26, 26]]      [1, 6, 26, 26]           0       
   Linear-21        [[1, 26, 396]]        [1, 26, 400]         158,800    
  Attention-6       [[1, 26, 400]]        [1, 26, 400]            0       
  LayerNorm-11      [[1, 26, 400]]        [1, 26, 400]           800      
   Linear-22        [[1, 26, 400]]        [1, 26, 800]         320,800    
     GELU-5         [[1, 26, 800]]        [1, 26, 800]            0       
   Dropout-7        [[1, 26, 400]]        [1, 26, 400]            0       
   Linear-23        [[1, 26, 800]]        [1, 26, 400]         320,400    
     Mlp-5          [[1, 26, 400]]        [1, 26, 400]            0       
 EncoderLayer-5     [[1, 26, 400]]        [1, 26, 400]            0       
  LayerNorm-12      [[1, 26, 400]]        [1, 26, 400]           800      
   Linear-24        [[1, 26, 400]]       [1, 26, 1188]         475,200    
   Softmax-7       [[1, 6, 26, 26]]      [1, 6, 26, 26]           0       
   Linear-25        [[1, 26, 396]]        [1, 26, 400]         158,800    
  Attention-7       [[1, 26, 400]]        [1, 26, 400]            0       
  LayerNorm-13      [[1, 26, 400]]        [1, 26, 400]           800      
   Linear-26        [[1, 26, 400]]        [1, 26, 800]         320,800    
     GELU-6         [[1, 26, 800]]        [1, 26, 800]            0       
   Dropout-8        [[1, 26, 400]]        [1, 26, 400]            0       
   Linear-27        [[1, 26, 800]]        [1, 26, 400]         320,400    
     Mlp-6          [[1, 26, 400]]        [1, 26, 400]            0       
 EncoderLayer-6     [[1, 26, 400]]        [1, 26, 400]            0       
  LayerNorm-14      [[1, 26, 400]]        [1, 26, 400]           800      
   Encoder-2        [[1, 26, 400]]        [1, 26, 400]            0       
   Linear-28          [[1, 400]]             [1, 5]             2,005     
============================================================================
Total params: 3,850,805
Trainable params: 3,850,805
Non-trainable params: 0
----------------------------------------------------------------------------
Input size (MB): 0.00
Forward/backward pass size (MB): 3.97
Params size (MB): 14.69
Estimated Total Size (MB): 18.66
----------------------------------------------------------------------------






{'total_params': 3850805, 'trainable_params': 3850805}

4.6 模型训练

EPOCH_NUM = 1000

model = paddle.Model(vit)  # 模型封装

scheduler = paddle.optimizer.lr.CosineAnnealingDecay(learning_rate=8e-3, T_max=int(EPOCH_NUM//2), verbose=True)

# 配置优化器、损失函数、评估指标
model.prepare(paddle.optimizer.Adam(learning_rate=scheduler, parameters=model.parameters()), 
              paddle.nn.CrossEntropyLoss(), 
              paddle.metric.Accuracy())

# 训练可视化VisualDL工具的回调函数
visualdl = paddle.callbacks.VisualDL(log_dir='visualdl_log')   

# 启动模型全流程训练
model.fit(train_dataset,    # 训练数据集
          eval_dataset,     # 评估数据集
          epochs=EPOCH_NUM, # 训练的总轮次
          batch_size=256,    # 训练使用的批大小
          verbose=1,        # 日志展示形式
          save_dir='output-model',
          save_freq=20,
          callbacks=[visualdl])  # 设置可视化

4.7 训练日志可视化

经典版环境不支持看可视化，这里在本地进行可视化的查看。

在本地打开anaconda的终端环境：

(paddle-cpu) C:\Users\Administrator>pip install visualdl
(paddle-cpu) C:\Users\Administrator>d:

执行可视化命令会出现：

(paddle-cpu) D:\>visualdl --logdir d:Desktop --port 8080
VisualDL 2.2.3
Running VisualDL at http://localhost:8080/ (Press CTRL+C to quit)
Serving VisualDL on localhost; to expose to the network, use a proxy or pass --host 0.0.0.0

在本地打开上面的网址即可：

可以看到，结果没有什么长进。

5 模型保存及简单验证

5.1 模型保存

model.save('output/model')  # 保存模型

5.2 模型验证

其实上面的可视化图已经是对结果的一个很好的验证了，这里我们通过实际的数据的展示，来看一下效果。

import paddle.nn.functional as F

label_txt = ['Adware','Banking','SMS malware','Riskware','Benign']

model = vit

para_dict = paddle.load('output-model/final.pdparams')
model.load_dict(para_dict)

model.eval()

# 找一张验证集的数据进行展示 mal_data/3/3684.jpg 2
image_path = 'mal_data/3/3684.jpg'
im = Image.open(image_path).convert('L')
im = im.resize((20, 20), Image.ANTIALIAS)
im = np.array(im).reshape(1, 1, 20, 20).astype(np.float32)
x = paddle.to_tensor(im)
predicts = model(x)
p = predicts.numpy().argmax() # 最大索引
print("实际标签为:",label_txt[int(image_path.split('/')[1])-1])
print("预测标签为:",label_txt[p])
al_data/3/3684.jpg 2
image_path = 'mal_data/3/3684.jpg'
im = Image.open(image_path).convert('L')
im = im.resize((20, 20), Image.ANTIALIAS)
im = np.array(im).reshape(1, 1, 20, 20).astype(np.float32)
x = paddle.to_tensor(im)
predicts = model(x)
p = predicts.numpy().argmax() # 最大索引
print("实际标签为:",label_txt[int(image_path.split('/')[1])-1])
print("预测标签为:",label_txt[p])
print(F.softmax(predicts)) # 0 对应的就是标签1的 预测正确.

实际标签为: SMS malware
预测标签为: SMS malware
Tensor(shape=[1, 5], dtype=float32, place=Place(gpu:0), stop_gradient=False,
       [[0.00000176, 0.00021302, 0.99953210, 0.00011784, 0.00013525]])

项目总结

本项目用图像的思想去解决恶意软件的分类问题，且达到了一个还不错的效果。

当然，后续还是要大量优化的，因为分类问题需求的精度还是很高的。

个人总结

全网同名：

iterhui

我在AI Studio上获得至尊等级，点亮10个徽章，来互关呀~

https://aistudio.baidu.com/aistudio/personalcenter/thirdview/643467

此文仅为搬运，原作链接：https://aistudio.baidu.com/aistudio/projectdetail/4356373