模型压缩技术在边缘设备部署中的安全基线验证方法

edge_model_security/
├── configs/                    # 配置文件
│   └── default_config.yaml
├── core/                       # 核心逻辑模块
│   ├── compression.py          # 模型压缩器
│   ├── attack.py               # 对抗攻击器
│   ├── detector.py             # 运行时异常检测器
│   └── evaluator.py            # 综合评估器
├── models/                     # 模型定义与加载
│   └── model_loader.py
├── utils/                      # 工具函数
│   ├── data_loader.py
│   └── logger.py
├── scripts/                    # 可执行脚本
│   ├── run_pipeline.py         # 主流程脚本
│   └── visualize_results.py    # 可视化脚本
├── outputs/                    # 运行输出目录（自动生成）
│   ├── compressed_models/
│   ├── attack_results/
│   └── reports/
├── requirements.txt
└── README.md                   # 项目说明（此处仅为示意，正文不展开）

# 数据配置
data:
  name: "CIFAR10"
  root: "./data"
  batch_size: 128

# 模型配置
model:
  base_model: "resnet18"
  pretrained: true
  num_classes: 10

# 压缩配置
compression:
  quantize:
    enabled: true
    bits: 8
  prune:
    enabled: false
    amount: 0.3 # 剪枝比例

# 攻击配置 (用于鲁棒性评估)
attack:
  methods: ["fgsm", "pgd"]
  fgsm:
    eps: 0.03
  pgd:
    eps: 0.03
    alpha: 0.01
    steps: 7

# 评估配置
evaluation:
  metrics: ["clean_accuracy", "adversarial_accuracy", "model_size", "inference_time"]
  detector:
    enabled: true
    feature_layer: "layer4" # 用于统计特征图的层名

import torch
import torch.nn as nn
import torch.quantization
import torch.nn.utils.prune as prune
from copy import deepcopy
import os

class ModelCompressor:
    """
    模型压缩器，集成量化与剪枝。
    """
    def __init__(self, config):
        self.config = config
        self.compressed_model = None

    def quantize(self, model, calibration_loader=None):
        """执行训练后静态量化"""
        model_to_quantize = deepcopy(model).eval()
        model_to_quantize.qconfig = torch.quantization.get_default_qconfig('fbgemm') # 针对服务器/CPU
        # 针对ARM，可使用 'qnnpack'

        # 准备模型，插入观察节点
        torch.quantization.prepare(model_to_quantize, inplace=True)
        # 使用校准数据校准量化参数
        if calibration_loader:
            self._calibrate(model_to_quantize, calibration_loader)
        # 转换模型
        quantized_model = torch.quantization.convert(model_to_quantize, inplace=False)
        print("[INFO] 模型量化完成。")
        return quantized_model

    def _calibrate(self, model, calibration_loader, num_batches=32):
        """使用部分数据校准量化参数"""
        model.eval()
        with torch.no_grad():
            for i, (data, _) in enumerate(calibration_loader):
                model(data)
                if i >= num_batches:
                    break

    def prune_structured(self, model, amount=0.3):
        """对模型的Conv2d层进行L1非结构化剪枝"""
        model_to_prune = deepcopy(model)
        parameters_to_prune = []
        for name, module in model_to_prune.named_modules():
            if isinstance(module, nn.Conv2d):
                parameters_to_prune.append((module, 'weight'))

        prune.global_unstructured(
            parameters_to_prune,
            pruning_method=prune.L1Unstructured,
            amount=amount,
        )
        # 永久移除被剪枝的权重
        for module, _ in parameters_to_prune:
            prune.remove(module, 'weight')
        print(f"[INFO] 模型结构化剪枝完成，比例: {amount}。")
        return model_to_prune

    def run(self, model, data_loader):
        """执行配置中启用的压缩流程"""
        self.compressed_model = deepcopy(model)
        compress_cfg = self.config['compression']

        if compress_cfg['prune']['enabled']:
            self.compressed_model = self.prune_structured(
                self.compressed_model,
                amount=compress_cfg['prune']['amount']
            )
        if compress_cfg['quantize']['enabled']:
            # 量化需要校准数据
            self.compressed_model = self.quantize(
                self.compressed_model,
                calibration_loader=data_loader
            )
        # 计算并保存压缩模型大小
        self._save_model_info(model, self.compressed_model)
        return self.compressed_model

    def _save_model_info(self, original_model, compressed_model):
        """保存模型大小信息（示例）"""
        # 实际项目中可保存模型文件并计算大小
        pass

import torch
import torch.nn as nn
import torch.nn.functional as F

class AdversarialAttacker:
    """
    生成对抗样本以评估模型鲁棒性。
    """
    def __init__(self, config):
        self.config = config['attack']

    def fgsm_attack(self, model, data, target, eps):
        """快速梯度符号攻击"""
        data.requires_grad = True
        output = model(data)
        loss = F.cross_entropy(output, target)
        model.zero_grad()
        loss.backward()
        data_grad = data.grad.data
        # 创建对抗样本
        perturbed_data = data + eps * data_grad.sign()
        # 将像素值裁剪到原始范围 [0,1]
        perturbed_data = torch.clamp(perturbed_data, 0, 1)
        return perturbed_data.detach()

    def pgd_attack(self, model, data, target, eps, alpha, steps):
        """投影梯度下降攻击（更强）"""
        original_data = data.clone().detach()
        perturbed_data = data.clone().detach()
        # 随机初始化扰动
        perturbed_data = perturbed_data + torch.empty_like(perturbed_data).uniform_(-eps, eps)
        perturbed_data = torch.clamp(perturbed_data, 0, 1).detach()

        for _ in range(steps):
            perturbed_data.requires_grad = True
            output = model(perturbed_data)
            loss = F.cross_entropy(output, target)
            model.zero_grad()
            loss.backward()
            grad = perturbed_data.grad.data
            # 梯度上升步
            perturbed_data = perturbed_data.detach() + alpha * grad.sign()
            # 投影到 eps 球内
            delta = torch.clamp(perturbed_data - original_data, min=-eps, max=eps)
            perturbed_data = torch.clamp(original_data + delta, 0, 1).detach()
        return perturbed_data

    def evaluate_robustness(self, model, data_loader, device='cpu'):
        """
        计算模型在干净样本和对抗样本上的精度。
        返回一个包含结果的字典。
        """
        model.eval()
        results = {}
        attack_methods = self.config.get('methods', [])

        for attack_name in attack_methods:
            correct_clean = 0
            correct_adv = 0
            total = 0
            for data, target in data_loader:
                data, target = data.to(device), target.to(device)
                total += target.size(0)

                # 干净样本精度
                with torch.no_grad():
                    output_clean = model(data)
                    pred_clean = output_clean.argmax(dim=1)
                    correct_clean += (pred_clean == target).sum().item()

                # 生成对抗样本
                if attack_name == 'fgsm':
                    eps = self.config['fgsm']['eps']
                    adv_data = self.fgsm_attack(model, data, target, eps)
                elif attack_name == 'pgd':
                    eps = self.config['pgd']['eps']
                    alpha = self.config['pgd']['alpha']
                    steps = self.config['pgd']['steps']
                    adv_data = self.pgd_attack(model, data, target, eps, alpha, steps)
                else:
                    continue

                # 对抗样本精度
                with torch.no_grad():
                    output_adv = model(adv_data)
                    pred_adv = output_adv.argmax(dim=1)
                    correct_adv += (pred_adv == target).sum().item()

            clean_acc = 100. * correct_clean / total
            adv_acc = 100. * correct_adv / total
            results[attack_name] = {
                'clean_accuracy': clean_acc,
                'adversarial_accuracy': adv_acc,
                'robustness_drop': clean_acc - adv_acc
            }
            print(f"[ATTACK] {attack_name.upper()} - 干净精度: {clean_acc:.2f}%, 对抗精度: {adv_acc:.2f}%, 下降: {clean_acc-adv_acc:.2f}%")
        return results

import torch
import numpy as np
from sklearn.covariance import EllipticEnvelope
from scipy import stats

class RuntimeAnomalyDetector:
    """
    基于模型内部特征统计的轻量级运行时异常检测器。
    在部署阶段，可以将其置于模型旁边，对输入进行筛查。
    """
    def __init__(self, model, feature_layer_name='layer4'):
        self.model = model
        self.feature_layer_name = feature_layer_name
        self.reference_features = [] # 用于拟合的参考特征
        self.detector = None # 如EllipticEnvelope
        self.feature_mean = None
        self.feature_std = None
        self._register_hook()

    def _get_features_hook(self, module, input, output):
        """钩子函数，用于提取指定层的特征"""
        # 提取特征并全局平均池化，得到一个特征向量
        feats = output.detach().cpu().mean(dim=[2,3]).view(output.size(0), -1).numpy()
        self.last_features = feats # 暂存最后一次前向的特征

    def _register_hook(self):
        """在指定层注册前向钩子"""
        for name, module in self.model.named_modules():
            if name == self.feature_layer_name:
                module.register_forward_hook(self._get_features_hook)
                print(f"[DETECTOR] 钩子注册在层: {name}")
                return
        raise ValueError(f"未找到层: {self.feature_layer_name}")

    def fit(self, data_loader, use_method='mahalanobis'):
        """
        使用干净训练数据拟合检测器。
        use_method: 'mahalanobis' 或 'zscore'
        """
        feature_list = []
        self.model.eval()
        with torch.no_grad():
            for data, _ in data_loader:
                _ = self.model(data) # 前向传播，触发钩子
                if hasattr(self, 'last_features'):
                    feature_list.append(self.last_features)
        self.reference_features = np.vstack(feature_list)
        print(f"[DETECTOR] 收集到 {self.reference_features.shape[0]} 个参考特征样本。")

        if use_method == 'mahalanobis':
            # 使用协方差椭圆拟合（对多变量异常值敏感）
            self.detector = EllipticEnvelope(contamination=0.1) # 假设10%异常污染率
            self.detector.fit(self.reference_features)
        elif use_method == 'zscore':
            # 简单标准化，假设特征独立
            self.feature_mean = self.reference_features.mean(axis=0)
            self.feature_std = self.reference_features.std(axis=0) + 1e-8

    def predict(self, input_tensor, threshold=3.0):
        """
        对单个输入进行预测，返回是否为异常的布尔值及异常分数。
        仅适用于 'zscore' 方法示例。
        """
        if self.feature_mean is None:
            raise RuntimeError("检测器尚未拟合。请先调用 fit() 方法。")
        with torch.no_grad():
            _ = self.model(input_tensor)
            if not hasattr(self, 'last_features'):
                return False, 0.0
            test_feat = self.last_features[0] # 假设批大小为1
        # 计算马氏距离 (简化版，假设特征独立，即为标准化欧氏距离)
        z_scores = np.abs((test_feat - self.feature_mean) / self.feature_std)
        anomaly_score = z_scores.max() # 取最大Z分数作为异常分数
        is_anomaly = anomaly_score > threshold
        return is_anomaly, anomaly_score

import torch
import time
import json
import os
from pathlib import Path

class SecurityBaselineEvaluator:
    """
    综合评估器：计算精度、效率、鲁棒性等基线指标。
    """
    def __init__(self, config, output_dir='./outputs'):
        self.config = config
        self.output_dir = Path(output_dir)
        self.output_dir.mkdir(parents=True, exist_ok=True)

    def evaluate_model(self, model, model_name, data_loader, attacker, device='cpu'):
        """对单个模型进行全面评估"""
        model.to(device).eval()
        results = {'model_name': model_name}

        # 1. 评估干净精度
        print(f"[EVAL] 评估 {model_name} 的干净精度...")
        clean_acc = self._compute_accuracy(model, data_loader, device)
        results['clean_accuracy'] = clean_acc

        # 2. 评估对抗鲁棒性
        print(f"[EVAL] 评估 {model_name} 的对抗鲁棒性...")
        robustness_results = attacker.evaluate_robustness(model, data_loader, device)
        results.update(robustness_results)

        # 3. 评估推理速度 (示例，使用CPU)
        print(f"[EVAL] 评估 {model_name} 的推理延迟...")
        latency = self._compute_latency(model, data_loader, device, num_samples=100)
        results['avg_inference_latency_ms'] = latency

        # 4. 评估模型大小 (参数数量)
        num_params = sum(p.numel() for p in model.parameters())
        results['num_parameters'] = num_params
        # 估算模型文件大小 (假设FP32，单位MB)
        estimated_size_mb = (num_params * 4) / (1024 ** 2)
        if hasattr(model, 'is_quantized') and model.is_quantized:
            # 量化模型 (INT8) 估算
            estimated_size_mb = (num_params * 1) / (1024 ** 2)
        results['estimated_size_mb'] = round(estimated_size_mb, 2)

        return results

    def _compute_accuracy(self, model, data_loader, device):
        correct = 0
        total = 0
        with torch.no_grad():
            for data, target in data_loader:
                data, target = data.to(device), target.to(device)
                outputs = model(data)
                _, predicted = outputs.max(1)
                total += target.size(0)
                correct += predicted.eq(target).sum().item()
        return 100. * correct / total

    def _compute_latency(self, model, data_loader, device, num_samples=100):
        model.eval()
        latencies = []
        with torch.no_grad():
            for i, (data, _) in enumerate(data_loader):
                if i * data.size(0) >= num_samples:
                    break
                data = data.to(device)
                start = time.perf_counter()
                _ = model(data)
                end = time.perf_counter()
                latencies.append((end - start) * 1000) # 转换为毫秒
        return sum(latencies) / len(latencies) if latencies else 0

    def generate_report(self, original_results, compressed_results):
        """生成对比报告并保存为JSON"""
        report = {
            'original_model': original_results,
            'compressed_model': compressed_results,
            'comparison': {
                'accuracy_drop': original_results['clean_accuracy'] - compressed_results['clean_accuracy'],
                'robustness_drop_fgsm': original_results.get('fgsm', {}).get('robustness_drop', 0) -
                                         compressed_results.get('fgsm', {}).get('robustness_drop', 0),
                'size_reduction_ratio': 1 - (compressed_results['estimated_size_mb'] / original_results['estimated_size_mb']),
                'speedup_ratio': original_results['avg_inference_latency_ms'] / max(compressed_results['avg_inference_latency_ms'], 1e-5)
            }
        }
        # 安全基线判断 (示例规则)
        baseline_passed = True
        issues = []
        if report['comparison']['accuracy_drop'] > 5.0:
            issues.append(f"精度下降过多: {report['comparison']['accuracy_drop']:.2f}%")
            baseline_passed = False
        if report['comparison']['robustness_drop_fgsm'] > 10.0:
            issues.append(f"FGSM鲁棒性下降过多: {report['comparison']['robustness_drop_fgsm']:.2f}%")
            baseline_passed = False

        report['security_baseline_verdict'] = 'PASS' if baseline_passed else 'FAIL'
        report['issues'] = issues

        report_path = self.output_dir / 'reports' / 'security_report.json'
        report_path.parent.mkdir(exist_ok=True)
        with open(report_path, 'w') as f:
            json.dump(report, f, indent=4)
        print(f"[REPORT] 安全评估报告已生成: {report_path}")
        return report

#!/usr/bin/env python3
"""
主流程脚本：串联模型加载、压缩、攻击评估、基线验证全流程。
"""
import sys
import os
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

import yaml
import torch
from torchvision import datasets, transforms
from models.model_loader import load_pretrained_model
from core.compression import ModelCompressor
from core.attack import AdversarialAttacker
from core.detector import RuntimeAnomalyDetector
from core.evaluator import SecurityBaselineEvaluator
from utils.data_loader import get_data_loader

def main():
    # 1. 加载配置
    with open('configs/default_config.yaml', 'r') as f:
        config = yaml.safe_load(f)
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    print(f"使用设备: {device}")

    # 2. 加载数据
    train_loader, test_loader = get_data_loader(config['data'])

    # 3. 加载原始模型
    original_model = load_pretrained_model(config['model'])
    original_model.to(device).eval()
    print(f"加载原始模型: {config['model']['base_model']}")

    # 4. 初始化各模块
    compressor = ModelCompressor(config)
    attacker = AdversarialAttacker(config)
    evaluator = SecurityBaselineEvaluator(config)

    # 5. 评估原始模型
    print("\n" + "="*50)
    print("开始评估原始模型安全基线...")
    original_results = evaluator.evaluate_model(
        original_model, "Original_Model", test_loader, attacker, device
    )

    # 6. 压缩模型
    print("\n" + "="*50)
    print("开始模型压缩...")
    compressed_model = compressor.run(original_model, train_loader) # 使用训练集进行量化校准
    compressed_model.to(device).eval()

    # 7. 评估压缩后模型
    print("\n" + "="*50)
    print("开始评估压缩模型安全基线...")
    compressed_results = evaluator.evaluate_model(
        compressed_model, "Compressed_Model", test_loader, attacker, device
    )

    # 8. 生成对比报告与安全基线验证结论
    print("\n" + "="*50)
    print("生成安全基线验证报告...")
    final_report = evaluator.generate_report(original_results, compressed_results)
    print(f"\n安全基线验证结果: {final_report['security_baseline_verdict']}")
    if final_report['issues']:
        print("发现的问题:")
        for issue in final_report['issues']:
            print(f"  - {issue}")

    # 9. (可选) 训练并测试运行时异常检测器
    if config['evaluation']['detector']['enabled']:
        print("\n" + "="*50)
        print("初始化并训练运行时异常检测器...")
        detector = RuntimeAnomalyDetector(
            compressed_model,
            feature_layer_name=config['evaluation']['detector']['feature_layer']
        )
        # 使用部分训练数据拟合检测器
        detector.fit(train_loader, use_method='zscore')
        # 简单测试：用测试集中的第一个样本
        test_iter = iter(test_loader)
        test_data, test_target = next(test_iter)
        test_sample = test_data[0:1].to(device)
        is_anomaly, score = detector.predict(test_sample, threshold=3.0)
        print(f"测试样本异常检测结果: 异常={is_anomaly}, 分数={score:.4f}")

    print("\n全流程执行完毕。")

if __name__ == '__main__':
    main()

torch>=1.9.0
torchvision>=0.10.0
numpy>=1.19.5
scikit-learn>=0.24.2
scipy>=1.7.1
pyyaml>=5.4.1
matplotlib>=3.3.4  # 用于可视化脚本
tqdm>=4.62.0       # 可选，用于进度条

# 创建虚拟环境 (推荐)
    python -m venv venv
    source venv/bin/activate  # Linux/macOS
    # 或 venv\Scripts\activate  # Windows

    # 安装依赖
    pip install -r requirements.txt

cd edge_model_security
    python scripts/run_pipeline.py

import sys
import os
sys.path.append(os.path.join(os.path.dirname(__file__), '..'))

import torch
import torch.nn as nn
from core.attack import AdversarialAttacker

def test_fgsm_attack():
    """测试FGSM攻击是否能成功降低模型准确率"""
    # 1. 创建一个简单的模型和虚拟数据
    class SimpleModel(nn.Module):
        def __init__(self):
            super().__init__()
            self.linear = nn.Linear(10, 2)
        def forward(self, x):
            return self.linear(x)

    model = SimpleModel()
    # 2. 配置攻击参数
    config = {
        'attack': {
            'methods': ['fgsm'],
            'fgsm': {'eps': 0.1}
        }
    }
    attacker = AdversarialAttacker(config)
    # 3. 生成虚拟数据
    torch.manual_seed(42)
    batch_size = 4
    data = torch.randn(batch_size, 10)
    target = torch.randint(0, 2, (batch_size,))
    # 4. 执行攻击
    perturbed_data = attacker.fgsm_attack(model, data, target, eps=0.1)
    # 5. 验证扰动数据与原始数据不同，且在允许范围内
    assert not torch.allclose(data, perturbed_data), "扰动数据应与原始数据不同"
    assert torch.all(torch.abs(perturbed_data - data) <= 0.1 + 1e-5), "扰动超出epsilon范围"
    assert torch.all(perturbed_data >= 0) and torch.all(perturbed_data <= 1), "数据超出[0,1]范围 (经过clamp后)"
    print("[PASS] test_fgsm_attack 通过。")

if __name__ == '__main__':
    test_fgsm_attack()

python tests/test_attack.py

import json
import matplotlib.pyplot as plt
import numpy as np

def plot_comparison_bars(report_path):
    with open(report_path, 'r') as f:
        report = json.load(f)
    orig = report['original_model']
    comp = report['compressed_model']
    labels = ['Clean Acc.(%)', 'Adv. Acc.(FGSM,%)', 'Size(MB)', 'Latency(ms)']
    orig_vals = [orig['clean_accuracy'],
                 orig.get('fgsm', {}).get('adversarial_accuracy', 0),
                 orig['estimated_size_mb'],
                 orig['avg_inference_latency_ms']]
    comp_vals = [comp['clean_accuracy'],
                 comp.get('fgsm', {}).get('adversarial_accuracy', 0),
                 comp['estimated_size_mb'],
                 comp['avg_inference_latency_ms']]
    x = np.arange(len(labels))
    width = 0.35
    fig, ax = plt.subplots(figsize=(10,6))
    ax.bar(x - width/2, orig_vals, width, label='Original', color='skyblue')
    ax.bar(x + width/2, comp_vals, width, label='Compressed', color='lightcoral')
    ax.set_ylabel('Score')
    ax.set_title('Model Performance & Security Comparison')
    ax.set_xticks(x)
    ax.set_xticklabels(labels)
    ax.legend()
    fig.tight_layout()
    plt.savefig('./outputs/reports/comparison_chart.png', dpi=150)
    print("可视化图表已保存。")
    # plt.show() # 在非headless环境下可显示

python scripts/visualize_results.py