def process_adv_examples(evade_method='dfgsm_k', mode='gen'):
"""
This function is used for the `attack` track challenge for two purposes
With mode='gen', it is meant to craft transferable adversarial examples and store them to a numpy array
With mode='eval', it loads up the examples from the numpy array and evaluates them on the tested model
Note, ADV Examples are only crafted for malicious files
:param evade_method: evasion method (participants can implement their own), here we use `dfgsm_k` as an example
:param mode: 'gen' to generate and store the adv examples or 'eval' to load them and evaluate
:return:
"""
model.eval()
category = "malicious"
total_loss = 0
total_correct = 0
total = 0
adv_file = os.path.join(adv_example_filepath, 'aes.npy')
xs_adv = [] if mode == 'gen' else np.load(adv_file)
# to be inline with the code base, the attack dataset will also be
# decomposed into train, val, and test. However, all these subsets will be
# used as part of the attack challenge.
xs_adv_offset = 0
for dset_type in ['train', 'val', 'test']:
if dset_type == 'train':
dataloader = train_dataloader_dict[category]
elif dset_type == 'val':
dataloader = valid_dataloader_dict[category]
elif dset_type == 'test':
dataloader = test_dataloader_dict[category]
# to impose the order of the crafted examples, we manually loop over the dataset
# instead of using the dataloader' sampler
batch_size = dataloader.batch_size
num_pts = len(dataloader.dataset)
num_batches = (num_pts + batch_size - 1) // batch_size
for batch_idx in range(num_batches):
# get the batch data
bstart = batch_idx * batch_size
bend = min(num_pts, bstart + batch_size)
x, y = zip(
*[dataloader.dataset[i] for i in range(bstart, bend)])
x = torch.stack(x, dim=0)
y = torch.LongTensor(y)
if mode == 'gen':
# put your method here
# ---------------------------
x_adv = inner_maximizer(x,
y,
model,
loss_fct,
iterations=evasion_iterations,
method=evade_method)
# ---------------------------
else:
x_adv = torch.from_numpy(
xs_adv[xs_adv_offset +
batch_idx * batch_size:xs_adv_offset +
(batch_idx + 1) * batch_size, :])
# stack input
if is_cuda:
x_adv = Variable(x_adv.cuda())
y = Variable(y.cuda())
else:
x_adv = Variable(x_adv)
y = Variable(y)
# forward pass
y_model = model(x_adv)
# loss pass
loss = loss_fct(y_model, y).mean()
# predict pass
_, predicted = torch.topk(y_model, k=1)
correct = predicted.data.eq(y.data.view_as(
predicted.data)).cpu().sum()
# metrics
total_loss += loss.data[0] * len(y)
total_correct += correct
total += len(y)
# let's save the adversarial examples
_x = x.numpy()
_x_adv = x_adv.cpu().data.numpy(
) if is_cuda else x_adv.data.numpy()
assert np.allclose(np.logical_and(_x, _x_adv),
_x), "perturbation constraint violated"
if mode == 'gen':
xs_adv = xs_adv + [_x_adv]
xs_adv_offset += num_pts
if mode == 'gen':
np.save(adv_file, np.concatenate(xs_adv, axis=0))
# we keep the same structure of metrics for compatibility
metrics = {
"bscn_num_pts": 1,
"bscn_exp_pts": 1,
"mal": {
"total_loss": 1,
"total_correct": 1,
"total": 1,
"evasion": {
"total_loss": total_loss,
"total_correct": total_correct,
"total": total
}
},
"bon": {
"total_loss": 1,
"total_correct": 1,
"total": 1
}
}
return metrics
def train(epoch):
model.train()
total_correct = 0.
total_loss = 0.
total = 0.
current_time = time.time()
if is_synthetic_dataset:
# since generation of synthetic data set is random, we'd like them to be the same over epochs
torch.manual_seed(seed_val)
random.seed(seed_val)
for batch_idx, ((bon_x, bon_y), (mal_x, mal_y)) in enumerate(
zip(train_dataloader_dict["benign"],
train_dataloader_dict["malicious"])):
# Check for adversarial learning
mal_x = inner_maximizer(
mal_x,
mal_y,
model,
loss_fct,
iterations=evasion_iterations,
method=training_method,
mal_index=batch_idx *
int(parameters["hyperparam"]["training_batch_size"]),
dataset=train_dataloader_dict)
# stack input
if is_cuda:
x = Variable(stack_tensors(bon_x, mal_x).cuda())
y = Variable(stack_tensors(bon_y, mal_y).cuda())
else:
x = Variable(stack_tensors(bon_x, mal_x))
y = Variable(stack_tensors(bon_y, mal_y))
# forward pass
y_model = model(x)
# backward pass
optimizer.zero_grad()
loss = loss_fct(y_model, y).mean()
loss.backward()
optimizer.step()
# predict pass
_, predicted = torch.topk(y_model, k=1)
correct = predicted.data.eq(y.data.view_as(
predicted.data)).cpu().sum()
# metrics
total_loss += loss.data[0] * len(y)
total_correct += correct
total += len(y)
bscn.update_numerator_batch(batch_idx, mal_x)
if batch_idx % log_interval == 0:
print("Time Taken:", time.time() - current_time)
current_time = time.time()
print(
"Train Epoch ({}) | Batch ({}) | [{}/{} ({:.0f}%)]\tBatch Loss: {:.6f}\tBatch Accuracy: {:.1f}%\t BSCN: {:.12f}"
.format(
epoch, batch_idx, batch_idx * len(x),
len(train_dataloader_dict["malicious"].dataset) +
len(train_dataloader_dict["benign"].dataset), 100. *
batch_idx / len(train_dataloader_dict["benign"]),
loss.data[0], 100. * correct / len(y), bscn.ratio()))
model_filename = "{name}_epoch_{e}".format(name=experiment_name,
e=epoch)
if save_every_epoch:
torch.save(model, os.path.join("model_weights", model_filename))
def check_one_category(category="benign",
dset_type='test',
is_evade=False,
evade_method='dfgsm_k'):
"""
test the model in terms of loss and accuracy on category, this function also allows to perform perturbation
with respect to loss to evade
:param category: benign or malicious dataset
:param dset_type: 'val', 'test', or 'train' dataset
:param is_evade: to perform evasion or not
:param evade_method: evasion method (we can use on of the inner maximier methods), it is only relevant if is_evade
is True
:return:
"""
model.eval()
total_loss = 0
total_correct = 0
total = 0
evasion_mode = ""
if is_synthetic_dataset:
# since generation of synthetic data set is random, we'd like them to be the same over epochs
torch.manual_seed(seed_val)
random.seed(seed_val)
if dset_type == 'train':
dataloader = train_dataloader_dict[category]
elif dset_type == 'val':
dataloader = valid_dataloader_dict[category]
elif dset_type == 'test':
dataloader = test_dataloader_dict[category]
else:
raise Exception("Invalid Dataset type")
for batch_idx, (x, y) in enumerate(dataloader):
#
if is_evade:
x = inner_maximizer(x,
y,
model,
loss_fct,
iterations=evasion_iterations,
method=evade_method)
evasion_mode = "(evasion using %s)" % evade_method
# stack input
if is_cuda:
x = Variable(x.cuda())
y = Variable(y.cuda())
else:
x = Variable(x)
y = Variable(y)
# forward pass
y_model = model(x)
# loss pass
loss = loss_fct(y_model, y).mean()
# predict pass
_, predicted = torch.topk(y_model, k=1)
correct = predicted.data.eq(y.data.view_as(
predicted.data)).cpu().sum()
# metrics
total_loss += loss.data[0] * len(y)
total_correct += correct
total += len(y)
print(
"{} set for {} {}: Average Loss: {:.4f}, Accuracy: {:.2f}%".format(
dset_type, category, evasion_mode, total_loss / total,
total_correct * 100. / total))
return total_loss, total_correct, total
def train_single_adversarial():
"""
With a single adversarial point, use an evasion method to train it and points in its feasible ball region
over a fixed number of iterations, keeping track of loss
"""
print("Train Single Adversarial")
base_figure_directory = "loss_progressions"
if not os.path.exists(base_figure_directory):
os.mkdir(base_figure_directory)
# Output will be in a folder with the current time
figure_directory = os.path.join(base_figure_directory, exp_time)
if not os.path.exists(figure_directory):
os.mkdir(figure_directory)
model.eval()
fig = plt.figure()
for i, (mal_x, mal_y) in enumerate(malicious_dataloader):
print(i, train_method, evasion_method)
if i == num_malware_samples_to_use:
break
# Axis labels and title
plt.xlabel("Inner Maximization Iterations",
fontsize=14,
fontweight='bold')
plt.ylabel("Loss Value", fontsize=14, fontweight='bold')
all_losses = []
# Use the original sample - solid line
_, losses = inner_maximizer(mal_x,
mal_y,
model,
loss_fct,
epsilon=epsilon,
iterations=evasion_iterations,
method=evasion_method,
return_loss=True,
report_loss_diff=False)
plt.plot(losses, linestyle='solid', linewidth=2.0)
all_losses.append(losses)
filename = "training_{t_method}_evasion_{e_method}_{num}_evasion_iterations_{mal_num}_sample.png".format(
t_method=train_method,
e_method=evasion_method,
num=evasion_iterations,
mal_num=i)
pickle.dump(
all_losses,
open(os.path.join(figure_directory, filename + ".p"), "wb"))
fig.savefig(os.path.join(figure_directory, filename))
if matplotlib.get_backend() != 'Agg':
plt.show()
plt.clf()
def plot_histogram():
print("Plotting All Methods Histogram")
base_figure_directory = "histograms"
if not os.path.exists(base_figure_directory):
os.mkdir(base_figure_directory)
# Output will be in a folder with the current time
figure_directory = os.path.join(base_figure_directory, exp_time)
if not os.path.exists(figure_directory):
os.mkdir(figure_directory)
model.eval()
# Axis labels and title
for i, (mal_x, mal_y) in enumerate(malicious_dataloader):
if i == num_malware_samples_to_use:
break
fig, ax = plt.subplots()
plt.xlabel("Loss Values")
plt.ylabel("Counts")
final_loss_values = []
for eva_method, c in zip(EVASION_METHODS, COLORS):
loss_values = []
for j in range(extra_points_for_each_sample):
_, losses = inner_maximizer(mal_x,
mal_y,
model,
loss_fct,
epsilon=epsilon,
iterations=evasion_iterations,
method=eva_method,
use_sample=True,
return_loss=True,
report_loss_diff=False)
loss_values.append(losses[-1][0])
final_loss_values.append(loss_values)
max_loss = max(list(itertools.chain.from_iterable(final_loss_values)))
num_bins = 50
bin_width = float(max_loss / num_bins)
plt.hist(final_loss_values,
bins=np.arange(0, 1.05 * max_loss, float(bin_width)),
color=COLORS,
stacked=True)
plt.legend(handles=legend_handles, labels=legend_labels)
filename = "histogram_{t_method}-training_all_evasions_{mal_num}_sample_{extra}_extra_points".format(
t_method=train_method,
mal_num=i,
extra=extra_points_for_each_sample)
fig.savefig(os.path.join(figure_directory, filename))
if matplotlib.get_backend() != 'Agg':
plt.show()
plt.clf()
pickle.dump(
final_loss_values,
open(os.path.join(figure_directory, filename + ".p"), "wb"))
def train(epoch):
model.train()
total_correct = 0.
total_loss = 0.
total = 0.
current_time = time.time()
# if is_synthetic_dataset:
# # since generation of synthetic data set is random, we'd like them to be the same over epochs
# torch.manual_seed(seed_val)
# random.seed(seed_val)
print("dataloaderr",train_dataloader_dict["benign"].dataset, train_dataloader_dict["malicious"].dataset)
for batch_idx, ((bon_x, bon_y), (mal_x, mal_y)) in enumerate(
zip(train_dataloader_dict["benign"], train_dataloader_dict["malicious"])):
# Check for adversarial learning
print("This is bacth_idx",batch_idx)
print(bon_x)
print(bon_y)
print(mal_x)
print(mal_y)
mal_x1 = inner_maximizer(
mal_x, mal_y, model, loss_fct, iterations=evasion_iterations, method=training_method)
print("THis is mal_X",mal_x1)
# stack input
if is_cuda:
x = Variable(stack_tensors(bon_x, mal_x, mal_x1).cuda())
y = Variable(stack_tensors(bon_y, mal_y, mal_y).cuda())
else:
x = Variable(stack_tensors(bon_x, mal_x, mal_x1))
y = Variable(stack_tensors(bon_y, mal_y, mal_y))
print("This is Train Forward Pass",x, y, model )
# forward pass
print("FORWARD X")
print(x)
y_model = model(x)
# backward pass
optimizer.zero_grad()
loss = loss_fct(y_model, y).mean()
loss.backward()
optimizer.step()
# predict pass
_, predicted = torch.topk(y_model, k=1)
correct = predicted.data.eq(y.data.view_as(predicted.data)).cpu().sum()
# metrics
print("LOSS LINE 148")
print(len(y))
print(y)
print(total_loss)
print(loss.data)
print(loss.data.item())
# total_loss += loss.data[0] * len(y)
total_loss += loss.data.item() * len(y)
total_correct += correct
total += len(y)
# bscn.update_numerator_batch(batch_idx, mal_x)
print("COVERING UPDATES",mal_x.size(0))
for i in range(mal_x.size(0)):
print("UPDATING NORMAL MALS")
a = bscn.update(mal_x[i])
print("UPDATINGS ADVS")
b = bscn.update(mal_x1[i])
print("UPDATING GOOD SPOTS")
c = gscn.update(bon_x[i])
print(a,b,c)
if batch_idx % log_interval == 0:
print("Time Taken:", time.time() - current_time)
current_time = time.time()
print(
"Train Epoch ({}) | Batch ({}) | [{}/{} ({:.0f}%)]\tBatch Loss: {:.6f}\tBatch Accuracy: {:.1f}%\t BSCN: {:.12f}".
# format(epoch, batch_idx, batch_idx * len(x),
# len(train_dataloader_dict["malicious"].dataset) +
# len(train_dataloader_dict["benign"].dataset),
# 100. * batch_idx / len(train_dataloader_dict["benign"]), loss.data[0],
# 100. * correct / len(y), bscn.ratio()))
format(epoch, batch_idx, batch_idx * len(x),
len(train_dataloader_dict["malicious"].dataset) +
len(train_dataloader_dict["benign"].dataset),
100. * batch_idx / len(train_dataloader_dict["benign"]), loss.data.item(),
100. * correct / len(y), bscn.Covering_value()))
# if is_losswise:
# graph_accuracy.append(epoch, {
# "train_accuracy_%s" % experiment_name: 100. * total_correct / total
# })
# graph_loss.append(epoch, {"train_loss_%s" % experiment_name: total_loss / total})
# graph_coverage.append(epoch, {"train_coverage_%s" % experiment_name: bscn.ratio()})
model_filename = "{name}_epoch_{e}".format(name=experiment_name, e=epoch)
def main():
'''
Generates adversarial samples using each training method and evasion method combination
'''
if len(sys.argv) == 1:
parameters = load_parameters("generate_adversarial_parameters.ini")
else:
parameters = load_parameters(sys.argv[1])
# Seed so that synthetic data is the same
use_seed = eval(parameters['general']['use_seed'])
if use_seed:
seed_val = int(parameters["general"]["seed"])
else:
seed_val = random.randint(1, 10000)
random.seed(seed_val)
torch.manual_seed(seed_val)
malicious_dataloader = load_malicious_data(parameters)
output_directory = parameters['general']['output_directory_for_adv_vecs']
if not os.path.exists(output_directory):
os.mkdir(output_directory)
print("Generating adversarial samples for each method")
loss_fct = nn.NLLLoss(reduce=False)
experiment_name = parameters['general']['experiment_name']
evasion_iterations = int(parameters['hyperparam']['evasion_iterations'])
print("Using experiment models: ", experiment_name)
train_methods = ['natural']
evasion_methods = ['rfgsm_k', 'dfgsm_k', 'bga_k', 'bca_k']
saved_model_directory = parameters['general']['saved_model_directory']
is_synthetic = eval(parameters['general']['is_synthetic_dataset'])
for train_method in train_methods:
train_directory = os.path.join(output_directory, train_method)
if not os.path.exists(train_directory):
os.mkdir(train_directory)
model_filepath = os.path.join(
saved_model_directory,
"[training:{train_meth}|evasion:{train_meth}]_{exp_name}-model.pt".
format(train_meth=train_method, exp_name=experiment_name))
model = torch.load(model_filepath)
for evasion_method in evasion_methods:
print(train_method, evasion_method)
evasion_subdirectory = os.path.join(train_directory,
evasion_method)
if not os.path.exists(evasion_subdirectory):
os.mkdir(evasion_subdirectory)
if is_synthetic:
for i, (mal_x, mal_y) in enumerate(malicious_dataloader):
actual_filename = str(i) + ".p"
mal_x = inner_maximizer(mal_x,
mal_y,
model,
loss_fct,
iterations=evasion_iterations,
method=evasion_method,
report_loss_diff=False)
# Save the adversarial vector as a pickle file
pickle.dump(
mal_x,
open(
os.path.join(evasion_subdirectory,
actual_filename), "wb"))
else:
for i, (mal_x, mal_y,
filepath) in enumerate(malicious_dataloader):
actual_filename = filepath[0].split('/')[-1]
mal_x = inner_maximizer(mal_x,
mal_y,
model,
loss_fct,
iterations=evasion_iterations,
method=evasion_method,
report_loss_diff=False)
# Save the adversarial vector as a pickle file
pickle.dump(
mal_x,
open(
os.path.join(evasion_subdirectory,
actual_filename), "wb"))
