DeepRobust: 强大的对抗性机器学习库

DeepRobust简介

DeepRobust是一个基于PyTorch开发的对抗性机器学习库,专注于为图像和图数据提供攻击和防御方法。该项目由密歇根州立大学数据科学与工程系(DSE-MSU)开发和维护,旨在帮助研究人员和开发者构建更加鲁棒的机器学习模型。

DeepRobust logo

DeepRobust的主要特点包括:

支持图像和图数据的对抗性攻击和防御
提供丰富的算法实现,包括最新的研究成果
易于使用和扩展的API设计
详细的文档和示例代码
持续更新和维护

安装和环境配置

DeepRobust支持通过pip安装或从源代码安装。基本环境要求如下:

Python >= 3.6
PyTorch >= 1.2.0

通过pip安装:

pip install deeprobust

从源代码安装:

git clone https://github.com/DSE-MSU/DeepRobust.git
cd DeepRobust
python setup.py install

如果安装依赖项遇到困难,可以尝试:

python setup_empty.py install

这将只安装DeepRobust核心库,而不安装其他依赖包。

主要功能

DeepRobust主要包含两个核心模块:图像模块和图模块。

图像模块

图像模块提供了针对图像数据的攻击和防御方法。主要功能包括:

训练模型
实例化攻击和防御方法
生成对抗样本
评估攻击效果

以下是一个使用PGD(Projected Gradient Descent)攻击的简单示例:

from deeprobust.image.attack.pgd import PGD
from deeprobust.image.config import attack_params
from deeprobust.image.utils import download_model
import torch
import deeprobust.image.netmodels.resnet as resnet
from torchvision import transforms, datasets

# 下载预训练模型
URL = "https://github.com/I-am-Bot/deeprobust_model/raw/master/CIFAR10_ResNet18_epoch_20.pt"
download_model(URL, "$MODEL_PATH$")

# 加载模型
model = resnet.ResNet18().to('cuda')
model.load_state_dict(torch.load("$MODEL_PATH$"))
model.eval()

# 准备数据
transform_val = transforms.Compose([transforms.ToTensor()])
test_loader = torch.utils.data.DataLoader(
datasets.CIFAR10('deeprobust/image/data', train=False, download=True,
transform=transform_val),
batch_size=10, shuffle=True)

x, y = next(iter(test_loader))
x = x.to('cuda').float()

# 执行PGD攻击
adversary = PGD(model, 'cuda')
Adv_img = adversary.generate(x, y, **attack_params['PGD_CIFAR10'])

图模块

图模块提供了针对图数据的攻击和防御方法。主要功能包括:

加载图数据集
设置代理模型
执行图攻击
应用图防御方法

以下是一个使用Metattack攻击图神经网络的示例:

import torch
import numpy as np
from deeprobust.graph.data import Dataset
from deeprobust.graph.defense import GCN
from deeprobust.graph.global_attack import Metattack

# 加载数据集
data = Dataset(root='/tmp/', name='cora', setting='nettack')
adj, features, labels = data.adj, data.features, data.labels
idx_train, idx_val, idx_test = data.idx_train, data.idx_val, data.idx_test
idx_unlabeled = np.union1d(idx_val, idx_test)

# 设置代理模型
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
surrogate = GCN(nfeat=features.shape[1], nclass=labels.max().item()+1, nhid=16,
with_relu=False, device=device)
surrogate = surrogate.to(device)
surrogate.fit(features, adj, labels, idx_train)

# 执行Metattack
model = Metattack(model=surrogate, nnodes=adj.shape[0], feature_shape=features.shape, device=device)
model = model.to(device)
perturbations = int(0.05 * (adj.sum() // 2))
model.attack(features, adj, labels, idx_train, idx_unlabeled, perturbations, ll_constraint=False)
modified_adj = model.modified_adj