简介：这是一份基于Python卷积神经网络的恶意软件检测项目源码，面向毕业设计、期末大作业及课程设计等场景，适合需要快速搭建完整系统的学生或开发者，也可作为网络安全与深度学习交叉方向的学习范例。资源共164个文件，含18个Python脚本、4个Jupyter Notebook（覆盖数据增强、数据分析等环节）、135张用于展示与验证的png/jpg图片，以及Markdown说明文档和Shell部署脚本，压缩包大小约34.63MB，整体轻量、易于部署。代码均附有详细注释，从数据预处理、数据增强、模型训练到检测流程形成完整闭环，个人手打98分项目，功能完善、操作简单，已有206人学习使用。下载后可直接运行，配合文档与Notebook分步解析，能够快速理解卷积神经网络在恶意软件识别中的实际应用，也方便在此基础上扩展或二次开发。
1. 为什么恶意软件检测要用卷积神经网络而不是特征码比对
内网有一台机器上报了异常行为，拉下来的 exe 放进杀毒软件没有任何反应：特征库还没收录这个家族。这类场景每天都发生在安全工程师手里。基于签名的检测对已知恶意软件有效，面对变种、加壳和混淆样本，覆盖率迅速下降。于是检测思路从“匹配字节”转向“学习内容”：把恶意软件本身转成灰度图，用卷积神经网络（CNN）提取纹理特征做家族分类。核心假设是同一家族的程序在汇编指令布局上存在一致性，这种局部相关模式恰好是卷积核擅长的东西。整个链路只用 Python 和 PyTorch 就能跑通，适合安全分析人员、数据工程师，也适合拿它做毕业设计的同学。下面按数据预处理、模型训练、服务封装的顺序，把一套可复现的做法讲清楚，最后补上模型上线前必须做的校准技巧。
2. 把二进制样本变成 CNN 能吃的灰度图：预处理与数据集准备
2.1 为什么是灰度图：卷积核、步长与填充在恶意软件检测里的作用
恶意软件本质上是一串字节，而 CNN 最成熟的输入是图像。把每个字节（0~255）映射成一个像素的灰度值，再按固定宽度排列成二维矩阵，就得到一张“程序图像”。Nataraj 等人在 2011 年的论文里用这种方法把恶意软件可视化，结果发现同一家族的样本在图像上呈现出肉眼可辨的纹理相似性。
二维卷积之所以适合这种数据，是因为指令流里存在空间局部性。一个 3×3 的卷积核扫描图像时，每个位置同时看到相邻 9 个像素的关系；一个函数的序言、跳转指令、常量池填充在局部区域会形成重复出现的模式。步长（stride）控制卷积核每次移动的距离，stride=1 时特征图分辨率不变，信息不丢失；填充（padding）用来控制边界像素的参与程度，padding=1 时输出尺寸不变，避免图像边缘模式只被扫到一次。池化层则把提取到的模式做降采样，stride=2 的最大池化让网络对指令位置的轻微偏移更鲁棒，这正是处理变种样本最需要的性质。
如果坚持用一维卷积处理原始字节序列，也不是不行，但序列长度动辄数万字节，模型要自己学会哪些偏移是无关的。转成灰度图之后可以复用图像分类领域几十年的工程经验，包括数据增强、预训练权重和部署工具链，这是实践中更常见的选择。
2.2 最小预处理脚本：把 PE/ELF 文件的字节流转成固定尺寸灰度图
下面这段代码是整套流程的地基。它接收一个文件的原始字节，返回一个 64×64 的归一化浮点数组。
import numpy as np IMAGE_WIDTH = 64 # 图像宽高，最终形状 (64, 64)，对应 4096 字节 def resize_bytes(data: bytes, length: int) -> bytes: """把任意长度的字节流均匀采样/补零到固定长度。""" arr = np.frombuffer(data, dtype=np.uint8) if len(arr) >= length: idx = np.linspace(0, len(arr) - 1, length).astype(int) return arr[idx].tobytes() return data + b"\x00" * (length - len(arr)) def bytes_to_gray(data: bytes, width: int = IMAGE_WIDTH, sample: str = "pad") -> np.ndarray: """字节流转灰度图。 sample="pad"：截取文件最前面 width*width 个字节，不足补 0x00； sample="resize"：把整个文件均匀采样到 width*width 字节。 """ size = width * width if sample == "pad": raw = data[:size] + b"\x00" * max(0, size - len(data)) else: raw = resize_bytes(data, size) arr = np.frombuffer(raw, dtype=np.uint8).reshape(width, width) return arr.astype(np.float32) / 255.0bytes_to_gray的两个采样策略分别对应两种检测思路。pad保留文件最前面的 4096 字节，PE 文件的 DOS 头、NT 头以及入口点附近的代码都集中在文件前部，这部分是判别家族的最强信号，而且读取速度快，适合构建大规模扫描管线。resize则把整个文件的纹理压缩到固定尺寸，文件长度不影响输入形状，对加壳后在尾部追加数据的样本更鲁棒。实际项目中我通常两种都生成，分别训练两个模型，推理时融合置信度。
参数调整上，IMAGE_WIDTH=64是时间和精度的平衡点。调到 32 时单样本计算量降到四分之一，但纹理细节损失严重，变种之间差异变小；调到 128 时输入变成 16384 字节，模型参数和训练时间明显上涨，而公开数据集上的准确率提升通常不超过一个点。另外注意，reshape要求字节数恰好等于width * width，改宽度时必须同步调整截断长度，这是预处理环节最常见的报错来源。
2.3 用公开数据集做基准：Malimg 的结构与样本划分
自己的样本往往数量少且类别偏，第一步先在公开数据集上跑通流程。Malimg 是恶意软件图像识别最常用的基准：从 25 个恶意软件家族提取了 9000 多张 64×64 灰度图，每张图由一个真实样本转换而来。数据量不大，普通笔记本 CPU 就能完成一整个训练实验，非常适合验证网络设计是否合理。
拿到数据后，先把文件列表和标签整理成表格，再做分层划分。下面用 scikit-learn 完成训练、验证、测试三段切分：
from pathlib import Path import pandas as pd from sklearn.model_selection import train_test_split # 假设目录结构为 data/malimg/<family>/<sample>.png rows = [] for img_path in Path("data/malimg").glob("*/*.png"): rows.append({"path": str(img_path), "family": img_path.parent.name}) df = pd.DataFrame(rows) train_df, val_test_df = train_test_split( df, test_size=0.30, stratify=df["family"], random_state=42 ) val_df, test_df = train_test_split( val_test_df, test_size=0.50, stratify=val_test_df["family"], random_state=42 ) print(train_df["family"].value_counts())stratify=df["family"]保证每个家族在训练集和验证集中的比例与原数据集一致。恶意软件数据集的类别分布通常很不均匀，某些家族有上千个样本，某些只有一两百个，不做分层划分就会出现在验证集里完全看不到某个家族的情况。random_state=42固定随机种子，保证每次跑实验的划分一致，这条习惯在调参时能省掉大量“这轮精度为什么差了两个点”的排查时间。
预处理和目标划分就绪后，下一步是设计模型。对这类小图、中等类别数的任务，CNN 的层数选择有明确边界，不是越深越好。
3. 搭建 CNN 模型并在 PyTorch 里完成训练
3.1 一个不堆层数的 CNN 结构：三层卷积块加两层全连接
针对 64×64 的灰度图，三层卷积块已经足够提取可分的纹理特征。模型太深反而容易在小数据集上过拟合，训练时间和显存开销却成倍增长，这是第一次做这类项目最容易踩的坑。
import torch.nn as nn class MalConvNet(nn.Module): def __init__(self, num_classes: int = 25): super().__init__() self.features = nn.Sequential( # block1: 64x64 -> 32x32 nn.Conv2d(1, 32, kernel_size=3, padding=1), nn.BatchNorm2d(32), nn.ReLU(inplace=True), nn.Conv2d(32, 32, kernel_size=3, padding=1), nn.BatchNorm2d(32), nn.ReLU(inplace=True), nn.MaxPool2d(kernel_size=2, stride=2), # block2: 32x32 -> 16x16 nn.Conv2d(32, 64, kernel_size=3, padding=1), nn.BatchNorm2d(64), nn.ReLU(inplace=True), nn.Conv2d(64, 64, kernel_size=3, padding=1), nn.BatchNorm2d(64), nn.ReLU(inplace=True), nn.MaxPool2d(kernel_size=2, stride=2), # block3: 16x16 -> 8x8 nn.Conv2d(64, 128, kernel_size=3, padding=1), nn.BatchNorm2d(128), nn.ReLU(inplace=True), nn.Conv2d(128, 128, kernel_size=3, padding=1), nn.BatchNorm2d(128), nn.ReLU(inplace=True), nn.MaxPool2d(kernel_size=2, stride=2), ) self.classifier = nn.Sequential( nn.Flatten(), nn.Dropout(p=0.5), nn.Linear(128 * 8 * 8, 256), nn.ReLU(inplace=True), nn.Dropout(p=0.3), nn.Linear(256, num_classes), ) def forward(self, x): return self.classifier(self.features(x))每个卷积块里放两个 3×3 卷积，感受野等效于一个 5×5 卷积，但参数量更小，中间还夹了一次 ReLU，非线性表达能力更强。BatchNorm2d放在卷积和激活之间，作用是稳定每层输入的分布，允许使用更大的学习率。padding=1让 3×3 卷积不改变特征图尺寸，空间信息由池化层统一压缩。分类头的两个 Dropout 是专门用来对抗过拟合的——恶意软件纹理中有大量与标签无关的噪声，比如编译器版本带来的字节差异，Dropout 迫使网络不依赖单一神经元。
不用 ResNet 这类深层网络的原因很简单：Malimg 的训练集只有六千多张 64×64 小图，而恶意软件图像是人工构造的表示，信息密度远低于自然图像。深度网络在小数据上没有任何优势，三层卷积块加全连接已经能在验证集上逼近 96% 的加权 F1。
3.2 训练超参与类别不平衡处理
训练代码的核心超参数如下表所示。这里给出的数值是经过多轮实验的稳定起点，不是拍脑袋的默认值。
| 超参数 | 数值 | 设置理由 |
|---|---|---|
| optimizer | AdamW | 解耦权重衰减，比 Adam 更稳 |
| learning_rate | 1e-3 | 配合 BatchNorm，足够快且不震荡 |
| weight_decay | 1e-2 | 约束权重范数，抑制过拟合 |
| batch_size | 64 | 平衡训练速度与梯度稳定性 |
| label_smoothing | 0.1 | 软化 one-hot 标签，缓解过度自信 |
| clip_grad_norm | 1.0 | 防止个别样本产生异常梯度 |
| scheduler | CosineAnnealingLR | 周期衰减学习率，收敛更平滑 |
类别不平衡的问题用 WeightedRandomSampler 处理。它的原理是给样本分配权重，类别样本数越少，单条样本被抽中的概率越高。
import numpy as np import torch from torch.utils.data import DataLoader, WeightedRandomSampler labels = train_df["family"].factorize()[0] counts = np.bincount(labels) weights = 1.0 / counts[labels] sampler = WeightedRandomSampler( weights=torch.from_numpy(weights).double(), num_samples=len(labels), replacement=True, ) train_loader = DataLoader(train_dataset, batch_size=64, sampler=sampler) criterion = nn.CrossEntropyLoss(label_smoothing=0.1)注意一个隐藏的坑：DataLoader一旦传入sampler，就不能再设置shuffle=True，两者互斥。replacement=True表示允许同一轮训练中重复抽取少数类样本，这是加权采样的常规做法。
训练循环里需要同时做梯度裁剪、学习率调度和模型保存：
EPOCHS, best_acc = 30, 0.0 optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3, weight_decay=1e-2) scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=EPOCHS) for epoch in range(1, EPOCHS + 1): model.train() for x, y in train_loader: x, y = x.to(device), y.to(device) optimizer.zero_grad() loss = criterion(model(x), y) loss.backward() torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) optimizer.step() scheduler.step() val_acc = evaluate(model, val_loader) if val_acc > best_acc: best_acc = val_acc torch.save(model.state_dict(), "best_model.pt")clip_grad_norm_把整个模型的梯度范数限制在 1.0 以内，作用是防止某个离群样本把参数一步推飞。CosineAnnealingLR让学习率从 1e-3 按余弦曲线下降到接近 0，前几个 epoch 大步探索，后期小步精调。保存模型时只存state_dict()而不是整个模型对象，这样换 Python 版本或者改模型结构时不会因为序列化兼容问题打不开权重。
3.3 别只看准确率：混淆矩阵与加权 F1
恶意软件家族分类的评估指标，准确率是最没有信息量的一个。当数据集中有 10% 的大类样本时，模型把一切预测成这个大类就能拿到 90% 准确率。正确做法是同时看加权 F1 和混淆矩阵。
from sklearn.metrics import classification_report, confusion_matrix import seaborn as sns import matplotlib.pyplot as plt preds, trues = predict_all(model, test_loader) report = classification_report(trues, preds, zero_division=0) print(report) cm = confusion_matrix(trues, preds) sns.heatmap(cm, cmap="Blues", xticklabels=families, yticklabels=families) plt.savefig("confusion_matrix.png", dpi=150)classification_report会输出每个家族的精确率、召回率和 F1。查看这张报告时重点看两类家族：样本量最少的那几个，以及形态上长得像的那几个。Malimg 里 Allaple 系列经常互相混淆，因为它们的加壳方式和指令序列高度相似，人眼都难以分辨。如果某个家族的 F1 明显低于平均水平，通常不是模型问题，而是这个家族本身的区分度不够，需要回到预处理环节考虑补充特征。
到这里模型已经能在测试集上达到可用的精度，距离真实的检测工具还差一层封装。无论是对目录批量扫描，对外提供接口，还是导出成其他框架可加载的格式，都需要把模型从训练脚本里抽离出来，单独组织成一个可交付的检测模块。
4. 从模型到可交付物：批量扫描脚本与 API 服务
4.1 离线批量扫描：os.walk 遍历目录并输出 CSV
真实使用场景不是单张图片预测，而是给一个目录，让它扫出哪些文件有风险。扫描脚本需要处理两个问题：避免加载整个大文件，以及跳过非可执行文件。
import csv import os from pathlib import Path def predict_file(path: Path, model, transform, idx_to_family, threshold=0.5): """对单个文件做预测，非 PE/ELF 返回 None。""" with open(path, "rb") as f: head = f.read(2) f.seek(0) if head == b"MZ": data = f.read(4096) # 只读前 4096 字节 elif head == b"\x7fE": data = f.read(4096) else: return None img = transform(data, sample="pad") with torch.no_grad(): logits = model(torch.tensor(img).unsqueeze(0).unsqueeze(0).to(device)) prob = torch.softmax(logits, dim=1)[0] idx = int(prob.argmax()) return { "path": str(path), "family": idx_to_family[idx], "confidence": float(prob[idx]), } results = [] for root, _, files in os.walk("/path/to/scan"): for name in files: path = Path(root) / name pred = predict_file(path, model, transform, idx_to_family) if pred and pred["confidence"] >= threshold: results.append(pred) with open("scan_result.csv", "w", newline="") as f: writer = csv.DictWriter(f, fieldnames=["path", "family", "confidence"]) writer.writeheader() writer.writerows(results)扫描时只读文件前 4096 字节，对几十 MB 的文件也能瞬间完成。PE 文件以MZ开头，ELF 文件以0x7f 0x45开头，这两个魔数判断把绝大部分无关文件挡在预测之外。threshold=0.5是保守阈值，扫描场景宁可多报一些候选交给人工分析，也不要漏掉真正的恶意样本。
4.2 把模型包成 FastAPI 接口：上传 exe 直接返回家族
离线扫描适合批量检查，实时接口适合嵌入到邮件网关、文件上传服务或沙箱调度器里。用 FastAPI 封装一个预测接口，整个服务不到六十行代码。
from fastapi import FastAPI, UploadFile import torch app = FastAPI() model = MalConvNet(num_classes=len(idx_to_family)) model.load_state_dict(torch.load("best_model.pt", map_location="cpu")) model.eval() @app.post("/predict") async def predict(file: UploadFile): data = await file.read(4096) if len(data) < 2 or data[:2] not in (b"MZ", b"\x7fE"): return {"filename": file.filename, "error": "not a PE/ELF file"} img = bytes_to_gray(data, sample="pad") with torch.no_grad(): logits = model(torch.tensor(img).unsqueeze(0).unsqueeze(0).float()) prob = torch.softmax(logits, dim=1)[0] idx = int(prob.argmax()) return { "filename": file.filename, "family": idx_to_family[idx], "confidence": round(float(prob[idx]), 4), "status": "positive" if prob[idx] >= 0.8 else "low_confidence", }启动服务并调用一次，命令如下：
uvicorn api_server:app --host 0.0.0.0 --port 8000 curl -F "file=@sample.exe" http://127.0.0.1:8000/predictmap_location="cpu"让服务端不依赖 GPU 也能运行，推理一个样本通常在 10 毫秒以内。响应里单独给出confidence让调用方自己决定处理策略。status字段区分高置信和低置信结果——在安全场景里，低置信结果不能直接放行，通常会转给沙箱或人工复核。这个分层设计是好封装的细节，直接复用训练好的模型，不碰训练脚本里的任何逻辑。
5. 让检测器在真实主机上少乱报：ONNX 导出与阈值校准
5.1 ONNX 导出三步
PyTorch 模型部署到生产环境，最省心的方式是导出为 ONNX。导出后可以脱离 Python 环境，用 ONNX Runtime 加载，CPU 推理速度比 PyTorch 快一倍左右。
dummy = torch.randn(1, 1, 64, 64) torch.onnx.export( model, dummy, "malware_cnn.onnx", input_names=["input"], output_names=["logits"], dynamic_axes={"input": {0: "batch"}}, opset_version=17, )dynamic_axes把 batch 维度设为动态，这样接口一次可以喂多张图做批量推理。opset_version保持与 ONNX Runtime 安装版本兼容即可，不必刻意追新。
5.2 用良性样本集校准置信度阈值
模型在 Malimg 测试集上 96% 的 F1 没有任何实际意义：真实主机上 99.9% 的可执行文件是良性的，误报率 2% 就意味着每扫描一百个软件就会报两个误报。需要一个良性样本集来校准阈值。做法是找一台干净的主机，把系统目录和常用软件目录里的一批可执行文件收集起来，跑一遍推理，统计预测为正类时的置信度分布。
benign_scores = [] for path in Path("/opt/clean_binaries").rglob("*"): pred = predict_file(path, model, transform, idx_to_family) if pred: benign_scores.append(pred["confidence"]) threshold = np.quantile(benign_scores, 0.95) print(f"建议置信度阈值: {threshold:.3f}")取置信度分布的 95 分位作为阈值，含义是：在这个良性样本集上，最多只有 5% 的样本会被判为正类。想要更严格的误报控制，就取 99 分位；想要更高的检出率，就取 90 分位。校准阈值是一锤子的事，但每次重新训练模型后都必须重跑一遍。
5.3 加壳样本回归测试
UPX 加壳是恶意软件最常见的变形方式，模型对加壳样本的稳定性需要单独验证。找几个训练时没见过的加壳样本，预测后的家族标签如果与原始样本一致，说明模型学到的是逻辑层特征而不是壳的字节特征；如果标签漂移了，说明预处理协议里pad策略截断了壳的展开代码。把这段判断写成一个断言脚本，加进 CI：任何新模型，只有加壳样本上的预测一致率超过 90%，才允许发布。
for packed in upx_test_files: pred = predict_file(packed, model, transform, idx_to_family) assert pred["confidence"] >= threshold - 0.02, packed校准完阈值、跑通加壳回归，模型才算真正具备落地条件。把这两套脚本和模型检查点一起收进仓库，后续每次数据更新后重跑一遍即可。
本文还有配套的精品资源，点击获取
