def eval_fgsm_bnn(model,
data,
inv_factors,
estimator='kfac',
samples=30,
epsilon=0.1,
stats=True,
device=torch.device('cuda'),
verbose=True):
model.eval()
mean_state = copy.deepcopy(model.state_dict())
mean_predictions = 0
samples = tqdm.tqdm(range(samples), disable=not verbose)
for sample in samples:
samples.set_postfix({'RAM': ram(), 'VRAM': vram()})
sample_and_replace_weights(model, inv_factors, estimator)
predictions, labels, _ = eval_fgsm(model, data, epsilon, stats=False, device=device, verbose=False)
mean_predictions += predictions
model.load_state_dict(mean_state)
mean_predictions /= len(samples)
if stats:
acc = accuracy(mean_predictions, labels)
ece1 = 100 * expected_calibration_error(mean_predictions, labels)[0]
ece2 = 100 * calibration_curve(mean_predictions, labels)[0]
nll = negative_log_likelihood(mean_predictions, labels)
ent = predictive_entropy(mean_predictions, mean=True)
stats_dict = {"eps": epsilon, "acc": acc, "ece1": ece1, "ece2": ece2, "nll": nll, "ent": ent}
if verbose:
print(f"Step: {epsilon:.2f} | Adv. Entropy: {stats_dict['ent']:.2f} | Adv. Accuracy: {stats_dict['acc']:.2f}%")
return mean_predictions, labels, stats_dict
def compute_factors(args: Any,
model: Union[torch.nn.Module, torch.nn.Sequential],
data: iter,
factors=None):
model.train()
criterion = torch.nn.CrossEntropyLoss().to(args.device)
est_base = getattr(fisher, args.estimator.upper())
if args.estimator == 'efb':
est = est_base(model, factors)
else:
est = est_base(model)
for epoch in range(args.epochs):
data = tqdm.tqdm(data,
desc=f"Epoch [{epoch + 1}/{args.epochs}]",
disable=not args.verbose)
for batch, (images, labels) in enumerate(data):
data.set_postfix({'RAM': ram(), 'VRAM': vram()})
logits = model(images.to(args.device, non_blocking=True))
dist = torch.distributions.Categorical(logits=logits)
for sample in range(args.samples):
labels = dist.sample()
loss = criterion(logits, labels)
model.zero_grad()
loss.backward(retain_graph=True)
est.update(images.size(0))
return est
def eval_bnn(model,
dataset,
inv_factors,
estimator='kfac',
samples=30, stats=False,
device=torch.device('cuda'),
verbose=True):
model.eval()
mean_state = copy.deepcopy(model.state_dict())
mean_predictions = 0
stats_list = {"acc": [], "ece": [], "nll": [], "ent": []}
with torch.no_grad():
samples = tqdm.tqdm(range(samples), disable=not verbose)
for sample in samples:
samples.set_postfix({'RAM': ram(), 'VRAM': vram()})
sample_and_replace_weights(model, inv_factors, estimator)
predictions, labels = eval_nn(model, dataset, device)
mean_predictions += predictions
model.load_state_dict(mean_state)
if stats:
running_mean = mean_predictions / (sample + 1)
stats_list["acc"].append(accuracy(running_mean, labels))
stats_list["ece"].append(100 * expected_calibration_error(running_mean, labels)[0])
stats_list["nll"].append(negative_log_likelihood(predictions, labels))
stats_list["ent"].append(predictive_entropy(running_mean, mean=True))
mean_predictions /= len(samples)
if verbose:
print(f"Accuracy: {accuracy(mean_predictions, labels):.2f}% | ECE: {100 * expected_calibration_error(mean_predictions, labels)[0]:.2f}%")
return mean_predictions, labels, stats_list
def test_loader():
args = setup()
print("Loading data")
if args.data == 'cifar10':
data = cifar10(args.torch_data, splits='test')
if args.data == 'mnist':
data = mnist(args.torch_data, splits='test')
elif args.data == 'tiny':
data = imagenet(args.data_dir,
img_size=64,
batch_size=args.batch_size,
splits='test',
tiny=True)
elif args.data == 'imagenet':
img_size = 224
if args.model in ['googlenet', 'inception_v3']:
img_size = 299
data = imagenet(args.data_dir,
img_size,
args.batch_size,
workers=args.workers,
splits='train')
elif args.data == 'gtsrb':
data = gtsrb(args.data_dir,
batch_size=args.batch_size,
workers=args.workers,
splits='train')
for epoch in range(args.epochs):
data = tqdm.tqdm(data, desc=f"Epoch [{epoch + 1}/{args.epochs}]")
for batch, (images, labels) in enumerate(data):
data.set_postfix({'RAM': ram(), 'VRAM': vram()})
fig, ax = plt.subplots(5, 12)
i = 0
j = 0
for img in images:
img = np.moveaxis(img.numpy(), 0, -1)
img = (img - np.min(img)) / np.ptp(img)
ax[i, j].imshow(img)
ax[i, j].axis('off')
j += 1
if j > 0 and j % 12 == 0:
j = 0
i += 1
plt.show()
plt.close()
def eval_nn(model, dataset, device=torch.device('cuda'), verbose=False):
model.eval()
with torch.no_grad():
logits_list = torch.Tensor().to(device)
labels_list = torch.LongTensor()
dataset = tqdm.tqdm(dataset, disable=not verbose or len(dataset) == 1)
for images, labels in dataset:
dataset.set_postfix({'RAM': ram(), 'VRAM': vram()})
logits = model(images.to(device, non_blocking=True))
logits_list = torch.cat([logits_list, logits])
labels_list = torch.cat([labels_list, labels])
predictions = torch.nn.functional.softmax(logits_list, dim=1).cpu().numpy()
labels = labels_list.numpy()
if verbose:
print(f"Accuracy: {accuracy(predictions, labels):.2f}% | ECE: {100 * expected_calibration_error(predictions, labels)[0]:.2f}%")
return predictions, labels
def eval_fgsm(model,
data,
epsilon=0.1,
stats=True,
device=torch.device('cuda'),
verbose=True):
model.eval()
logits_list = torch.Tensor().to(device)
labels_list = torch.LongTensor()
stats_dict = None
data = tqdm.tqdm(data, disable=not verbose or len(data) == 1)
for images, labels in data:
data.set_postfix({'RAM': ram(), 'VRAM': vram()})
adv_images = datasets.fgsm(model, images.to(device, non_blocking=True), labels.to(device, non_blocking=True),
epsilon=epsilon)
with torch.no_grad():
adv_logits = model(adv_images)
logits_list = torch.cat([logits_list, adv_logits])
labels_list = torch.cat([labels_list, labels])
adv_predictions = torch.nn.functional.softmax(logits_list, dim=1).detach().cpu().numpy()
labels = labels_list.numpy()
if stats:
acc = accuracy(adv_predictions, labels)
ece1 = 100 * expected_calibration_error(adv_predictions, labels)[0]
ece2 = 100 * calibration_curve(adv_predictions, labels)[0]
nll = negative_log_likelihood(adv_predictions, labels)
ent = predictive_entropy(adv_predictions, mean=True)
stats_dict = {"eps": epsilon, "acc": acc, "ece1": ece1, "ece2": ece2, "nll": nll, "ent": ent}
if verbose:
print(f"Step: {epsilon:.2f} | Adv. Entropy: {stats_dict['ent']:.2f} | Adv. Accuracy: {stats_dict['acc']:.2f}%")
return adv_predictions, labels, stats_dict
def train(model, train_loader, val_loader, optimizer, criterion, epochs, learning_rate, device):
train_loss = 0
for epoch in range(epochs):
model.train()
train_loader = tqdm.tqdm(train_loader, desc=f"Epoch [{epoch + 1}/{epochs}]")
for batch, (images, labels) in enumerate(train_loader):
train_loader.set_postfix({'Train loss': train_loss / ((batch + 1) + (epoch * len(train_loader))),
'Train acc.': train_acc if batch > 10 else 0,
'RAM': ram(),
'VRAM': vram()})
logits = model(images.to(device, non_blocking=True))
loss = criterion(logits, labels.to(device, non_blocking=True))
optimizer.zero_grad()
loss.backward()
optimizer.step()
if batch % (len(train_loader) // 10) == 0:
train_loss += loss.detach().cpu().numpy()
train_acc = accuracy(logits.detach().cpu().numpy(), labels.numpy())
eval_nn(model, val_loader, device, verbose=True)