def deepfool_choose(target_model: Module, blackbox: Blackbox, queryset,
testset, selection: set, transferset: List,
indices_list: List, device, penalty: ndarray, **params):
model_dir = os.path.join(params['model_dir'],
'deepfool.{}'.format(params['deepfool_budget']))
if not os.path.exists(model_dir):
os.mkdir(model_dir)
# transferset, num_classes = load_transferset(os.path.join(model_dir, 'transferset.pickle'))
surrogate = target_model
evalutation_set = query(blackbox, [data[0] for data in testset],
len(testset),
device=device,
argmax=True)
optimizer = get_optimizer(surrogate.parameters(),
params['optimizer_choice'], **params)
reverse = True
budget = params['deepfool_budget']
ispenalty = params['ispenalty']
# batch_samples = []
batch_permutation = []
total = set([i for i in range(len(queryset))])
unselected = list(total - selection)
num_classes = len(testset.classes)
print('{} object to calculate.'.format(len(unselected)))
dm = DeepfoolMappable(surrogate, num_classes, queryset)
# with multiprocessing.Pool(4) as pool:
# results = pool.map(dm, unselected)
# batch_permutation = np.array([item[1] for item in results])
for index in tqdm(unselected):
permutation, _, _, _, _ = deepfool(queryset[index][0], surrogate,
penalty, ispenalty, num_classes)
# batch_samples.append(result.squeeze(0))
batch_permutation.append(np.linalg.norm(permutation))
batch_permutation = np.array(batch_permutation)
current_selection = batch_permutation.argsort(
0)[:budget] if not reverse else batch_permutation.argsort(
0)[len(unselected) - budget:]
assert len(current_selection) == budget
training_batch = [queryset[unselected[i]][0] for i in current_selection]
current_selection = [unselected[i] for i in current_selection]
indices_list.extend(current_selection)
selection.update(current_selection)
transferset.extend(query(blackbox, training_batch, budget, device=device))
model_utils.train_model(surrogate,
transferset,
model_dir,
testset=evalutation_set,
criterion_train=model_utils.soft_cross_entropy,
optimizer=optimizer,
checkpoint_suffix='.deepfool.{}'.format(
len(transferset)),
**params)
save_selection_state(transferset, selection, indices_list, model_dir)
def get_transferset(self):
"""
:return:
"""
# for rho_current in range(self.rho):
rho_current = 0
while self.blackbox.call_count < self.budget:
print('=> Beginning substitute epoch {} (|D| = {})'.format(rho_current, len(self.D)))
# -------------------------- 0. Initialize Model
model_adv = zoo.get_net(self.model_adv_name, self.modelfamily, self.model_adv_pretrained,
num_classes=self.num_classes)
model_adv = model_adv.to(self.device)
# -------------------------- 1. Train model on D
model_adv = model_utils.train_model(model_adv, self.D, self.out_dir, num_workers=10,
checkpoint_suffix='.{}'.format(self.blackbox.call_count),
device=self.device, epochs=self.train_epochs, log_interval=500, lr=0.1,
momentum=0.9, batch_size=self.batch_size, lr_gamma=0.1,
testset=self.testset, criterion_train=model_utils.soft_cross_entropy)
# -------------------------- 2. Evaluate model
# _, acc = model_utils.test_step(model_adv, self.testloader, nn.CrossEntropyLoss(reduction='mean'),
# device=self.device, epoch=rho_current)
# self.accuracies.append(acc)
# -------------------------- 3. Jacobian-based data augmentation
if self.aug_strategy in ['jbda', 'jbself']:
self.D = self.jacobian_augmentation(model_adv, rho_current)
elif self.aug_strategy == 'jbtop{}'.format(self.topk):
self.D = self.jacobian_augmentation_topk(model_adv, rho_current)
else:
raise ValueError('Unrecognized augmentation strategy: "{}"'.format(self.aug_strategy))
# -------------------------- 4. End if necessary
rho_current += 1
if (self.blackbox.call_count >= self.budget) or ((self.rho is not None) and (rho_current >= self.rho)):
print('=> # BB Queries ({}) >= budget ({}). Ending attack.'.format(self.blackbox.call_count,
self.budget))
model_adv = zoo.get_net(self.model_adv_name, self.modelfamily, self.model_adv_pretrained,
num_classes=self.num_classes)
model_adv = model_adv.to(self.device)
model_adv = model_utils.train_model(model_adv, self.D, self.out_dir, num_workers=10,
checkpoint_suffix='.{}'.format(self.blackbox.call_count),
device=self.device, epochs=self.final_train_epochs,
log_interval=500, lr=0.01, momentum=0.9, batch_size=self.batch_size,
lr_gamma=0.1, testset=self.testset,
criterion_train=model_utils.soft_cross_entropy)
break
print()
return self.D, model_adv
def deepfool_active():
params = parser_dealer({
'transfer': False,
'active': False,
'sampling': False,
'synthetic': False,
'black_box': True,
'train': True
})
model_dir = params['model_dir']
transferset, num_classes = load_transferset(
os.path.join(model_dir, 'transferset.pickle'))
surrogate = params['surrogate']
blackbox = params['blackbox']
device = params['device']
testset = params['testset']
# todo: make the iteration batch and rounds part of parameters.
iter_batch = 1000
rounds = 20
remnant = set(range(len(transferset)))
selected = set()
params['testset'] = query(blackbox, [data[0] for data in testset],
len(testset),
device=device,
argmax=True)
optimizer = get_optimizer(surrogate.parameters(),
params['optimizer_choice'], **params)
for i in range(rounds):
batch_samples = []
batch_permutation = []
current_round = list(remnant)
print('round {}: {} object to calculate.'.format(i + 1, len(remnant)))
for index in current_round:
permutation, _, _, _, result = deepfool(transferset[index][0],
surrogate, num_classes)
batch_samples.append(result.squeeze(0))
batch_permutation.append(np.linalg.norm(permutation))
batch_permutation = np.array(batch_permutation)
selection = batch_permutation.argsort(0)[:iter_batch]
training_batch = [batch_samples[i] for i in selection]
selection = [current_round[i] for i in selection]
selected.update(selection)
remnant.difference_update(selection)
transferset.extend(
query(blackbox, training_batch, iter_batch, device=device))
model_utils.train_model(surrogate,
transferset,
model_dir,
criterion_train=model_utils.soft_cross_entropy,
optimizer=optimizer,
**params)
def train(self):
# self.surrogate = zoo.get_net(self.kwargs["model_arch"], 'custom_cnn', None, num_classes=43).to(self.device)
# self.optim = get_optimizer(self.surrogate.parameters(), self.optimizer_choice, **self.kwargs)
model_utils.train_model(self.surrogate,
self.selected,
self.path,
batch_size=self.batch_size,
testset=self.evaluation_set,
criterion_train=self.criterion,
checkpoint_suffix='.active.{}'.format(
len(self.selected)),
device=self.device,
optimizer=self.optim,
**self.kwargs)
def train(model, transferset_samples, budget, round):
transferset = samples_to_transferset(transferset_samples, budget=budget)
print()
print('=> Training at budget = {}'.format(len(transferset)))
checkpoint_suffix = '.{}.{}'.format(budget, round)
model_utils.train_model(model,
transferset,
train.model_dir,
testset=train.testset,
criterion_train=train.criterion_train,
checkpoint_suffix=checkpoint_suffix,
device=train.device,
optimizer=train.optimizer,
**train.params)
def train(self,
trainset: Union[Dataset, Iterable[Tuple[Tensor, Tensor]],
Iterable[Tensor], Iterable[int]],
to_query: bool = False):
if self.blackbox is None and to_query:
raise Exception("Blackbox didn't exists, couldn't query")
elif to_query:
if isinstance(trainset[0], Tuple):
# This means trainset is either Dataset or Iterable[Tuple[Tensor, Tensor]], which needs to unpack.
data = self.query(unpack(trainset))
elif isinstance(trainset[0], Tensor):
# This means trainset has already unpack.
data = self.query(trainset)
elif isinstance(trainset[0], int):
assert self.sampleset is not None
duplication = self.selection.intersection(trainset)
if len(duplication) > 0:
print('{} samples duplicated.'.format(len(duplication)))
difference = set(trainset) - duplication
data = self.query([self.sampleset[i][0] for i in difference])
self.selection.update(difference)
else:
# dataset is baked
data = trainset
self.transfer.extend(data)
train_model(self.target_model,
self.transfer,
self.state_dir,
self.batch_size,
self.criterion,
testset=self.evaluation_set,
device=self.device,
num_workers=self.num_workers,
optimizer=self.optim,
**self.kwargs)
def main():
#torch.backends.cudnn.enabled = False
parser = argparse.ArgumentParser(description='Train a model')
# Required arguments
parser.add_argument('dataset',
metavar='DS_NAME',
type=str,
help='Dataset name')
parser.add_argument('model_arch',
metavar='MODEL_ARCH',
type=str,
help='Model name')
# Optional arguments
parser.add_argument('-x',
'--complexity',
metavar='X',
type=int,
help='Complexity of conv layer channel.',
default=64)
parser.add_argument('-o',
'--out_path',
metavar='PATH',
type=str,
help='Output path for model',
default=cfg.MODEL_DIR)
parser.add_argument('-d',
'--device_id',
metavar='D',
type=int,
help='Device id. -1 for CPU.',
default=0)
parser.add_argument('-b',
'--batch-size',
type=int,
default=64,
metavar='N',
help='input batch size for training (default: 64)')
parser.add_argument('-e',
'--epochs',
type=int,
default=100,
metavar='N',
help='number of epochs to train (default: 100)')
parser.add_argument('--lr',
type=float,
default=0.1,
metavar='LR',
help='learning rate (default: 0.1)')
parser.add_argument('--momentum',
type=float,
default=0.5,
metavar='M',
help='SGD momentum (default: 0.5)')
parser.add_argument(
'--log-interval',
type=int,
default=100,
metavar='N',
help='how many batches to wait before logging training status')
parser.add_argument('--resume',
default=None,
type=str,
metavar='PATH',
help='path to latest checkpoint (default: none)')
parser.add_argument('--lr-step',
type=int,
default=30,
metavar='N',
help='Step sizes for LR')
parser.add_argument('--lr-gamma',
type=float,
default=0.1,
metavar='N',
help='LR Decay Rate')
parser.add_argument('-w',
'--num_workers',
metavar='N',
type=int,
help='# Worker threads to load data',
default=10)
parser.add_argument('--train_subset',
type=int,
help='Use a subset of train set',
default=None)
parser.add_argument('--pretrained',
type=str,
help='Use pretrained network',
default=None)
parser.add_argument('--weighted-loss',
action='store_true',
help='Use a weighted loss',
default=None)
parser.add_argument('--optimizer-choice',
type=str,
help='Optimizer',
default='sgdm',
choices=('sgd', 'sgdm', 'adam', 'adagrad'))
args = parser.parse_args()
params = vars(args)
# torch.manual_seed(cfg.DEFAULT_SEED)
if params['device_id'] >= 0:
os.environ["CUDA_VISIBLE_DEVICES"] = str(params['device_id'])
device = torch.device('cuda')
#torch.cuda.set_device(str(params['device_id']))
else:
device = torch.device('cpu')
# ----------- Set up dataset
dataset_name = params['dataset']
valid_datasets = datasets.__dict__.keys()
if dataset_name not in valid_datasets:
raise ValueError(
'Dataset not found. Valid arguments = {}'.format(valid_datasets))
dataset = datasets.__dict__[dataset_name]
modelfamily = datasets.dataset_to_modelfamily[dataset_name]
train_transform = datasets.modelfamily_to_transforms[modelfamily]['train']
test_transform = datasets.modelfamily_to_transforms[modelfamily]['test']
trainset = dataset(train=True, transform=train_transform)
testset = dataset(train=False, transform=test_transform)
num_classes = len(trainset.classes)
sample = testset[0][0]
if len(sample.shape) <= 2:
# 2 dimensional images
channel = 1
else:
channel = sample.shape[0]
params['channel'] = channel
params['num_classes'] = num_classes
if params['train_subset'] is not None:
idxs = np.arange(len(trainset))
ntrainsubset = params['train_subset']
idxs = np.random.choice(idxs, size=ntrainsubset, replace=False)
trainset = Subset(trainset, idxs)
# ----------- Set up model
model_name = params['model_arch']
pretrained = params['pretrained']
complexity = params['complexity']
# model = model_utils.get_net(model_name, n_output_classes=num_classes, pretrained=pretrained)
model = zoo.get_net(model_name,
modelfamily,
pretrained,
num_classes=num_classes,
channel=channel,
complexity=complexity)
model = model.to(device)
optimizer = get_optimizer(model.parameters(), params['optimizer_choice'],
**params)
# ----------- Train
out_path = params['out_path']
model_utils.train_model(model,
trainset,
testset=testset,
device=device,
optimizer=optimizer,
**params)
# Store arguments
params['created_on'] = str(datetime.now())
params_out_path = osp.join(out_path, 'params.json')
with open(params_out_path, 'w') as jf:
json.dump(params, jf, indent=True)
def main():
parser = argparse.ArgumentParser(description='Train a model')
# Required arguments
parser.add_argument('model_dir', metavar='DIR', type=str, help='Directory containing transferset.pickle')
parser.add_argument('model_arch', metavar='MODEL_ARCH', type=str, help='Model name')
parser.add_argument('testdataset', metavar='DS_NAME', type=str, help='Name of test')
parser.add_argument('--budgets', metavar='B', type=str,
help='Comma separated values of budgets. Knockoffs will be trained for each budget.')
# Optional arguments
parser.add_argument('-d', '--device_id', metavar='D', type=int, help='Device id. -1 for CPU.', default=0)
parser.add_argument('-b', '--batch-size', type=int, default=64, metavar='N',
help='input batch size for training (default: 64)')
parser.add_argument('-e', '--epochs', type=int, default=100, metavar='N',
help='number of epochs to train (default: 100)')
parser.add_argument('--lr', type=float, default=0.01, metavar='LR',
help='learning rate (default: 0.01)')
parser.add_argument('--momentum', type=float, default=0.5, metavar='M',
help='SGD momentum (default: 0.5)')
parser.add_argument('--log-interval', type=int, default=50, metavar='N',
help='how many batches to wait before logging training status')
parser.add_argument('--resume', default=None, type=str, metavar='PATH',
help='path to latest checkpoint (default: none)')
parser.add_argument('--lr-step', type=int, default=60, metavar='N',
help='Step sizes for LR')
parser.add_argument('--lr-gamma', type=float, default=0.1, metavar='N',
help='LR Decay Rate')
parser.add_argument('-w', '--num_workers', metavar='N', type=int, help='# Worker threads to load data', default=10)
parser.add_argument('--pretrained', type=str, help='Use pretrained network', default=None)
parser.add_argument('--weighted-loss', action='store_true', help='Use a weighted loss', default=False)
# Attacker's defense
parser.add_argument('--argmaxed', action='store_true', help='Only consider argmax labels', default=False)
parser.add_argument('--optimizer_choice', type=str, help='Optimizer', default='sgdm', choices=('sgd', 'sgdm', 'adam', 'adagrad'))
args = parser.parse_args()
params = vars(args)
torch.manual_seed(cfg.DEFAULT_SEED)
if params['device_id'] >= 0:
os.environ["CUDA_VISIBLE_DEVICES"] = str(params['device_id'])
device = torch.device('cuda')
else:
device = torch.device('cpu')
model_dir = params['model_dir']
# ----------- Set up transferset
transferset_path = osp.join(model_dir, 'transferset.pickle')
with open(transferset_path, 'rb') as rf:
transferset_samples = pickle.load(rf)
num_classes = transferset_samples[0][1].size(0)
print('=> found transfer set with {} samples, {} classes'.format(len(transferset_samples), num_classes))
# ----------- Clean up transfer (if necessary)
if params['argmaxed']:
new_transferset_samples = []
print('=> Using argmax labels (instead of posterior probabilities)')
for i in range(len(transferset_samples)):
x_i, y_i = transferset_samples[i]
argmax_k = y_i.argmax()
y_i_1hot = torch.zeros_like(y_i)
y_i_1hot[argmax_k] = 1.
new_transferset_samples.append((x_i, y_i_1hot))
transferset_samples = new_transferset_samples
# ----------- Set up testset
dataset_name = params['testdataset']
valid_datasets = datasets.__dict__.keys()
modelfamily = datasets.dataset_to_modelfamily[dataset_name]
transform = datasets.modelfamily_to_transforms[modelfamily]['test']
if dataset_name not in valid_datasets:
raise ValueError('Dataset not found. Valid arguments = {}'.format(valid_datasets))
dataset = datasets.__dict__[dataset_name]
testset = dataset(train=False, transform=transform)
if len(testset.classes) != num_classes:
raise ValueError('# Transfer classes ({}) != # Testset classes ({})'.format(num_classes, len(testset.classes)))
# ----------- Set up model
model_name = params['model_arch']
pretrained = params['pretrained']
# model = model_utils.get_net(model_name, n_output_classes=num_classes, pretrained=pretrained)
model = zoo.get_net(model_name, modelfamily, pretrained, num_classes=num_classes)
model = model.to(device)
# ----------- Train
budgets = [int(b) for b in params['budgets'].split(',')]
for b in budgets:
np.random.seed(cfg.DEFAULT_SEED)
torch.manual_seed(cfg.DEFAULT_SEED)
torch.cuda.manual_seed(cfg.DEFAULT_SEED)
transferset = samples_to_transferset(transferset_samples, budget=b, transform=transform)
print()
print('=> Training at budget = {}'.format(len(transferset)))
optimizer = get_optimizer(model.parameters(), params['optimizer_choice'], **params)
print(params)
checkpoint_suffix = '.{}'.format(b)
criterion_train = model_utils.soft_cross_entropy
model_utils.train_model(model, transferset, model_dir, testset=testset, criterion_train=criterion_train,
checkpoint_suffix=checkpoint_suffix, device=device, optimizer=optimizer, **params)
# Store arguments
params['created_on'] = str(datetime.now())
params_out_path = osp.join(model_dir, 'params_train.json')
with open(params_out_path, 'w') as jf:
json.dump(params, jf, indent=True)
def __init__(self, blackbox, budget, model_adv_name, model_adv_pretrained, modelfamily, seedset, testset, device,
out_dir, batch_size=cfg.DEFAULT_BATCH_SIZE, train_epochs=20, kappa=400, tau=None, rho=6, sigma=-1,
query_batch_size=1, aug_strategy='jbda', useprobs=True, final_train_epochs=100):
self.blackbox = blackbox
self.budget = budget
self.model_adv_name = model_adv_name
self.model_adv_pretrained = model_adv_pretrained
self.model_adv = None
self.modelfamily = modelfamily
self.seedset = seedset
self.testset = testset
self.batch_size = batch_size
self.query_batch_size = query_batch_size
self.testloader = DataLoader(self.testset, batch_size=self.batch_size, pin_memory=True)
self.train_epochs = train_epochs
self.final_train_epochs = final_train_epochs
self.kappa = kappa
self.tau = tau
self.rho = rho
self.sigma = sigma
self.device = device
self.out_dir = out_dir
self.num_classes = len(self.testset.classes)
assert (aug_strategy in ['jbda', 'jbself']) or 'jbtop' in aug_strategy
self.aug_strategy = aug_strategy
self.topk = 0
if 'jbtop' in aug_strategy:
# extract k from "jbtop<k>"
self.topk = int(aug_strategy.replace('jbtop', ''))
self.accuracies = [] # Track test accuracies over time
self.useprobs = useprobs
# -------------------------- Initialize seed data
print('=> Obtaining predictions over {} seed samples using strategy {}'.format(len(self.seedset),
self.aug_strategy))
Dx = torch.cat([self.seedset[i][0].unsqueeze(0) for i in range(len(self.seedset))])
Dy = []
# Populate Dy
with torch.no_grad():
for inputs, in DataLoader(TensorDataset(Dx), batch_size=self.query_batch_size):
inputs = inputs.to(self.device)
outputs = blackbox(inputs).cpu()
if not self.useprobs:
labels = torch.argmax(outputs, dim=1)
labels_onehot = make_one_hot(labels, outputs.shape[1])
outputs = labels_onehot
Dy.append(outputs)
# Dy = torch.tensor(Dy)
Dy = torch.cat(Dy)
# TensorDataset D
self.D = TensorDataset(Dx, Dy)
### Block memory required for training later on
model_adv = zoo.get_net(self.model_adv_name, self.modelfamily, self.model_adv_pretrained,
num_classes=self.num_classes)
model_adv = model_adv.to(self.device)
model_adv = model_utils.train_model(model_adv, self.D, self.out_dir, num_workers=10,
checkpoint_suffix='.{}'.format(self.blackbox.call_count),
device=self.device, epochs=1,
log_interval=500, lr=0.01, momentum=0.9, batch_size=self.batch_size,
lr_gamma=0.1, testset=self.testset,
criterion_train=model_utils.soft_cross_entropy)
def main():
parser = argparse.ArgumentParser(description='Train a model')
# Required arguments
parser.add_argument('dataset',
metavar='DS_NAME',
type=str,
help='Dataset name')
parser.add_argument('model_arch',
metavar='MODEL_ARCH',
type=str,
help='Model name')
# Optional arguments
parser.add_argument('-o',
'--out_path',
metavar='PATH',
type=str,
help='Output path for model',
default=cfg.MODEL_DIR)
parser.add_argument('-d',
'--device_id',
metavar='D',
type=int,
help='Device id. -1 for CPU.',
default=0)
parser.add_argument('-b',
'--batch-size',
type=int,
default=64,
metavar='N',
help='input batch size for training (default: 64)')
parser.add_argument('-e',
'--epochs',
type=int,
default=100,
metavar='N',
help='number of epochs to train (default: 100)')
parser.add_argument('--lr',
type=float,
default=0.1,
metavar='LR',
help='learning rate (default: 0.1)')
parser.add_argument('--momentum',
type=float,
default=0.5,
metavar='M',
help='SGD momentum (default: 0.5)')
parser.add_argument(
'--log-interval',
type=int,
default=100,
metavar='N',
help='how many batches to wait before logging training status')
parser.add_argument('--resume',
default=None,
type=str,
metavar='PATH',
help='path to latest checkpoint (default: none)')
parser.add_argument('--lr-step',
type=int,
default=30,
metavar='N',
help='Step sizes for LR')
parser.add_argument('--lr-gamma',
type=float,
default=0.1,
metavar='N',
help='LR Decay Rate')
parser.add_argument('-w',
'--num_workers',
metavar='N',
type=int,
help='# Worker threads to load data',
default=10)
parser.add_argument('--pretrained',
type=str,
help='Use pretrained network',
default=None)
parser.add_argument('--weighted-loss',
action='store_true',
help='Use a weighted loss',
default=None)
args = parser.parse_args()
params = vars(args)
torch.manual_seed(cfg.DEFAULT_SEED)
if params['device_id'] >= 0:
os.environ["CUDA_VISIBLE_DEVICES"] = str(params['device_id'])
device = torch.device('cuda')
else:
device = torch.device('cpu')
# ----------- Set up dataset
dataset_name = params['dataset']
valid_datasets = datasets.__dict__.keys()
if dataset_name not in valid_datasets:
raise ValueError(
'Dataset not found. Valid arguments = {}'.format(valid_datasets))
dataset = datasets.__dict__[dataset_name]
modelfamily = datasets.dataset_to_modelfamily[dataset_name]
train_transform = datasets.modelfamily_to_transforms[modelfamily]['train']
test_transform = datasets.modelfamily_to_transforms[modelfamily]['test']
trainset = dataset(train=True, transform=train_transform)
testset = dataset(train=False, transform=test_transform)
num_classes = len(trainset.classes)
params['num_classes'] = num_classes
# ----------- Set up model
model_name = params['model_arch']
pretrained = params['pretrained']
# model = model_utils.get_net(model_name, n_output_classes=num_classes, pretrained=pretrained)
model = zoo.get_net(model_name,
modelfamily,
pretrained,
num_classes=num_classes)
model = model.to(device)
# ----------- Train
out_path = params['out_path']
model_utils.train_model(model,
trainset,
testset=testset,
device=device,
**params)
# Store arguments
params['created_on'] = str(datetime.now())
params_out_path = osp.join(out_path, 'params.json')
with open(params_out_path, 'w') as jf:
json.dump(params, jf, indent=True)
def main():
parser = argparse.ArgumentParser(description='Train a model')
# Required arguments
parser.add_argument('model_dir',
metavar='DIR',
type=str,
help='Directory containing transferset.pickle')
parser.add_argument('model_arch',
metavar='MODEL_ARCH',
type=str,
help='Model name')
parser.add_argument('testdataset',
metavar='DS_NAME',
type=str,
help='Name of test')
parser.add_argument('--budgets',
metavar='B',
type=int,
help='Knockoffs will be trained for budget.')
# Optional arguments
parser.add_argument('-d',
'--device_id',
metavar='D',
type=int,
help='Device id. -1 for CPU.',
default=0)
parser.add_argument('-b',
'--batch-size',
type=int,
default=64,
metavar='N',
help='input batch size for training (default: 64)')
parser.add_argument('-e',
'--epochs',
type=int,
default=10,
metavar='N',
help='number of epochs to train (default: 10)')
parser.add_argument('--lr',
type=float,
default=0.01,
metavar='LR',
help='learning rate (default: 0.01)')
parser.add_argument('--momentum',
type=float,
default=0.5,
metavar='M',
help='SGD momentum (default: 0.5)')
parser.add_argument(
'--log-interval',
type=int,
default=50,
metavar='N',
help='how many batches to wait before logging training status')
parser.add_argument('--resume',
default=None,
type=str,
metavar='PATH',
help='path to latest checkpoint (default: none)')
parser.add_argument('--lr-step',
type=int,
default=60,
metavar='N',
help='Step sizes for LR')
parser.add_argument('--lr-gamma',
type=float,
default=0.1,
metavar='N',
help='LR Decay Rate')
parser.add_argument('-w',
'--num_workers',
metavar='N',
type=int,
help='# Worker threads to load data',
default=10)
parser.add_argument('--pretrained',
type=str,
help='Use pretrained network',
default=None)
parser.add_argument('--weighted-loss',
action='store_true',
help='Use a weighted loss',
default=False)
# RL arguments
parser.add_argument('--traj_length',
metavar='N',
type=int,
help='# Step in one trajactory',
default=10)
parser.add_argument('--num_each_class',
metavar='N',
type=int,
help='# sample in each class',
default=1)
parser.add_argument('--n_iter',
metavar='N',
type=int,
help='# iterations of RL training',
default=10)
parser.add_argument('--n_traj_each_iter',
metavar='N',
type=int,
help='# trajactories / iter',
default=10)
parser.add_argument('--queryset',
metavar='DS_NAME',
type=str,
help='Name of test')
parser.add_argument('--victim_model_dir',
default=None,
type=str,
metavar='PATH',
help='path to latest checkpoint (default: none)')
parser.add_argument('--n_layers',
metavar='N',
type=int,
help='# layers in policy',
default=4)
parser.add_argument('--size',
metavar='N',
type=int,
help='size of layer in policy',
default=64)
parser.add_argument('--policy_lr',
type=float,
default=1e-4,
metavar='N',
help='Policy learning rate')
parser.add_argument('--num_agent_train_steps_per_iter',
metavar='N',
type=int,
help='num_agent_train_steps_per_iter',
default=10)
parser.add_argument('--agent_train_batch_size',
metavar='N',
type=int,
help='num_agent_train_steps_per_iter',
default=990)
parser.add_argument('--policy_gamma',
type=float,
default=0.9,
metavar='N',
help='reward discounting')
parser.add_argument('--eps_random',
type=float,
default=-1,
metavar='N',
help='eps random exploration')
parser.add_argument('--nn_baseline',
action='store_true',
help='Use nn baseline',
default=False)
# Attacker's defense
parser.add_argument('--argmaxed',
action='store_true',
help='Only consider argmax labels',
default=False)
parser.add_argument('--optimizer_choice',
type=str,
help='Optimizer',
default='sgdm',
choices=('sgd', 'sgdm', 'adam', 'adagrad'))
args = parser.parse_args()
params = vars(args)
torch.manual_seed(cfg.DEFAULT_SEED)
if params['device_id'] >= 0:
os.environ["CUDA_VISIBLE_DEVICES"] = str(params['device_id'])
device = torch.device('cuda')
ptu.init_gpu()
else:
device = torch.device('cpu')
model_dir = params['model_dir']
# ----------- Set up testset
dataset_name = params['testdataset']
valid_datasets = datasets.__dict__.keys()
modelfamily = datasets.dataset_to_modelfamily[dataset_name]
transform = datasets.modelfamily_to_transforms[modelfamily]['test']
if dataset_name not in valid_datasets:
raise ValueError(
'Dataset not found. Valid arguments = {}'.format(valid_datasets))
dataset = datasets.__dict__[dataset_name]
testset = dataset(train=False, transform=transform)
#if len(testset.classes) != num_classes:
# raise ValueError('# Transfer classes ({}) != # Testset classes ({})'.format(num_classes, len(testset.classes)))
# ----------- Set up queryset
queryset_name = params['queryset']
valid_datasets = datasets.__dict__.keys()
if queryset_name not in valid_datasets:
raise ValueError(
'Dataset not found. Valid arguments = {}'.format(valid_datasets))
modelfamily = datasets.dataset_to_modelfamily[queryset_name]
transform = datasets.modelfamily_to_transforms[modelfamily]['train']
try:
queryset = datasets.__dict__[queryset_name](train=True,
transform=transform)
except:
queryset = datasets.__dict__[queryset_name](split="train",
transform=transform)
num_classes = len(queryset.classes)
# ----------- Initialize blackbox
blackbox_dir = params['victim_model_dir']
blackbox = Blackbox.from_modeldir(blackbox_dir, device)
# ----------- Set up adversary model
model_name = params['model_arch']
pretrained = params['pretrained']
# model = model_utils.get_net(model_name, n_output_classes=num_classes, pretrained=pretrained)
adv_model = zoo.get_net(model_name,
modelfamily,
pretrained,
num_classes=10)
adv_model = adv_model.to(device)
# ----------- Initialize adversary
num_each_class = params['num_each_class']
agent_params = {
"ac_dim": num_classes,
"ob_dim": len(testset.classes),
"n_layers": params["n_layers"],
"size": params["size"],
"discrete": True,
"learning_rate": params["policy_lr"],
"num_agent_train_steps_per_iter":
params["num_agent_train_steps_per_iter"],
"agent_train_batch_size": params["agent_train_batch_size"],
"gamma": params["policy_gamma"],
"reward_to_go": True,
"nn_baseline": params["nn_baseline"],
"standardize_advantages": True,
"eps_random": params["eps_random"]
}
adversary = PGAdversary(queryset, num_each_class, agent_params)
# ----------- Set up transferset
def collect_training_trajectories(length, n_traj=10):
nonlocal avg_rewards, avg_components
paths = []
mean_rew = 0
mean_cert = mean_L = mean_E = mean_div = 0
X_paths, Y_paths = [], []
for _ in range(n_traj):
obs, acs, rewards, next_obs = [], [], [], []
r_certs, r_Ls, r_Es, r_divs = [], [], [], []
X_path, Y_path = [], []
X, actions = adversary.init_sampling()
X_path.append(X)
ob = blackbox(X)
Y_path.append(ob)
ob = ob.numpy()
for t in range(length - 1):
with torch.no_grad():
# Observe and react
obs.append(ob)
X_new, actions = adversary.sample(ob)
X_path.append(X_new)
acs.append(actions)
# Env gives feedback, which is a new observation
X_new = X_new.to(device)
ob = blackbox(X_new)
Y_path.append(ob)
ob = ob.cpu().numpy()
next_obs.append(ob)
Y_adv = adv_model(X_new)
Y_adv = F.softmax(Y_adv, dim=1).cpu().numpy()
reward, r_cert, r_L, r_E, r_div = adversary.agent.calculate_reward(
ob, np.concatenate(acs), Y_adv)
rewards.append(reward)
r_certs.append(r_cert)
r_Ls.append(r_L)
r_Es.append(r_E)
r_divs.append(r_div)
obs = np.concatenate(obs)
acs = np.concatenate(acs)
rewards = np.concatenate(rewards)
mean_rew += np.mean(rewards)
mean_cert += np.mean(np.concatenate(r_certs))
mean_L += np.mean(np.concatenate(r_Ls))
mean_E += np.mean(np.array(r_Es))
mean_div += np.mean(np.array(r_divs))
next_obs = np.concatenate(next_obs)
path = {
"observation": obs,
"action": acs,
"reward": rewards,
"next_observation": next_obs
}
paths.append(path)
X_paths.append(torch.cat(X_path))
Y_paths.append(torch.cat(Y_path))
print(f"==> Avg reward: {mean_rew / n_traj}")
avg_rewards.append(mean_rew / n_traj)
avg_components["avg_cert"].append(mean_cert / n_traj)
avg_components["avg_L"].append(mean_L / n_traj)
avg_components["avg_E"].append(mean_E / n_traj)
avg_components["avg_div"].append(mean_div / n_traj)
return torch.cat(X_paths), torch.cat(Y_paths), paths
traj_length = params['traj_length']
num_each_class = params['num_each_class']
n_iter = params['n_iter']
X, Y = None, None
budgets = params['budgets']
n_traj = params['n_traj_each_iter']
criterion_train = model_utils.soft_cross_entropy
if traj_length > 0:
n_iter = budgets // (traj_length * n_traj)
print(f"==> Budget = {n_iter} x {traj_length} x {n_traj}")
best_test_acc = []
best_acc = -1
avg_rewards = []
avg_components = collections.defaultdict(list)
for iter in range(1, n_iter + 1):
# n_iter * traj_length = budget
print(f"==> Iteration: {iter}/{n_iter}")
X_path, Y_path, paths = collect_training_trajectories(traj_length,
n_traj=n_traj)
adversary.add_to_replay_buffer(paths)
adversary.train_agent()
if X is None:
X, Y = X_path, Y_path
else:
X = torch.cat((X, X_path))
Y = torch.cat((Y, Y_path))
transferset = ImageTensorSet((X, Y))
# ----------- Train
#np.random.seed(cfg.DEFAULT_SEED)
#torch.manual_seed(cfg.DEFAULT_SEED)
#torch.cuda.manual_seed(cfg.DEFAULT_SEED)
optimizer = get_optimizer(adv_model.parameters(),
params['optimizer_choice'], **params)
print(f"Train on {len(transferset)} samples")
checkpoint_suffix = '.{extraction}'
best_acc = model_utils.train_model(adv_model,
transferset,
model_dir,
testset=testset,
criterion_train=criterion_train,
checkpoint_suffix=checkpoint_suffix,
device=device,
optimizer=optimizer,
benchmark=best_acc,
**params)
best_test_acc.append(best_acc)
adversary.agent.actor.save(
osp.join(model_dir, "checkpoint.agent.state_dict"))
# ----------- Log
torch.save(best_test_acc, osp.join(model_dir, "best_acc.pylist"))
torch.save(avg_rewards, osp.join(model_dir, "avg_rewards.pylist"))
torch.save(avg_components, osp.join(model_dir,
"avg_components.pydict"))
torch.save(adversary.idx_counter,
osp.join(model_dir, "idx_counter.pydict"))
torch.save(transferset, osp.join(model_dir, "transferset.pt"))
# Store arguments
params['created_on'] = str(datetime.now())
params_out_path = osp.join(model_dir, 'params_train.json')
with open(params_out_path, 'w') as jf:
json.dump(params, jf, indent=True)
agent_params_out_path = osp.join(model_dir, 'agent_params_train.json')
with open(agent_params_out_path, 'w') as jf:
json.dump(agent_params, jf, indent=True)
