def load_from_file(dirname):
model = fc_100_100_10()
model.load_weights(f"{dirname}/weights.h5")
X_train, _, _, _ = mnist()
if exists(f"{dirname}/pca.pkl"):
sklearn_transformer = load_pickle_from_file(f"{dirname}/pca.pkl")
model = filtered_model(model, X_train, sklearn_transformer)
elif exists(f"{dirname}/fastica.pkl"):
sklearn_transformer = load_pickle_from_file(f"{dirname}/fastica.pkl")
model = filtered_model(model, X_train, sklearn_transformer)
elif exists(f"{dirname}/nmf.pkl"):
sklearn_transformer = load_pickle_from_file(f"{dirname}/nmf.pkl")
model = filtered_model(model, X_train, sklearn_transformer)
elif exists(f"{dirname}/kernelpca.pkl"):
sklearn_transformer = load_pickle_from_file(f"{dirname}/kernelpca.pkl")
model = filtered_model(model, X_train, sklearn_transformer)
elif exists(f"{dirname}/truncatedsvd.pkl"):
sklearn_transformer = load_pickle_from_file(
f"{dirname}/truncatedsvd.pkl")
model = filtered_model(model, X_train, sklearn_transformer)
elif exists(f"{dirname}/incrementalpca.pkl"):
sklearn_transformer = load_pickle_from_file(
f"{dirname}/incrementalpca.pkl")
model = filtered_model(model, X_train, sklearn_transformer)
return model
def adversarial_attack(args, model, inv_factors, results_path, fig_path):
print("Loading data")
if args.data == 'cifar10':
test_loader = datasets.cifar10(args.torch_data, splits='test')
elif args.data == 'gtsrb':
test_loader = datasets.gtsrb(args.data_dir, batch_size=args.batch_size, splits='test')
if args.data == 'mnist':
test_loader = datasets.mnist(args.torch_data, splits='test')
elif args.data == 'tiny':
test_loader = datasets.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
test_loader = datasets.imagenet(args.data_dir, img_size, args.batch_size, workers=args.workers, splits='test')
if args.epsilon > 0:
print(eval_fgsm(model, test_loader, args.epsilon, args.device)[-1])
else:
stats_dict = {"eps": [], "acc": [], "ece1": [], "ece2": [], "nll": [], "ent": []}
bnn_stats_dict = {"eps": [], "acc": [], "ece1": [], "ece2": [], "nll": [], "ent": []}
steps = np.concatenate([np.linspace(0, 0.2, 11), np.linspace(0.3, 1, 8)])
for step in steps:
stats = eval_fgsm(model, test_loader, step, args.device, verbose=False)[-1]
bnn_stats = eval_fgsm_bnn(model, test_loader, inv_factors, args.estimator, args.samples, step,
device=args.device)[-1]
for (k1, v1), (k2, v2) in zip(stats.items(), bnn_stats.items()):
stats_dict[k1].append(v1)
bnn_stats_dict[k2].append(v2)
np.savez(results_path + "_fgsm.npz", stats=stats_dict, bnn_stats=bnn_stats_dict)
print(tabulate.tabulate(stats_dict, headers="keys"))
print(tabulate.tabulate(bnn_stats_dict, headers="keys"))
plot.adversarial_results(steps, stats_dict, bnn_stats_dict, fig_path)
def save_models():
X_train, y_train, X_test, y_test = mnist()
prefix = "/tmp"
model = pca_filtered_model(fc_100_100_10(), X_train, 10)
save_to_file(model, prefix)
model = fastica_filtered_model(fc_100_100_10(), X_train, 10)
save_to_file(model, prefix)
model = truncatedsvd_filtered_model(fc_100_100_10(), X_train, 10)
save_to_file(model, prefix)
model = kernelpca_filtered_model(fc_100_100_10(), X_train[:1000], 10)
save_to_file(model, prefix)
model = incrementalpca_filtered_model(fc_100_100_10(), X_train, 10)
save_to_file(model, prefix)
model = nmf_filtered_model(fc_100_100_10(), X_train, 10)
save_to_file(model, prefix)
yield
shutil.rmtree("/tmp/model/pca-filtered-model-10-components")
shutil.rmtree("/tmp/model/fastica-filtered-model-10-components")
shutil.rmtree("/tmp/model/truncatedsvd-filtered-model-10-components")
shutil.rmtree("/tmp/model/kernelpca-filtered-model-10-components")
shutil.rmtree("/tmp/model/incrementalpca-filtered-model-10-components")
shutil.rmtree("/tmp/model/nmf-filtered-model-10-components")
def main(unused_argv):
# Build data and .
print('Loading data.')
x_train, y_train, x_test, y_test = datasets.mnist(permute_train=True)
# Build the network
init_fn, f = stax.serial(
stax.Dense(2048),
stax.Tanh,
stax.Dense(10))
key = random.PRNGKey(0)
_, params = init_fn(key, (-1, 784))
# Linearize the network about its initial parameters.
f_lin = linearize(f, params)
# Create and initialize an optimizer for both f and f_lin.
opt_init, opt_apply, get_params = optimizers.momentum(FLAGS.learning_rate,
0.9)
opt_apply = jit(opt_apply)
state = opt_init(params)
state_lin = opt_init(params)
# Create a cross-entropy loss function.
loss = lambda fx, y_hat: -np.mean(stax.logsoftmax(fx) * y_hat)
# Specialize the loss function to compute gradients for both linearized and
# full networks.
grad_loss = jit(grad(lambda params, x, y: loss(f(params, x), y)))
grad_loss_lin = jit(grad(lambda params, x, y: loss(f_lin(params, x), y)))
# Train the network.
print('Training.')
print('Epoch\tLoss\tLinearized Loss')
print('------------------------------------------')
epoch = 0
steps_per_epoch = 50000 // FLAGS.batch_size
for i, (x, y) in enumerate(datasets.minibatch(
x_train, y_train, FLAGS.batch_size, FLAGS.train_epochs)):
params = get_params(state)
state = opt_apply(i, grad_loss(params, x, y), state)
params_lin = get_params(state_lin)
state_lin = opt_apply(i, grad_loss_lin(params_lin, x, y), state_lin)
if i % steps_per_epoch == 0:
print('{}\t{:.4f}\t{:.4f}'.format(
epoch, loss(f(params, x), y), loss(f_lin(params_lin, x), y)))
epoch += 1
# Print out summary data comparing the linear / nonlinear model.
x, y = x_train[:10000], y_train[:10000]
util.print_summary('train', y, f(params, x), f_lin(params_lin, x), loss)
util.print_summary(
'test', y_test, f(params, x_test), f_lin(params_lin, x_test), loss)
def save_net_results(model_name, dataset_name):
if dataset_name == 'mnist':
train_images, train_labels, test_images, test_labels = datasets.mnist()
elif dataset_name == 'shapes':
train_images, train_labels, test_images, test_labels = datasets.shapes(
'img/shapes')
elif dataset_name == 'alienator':
train_images, train_labels, test_images, test_labels = datasets.alienator(
'img', 'train_keypoints.txt', 'test_keypoints.txt', rotated=False)
elif dataset_name == 'alienator_custom':
train_images, train_labels, test_images, test_labels = datasets.alienator(
'.',
'train_keypoints_custom.txt',
'test_keypoints_custom.txt',
rotated=False,
kp_size_multiplier=30)
elif dataset_name == 'alienator_custom_ns':
train_images, train_labels, test_images, test_labels = datasets.alienator(
'.',
'train_keypoints_custom_ns.txt',
'test_keypoints_custom_ns.txt',
rotated=False,
kp_size_multiplier=30)
elif dataset_name == 'alienator2':
train_images, train_labels, test_images, test_labels = datasets.alienator(
'img',
'train_keypoints2.txt',
'test_keypoints2.txt',
rotated=False)
else:
train_images, train_labels, test_images, test_labels = datasets.brown(
dataset_name)
train_dataset = Dataset(train_images, train_labels, size=64)
test_dataset = Dataset(test_images,
test_labels,
mean=train_dataset.mean,
std=train_dataset.std,
size=64)
network_desc = NetworkDesc(model_file=model_name + '.h5')
# network_desc = NetworkDescPN(model_file=model_name + '.h5')
batch = test_dataset.get_batch_triplets(100000)
positives_net, negatives_net = get_positives_negatives(
get_net_descriptors(network_desc, batch[0]),
get_net_descriptors(network_desc, batch[1]),
get_net_descriptors(network_desc, batch[2]))
results_dir = 'results/{}/'.format(model_name)
if not os.path.isdir(results_dir):
os.makedirs(results_dir)
np.save('{}{}.positives'.format(results_dir, dataset_name), positives_net)
np.save('{}{}.negatives'.format(results_dir, dataset_name), negatives_net)
def main():
# 解析命令行参数
arg_parser = ArgumentParser('Distributed Optimization Simulator',
usage='一个分布式优化模拟器')
arg_parser.add_argument('--direct',
action='store_true',
default=False,
help='如果加上该选项，则创建有向图')
arg_parser.add_argument('--graphfile', default='', help='网络定义文件')
arg_parser.add_argument('--nodenum', default=0, help='节点个数')
args = arg_parser.parse_args()
# 构建网络
g = nx.DiGraph() if args.direct else nx.Graph()
if args.graphfile:
# 从文件导入网络结构
with open(args.graphfile) as f:
pass
else:
# 节点
node_num = args.nodenum
while node_num <= 0:
node_num = int(input('请输入节点数：'))
node_list = [Node('Node{}'.format(i + 1)) for i in range(node_num)]
g.add_nodes_from([str(i) for i in list(range(1, node_num + 1))])
# 边
print('请输入网络的边，空格或换行符分割的数字对，如1,2 3,4')
edge_list = []
x = input('>')
while x != '':
edge_list.extend([t.split(',') for t in x.split()])
x = input('>')
g.add_edges_from(edge_list)
# 日志
print('网络生成成功')
print('节点：', list(g.nodes))
print('边：', list(g.edges))
# 绘图
plt.figure()
nx.draw(g, with_labels=True, pos=nx.spring_layout(g))
plt.show()
# 数据
x, y, _, _, _, _ = mnist(num_train=10000)
# 创建控制器
controller = Controller(g, node_list, {'x': x, 'y': y})
# 分发数据
controller.distribute_data()
def test_parsed_mnist_has_expected_shape():
X_train, y_train, X_test, y_test = mnist()
assert len(X_train) == 60000
assert len(y_train) == 60000
assert X_train.shape == (60000, 28, 28)
assert y_train.shape == (60000, )
assert len(X_test) == 10000
assert len(y_test) == 10000
assert X_test.shape == (10000, 28, 28)
assert y_test.shape == (10000, )
def main():
rng = random.PRNGKey(0)
batch_size = 128
step_size = 0.001
num_epochs = 10
momentum_mass = 0.9
train_images, train_labels, test_images, test_labels = datasets.mnist()
num_train = train_images.shape[0]
num_complete_batches, leftover = divmod(num_train, batch_size)
num_batches = num_complete_batches + bool(leftover)
# define data stream
def data_stream():
rng = npr.RandomState(0)
while True:
perm = rng.permutation(num_train)
for i in range(num_batches):
batch_indices = perm[i*batch_size:(i+1)*batch_size]
yield train_images[batch_size], train_labels[batch_indices]
batches = data_stream()
# define optimizer
opt_init, opt_update, get_params = optimizers.momentum(step_size, mass=momentum_mass)
@jit
def update(i, opt_state, batch):
params = get_params(opt_state)
return opt_update(i, grad(loss)(params, batch), opt_state)
_, init_params = init_random_params(rng, (-1, 28*28))
opt_state = opt_init(init_params)
itercount = itertools.count()
print('\nStarting training...')
for epoch in range(num_epochs):
start_tm = time.time()
for _ in range(num_epochs):
opt_state = update(next(itercount), opt_state, next(batches))
epoch_tm = time.time() - start_tm
params = get_params(opt_state)
train_acc = accuracy(params, (train_images, train_labels))
test_acc = accuracy(params, (test_images, test_labels))
print(f'Epoch {epoch} in {epoch_tm:0.2f} sec')
print(f'Training set accuracy {train_acc}')
print(f'Test set accuracy {test_acc}')
print('DONE')
def main(unused_argv):
# Build data pipelines.
print('Loading data.')
x_train, y_train, x_test, y_test = \
datasets.mnist(FLAGS.train_size, FLAGS.test_size)
# Build the network
init_fn, f, _ = stax.serial(
stax.Dense(4096, 1., 0.),
stax.Erf(),
stax.Dense(10, 1., 0.))
key = random.PRNGKey(0)
_, params = init_fn(key, (-1, 784))
# Create and initialize an optimizer.
opt_init, opt_apply, get_params = optimizers.sgd(FLAGS.learning_rate)
state = opt_init(params)
# Create an mse loss function and a gradient function.
loss = lambda fx, y_hat: 0.5 * np.mean((fx - y_hat) ** 2)
grad_loss = jit(grad(lambda params, x, y: loss(f(params, x), y)))
# Create an MSE predictor to solve the NTK equation in function space.
ntk = batch(get_ntk_fun_empirical(f), batch_size=32, device_count=1)
g_dd = ntk(x_train, None, params)
g_td = ntk(x_test, x_train, params)
predictor = predict.analytic_mse(g_dd, y_train, g_td)
# Get initial values of the network in function space.
fx_train = f(params, x_train)
fx_test = f(params, x_test)
# Train the network.
train_steps = int(FLAGS.train_time // FLAGS.learning_rate)
print('Training for {} steps'.format(train_steps))
for i in range(train_steps):
params = get_params(state)
state = opt_apply(i, grad_loss(params, x_train, y_train), state)
# Get predictions from analytic computation.
print('Computing analytic prediction.')
fx_train, fx_test = predictor(FLAGS.train_time, fx_train, fx_test)
# Print out summary data comparing the linear / nonlinear model.
util.print_summary('train', y_train, f(params, x_train), fx_train, loss)
util.print_summary('test', y_test, f(params, x_test), fx_test, loss)
def load_datasets(dataset_name):
if dataset_name == 'mnist':
train_images, train_labels, test_images, test_labels = datasets.mnist()
elif dataset_name == 'shapes':
train_images, train_labels, test_images, test_labels = datasets.shapes(
'img/shapes')
elif dataset_name == 'alienator':
train_images, train_labels, test_images, test_labels = datasets.alienator(
'img', 'train_keypoints.txt', 'test_keypoints.txt', rotated=False)
elif dataset_name == 'alienator_custom':
train_images, train_labels, test_images, test_labels = datasets.alienator(
'.',
'train_keypoints_custom.txt',
'test_keypoints_custom.txt',
rotated=False,
kp_size_multiplier=30)
elif dataset_name == 'alienator_custom_ns':
train_images, train_labels, test_images, test_labels = datasets.alienator(
'.',
'train_keypoints_custom_ns.txt',
'test_keypoints_custom_ns.txt',
rotated=False,
kp_size_multiplier=30)
elif dataset_name == 'alienator2':
train_images, train_labels, test_images, test_labels = datasets.alienator(
'img',
'train_keypoints2.txt',
'test_keypoints2.txt',
rotated=False)
else:
train_images, train_labels, test_images, test_labels = datasets.brown(
dataset_name)
train = Dataset(train_images, train_labels, size=64)
test = Dataset(test_images,
test_labels,
mean=train.mean,
std=train.std,
size=64)
return train, test
def test(args, model, fig_path=""):
print("Loading data")
if args.data == 'cifar10':
test_loader = datasets.cifar10(args.torch_data, splits='test')
elif args.data == 'gtsrb':
test_loader = datasets.gtsrb(args.data_dir, batch_size=args.batch_size, splits='test')
if args.data == 'mnist':
test_loader = datasets.mnist(args.torch_data, splits='test')
elif args.data == 'tiny':
test_loader = datasets.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
test_loader = datasets.imagenet(args.data_dir, img_size, args.batch_size, workers=args.workers, splits='test')
predictions, labels = eval_nn(model, test_loader, args.device, args.verbose)
print("Plotting results")
plot.reliability_diagram(predictions, labels, path=fig_path + "_reliability.pdf")
def save_sift_results(dataset_name):
if dataset_name == 'mnist':
train_images, train_labels, test_images, test_labels = datasets.mnist()
elif dataset_name == 'shapes':
train_images, train_labels, test_images, test_labels = datasets.shapes(
'img/shapes')
elif dataset_name == 'alienator':
train_images, train_labels, test_images, test_labels = datasets.alienator(
'img', 'train_keypoints.txt', 'test_keypoints.txt', rotated=False)
elif dataset_name == 'alienator2':
train_images, train_labels, test_images, test_labels = datasets.alienator(
'img',
'train_keypoints2.txt',
'test_keypoints2.txt',
rotated=False)
else:
train_images, train_labels, test_images, test_labels = datasets.brown(
dataset_name)
train_dataset = Dataset(train_images, train_labels)
test_dataset = Dataset(test_images,
test_labels,
mean=train_dataset.mean,
std=train_dataset.std)
batch = test_dataset.get_batch_triplets(100000)
positives_sift, negatives_sift = get_positives_negatives(
get_sift_descriptors(batch[0]), get_sift_descriptors(batch[1]),
get_sift_descriptors(batch[2]))
results_dir = 'results/sift/'
if not os.path.isdir(results_dir):
os.makedirs(results_dir)
np.save('{}{}.positives'.format(results_dir, dataset_name), positives_sift)
np.save('{}{}.negatives'.format(results_dir, dataset_name), negatives_sift)
predicted_class = jnp.argmax(predict(params, inputs), axis=1)
return jnp.mean(predicted_class == target_class)
init_random_params, predict = stax.serial(Dense(1024), Relu, Dense(1024), Relu,
Dense(10), LogSoftmax)
if __name__ == "__main__":
rng = random.PRNGKey(0)
step_size = 0.001
num_epochs = 10
batch_size = 128
momentum_mass = 0.9
train_images, train_labels, test_images, test_labels = datasets.mnist()
num_train = train_images.shape[0]
num_complete_batches, leftover = divmod(num_train, batch_size)
num_batches = num_complete_batches + bool(leftover)
def data_stream():
rng = npr.RandomState(0)
while True:
perm = rng.permutation(num_train)
for i in range(num_batches):
batch_idx = perm[i * batch_size:(i + 1) * batch_size]
yield train_images[batch_idx], train_labels[batch_idx]
batches = data_stream()
opt_init, opt_update, get_params = adahessian()
import pytest
import random
import numpy.matlib
from utils import *
random.seed(1)
np.random.seed(1)
train_samples = 1500
val_samples = 500
test_samples = 1000
#loading train and test data
digits = list(range(10))
trX, trY, tsX, tsY, valX, valY = mnist(train_samples,
val_samples,
test_samples,
digits=digits)
#visualizing few samples
plt.imshow(trX[:, 0].reshape(28, 28))
plt.imshow(tsX[:, 100].reshape(28, 28))
plt.imshow(valX[:, 0].reshape(28, 28))
#Testing parameter initialization
net_dims_tst = [5, 4, 1]
parameters_tst = initialize(net_dims_tst)
assert parameters_tst['W1'].shape == (4, 5)
assert parameters_tst['b1'].shape == (4, 1)
#Testing relu activation function
z_tst = np.array([[-1, 2], [3, -6]])
def out_of_domain(args, model, inv_factors, results_path="", fig_path=""):
"""Evaluates the model on in- and out-of-domain data.
Each dataset has its own out-of-domain dataset which is loaded automatically alongside the in-domain dataset
specified in `args.data`. For each image (batch) in the in- and out-of-domain data a forward pass through the
provided `model` is performed and the predictions are stored under `results_path`. This is repeated for the Bayesian
variant of the model (Laplace approximation).
Parameters
----------
args : Todo: Check type
The arguments provided to the script on execution.
model : torch.nn.Module Todo: Verify
A `torchvision` or custom neural network (a `torch.nn.Module` or `torch.nn.Sequential` instance)
inv_factors : list
A list KFAC factors, Eigenvectors of KFAC factors or diagonal terms. Todo: INF
results_path : string, optional
The path where results (in- and out-of-domain predictions) should be stored. Results are not stored if
argument `args.no_results` is provided.
fig_path : string, optional
The path where figures should be stored. Figures are only generated if argument `args.plot` is provided.
"""
print("Loading data")
if args.data == 'cifar10':
in_data = datasets.cifar10(args.torch_data, splits='test')
out_data = datasets.svhn(args.torch_data, splits='test')
elif args.data == 'mnist':
in_data = datasets.mnist(args.torch_data, splits='test')
out_data = datasets.kmnist(args.torch_data, splits='test')
elif args.data == 'gtsrb':
in_data = datasets.gtsrb(args.data_dir, batch_size=args.batch_size, splits='test')
out_data = datasets.cifar10(args.torch_data, splits='test')
elif args.data == 'tiny':
in_data = datasets.imagenet(args.data_dir, img_size=64, batch_size=args.batch_size, splits='test', tiny=True,
use_cache=True)
out_data = datasets.art(args.data_dir, img_size=64, batch_size=args.batch_size, use_cache=True)
elif args.data == 'imagenet':
img_size = 224
if args.model in ['googlenet', 'inception_v3']:
img_size = 299
in_data = datasets.imagenet(args.data_dir, img_size, args.batch_size, workers=args.workers, splits='test',
use_cache=True)
out_data = datasets.art(args.data_dir, img_size, args.batch_size, workers=args.workers, use_cache=True)
# Compute NN and BNN predictions on validation set of training data
predictions, bnn_predictions, labels, stats = eval_nn_and_bnn(model, in_data, inv_factors, args.estimator,
args.samples, args.stats, args.device, verbose=True)
# Compute NN and BNN predictions on out-of-distribution data
ood_predictions, bnn_ood_predictions, _, _ = eval_nn_and_bnn(model, out_data, inv_factors, args.estimator,
args.samples, False, args.device, verbose=True)
if not args.no_results:
print("Saving results")
np.savez_compressed(results_path,
stats=stats,
labels=labels,
predictions=predictions,
bnn_predictions=bnn_predictions,
ood_predictions=ood_predictions,
bnn_ood_predictions=bnn_ood_predictions)
if args.plot:
print("Plotting results")
fig, ax = plt.subplots(figsize=(12, 7), tight_layout=True)
plot.inv_ecdf_vs_pred_entropy(predictions, color='dodgerblue', linestyle='--', axis=ax)
plot.inv_ecdf_vs_pred_entropy(ood_predictions, color='crimson', linestyle='--', axis=ax)
plot.inv_ecdf_vs_pred_entropy(bnn_predictions, color='dodgerblue', axis=ax)
plot.inv_ecdf_vs_pred_entropy(bnn_ood_predictions, color='crimson', axis=ax)
ax.legend([f"NN {args.data.upper()} | Acc.: {accuracy(predictions, labels):.2f}%",
f"NN OOD",
f"BNN {args.data.upper()} | Acc.: {accuracy(bnn_predictions, labels):.2f}%",
f"BNN OOD"], fontsize=16, frameon=False)
plt.savefig(fig_path + "_ecdf.pdf", forma='pdf', dpi=1200)
plot.reliability_diagram(predictions, labels, path=fig_path + "_reliability.pdf")
plot.reliability_diagram(bnn_predictions, labels, path=fig_path + "_bnn_reliability.pdf")
plot.entropy_hist(predictions, ood_predictions, path=fig_path + "_entropy.pdf")
plot.entropy_hist(bnn_predictions, bnn_ood_predictions, path=fig_path + "_bnn_entropy.pdf")
def getAccuracy(out, labels):
tot = out.size()[0]
return torch.sum(torch.eq(torch.argmax(out, dim=1).int(),
labels.int())).float() / tot
def criterion(out, labels):
t = nn.CrossEntropyLoss()
return t(out, labels.long())
Demo = True
if Demo: train_size = 0
else: train_size = 60000
test_size = 10000
mnist_data = mnist(train_size, test_size)
learning_rate = 0.01
generation = 1000
batch_size = 100
d = (learning_rate - 0.0001) / 100.0
net = Net()
print(net)
out = None
in_labels = None
rand_x = None
if not Demo:
for i in tqdm.tqdm(range(generation)):
Dense(512), Relu,
Dense(28 * 28),
)
if __name__ == "__main__":
step_size = 0.001
num_epochs = 100
batch_size = 32
nrow, ncol = 10, 10 # sampled image grid size
rng = random.PRNGKey(0)
test_rng = random.PRNGKey(1) # fixed prng key for evaluation
imfile = os.path.join(os.getenv("TMPDIR", "/tmp/"), "mnist_vae_{:03d}.png")
train_images, _, test_images, _ = datasets.mnist(permute_train=True)
num_complete_batches, leftover = divmod(train_images.shape[0], batch_size)
num_batches = num_complete_batches + bool(leftover)
_, init_encoder_params = encoder_init((batch_size, 28 * 28))
_, init_decoder_params = decoder_init((batch_size, 10))
init_params = init_encoder_params, init_decoder_params
opt_init, opt_update = minmax.momentum(step_size, mass=0.9)
def binarize_batch(rng, i, images):
i = i % num_batches
batch = lax.dynamic_slice_in_dim(images, i * batch_size, batch_size)
return random.bernoulli(rng, batch)
@jit
def main(_):
rng = random.PRNGKey(0)
# Load MNIST dataset
train_images, train_labels, test_images, test_labels = datasets.mnist()
batch_size = 128
batch_shape = (-1, 28, 28, 1)
num_train = train_images.shape[0]
num_complete_batches, leftover = divmod(num_train, batch_size)
num_batches = num_complete_batches + bool(leftover)
train_images = np.reshape(train_images, batch_shape)
test_images = np.reshape(test_images, batch_shape)
def data_stream():
rng = npr.RandomState(0)
while True:
perm = rng.permutation(num_train)
for i in range(num_batches):
batch_idx = perm[i * batch_size:(i + 1) * batch_size]
yield train_images[batch_idx], train_labels[batch_idx]
def save(fn, opt_state):
params = deepcopy(get_params(opt_state))
save_dict = {}
for idx, p in enumerate(params):
if (p != ()):
pp = (p[0].tolist(), p[1].tolist())
params[idx] = pp
save_dict["params"] = params
with open(fn, "w") as f:
json.dump(save_dict, f)
def load(fn):
with open(fn, "r") as f:
params = json.load(f)
params = params["params"]
for idx, p in enumerate(params):
if (p != []):
pp = (np.array(p[0]), np.array(p[1]))
params[idx] = pp
else:
params[idx] = ()
return opt_init(params)
batches = data_stream()
# Model, loss, and accuracy functions
init_random_params, predict = stax.serial(
stax.Conv(32, (8, 8), strides=(2, 2), padding="SAME"),
stax.Relu,
stax.Conv(128, (6, 6), strides=(2, 2), padding="VALID"),
stax.Relu,
stax.Conv(128, (5, 5), strides=(1, 1), padding="VALID"),
stax.Flatten,
stax.Dense(128),
stax.Relu,
stax.Dense(10),
)
def loss(params, batch):
inputs, targets = batch
preds = predict(params, inputs)
return -np.mean(logsoftmax(preds) * targets)
def accuracy(params, batch):
inputs, targets = batch
target_class = np.argmax(targets, axis=1)
predicted_class = np.argmax(predict(params, inputs), axis=1)
return np.mean(predicted_class == target_class)
def gen_ellipsoid(X, zeta_rel, zeta_const, alpha, N_steps):
zeta = (np.abs(X).T * zeta_rel).T + zeta_const
if (alpha is None):
alpha = 1 / N_steps * zeta
else:
assert isinstance(alpha, float), "Alpha must be float"
alpha = alpha * np.ones_like(X)
return zeta, alpha
def gen_ellipsoid_match_volume(X, zeta_const, eps, alpha, N_steps):
x_norms = np.linalg.norm(np.reshape(X, (X.shape[0], -1)),
ord=1,
axis=1)
N = np.prod(X.shape[1:])
zeta_rel = N * (eps - zeta_const) / x_norms
assert (zeta_rel <= 1.0).all(
), "Zeta rel cannot be larger than 1. Please increase zeta const or reduce eps"
zeta_rel = np.clip(0.0, zeta_rel, 1.0)
return gen_ellipsoid(X, zeta_rel, zeta_const, alpha, N_steps)
# Instantiate an optimizer
opt_init, opt_update, get_params = optimizers.adam(0.001)
@jit
def update(i, opt_state, batch):
params = get_params(opt_state)
return opt_update(i, grad(loss)(params, batch), opt_state)
# Initialize model
_, init_params = init_random_params(rng, batch_shape)
opt_state = opt_init(init_params)
itercount = itertools.count()
try:
opt_state = load("tutorials/jax/test_model.json")
except:
# Training loop
print("\nStarting training...")
for _ in range(num_batches):
opt_state = update(next(itercount), opt_state, next(batches))
epoch_time = time.time() - start_time
save("tutorials/jax/test_model.json", opt_state)
# Evaluate model on clean data
params = get_params(opt_state)
# Evaluate model on adversarial data
model_fn = lambda images: predict(params, images)
# Generate single attacking test image
idx = 0
plt.figure(figsize=(15, 6), constrained_layout=True)
zeta, alpha = gen_ellipsoid(X=test_images[idx].reshape((1, 28, 28, 1)),
zeta_rel=FLAGS.zeta_rel,
zeta_const=FLAGS.zeta_const,
alpha=None,
N_steps=40)
# zeta, alpha = gen_ellipsoid_match_volume(X=test_images[idx].reshape((1,28,28,1)), zeta_const=FLAGS.zeta_const, eps=FLAGS.eps, alpha=None, N_steps=40)
test_images_pgd_ellipsoid = projected_gradient_descent(
model_fn, test_images[idx].reshape((1, 28, 28, 1)), zeta, alpha, 40,
np.inf)
predict_pgd_ellipsoid = np.argmax(predict(params,
test_images_pgd_ellipsoid),
axis=1)
test_images_fgm = fast_gradient_method(
model_fn, test_images[idx].reshape((1, 28, 28, 1)), 0.075, np.inf)
predict_fgm = np.argmax(predict(params, test_images_fgm), axis=1)
test_images_pgd = projected_gradient_descent(
model_fn, test_images[idx].reshape((1, 28, 28, 1)), FLAGS.eps, 0.01,
40, 2)
predict_pgd = np.argmax(predict(params, test_images_pgd), axis=1)
base = 100
f_ = lambda x: np.log(x) / np.log(base)
a = base - 1
transform = 1 + a * test_images[idx].reshape((1, 28, 28, 1)) # [1,base]
# test_images_pgd_transform = projected_gradient_descent(model_fn, f_(np.where(transform > base,base,transform)), FLAGS.zeta_rel, 0.01, 40, np.inf)
test_images_pgd_transform = projected_gradient_descent(
model_fn, f_(np.where(transform > base, base, transform)), 1.8, 0.01,
40, 2)
test_images_pgd_transform = np.clip(test_images_pgd_transform, 0.0, 1.0)
test_images_pgd_transform = (base**test_images_pgd_transform - 1) / a
predict_transform = np.argmax(predict(params, test_images_pgd_transform),
axis=1)
plt.subplot(151)
plt.imshow(np.squeeze(test_images[idx]), cmap='gray')
plt.title("Original")
plt.subplot(152)
plt.imshow(np.squeeze(test_images_fgm), cmap='gray')
plt.title(f"FGM L-Inf Pred: {predict_fgm}")
plt.subplot(153)
plt.imshow(np.squeeze(test_images_pgd), cmap='gray')
plt.title(f"PGD L2 {predict_pgd}")
plt.subplot(154)
plt.imshow(np.squeeze(test_images_pgd_ellipsoid), cmap='gray')
plt.title(f"PGD Ellipsoid L-Inf Pred: {predict_pgd_ellipsoid}")
plt.subplot(155)
plt.imshow(np.squeeze(test_images_pgd_transform), cmap='gray')
plt.title(f"PGD log{base} L2 Pred: {predict_transform}")
plt.show()
transform = 1 + a * test_images
test_images_pgd_transform = projected_gradient_descent(
model_fn, f_(np.where(transform > base, base, transform)),
FLAGS.zeta_rel, 0.01, 40, np.inf)
test_images_pgd_transform = np.clip(test_images_pgd_transform, 0.0, 1.0)
test_images_pgd_transform = (base**test_images_pgd_transform - 1) / a
test_acc_pgd_transform = accuracy(params,
(test_images_pgd_transform, test_labels))
# Generate whole attacking test images
# zeta, alpha = gen_ellipsoid(X=test_images, zeta_rel=FLAGS.zeta_rel, zeta_const=FLAGS.zeta_const, alpha=None, N_steps=40)
zeta, alpha = gen_ellipsoid_match_volume(X=test_images,
zeta_const=FLAGS.zeta_const,
eps=FLAGS.eps,
alpha=None,
N_steps=40)
test_images_pgd_ellipsoid = projected_gradient_descent(
model_fn, test_images, zeta, alpha, 40, np.inf)
test_acc_pgd_ellipsoid = accuracy(params,
(test_images_pgd_ellipsoid, test_labels))
test_images_fgm = fast_gradient_method(model_fn, test_images, FLAGS.eps,
np.inf)
test_images_pgd = projected_gradient_descent(model_fn, test_images,
FLAGS.eps, 0.01, 40, np.inf)
test_acc_fgm = accuracy(params, (test_images_fgm, test_labels))
test_acc_pgd = accuracy(params, (test_images_pgd, test_labels))
train_acc = accuracy(params, (train_images, train_labels))
test_acc = accuracy(params, (test_images, test_labels))
print("Training set accuracy: {}".format(train_acc))
print("Test set accuracy on clean examples: {}".format(test_acc))
print("Test set accuracy on FGM adversarial examples: {}".format(
test_acc_fgm))
print("Test set accuracy on PGD adversarial examples: {}".format(
test_acc_pgd))
print("Test set accuracy on PGD Ellipsoid adversarial examples: {}".format(
test_acc_pgd_ellipsoid))
print(
"Test set accuracy on PGD Ellipsoid via transform adversarial examples: {}"
.format(test_acc_pgd_transform))
# Save config file in experiment directory
with open(directory + '/config.json', 'w') as config_file:
json.dump(config, config_file)
img_size = config["img_size"]
batch_size = config["batch_size"]
r_dim = config["r_dim"]
h_dim = config["h_dim"]
z_dim = config["z_dim"]
num_context_range = config["num_context_range"]
num_extra_target_range = config["num_extra_target_range"]
epochs = config["epochs"]
if config["dataset"] == "mnist":
data_loader, _ = mnist(batch_size=batch_size, size=img_size[1])
elif config["dataset"] == "celeba":
data_loader = celeba(batch_size=batch_size, size=img_size[1])
elif config["dataset"] == "calligraphy":
data_loader = calligraphy(batch_size=batch_size, size=img_size[1])
np_img = NeuralProcessImg(img_size, r_dim, z_dim, h_dim).to(device)
optimizer = torch.optim.Adam(np_img.parameters(), lr=config["lr"])
np_trainer = NeuralProcessTrainer(device,
np_img,
optimizer,
num_context_range,
num_extra_target_range,
print_freq=100)
import torch.optim as optim
import torch.nn as nn
import torch.nn.functional as F
from torch.autograd import Variable
from datasets import mnist, loader
batch_size = 64
ts_batch_size = 1000
epochs = 2
lr = 0.01
momentum = 0.5
seed = 1
log_freq = 100
tn_loader, ts_loader = loader(mnist(), batch_size, ts_batch_size)
class Net(nn.Module):
def __init__(self):
super(Net, self).__init__()
self.conv1 = nn.Conv2d(1, 10, kernel_size=5)
self.conv2 = nn.Conv2d(10, 20, kernel_size=5)
self.conv2_drop = nn.Dropout2d()
self.fc1 = nn.Linear(320, 50)
self.fc2 = nn.Linear(50, 10)
def forward(self, x):
x = F.relu(F.max_pool2d(self.conv1(x), 2))
x = F.relu(F.max_pool2d(self.conv2_drop(self.conv2(x)), 2))
x = x.view(-1, 320)
x = F.relu(self.fc1(x))
# TODO: fix these magic numbers (especially the 800)
def f(X):
layer0 = X.reshape((X.shape[0], 1, 28, 28))
layer1 = _build_conv_pool(P, 1, layer0, 20, 1, 5, 2)
layer2_= _build_conv_pool(P, 2, layer1, 50, 20, 5, 2)
layer2 = layer2_.flatten(2)
output = T.nnet.softmax(T.dot(layer2, P.W_hidden_output) + P.b_output)
return output
return f
def cost(P, Y_hat, Y, l2 = 0):
return (T.mean(T.nnet.categorical_crossentropy(Y_hat, Y)) +
l2 * sum(T.mean(p**2) for p in P.values()))
if __name__ == "__main__":
import datasets
x,y = datasets.mnist()
x,y = x[0:1000],y[0:1000]
P = Parameters()
X = T.matrix('X')
Y = T.ivector('Y')
net = build(P, 784, 800, 10)
Y_hat = net(X)
f = theano.function(inputs = [X], outputs = Y_hat)
J = cost(P, Y_hat, Y)
grad = T.grad(J, wrt=P.values())
def test_parsed_mnist_is_not_writeable():
X_train, y_train, X_test, y_test = mnist()
with pytest.raises(ValueError):
X_train[:100] = X_test[:100]
def main():
args = setup()
print("Preparing directories")
filename = f"{args.prefix}{args.model}_{args.data}_{args.estimator}{args.suffix}"
factors_path = os.path.join(args.root_dir, "factors", filename)
weights_path = os.path.join(args.root_dir, "weights",
f"{args.model}_{args.data}.pth")
if args.exp_id == -1:
if not args.no_results:
os.makedirs(os.path.join(args.results_dir, args.model, "data",
args.estimator, args.optimizer),
exist_ok=True)
if args.plot:
os.makedirs(os.path.join(args.results_dir, args.model, "figures",
args.estimator, args.optimizer),
exist_ok=True)
results_path = os.path.join(args.results_dir, args.model, "data",
args.estimator, args.optimizer, filename)
else:
if not args.no_results:
os.makedirs(os.path.join(args.results_dir, args.model, "data",
args.estimator, args.optimizer,
args.exp_id),
exist_ok=True)
if args.plot:
os.makedirs(os.path.join(args.results_dir, args.model, "figures",
args.estimator, args.optimizer,
args.exp_id),
exist_ok=True)
results_path = os.path.join(args.results_dir, args.model, "data",
args.estimator, args.optimizer,
args.exp_id, filename)
print("Loading model")
if args.model == 'lenet5':
model = lenet5(pretrained=args.data, device=args.device)
elif args.model == 'resnet18' and args.data != 'imagenet':
model = resnet18(pretrained=weights_path,
num_classes=43 if args.data == 'gtsrb' else 10,
device=args.device)
else:
model_class = getattr(torchvision.models, args.model)
if args.model in ['googlenet', 'inception_v3']:
model = model_class(pretrained=True, aux_logits=False)
else:
model = model_class(pretrained=True)
model.to(args.device).eval()
if args.parallel:
model = torch.nn.parallel.DataParallel(model)
print("Loading data")
if args.data == 'mnist':
val_loader = datasets.mnist(args.torch_data, splits='val')
elif args.data == 'cifar10':
val_loader = datasets.cifar10(args.torch_data, splits='val')
elif args.data == 'gtsrb':
val_loader = datasets.gtsrb(args.data_dir,
batch_size=args.batch_size,
splits='val')
elif args.data == 'imagenet':
img_size = 224
if args.model in ['googlenet', 'inception_v3']:
img_size = 299
val_loader = datasets.imagenet(args.data_dir,
img_size,
args.batch_size,
args.workers,
splits='val',
use_cache=True,
pre_cache=True)
print("Loading factors")
if args.estimator in ["diag", "kfac"]:
factors = torch.load(factors_path + '.pth')
elif args.estimator == 'efb':
kfac_factors = torch.load(factors_path.replace("efb", "kfac") + '.pth')
lambdas = torch.load(factors_path + '.pth')
factors = list()
eigvecs = get_eigenvectors(kfac_factors)
for eigvec, lambda_ in zip(eigvecs, lambdas):
factors.append((eigvec[0], eigvec[1], lambda_))
elif args.estimator == 'inf':
factors = torch.load(f"{factors_path}{args.rank}.pth")
torch.backends.cudnn.benchmark = True
norm_min = -10
norm_max = 10
scale_min = -10
scale_max = 10
if args.boundaries:
x0 = list()
boundaries = [[norm_min, scale_min], [norm_max, scale_max],
[norm_min, scale_max], [norm_max, scale_min],
[norm_min / 2., scale_min], [norm_max / 2., scale_max],
[norm_min, scale_max / 2.], [norm_max, scale_min / 2.],
[norm_min / 2., scale_min / 2.],
[norm_max / 2., scale_max / 2.],
[norm_min / 2., scale_max / 2.],
[norm_max / 2., scale_min / 2.]]
for b in boundaries:
tmp = list()
for _ in range(3 if args.layer else 1):
tmp.extend(b)
x0.append(tmp)
else:
x0 = None
f_norms = np.array([factor.norm().cpu().numpy() for factor in factors])
space = list()
for i in range(3 if args.layer else 1):
space.append(
skopt.space.Real(norm_min,
norm_max,
name=f"norm{i}",
prior='uniform'))
space.append(
skopt.space.Real(scale_min,
scale_max,
name=f"scale{i}",
prior='uniform'))
stats = {
"norms": [],
"scales": [],
"acc": [],
"ece": [],
"nll": [],
"ent": [],
"cost": []
}
@skopt.utils.use_named_args(dimensions=space)
def objective(**params):
norms = list()
scales = list()
for f in f_norms:
if args.layer:
# Closest to max
if abs(f_norms.max() - f) < abs(f_norms.min() - f) and abs(
f_norms.max() - f) < abs(f_norms.mean() - f):
norms.append(10**params['norm0'])
scales.append(10**params['scale0'])
# Closest to min
elif abs(f_norms.min() - f) < abs(f_norms.max() - f) and abs(
f_norms.min() - f) < abs(f_norms.mean() - f):
norms.append(10**params['norm1'])
scales.append(10**params['scale1'])
# Closest to mean
else:
norms.append(10**params['norm2'])
scales.append(10**params['scale2'])
else:
norms.append(10**params['norm0'])
scales.append(10**params['scale0'])
if args.layer:
print(
tabulate.tabulate(
{
'Layer': np.arange(len(factors)),
'F-Norm:': f_norms,
'Norms': norms,
'Scales': scales
},
headers='keys',
numalign='right'))
else:
print("Norm:", norms[0], "Scale:", scales[0])
try:
inv_factors = invert_factors(factors, norms,
args.pre_scale * scales,
args.estimator)
except (RuntimeError, np.linalg.LinAlgError):
print(f"Error: Singular matrix")
return 200
predictions, labels, _ = eval_bnn(model,
val_loader,
inv_factors,
args.estimator,
args.samples,
stats=False,
device=args.device,
verbose=False)
err = 100 - accuracy(predictions, labels)
ece = 100 * expected_calibration_error(predictions, labels)[0]
nll = negative_log_likelihood(predictions, labels)
ent = predictive_entropy(predictions, mean=True)
stats["norms"].append(norms)
stats["scales"].append(scales)
stats["acc"].append(100 - err)
stats["ece"].append(ece)
stats["nll"].append(nll)
stats["ent"].append(ent)
stats["cost"].append(err + ece)
print(
f"Err.: {err:.2f}% | ECE: {ece:.2f}% | NLL: {nll:.3f} | Ent.: {ent:.3f}"
)
return err + ece
with warnings.catch_warnings():
warnings.filterwarnings("ignore", category=FutureWarning)
if args.optimizer == "gbrt":
res = skopt.gbrt_minimize(func=objective,
dimensions=space,
n_calls=args.calls,
x0=x0,
verbose=True,
n_jobs=args.workers,
n_random_starts=0 if x0 else 10,
acq_func='EI')
# EI (neg. expected improvement)
# LCB (lower confidence bound)
# PI (neg. prob. of improvement): Usually favours exploitation over exploration
# gp_hedge (choose probabilistically between all)
if args.optimizer == "gp":
res = skopt.gp_minimize(func=objective,
dimensions=space,
n_calls=args.calls,
x0=x0,
verbose=True,
n_jobs=args.workers,
n_random_starts=0 if x0 else 1,
acq_func='gp_hedge')
# acq_func: EI (neg. expected improvement), LCB (lower confidence bound), PI (neg. prob. of improvement)
# xi: how much improvement one wants over the previous best values.
# kappa: Importance of variance of predicted values. High: exploration > exploitation
# base_estimator: RF (random forest), ET (extra trees)
elif args.optimizer == "forest":
res = skopt.forest_minimize(func=objective,
dimensions=space,
n_calls=args.calls,
x0=x0,
verbose=True,
n_jobs=args.workers,
n_random_starts=0 if x0 else 1,
acq_func='EI')
elif args.optimizer == "random":
res = skopt.dummy_minimize(func=objective,
dimensions=space,
n_calls=args.calls,
x0=x0,
verbose=True)
elif args.optimizer == "grid":
space = [
np.arange(norm_min, norm_max + 1, 10),
np.arange(scale_min, scale_max + 1, 10)
]
res = grid(func=objective, dimensions=space)
print(f"Minimal cost of {min(stats['cost'])} found at:")
if args.layer:
print(
tabulate.tabulate(
{
'Layer': np.arange(len(factors)),
'F-Norm:': f_norms,
'Norms': stats['norms'][np.argmin(stats['cost'])],
'Scales': stats['scales'][np.argmin(stats['cost'])]
},
headers='keys',
numalign='right'))
else:
print("Norm:", stats['norms'][np.argmin(stats['cost'])][0],
"Scale:", stats['scales'][np.argmin(stats['cost'])][0])
if not args.no_results:
print("Saving results")
del res.specs['args']['func']
np.save(
results_path +
f"_best_params{'_layer.npy' if args.layer else '.npy'}", [
stats['norms'][np.argmin(stats['cost'])],
stats['scales'][np.argmin(stats['cost'])]
])
np.save(
results_path +
f"_hyperopt_stats{'_layer.npy' if args.layer else '.npy'}", stats)
skopt.dump(
res, results_path +
f"_hyperopt_dump{'_layer.pkl' if args.layer else '.pkl'}")
if args.plot:
print("Plotting results")
hyperparameters(args)
type=int,
nargs="+",
default=[],
help="number of components for image filters")
argument_parser.add_argument("--epochs",
type=int,
default=-1,
help="default: let the model choose")
argument_parser.add_argument("--random-seed",
action="store_true",
help="initialize model with random seed")
args = argument_parser.parse_args()
PREFIX = os.environ.get('PREFIX', '.')
X_train, y_train, X_test, y_test = mnist()
if not args.random_seed:
K.clear_session()
tf.set_random_seed(1234)
np.random.seed(1234)
no_defense_model = fc_100_100_10()
print(f"Training {no_defense_model.name}...")
train(no_defense_model,
X_train,
y_train,
args.epochs,
verbose=True,
stop_on_stable_weights=True,
reduce_lr_on_plateau=True,
argument_parser = argparse.ArgumentParser()
argument_parser.add_argument('--model',
nargs="+",
dest="model_paths",
help="models to test the accuracy of")
args = argument_parser.parse_args()
if not args.model_paths:
model_paths = glob("model/*") # compute accuracy for all the found models
else:
model_paths = args.model_paths
result = []
_, _, X_test, y_test = mnist()
for path in model_paths:
if not os.path.exists(path):
continue
print(f"Computing for {path}...", end="")
sys.stdout.flush()
model = load_from_file(path)
test_set_accuracy = accuracy(model, X_test, y_test)
result.append((path, test_set_accuracy))
print()
result.sort(key=lambda pair: pair[1], reverse=True)
for pair in result:
path, test_set_accuracy = pair
print(f"{path} -> {test_set_accuracy}")
def main():
args = setup()
print("Preparing directories")
os.makedirs(os.path.join(args.root_dir, "factors"), exist_ok=True)
filename = f"{args.prefix}{args.model}_{args.data}_{args.estimator}{args.suffix}"
factors_path = os.path.join(args.root_dir, "factors", filename)
print("Loading model")
if args.model == 'lenet5':
model = lenet5.lenet5(pretrained=args.data, device=args.device)
elif args.model == 'resnet18' and args.data != 'imagenet':
model = resnet.resnet18(pretrained=os.path.join(
args.root_dir, 'weights', f"{args.model}_{args.data}.pth"),
num_classes=43 if args.data == 'gtsrb' else 10,
device=args.device)
else:
model_class = getattr(torchvision.models, args.model)
if args.model in ['googlenet', 'inception_v3']:
model = model_class(pretrained=True, aux_logits=False)
else:
model = model_class(pretrained=True)
model.to(args.device).train()
if args.parallel:
model = torch.nn.parallel.DataParallel(model)
if args.estimator != 'inf':
print(f"Loading data")
if args.data == 'cifar10':
data = datasets.cifar10(args.torch_data,
args.batch_size,
args.workers,
args.augment,
splits='train')
elif args.data == 'mnist':
data = datasets.mnist(args.torch_data,
args.batch_size,
args.workers,
args.augment,
splits='train')
elif args.data == 'gtsrb':
data = datasets.gtsrb(args.data_dir,
batch_size=args.batch_size,
workers=args.workers,
splits='train')
elif args.data == 'tiny':
img_size = 64
data = datasets.imagenet(args.data_dir,
img_size,
args.batch_size,
splits='train',
tiny=True)
elif args.data == 'imagenet':
img_size = 224
if args.model in ['googlenet', 'inception_v3']:
img_size = 299
data = datasets.imagenet(args.data_dir,
img_size,
args.batch_size,
workers=args.workers,
splits='train')
torch.backends.cudnn.benchmark = True
print("Computing factors")
if args.estimator == 'inf':
est = compute_inf(args)
elif args.estimator == 'efb':
factors = torch.load(factors_path.replace("efb", "kfac") + '.pth')
est = compute_factors(args, model, data, factors)
else:
est = compute_factors(args, model, data)
print("Saving factors")
if args.estimator == "inf":
torch.save(est.state, f"{factors_path}{args.rank}.pth")
elif args.estimator == "efb":
torch.save(list(est.state.values()), factors_path + '.pth')
torch.save(list(est.diags.values()),
factors_path.replace("efb", "diag") + '.pth')
else:
torch.save(list(est.state.values()), factors_path + '.pth')
def main(_):
rng = random.PRNGKey(0)
# Load MNIST dataset
train_images, train_labels, test_images, test_labels = datasets.mnist()
batch_size = 128
batch_shape = (-1, 28, 28, 1)
num_train = train_images.shape[0]
num_complete_batches, leftover = divmod(num_train, batch_size)
num_batches = num_complete_batches + bool(leftover)
train_images = np.reshape(train_images, batch_shape)
test_images = np.reshape(test_images, batch_shape)
def data_stream():
rng = npr.RandomState(0)
while True:
perm = rng.permutation(num_train)
for i in range(num_batches):
batch_idx = perm[i * batch_size:(i + 1) * batch_size]
yield train_images[batch_idx], train_labels[batch_idx]
batches = data_stream()
# Model, loss, and accuracy functions
init_random_params, predict = stax.serial(
stax.Conv(32, (8, 8), strides=(2, 2), padding='SAME'), stax.Relu,
stax.Conv(128, (6, 6), strides=(2, 2), padding='VALID'), stax.Relu,
stax.Conv(128, (5, 5), strides=(1, 1), padding='VALID'), stax.Flatten,
stax.Dense(128), stax.Relu, stax.Dense(10))
def loss(params, batch):
inputs, targets = batch
preds = predict(params, inputs)
return -np.mean(logsoftmax(preds) * targets)
def accuracy(params, batch):
inputs, targets = batch
target_class = np.argmax(targets, axis=1)
predicted_class = np.argmax(predict(params, inputs), axis=1)
return np.mean(predicted_class == target_class)
# Instantiate an optimizer
opt_init, opt_update, get_params = optimizers.adam(0.001)
@jit
def update(i, opt_state, batch):
params = get_params(opt_state)
return opt_update(i, grad(loss)(params, batch), opt_state)
# Initialize model
_, init_params = init_random_params(rng, batch_shape)
opt_state = opt_init(init_params)
itercount = itertools.count()
# Training loop
print("\nStarting training...")
for epoch in range(FLAGS.nb_epochs):
start_time = time.time()
for _ in range(num_batches):
opt_state = update(next(itercount), opt_state, next(batches))
epoch_time = time.time() - start_time
# Evaluate model on clean data
params = get_params(opt_state)
train_acc = accuracy(params, (train_images, train_labels))
test_acc = accuracy(params, (test_images, test_labels))
# Evaluate model on adversarial data
model_fn = lambda images: predict(params, images)
test_images_fgm = fast_gradient_method(model_fn, test_images,
FLAGS.eps, np.inf)
test_images_pgd = projected_gradient_descent(model_fn, test_images,
FLAGS.eps, 0.01, 40,
np.inf)
test_acc_fgm = accuracy(params, (test_images_fgm, test_labels))
test_acc_pgd = accuracy(params, (test_images_pgd, test_labels))
print("Epoch {} in {:0.2f} sec".format(epoch, epoch_time))
print("Training set accuracy: {}".format(train_acc))
print("Test set accuracy on clean examples: {}".format(test_acc))
print("Test set accuracy on FGM adversarial examples: {}".format(
test_acc_fgm))
print("Test set accuracy on PGD adversarial examples: {}".format(
test_acc_pgd))
batch_size=64,
val_batch_size=1000,
epochs=10,
lr=0.01,
caps=[10, 20, 50, 70, 100, 200, 500],
momentum=0.5,
)
S_proto = struct(epoch=0, bn=0)
log = Logger('mnist_capacity_nodrop',
H_proto,
S_proto,
load=True,
metric_show_freq=1)
tn_loader, val_loader = loader(mnist(), H_proto.batch_size,
H_proto.val_batch_size)
for epoch in range(S_proto.epoch, len(H_proto.caps)):
S_proto.epoch = epoch
log.flush()
cap = H_proto.caps[epoch]
H = H_proto.copy()
H.cap = cap
S = struct(epoch=1, bn=1)
inner_log = Logger(
'cap{}'.format(cap),
H,
S,
# overwrite=True,
load=True,
x = F.relu(F.max_pool2d(self.conv1(x), 2))
x = F.relu(F.max_pool2d(self.conv2_drop(self.conv2(x)), 2))
x = x.view(-1, 320)
x = F.relu(self.fc1(x))
x = F.dropout(x, training=self.training)
x = self.fc2(x)
return F.log_softmax(x)
return Net()
def optimizer(H, params):
return optim.SGD(params, lr=H.lr, momentum=H.momentum)
loader_fn = lambda H: loader(mnist(), H.batch_size, H.val_batch_size)
metrics = lambda model, tn_loader, val_loader: struct(
val_acc=acc_metric(model, val_loader),
val_loss=nll_metric(model, val_loader),
tn_loss=nll_metric(model, tn_loader),
tn_acc=acc_metric(model, tn_loader))
meta_metrics = struct(best_acc=max_metametric('val_acc'),
best_acc_idx=argmax_metametric('val_acc'),
best_loss=min_metametric('val_loss'),
best_loss_idx=argmin_metametric('val_loss'),
best_tn_acc=max_metametric('tn_acc'),
best_tn_acc_idx=argmax_metametric('tn_acc'),
best_tn_loss=min_metametric('tn_loss'),
best_tn_loss_idx=argmin_metametric('tn_loss'))
"dec_im_shape": dec_im_shape,
"dec_filters": dec_filters,
"disc_neurons": disc_neurons,
"epochs": epochs,
"batch_size": batch_size,
"learning_rate": learning_rate,
"num_classes": num_classes,
"N_plot": N_plot,
"log_path": log_path
}
# Create Tensorboard Filewriter & Save Results
[fw, log_path_dt] = utils.create_tensorboard(sess, log_path)
# Load data for training
data = datasets.mnist()
#data = datasets.fashion_mnist()
[train_data, train_labels] = data.get_train_samples()
train_data = train_data / 255.0
sess.run(iterator.initializer,
feed_dict={
data_ph: train_data,
labels_ph: train_labels,
batch_size_ph: batch_size,
shufflebuffer_ph: train_data.shape[0],
epochs_ph: epochs
})
# Train model
i = 0
while True:
from datasets import mnist, loader
from logger import struct, Logger
from train import train
H = struct(
batch_size=64,
val_batch_size=1000,
epochs=2,
lr=0.01,
momentum=0.5,
)
S = struct(epoch=1, bn=1)
log = Logger('mnist_scratch', H, S, overwrite=True, metric_show_freq=100)
tn_loader, val_loader = loader(mnist(), H.batch_size, H.val_batch_size)
class Net(nn.Module):
def __init__(self):
super(Net, self).__init__()
self.conv1 = nn.Conv2d(1, 10, kernel_size=5)
self.conv2 = nn.Conv2d(10, 20, kernel_size=5)
self.conv2_drop = nn.Dropout2d()
self.fc1 = nn.Linear(320, 50)
self.fc2 = nn.Linear(50, 10)
def forward(self, x):
x = F.relu(F.max_pool2d(self.conv1(x), 2))
x = F.relu(F.max_pool2d(self.conv2_drop(self.conv2(x)), 2))
x = x.view(-1, 320)
