def main():
args = parser.parse_args()
if not os.path.isdir('CMDs'):
os.mkdir('CMDs')
with open('CMDs/setup_dpn.cmd', 'a') as f:
f.write(' '.join(sys.argv) + '\n')
f.write('--------------------------------\n')
if os.path.isdir(args.destination_path) and not args.override_directory:
raise EnvironmentError(
'Destination directory exists. To override the directory run with '
'the --override_directory flag.')
else:
os.makedirs(args.destination_path)
# Link and and create directories
os.chdir(args.destination_path)
os.mkdir('model')
model = ModelFactory.create_model(args.arch,
num_classes=args.num_classes,
small_inputs=args.small_inputs,
dropout_rate=args.drop_rate,
pretrained=False,
n_channels=args.n_channels)
ModelFactory.checkpoint_model(path='model/model.tar',
model=model,
arch=args.arch,
n_channels=args.n_channels,
num_classes=args.num_classes,
dropout_rate=args.drop_rate,
small_inputs=args.small_inputs,
n_in=args.n_in)
def main():
args = parser.parse_args()
if not os.path.isdir('CMDs'):
os.mkdir('CMDs')
with open('CMDs/setup_dpn.cmd', 'a') as f:
f.write(' '.join(sys.argv) + '\n')
f.write('--------------------------------\n')
ensemble_dir = Path(args.destination_path)
if os.path.isdir(ensemble_dir):
if args.override_directory:
shutil.rmtree(ensemble_dir)
else:
raise EnvironmentError(
'Destination directory exists. To override the directory run with '
'the --override_directory flag.')
os.makedirs(ensemble_dir)
for i in range(args.num_models):
model_dir = ensemble_dir / f'model{i}'
os.mkdir(model_dir)
set_random_seeds(i)
model = ModelFactory.create_model(args.arch,
num_classes=args.num_classes,
small_inputs=args.small_inputs,
pretrained=False)
ModelFactory.checkpoint_model(path=model_dir / 'model.tar',
model=model,
arch=args.arch,
n_channels=args.n_channels,
num_classes=args.num_classes,
small_inputs=args.small_inputs,
n_in=args.n_in)
def main():
args = parser.parse_args()
if not os.path.isdir('CMDs'):
os.mkdir('CMDs')
with open('CMDs/eval_ID.cmd', 'a') as f:
f.write(' '.join(sys.argv) + '\n')
f.write('--------------------------------\n')
if os.path.isdir(args.output_path) and not args.overwrite:
print(f'Directory {args.output_path} exists. Exiting...')
sys.exit()
elif os.path.isdir(args.output_path) and args.overwrite:
os.remove(args.output_path + '/*')
else:
os.makedirs(args.output_path)
# Check that we are using a sensible GPU
device = select_gpu(args.gpu)
# Load up the model
model_dir = Path(args.model_dir)
ckpt = torch.load(model_dir / 'model/model.tar', map_location=device)
model = ModelFactory.model_from_checkpoint(ckpt)
model.to(device)
model.eval()
# Load the in-domain evaluation data
if args.train:
dataset = DATASET_DICT[args.dataset](
root=args.data_path,
transform=construct_transforms(
n_in=ckpt['n_in'],
mean=DATASET_DICT[args.dataset].mean,
std=DATASET_DICT[args.dataset].std,
mode='train'),
target_transform=None,
download=True,
split='train')
else:
dataset = DATASET_DICT[args.dataset](
root=args.data_path,
transform=construct_transforms(
n_in=ckpt['n_in'],
mean=DATASET_DICT[args.dataset].mean,
std=DATASET_DICT[args.dataset].std,
mode='eval'),
target_transform=None,
download=True,
split='test')
# Evaluate the model
logits, labels = eval_logits_on_dataset(model=model,
dataset=dataset,
batch_size=args.batch_size,
device=device)
labels, probs, logits = labels.numpy(), F.softmax(
logits, dim=1).numpy(), logits.numpy()
# Save model outputs
np.savetxt(os.path.join(args.output_path, 'labels.txt'), labels)
np.savetxt(os.path.join(args.output_path, 'probs.txt'), probs)
np.savetxt(os.path.join(args.output_path, 'logits.txt'), logits)
# Get dictionary of uncertainties.
uncertainties = dirichlet_prior_network_uncertainty(logits)
# Save uncertainties
for key in uncertainties.keys():
np.savetxt(os.path.join(args.output_path, key + '.txt'),
uncertainties[key])
if args.ood:
sys.exit()
nll = -np.mean(
np.log(probs[np.arange(probs.shape[0]),
np.squeeze(labels)] + 1e-10))
accuracy = np.mean(
np.asarray(labels == np.argmax(probs, axis=1), dtype=np.float32))
with open(os.path.join(args.output_path, 'results.txt'), 'a') as f:
f.write(
f'Classification Error: {np.round(100 * (1.0 - accuracy), 1)} \n')
f.write(f'NLL: {np.round(nll, 3)} \n')
# TODO: Have different results files? Or maybedifferent folders
# Assess Misclassification Detection
eval_misc_detect(labels,
probs,
uncertainties,
save_path=args.output_path,
misc_positive=True)
# Assess Calibration
classification_calibration(labels=labels,
probs=probs,
save_path=args.output_path)
# Assess Rejection Performance
eval_rejection_ratio_class(labels=labels,
probs=probs,
uncertainties=uncertainties,
save_path=args.output_path)
def main():
args = parser.parse_args()
if not os.path.isdir('CMDs'):
os.mkdir('CMDs')
with open('CMDs/step_train_distillation.cmd', 'a') as f:
f.write(' '.join(sys.argv) + '\n')
f.write('--------------------------------\n')
model_dir = Path(args.model_dir)
# Check that we are training on a sensible GPU
assert max(args.gpu) <= torch.cuda.device_count() - 1
device = select_gpu(args.gpu)
# Load up the model
ckpt = torch.load(model_dir / 'model/model.tar', map_location=device)
model = ModelFactory.model_from_checkpoint(ckpt)
if len(args.gpu) > 1 and torch.cuda.device_count() > 1:
model = torch.nn.DataParallel(model, device_ids=args.gpu)
print('Using Multi-GPU training.')
model.to(device)
# Load the in-domain training and validation data
train_dataset_class = DATASET_DICT[args.dataset]
train_dataset_parameters = {
'root':
args.data_path,
'transform':
construct_transforms(n_in=ckpt['n_in'],
mode='train',
mean=DATASET_DICT[args.dataset].mean,
std=DATASET_DICT[args.dataset].std,
augment=args.augment,
rotation=args.rotate,
jitter=args.jitter),
'target_transform':
None,
'download':
True,
'split':
'train'
}
train_dataset = EnsembleDataset(
dataset=train_dataset_class,
dataset_parameters=train_dataset_parameters,
ensemble_path=args.ensemble_path,
model_dirs=args.model,
n_models=args.n_models,
folder='train')
val_dataset_class = DATASET_DICT[args.dataset]
val_dataset_parameters = {
'root':
args.data_path,
'transform':
construct_transforms(n_in=ckpt['n_in'],
mean=DATASET_DICT[args.dataset].mean,
std=DATASET_DICT[args.dataset].std,
mode='eval'),
'target_transform':
None,
'download':
True,
'split':
'val'
}
val_dataset = EnsembleDataset(dataset=val_dataset_class,
dataset_parameters=val_dataset_parameters,
ensemble_path=args.ensemble_path,
model_dirs=args.model,
n_models=args.n_models,
folder='eval')
if args.ood:
assert args.ood_folder is not None
ood_dataset_class = DATASET_DICT[args.ood_dataset]
ood_dataset_parameters = {
'root':
args.data_path,
'transform':
construct_transforms(n_in=ckpt['n_in'],
mean=DATASET_DICT[args.dataset].mean,
std=DATASET_DICT[args.dataset].std,
mode='train'),
'target_transform':
None,
'download':
True,
'split':
'train'
}
ood_dataset = EnsembleDataset(
dataset=ood_dataset_class,
dataset_parameters=ood_dataset_parameters,
ensemble_path=args.ensemble_path,
model_dirs=args.model,
n_models=args.n_models,
folder=args.ood_folder)
train_dataset = data.ConcatDataset([train_dataset, ood_dataset])
# Set up training and test criteria
test_criterion = torch.nn.CrossEntropyLoss()
if args.endd:
train_criterion = DirichletEnDDLoss()
elif args.endd_entemp:
train_criterion = DirichletEnDDEnTempLoss()
elif args.endd_dirtemp:
train_criterion = DirichletEnDDDirTempLoss()
elif args.endd_revtemp:
train_criterion = DirichletEnDDRevLoss()
else:
train_criterion = EnDLoss()
# Select optimizer and optimizer params
optimizer, optimizer_params = choose_optimizer(args.optimizer, args.lr,
args.weight_decay)
# Setup model trainer and train model
trainer = TrainerDistillation(model=model,
criterion=train_criterion,
test_criterion=test_criterion,
train_dataset=train_dataset,
test_dataset=val_dataset,
optimizer=optimizer,
device=device,
checkpoint_path=model_dir / 'model',
scheduler=optim.lr_scheduler.MultiStepLR,
temp_scheduler=LRTempScheduler,
optimizer_params=optimizer_params,
scheduler_params={
'milestones': args.lrc,
'gamma': args.lr_decay
},
temp_scheduler_params={
'init_temp': args.temperature,
'min_temp': args.min_temperature,
'decay_epoch': args.tdecay_epoch,
'decay_length': args.tdecay_length
},
batch_size=args.batch_size,
clip_norm=args.clip_norm)
if args.resume:
trainer.load_checkpoint(model_dir / 'model/checkpoint.tar',
True,
True,
map_location=device)
trainer.train(args.n_epochs, resume=args.resume)
# Save final model
if len(args.gpu) > 1 and torch.cuda.device_count() > 1:
model = model.module
ModelFactory.checkpoint_model(path=model_dir / 'model/model.tar',
model=model,
arch=ckpt['arch'],
n_channels=ckpt['n_channels'],
num_classes=ckpt['num_classes'],
small_inputs=ckpt['small_inputs'],
n_in=ckpt['n_in'])
def main():
args = parser.parse_args()
if not os.path.isdir('CMDs'):
os.mkdir('CMDs')
with open('CMDs/ood_detect.cmd', 'a') as f:
f.write(' '.join(sys.argv) + '\n')
f.write('--------------------------------\n')
if os.path.isdir(args.output_path) and not args.overwrite:
print(f'Directory {args.output_path} exists. Exiting...')
sys.exit()
elif os.path.isdir(args.output_path) and args.overwrite:
os.remove(args.output_path + '/*')
else:
os.makedirs(args.output_path)
# Check that we are using a sensible GPU device
device = select_gpu(args.gpu)
model_dir = Path(args.model_dir)
# Load up the model
ckpt = torch.load(model_dir, map_location=device)
model = ModelFactory.model_from_checkpoint(ckpt)
model.to(device)
model.eval()
# Load the in-domain evaluation data
id_dataset = DATASET_DICT[args.id_dataset](
root=args.data_path,
transform=construct_transforms(n_in=ckpt['n_in'],
mean=DATASET_DICT[args.id_dataset].mean,
std=DATASET_DICT[args.id_dataset].std,
mode='eval'),
target_transform=None,
download=True,
split='test')
# Evaluate the model
id_logits, id_labels = eval_logits_on_dataset(model=model,
dataset=id_dataset,
batch_size=args.batch_size,
device=device)
id_labels, id_probs, id_logits = id_labels.numpy(), F.softmax(
id_logits, dim=1).numpy(), id_logits.numpy()
accuracy = np.mean(np.squeeze(id_labels) == np.argmax(id_probs, axis=-1))
ece, mce = classification_calibration(id_labels, id_probs)
brier = brier_score(id_labels, id_probs)
uncertainty, resolution, reliability = brier_decomposition(
id_labels, id_probs)
with open(os.path.join(args.output_path, 'results.txt'), 'a') as f:
f.write(f"Accuracy: {accuracy * 100:.02f}\n")
f.write(f"ECE: {ece * 100:.02f}\n")
f.write(f"MCE: {mce * 100:.02f}\n")
f.write(f"Brier: {brier * 100:.02f}\n")
f.write(f"Brier UNC: {uncertainty * 100:.02f}\n")
f.write(f"Brier RES: {resolution * 100:.02f}\n")
f.write(f"Brier REL: {reliability * 100:.02f}\n")
id_uncertainties = dirichlet_prior_network_uncertainty(id_logits)
for ood_dataset in args.ood_dataset:
save_dir = os.path.join(args.output_path, ood_dataset)
os.mkdir(save_dir)
with open(os.path.join(args.output_path, 'results.txt'), 'a') as f:
f.write(f"\n\n{ood_dataset}\n")
split = ood_dataset.split("_", maxsplit=1)
if len(split) == 2 and split[1] == "unscaled":
unscaled_ds = True
ood_dataset_name = split[0]
else:
unscaled_ds = False
ood_dataset_name = ood_dataset
ood_dataset = DATASET_DICT[ood_dataset_name](
root=args.data_path,
transform=construct_transforms(
n_in=ckpt['n_in'],
mean=DATASET_DICT[args.id_dataset].mean,
std=DATASET_DICT[args.id_dataset].std,
mode='eval',
unscaled=unscaled_ds),
target_transform=None,
download=True,
split='test')
print(
f"ID dataset length: {len(id_dataset)}, OOD dataset length: {len(ood_dataset)}"
)
ood_logits, ood_labels = eval_logits_on_dataset(
model=model,
dataset=ood_dataset,
batch_size=args.batch_size,
device=device)
ood_labels, ood_probs, ood_logits = ood_labels.numpy(), F.softmax(
ood_logits, dim=1).numpy(), ood_logits.numpy()
ood_uncertainties = dirichlet_prior_network_uncertainty(ood_logits)
# Compute Labels
in_domain = np.zeros_like(id_labels)
out_domain = np.ones((ood_labels.shape[0], ))
domain_labels = np.concatenate((in_domain, out_domain), axis=0)
eval_ood_detect(domain_labels=domain_labels,
in_uncertainties=id_uncertainties,
out_uncertainties=ood_uncertainties,
save_path=save_dir)
def main():
args = parser.parse_args()
if not os.path.isdir('CMDs'):
os.mkdir('CMDs')
with open('CMDs/construct_adversarial_dataset.cmd', 'a') as f:
f.write(' '.join(sys.argv) + '\n')
f.write('--------------------------------\n')
if os.path.isdir(args.output_path) and not args.overwrite:
print(f'Directory {args.output_path} exists. Exiting...')
sys.exit()
elif os.path.isdir(args.output_path) and args.overwrite:
os.remove(args.output_path + '/*')
else:
os.makedirs(args.output_path)
# Check that we are using a sensible GPU
device = select_gpu(args.gpu)
# Load up the model
model_dir = Path(args.model_dir)
ckpt = torch.load(model_dir / 'model/model.tar', map_location=device)
model = ModelFactory.model_from_checkpoint(ckpt)
model.to(device)
model.eval()
# Load the in-domain evaluation data
if args.train:
dataset = DATASET_DICT[args.dataset](
root=args.data_path,
transform=construct_transforms(
n_in=ckpt['n_in'],
mean=DATASET_DICT[args.dataset].mean,
std=DATASET_DICT[args.dataset].std,
num_channels=args.n_channels,
mode='train'),
target_transform=None,
download=True,
split='train')
else:
dataset = DATASET_DICT[args.dataset](
root=args.data_path,
transform=construct_transforms(
n_in=ckpt['n_in'],
mean=DATASET_DICT[args.dataset].mean,
std=DATASET_DICT[args.dataset].std,
num_channels=args.n_channels,
mode='eval'),
target_transform=None,
download=True,
split='test')
# dataset = DataSpliter.reduceSize(dataset, 16)
mean = np.array(DATASET_DICT[args.dataset].mean).reshape((3, 1, 1))
std = np.array(DATASET_DICT[args.dataset].std).reshape((3, 1, 1))
# dataset = torch.utils.data.Subset(dataset, image_indices)
print("dataset length:", len(dataset))
org_dataset_folder = os.path.join(args.output_path, "org-images")
os.makedirs(org_dataset_folder)
persist_dataset(dataset, mean, std, args.n_channels, org_dataset_folder)
# perform original evaluation on the model using unperturbed images
logits, labels = eval_logits_on_dataset(model,
dataset=dataset,
device=device,
batch_size=args.batch_size)
labels = labels.numpy()
# determine correct classifications without attack (original non perturbed images)
org_preds = np.argmax(F.softmax(logits, dim=1).numpy(), axis=1)
correct_classifications = np.asarray(org_preds == labels, dtype=np.int32)
correct_classified_indices = np.argwhere(correct_classifications == 1)
# perform attacks on the same dataset, using different epsilon values.
adv_success_rates = []
for epsilon in args.epsilon:
attack_folder = os.path.join(args.output_path, f"e{epsilon}-attack")
out_path = os.path.join(attack_folder, "adv-images")
os.makedirs(out_path)
adv_success = perform_epsilon_attack(model,
epsilon,
dataset,
correct_classified_indices,
args.batch_size,
device,
args.n_channels,
attack_folder,
mean,
std,
args.attack_type,
norm=args.norm,
step_size=args.step_size,
max_steps=args.max_steps)
adv_success_rates.append(adv_success / len(correct_classified_indices))
# plot the epsilon, adversarial success rate graph (line plot)
plt.figure(figsize=(5, 5))
plt.plot(args.epsilon, adv_success_rates, "*-")
plt.yticks(np.arange(0, 1.1, step=0.1))
plt.xticks(np.arange(np.min(args.epsilon), np.max(args.epsilon), step=0.1))
plt.title("Adversarial Success Rate vs Epsilon")
plt.xlabel("Epsilon")
plt.ylabel("Adversarial Success Rate")
plt.savefig(os.path.join(args.output_path, "epsilon-curve.png"))
def main():
args = parser.parse_args()
if not os.path.isdir('CMDs'):
os.mkdir('CMDs')
with open('CMDs/step_train_dpn.cmd', 'a') as f:
f.write(' '.join(sys.argv) + '\n')
f.write('--------------------------------\n')
model_dir = Path(args.model_dir)
checkpoint_path = args.checkpoint_path
if checkpoint_path is None:
checkpoint_path = model_dir / 'model'
# Check that we are training on a sensible GPU
assert max(args.gpu) <= torch.cuda.device_count() - 1
device = select_gpu(args.gpu)
# Load up the model
ckpt = torch.load(model_dir / 'model/model.tar', map_location=device)
model = ModelFactory.model_from_checkpoint(ckpt)
if len(args.gpu) > 1 and torch.cuda.device_count() > 1:
model = torch.nn.DataParallel(model, device_ids=args.gpu)
print('Using Multi-GPU training.')
model.to(device)
if args.normalize:
mean = DATASET_DICT[args.id_dataset].mean
std = DATASET_DICT[args.id_dataset].std
else:
mean = (0.5, 0.5, 0.5)
std = (0.5, 0.5, 0.5)
# Load the in-domain training and validation data
train_dataset = DATASET_DICT[args.id_dataset](root=args.data_path,
transform=construct_transforms(
n_in=ckpt['n_in'],
mode='train',
mean=mean,
std=std,
augment=args.augment,
rotation=args.rotate,
jitter=args.jitter),
target_transform=None,
download=True,
split='train')
val_dataset = DATASET_DICT[args.id_dataset](root=args.data_path,
transform=construct_transforms(
n_in=ckpt['n_in'],
mean=mean,
std=std,
mode='eval',
rotation=args.rotate,
jitter=args.jitter),
target_transform=None,
download=True,
split='val')
# Load the out-of-domain training dataset
ood_dataset = DATASET_DICT[args.ood_dataset](root=args.data_path,
transform=construct_transforms(
n_in=ckpt['n_in'],
mean=mean,
std=std,
mode='ood'),
target_transform=None,
download=True,
split='train')
ood_val_dataset = DATASET_DICT[args.ood_dataset](root=args.data_path,
transform=construct_transforms(
n_in=ckpt['n_in'],
mean=mean,
std=std,
mode='eval'),
target_transform=None,
download=True,
split='val')
# Combine ID and OOD training datasets into a single dataset for
# training (necessary for DataParallel training)
assert len(val_dataset) == len(ood_val_dataset)
# Even out dataset lengths.
id_ratio = 1.0
if len(train_dataset) < len(ood_dataset):
id_ratio = np.ceil(float(len(ood_dataset)) / float(len(train_dataset)))
assert id_ratio.is_integer()
dataset_list = [train_dataset, ] * (int(id_ratio))
train_dataset = data.ConcatDataset(dataset_list)
if len(train_dataset) > len(ood_dataset):
ratio = np.ceil(float(len(train_dataset)) / float(len(ood_dataset)))
assert ratio.is_integer()
dataset_list = [ood_dataset, ] * int(ratio)
ood_dataset = data.ConcatDataset(dataset_list)
if len(ood_dataset) > len(train_dataset):
ood_dataset = data.Subset(ood_dataset, np.arange(0, len(train_dataset)))
assert len(train_dataset) == len(ood_dataset)
print(f"Validation dataset length: {len(val_dataset)}")
print(f"Train dataset length: {len(train_dataset)}")
# Set up training and test criteria
id_criterion = DirichletKLLoss(target_concentration=args.target_concentration,
concentration=args.concentration,
reverse=args.reverse_KL)
ood_criterion = DirichletKLLoss(target_concentration=0.0,
concentration=args.concentration,
reverse=args.reverse_KL)
criterion = PriorNetMixedLoss([id_criterion, ood_criterion], mixing_params=[1.0, args.gamma])
# Select optimizer and optimizer params
optimizer, optimizer_params = choose_optimizer(args.optimizer,
args.lr,
args.weight_decay)
# Setup model trainer and train model
lrc = [int(lrc / id_ratio) for lrc in args.lrc]
trainer = TrainerWithOOD(model=model,
criterion=criterion,
id_criterion=id_criterion,
ood_criterion=ood_criterion,
test_criterion=criterion,
ood_dataset=ood_dataset,
test_ood_dataset=ood_val_dataset,
train_dataset=train_dataset,
test_dataset=val_dataset,
optimizer=optimizer,
device=device,
checkpoint_path=checkpoint_path,
scheduler=optim.lr_scheduler.MultiStepLR,
optimizer_params=optimizer_params,
scheduler_params={'milestones': lrc, 'gamma': args.lr_decay},
batch_size=args.batch_size,
clip_norm=args.clip_norm)
if args.resume:
try:
trainer.load_checkpoint(True, True, map_location=device)
except:
print('No checkpoint found, training from empty model.')
pass
trainer.train(int(args.n_epochs / id_ratio), resume=args.resume)
# Save final model
if len(args.gpu) > 1 and torch.cuda.device_count() > 1:
model = model.module
ModelFactory.checkpoint_model(path=model_dir / 'model/model.tar',
model=model,
arch=ckpt['arch'],
dropout_rate=ckpt['dropout_rate'],
n_channels=ckpt['n_channels'],
num_classes=ckpt['num_classes'],
small_inputs=ckpt['small_inputs'],
n_in=ckpt['n_in'])
def main():
args = parser.parse_args()
if not os.path.isdir('CMDs'):
os.mkdir('CMDs')
with open('CMDs/ood_detect.cmd', 'a') as f:
f.write(' '.join(sys.argv) + '\n')
f.write('--------------------------------\n')
if os.path.isdir(args.output_path) and not args.overwrite:
print(f'Directory {args.output_path} exists. Exiting...')
sys.exit()
elif os.path.isdir(args.output_path) and args.overwrite:
os.remove(args.output_path + '/*')
else:
os.makedirs(args.output_path)
# Check that we are using a sensible GPU device
device = select_gpu(args.gpu)
model_dir = Path(args.model_dir)
# Load up the model
ckpt = torch.load(model_dir / 'model/model.tar', map_location=device)
model = ModelFactory.model_from_checkpoint(ckpt)
model.to(device)
model.eval()
# Load the in-domain evaluation data
id_dataset = DATASET_DICT[args.id_dataset](
root=args.data_path,
transform=construct_transforms(n_in=ckpt['n_in'],
mean=DATASET_DICT[args.id_dataset].mean,
std=DATASET_DICT[args.id_dataset].std,
num_channels=args.n_channels,
mode='eval'),
target_transform=None,
download=True,
split='test')
# id_dataset = DataSpliter.reduceSize(id_dataset, 10)
ood_dataset = DATASET_DICT[args.ood_dataset](
root=args.data_path,
transform=construct_transforms(n_in=ckpt['n_in'],
mean=DATASET_DICT[args.id_dataset].mean,
std=DATASET_DICT[args.id_dataset].std,
num_channels=args.n_channels,
mode='eval'),
target_transform=None,
download=True,
split='test')
# ood_dataset = DataSpliter.reduceSize(ood_dataset, 10)
print(
f"ID dataset length: {len(id_dataset)}, OOD dataset length: {len(ood_dataset)}"
)
# Evaluate the model
if args.use_attack_id_dataset:
id_logits, id_labels, _ = construct_adversarial_dataset(
model=model,
dataset=id_dataset,
epsilon=0.5,
batch_size=args.batch_size,
device=device,
attack_type='FGSM',
step_size=0.4,
norm='inf',
max_steps=10)
else:
id_logits, id_labels = eval_logits_on_dataset(
model=model,
dataset=id_dataset,
batch_size=args.batch_size,
device=device)
ood_logits, ood_labels = eval_logits_on_dataset(model=model,
dataset=ood_dataset,
batch_size=args.batch_size,
device=device)
id_labels, id_probs, id_logits = id_labels.numpy(), F.softmax(
id_logits, dim=1).numpy(), id_logits.numpy()
ood_labels, ood_probs, ood_logits = ood_labels.numpy(), F.softmax(
ood_logits, dim=1).numpy(), ood_logits.numpy()
# Save model outputs
np.savetxt(os.path.join(args.output_path, 'id_labels.txt'), id_labels)
np.savetxt(os.path.join(args.output_path, 'id_probs.txt'), id_probs)
np.savetxt(os.path.join(args.output_path, 'id_logits.txt'), id_logits)
np.savetxt(os.path.join(args.output_path, 'ood_labels.txt'), ood_labels)
np.savetxt(os.path.join(args.output_path, 'ood_probs.txt'), ood_probs)
np.savetxt(os.path.join(args.output_path, 'ood_logits.txt'), ood_logits)
# Get dictionary of uncertainties.
id_uncertainties = dirichlet_prior_network_uncertainty(id_logits)
ood_uncertainties = dirichlet_prior_network_uncertainty(ood_logits)
# Save uncertainties
for key in id_uncertainties.keys():
np.savetxt(os.path.join(args.output_path, key + '_id.txt'),
id_uncertainties[key])
np.savetxt(os.path.join(args.output_path, key + '_ood.txt'),
ood_uncertainties[key])
# Compute Labels
# For OOD detection task, we have a binary classification problem at hand, where label=0 indicates in-domain sample and
# label=1 indicates ood sample.
in_domain = np.zeros_like(id_labels)
out_domain = np.ones_like(ood_labels)
domain_labels = np.concatenate((in_domain, out_domain), axis=0)
eval_ood_detect(domain_labels=domain_labels,
in_uncertainties=id_uncertainties,
out_uncertainties=ood_uncertainties,
save_path=args.output_path)
def main():
args = parser.parse_args()
if not os.path.isdir('CMDs'):
os.mkdir('CMDs')
with open('CMDs/step_train_dnn.cmd', 'a') as f:
f.write(' '.join(sys.argv) + '\n')
f.write('--------------------------------\n')
model_dir = Path(args.model_dir)
checkpoint_path = args.checkpoint_path
if checkpoint_path is None:
checkpoint_path = model_dir / 'model'
# Load up the model
assert max(args.gpu) <= torch.cuda.device_count() - 1
device = select_gpu(args.gpu)
# Load up the model
ckpt = torch.load(model_dir / 'model/model.tar', map_location=device)
model = ModelFactory.model_from_checkpoint(ckpt)
if len(args.gpu) > 1 and torch.cuda.device_count() > 1:
model = torch.nn.DataParallel(model, device_ids=args.gpu)
print('Using Multi-GPU training.')
model.to(device)
if args.normalize:
mean = DATASET_DICT[args.dataset].mean
std = DATASET_DICT[args.dataset].std
else:
mean = (0.5, 0.5, 0.5)
std = (0.5, 0.5, 0.5)
# Load the in-domain training and validation data
train_dataset = DATASET_DICT[args.dataset](root=args.data_path,
transform=construct_transforms(
n_in=ckpt['n_in'],
mode='train',
mean=mean,
std=std,
augment=args.augment,
rotation=args.rotate,
jitter=args.jitter),
target_transform=None,
download=True,
split='train')
val_dataset = DATASET_DICT[args.dataset](root=args.data_path,
transform=construct_transforms(
n_in=ckpt['n_in'],
mean=mean,
std=std,
mode='eval'),
target_transform=None,
download=True,
split='val')
# Check that we are training on a sensible GPU
# Set up training and test criteria
criterion = torch.nn.CrossEntropyLoss()
# Select optimizer and optimizer params
optimizer, optimizer_params = choose_optimizer(args.optimizer, args.lr,
args.weight_decay)
# Setup model trainer and train model
trainer = Trainer(model=model,
criterion=criterion,
test_criterion=criterion,
train_dataset=train_dataset,
test_dataset=val_dataset,
optimizer=optimizer,
device=device,
checkpoint_path=checkpoint_path,
scheduler=optim.lr_scheduler.MultiStepLR,
optimizer_params=optimizer_params,
scheduler_params={
'milestones': args.lrc,
'gamma': args.lr_decay
},
batch_size=args.batch_size,
clip_norm=args.clip_norm)
if args.resume:
try:
trainer.load_checkpoint(True, True, map_location=device)
except:
print('No checkpoint found, training from empty model.')
pass
trainer.train(args.n_epochs, resume=args.resume)
# Save final model
if len(args.gpu) > 1 and torch.cuda.device_count() > 1:
model = model.module
ModelFactory.checkpoint_model(path=model_dir / 'model/model.tar',
model=model,
arch=ckpt['arch'],
dropout_rate=ckpt['dropout_rate'],
n_channels=ckpt['n_channels'],
num_classes=ckpt['num_classes'],
small_inputs=ckpt['small_inputs'],
n_in=ckpt['n_in'])
def main():
args = parser.parse_args()
if not os.path.isdir('CMDs'):
os.mkdir('CMDs')
with open('CMDs/construct_adversarial_attack.cmd', 'a') as f:
f.write(' '.join(sys.argv) + '\n')
f.write('--------------------------------\n')
if os.path.isdir(args.output_path) and not args.overwrite:
print(f'Directory {args.output_path} exists. Exiting...')
sys.exit()
elif os.path.isdir(args.output_path) and args.overwrite:
os.remove(args.output_path + '/*')
else:
os.makedirs(args.output_path)
os.makedirs(os.path.join(args.output_path, 'images'))
# Check that we are using a sensible GPU
device = select_gpu(args.gpu)
# Load up the model
model_dir = Path(args.model_dir)
ckpt = torch.load(os.path.join(model_dir, 'model/model.tar'),
map_location=device)
model = ModelFactory.model_from_checkpoint(ckpt)
model.to(device)
model.eval()
# Wrap model with a Foolbox wrapper.
mean = np.array([0.4914, 0.4823, 0.4465]).reshape((3, 1, 1))
std = np.array([0.247, 0.243, 0.261]).reshape((3, 1, 1))
fmodel = PyTorchModel(model,
bounds=(0, 1),
num_classes=ckpt['num_classes'],
preprocessing=(mean, std))
# Load the evaluation data
if args.train:
dataset = DATASET_DICT[args.dataset](root=args.data_path,
transform=construct_transforms(
n_in=ckpt['n_in'],
mode='train'),
target_transform=None,
download=True,
split='train')
else:
dataset = DATASET_DICT[args.dataset](root=args.data_path,
transform=construct_transforms(
n_in=ckpt['n_in'],
mode='eval'),
target_transform=None,
download=True,
split='test')
loader = DataLoader(dataset, batch_size=args.batch_size, num_workers=1)
# Construct adversarial attack
if args.attack == 'CWL2':
if args.adaptive:
attack = AdaptiveCarliniWagnerL2Attack(model=fmodel)
else:
attack = CarliniWagnerL2Attack(model=fmodel)
elif args.attack == 'EAD':
if args.adaptive:
attack = AdaptiveEADAttack(model=fmodel)
else:
attack = EADAttack(model=fmodel)
else:
raise NotImplementedError
adversarials = []
for i, data in enumerate(loader):
start = time.time()
images, labels = data
images = images.numpy()
labels = labels.numpy()
adversarials.extend(attack(inputs=images, labels=labels, unpack=False))
print(
f"Batch {i}/{len(loader)} took {np.round((time.time() - start) / 60.0, 1)} minutes."
)
adv_labels = np.stack(
[adversarial.adversarial_class for adversarial in adversarials],
axis=0)
labels = np.stack(
[adversarial.original_class for adversarial in adversarials], axis=0)
distances = np.stack(
[adversarial.distance for adversarial in adversarials], axis=0)
logits = np.stack([adversarial.output for adversarial in adversarials],
axis=0)
np.savetxt(os.path.join(args.output_path, 'labels.txt'),
labels,
dtype=np.int32)
np.savetxt(os.path.join(args.output_path, 'adv_labels.txt'),
adv_labels,
dtype=np.int32)
np.savetxt(os.path.join(args.output_path, 'logits.txt'),
logits,
dtype=np.float32)
np.savetxt(os.path.join(args.output_path, 'distances.txt'),
distances,
dtype=np.float32)
accuracy = np.mean(np.asarray(labels == adv_labels, dtype=np.float32))
sr = np.mean(np.asarray(labels != adv_labels, dtype=np.float32))
with open(os.path.join(args.output_path, 'results.txt'), 'a') as f:
f.write(
f'Classification Error: {np.round(100 * (1.0 - accuracy), 1)} \n')
f.write(f'Success Rate: {np.round(100 * sr, 1)} \n')
print("Saving images to folder...")
adversarial_images = np.stack(
[adversarial.perturbed for adversarial in adversarials], axis=0)
for i, image in enumerate([
np.asarray(255.0 * adversarial.perturbed, dtype=np.uint8)
for adversarial in adversarials
]):
print(np.max(adversarial_images), np.min(adversarial_images))
Image.fromarray(image).save(
os.path.join(args.output_path, f"images/{i}.png"))
