def main():
parser = argparse.ArgumentParser()
parser.add_argument('--time', type=int, default=1000)
parser.add_argument('--dt', type=float, default=1)
args = parser.parse_args()
dataset = MNIST(path=data_path, download=True, shuffle=True)
images, _ = dataset.get_train()
image = images[0]
spikes = rank_order(datum=image, time=args.time, dt=args.dt)
plt.matshow(spikes.view(-1, 784).t())
plt.matshow(spikes.sum(dim=0).view(28, 28))
plt.colorbar()
plt.matshow(image)
plt.show()
def test_mnist_pipeline(self):
network = DiehlAndCook2015(n_inpt=784, n_neurons=400, exc=22.5, inh=17.5, dt=1.0, norm=78.4)
environment = DatasetEnvironment(dataset=MNIST(path='../data/MNIST', download=True), train=True, intensity=0.25)
pipeline = Pipeline(network=network, environment=environment, encoding=poisson, time=350)
assert pipeline.network == network
assert pipeline.env == environment
assert pipeline.encoding == poisson
assert pipeline.time == 350
assert pipeline.history_length is None
def load(self, file_path):
self.network = load(file_path)
self.n_iter = 60000
dt = 1
intensity = 127.5
self.train_dataset = MNIST(
PoissonEncoder(time=self.time_max, dt=dt),
None,
"MNIST",
download=False,
train=True,
transform=transforms.Compose(
[transforms.ToTensor(), transforms.Lambda(lambda x: x * intensity)]
)
)
self.spikes = {}
for layer in set(self.network.layers):
self.spikes[layer] = Monitor(self.network.layers[layer], state_vars=["s"], time=self.time_max)
self.network.add_monitor(self.spikes[layer], name="%s_spikes" % layer)
#print('GlobalMonitor.state_vars:', self.GlobalMonitor.state_vars)
self.voltages = {}
for layer in set(self.network.layers) - {"X"}:
self.voltages[layer] = Monitor(self.network.layers[layer], state_vars=["v"], time=self.time_max)
self.network.add_monitor(self.voltages[layer], name="%s_voltages" % layer)
weights_XY = self.network.connections[('X', 'Y')].w
weights_XY = weights_XY.reshape(28, 28, -1)
weights_to_display = torch.zeros(0, 28*25)
i = 0
while i < 625:
for j in range(25):
weights_to_display_row = torch.zeros(28, 0)
for k in range(25):
weights_to_display_row = torch.cat((weights_to_display_row, weights_XY[:, :, i]), dim=1)
i += 1
weights_to_display = torch.cat((weights_to_display, weights_to_display_row), dim=0)
self.weights_XY = weights_to_display.numpy()
def _init_dataset(self, cfg):
if cfg['dataset'] not in ['MNIST']:
raise ValueError(name, 'is not a valid dataset in SNN training.')
if cfg['dataset'] == 'MNIST':
dataset = MNIST(
PoissonEncoder(time=cfg['time'], dt=cfg['dt']),
None,
root=os.path.join("data", "MNIST"),
download=True,
transform=tfs.Compose([
tfs.ToTensor(),
tfs.Lambda(lambda x: x * cfg['intensity'])
]),
)
n_classes = 10
return dataset, n_classes
network.add_connection(conv_conn, source="X", target="Y")
network.add_connection(recurrent_conn, source="Y", target="Y")
# Voltage recording for excitatory and inhibitory layers.
voltage_monitor = Monitor(network.layers["Y"], ["v"], time=time)
network.add_monitor(voltage_monitor, name="output_voltage")
if gpu:
network.to("cuda")
# Load MNIST data.
train_dataset = MNIST(
PoissonEncoder(time=time, dt=dt),
None,
"../../data/MNIST",
download=True,
train=True,
transform=transforms.Compose(
[transforms.ToTensor(),
transforms.Lambda(lambda x: x * intensity)]),
)
spikes = {}
for layer in set(network.layers):
spikes[layer] = Monitor(network.layers[layer], state_vars=["s"], time=time)
network.add_monitor(spikes[layer], name="%s_spikes" % layer)
voltages = {}
for layer in set(network.layers) - {"X"}:
voltages[layer] = Monitor(network.layers[layer],
state_vars=["v"],
time=time)
def main(seed=0,
n_train=60000,
n_test=10000,
kernel_size=(16, ),
stride=(4, ),
n_filters=25,
padding=0,
inhib=100,
time=25,
lr=1e-3,
lr_decay=0.99,
dt=1,
intensity=1,
progress_interval=10,
update_interval=250,
plot=False,
train=True,
gpu=False):
assert n_train % update_interval == 0 and n_test % update_interval == 0, \
'No. examples must be divisible by update_interval'
params = [
seed, n_train, kernel_size, stride, n_filters, padding, inhib, time,
lr, lr_decay, dt, intensity, update_interval
]
model_name = '_'.join([str(x) for x in params])
if not train:
test_params = [
seed, n_train, n_test, kernel_size, stride, n_filters, padding,
inhib, time, lr, lr_decay, dt, intensity, update_interval
]
np.random.seed(seed)
if gpu:
torch.set_default_tensor_type('torch.cuda.FloatTensor')
torch.cuda.manual_seed_all(seed)
else:
torch.manual_seed(seed)
n_examples = n_train if train else n_test
input_shape = [20, 20]
if kernel_size == input_shape:
conv_size = [1, 1]
else:
conv_size = (int((input_shape[0] - kernel_size[0]) / stride[0]) + 1,
int((input_shape[1] - kernel_size[1]) / stride[1]) + 1)
n_classes = 10
n_neurons = n_filters * np.prod(conv_size)
total_kernel_size = int(np.prod(kernel_size))
total_conv_size = int(np.prod(conv_size))
# Build network.
if train:
network = Network()
input_layer = Input(n=400, shape=(1, 1, 20, 20), traces=True)
conv_layer = DiehlAndCookNodes(n=n_filters * total_conv_size,
shape=(1, n_filters, *conv_size),
thresh=-64.0,
traces=True,
theta_plus=0.05 * (kernel_size[0] / 20),
refrac=0)
conv_layer2 = LIFNodes(n=n_filters * total_conv_size,
shape=(1, n_filters, *conv_size),
refrac=0)
conv_conn = Conv2dConnection(input_layer,
conv_layer,
kernel_size=kernel_size,
stride=stride,
update_rule=WeightDependentPostPre,
norm=0.05 * total_kernel_size,
nu=[0, lr],
wmin=0,
wmax=0.25)
conv_conn2 = Conv2dConnection(input_layer,
conv_layer2,
w=conv_conn.w,
kernel_size=kernel_size,
stride=stride,
update_rule=None,
wmax=0.25)
w = -inhib * torch.ones(n_filters, conv_size[0], conv_size[1],
n_filters, conv_size[0], conv_size[1])
for f in range(n_filters):
for f2 in range(n_filters):
if f != f2:
w[f, :, :f2, :, :] = 0
w = w.view(n_filters * conv_size[0] * conv_size[1],
n_filters * conv_size[0] * conv_size[1])
recurrent_conn = Connection(conv_layer, conv_layer, w=w)
network.add_layer(input_layer, name='X')
network.add_layer(conv_layer, name='Y')
network.add_layer(conv_layer2, name='Y_')
network.add_connection(conv_conn, source='X', target='Y')
network.add_connection(conv_conn2, source='X', target='Y_')
network.add_connection(recurrent_conn, source='Y', target='Y')
# Voltage recording for excitatory and inhibitory layers.
voltage_monitor = Monitor(network.layers['Y'], ['v'], time=time)
network.add_monitor(voltage_monitor, name='output_voltage')
else:
network = load_network(os.path.join(params_path, model_name + '.pt'))
network.connections['X', 'Y'].update_rule = NoOp(
connection=network.connections['X', 'Y'],
nu=network.connections['X', 'Y'].nu)
network.layers['Y'].theta_decay = 0
network.layers['Y'].theta_plus = 0
# Load MNIST data.
dataset = MNIST(data_path, download=True)
if train:
images, labels = dataset.get_train()
else:
images, labels = dataset.get_test()
images *= intensity
images = images[:, 4:-4, 4:-4].contiguous()
# Record spikes during the simulation.
spike_record = torch.zeros(update_interval, time, n_neurons)
full_spike_record = torch.zeros(n_examples, n_neurons)
# Neuron assignments and spike proportions.
if train:
logreg_model = LogisticRegression(warm_start=True,
n_jobs=-1,
solver='lbfgs',
max_iter=1000,
multi_class='multinomial')
logreg_model.coef_ = np.zeros([n_classes, n_neurons])
logreg_model.intercept_ = np.zeros(n_classes)
logreg_model.classes_ = np.arange(n_classes)
else:
path = os.path.join(params_path,
'_'.join(['auxiliary', model_name]) + '.pt')
logreg_coef, logreg_intercept = torch.load(open(path, 'rb'))
logreg_model = LogisticRegression(warm_start=True,
n_jobs=-1,
solver='lbfgs',
max_iter=1000,
multi_class='multinomial')
logreg_model.coef_ = logreg_coef
logreg_model.intercept_ = logreg_intercept
logreg_model.classes_ = np.arange(n_classes)
# Sequence of accuracy estimates.
curves = {'logreg': []}
predictions = {scheme: torch.Tensor().long() for scheme in curves.keys()}
if train:
best_accuracy = 0
spikes = {}
for layer in set(network.layers):
spikes[layer] = Monitor(network.layers[layer],
state_vars=['s'],
time=time)
network.add_monitor(spikes[layer], name='%s_spikes' % layer)
# Train the network.
if train:
print('\nBegin training.\n')
else:
print('\nBegin test.\n')
inpt_ims = None
inpt_axes = None
spike_ims = None
spike_axes = None
weights_im = None
plot_update_interval = 100
start = t()
for i in range(n_examples):
if i % progress_interval == 0:
print('Progress: %d / %d (%.4f seconds)' %
(i, n_examples, t() - start))
start = t()
if i % update_interval == 0 and i > 0:
if train:
network.connections['X', 'Y'].update_rule.nu[1] *= lr_decay
if i % len(labels) == 0:
current_labels = labels[-update_interval:]
current_record = full_spike_record[-update_interval:]
else:
current_labels = labels[i % len(labels) - update_interval:i %
len(labels)]
current_record = full_spike_record[i % len(labels) -
update_interval:i %
len(labels)]
# Update and print accuracy evaluations.
curves, preds = update_curves(curves,
current_labels,
n_classes,
full_spike_record=current_record,
logreg=logreg_model)
print_results(curves)
for scheme in preds:
predictions[scheme] = torch.cat(
[predictions[scheme], preds[scheme]], -1)
# Save accuracy curves to disk.
to_write = ['train'] + params if train else ['test'] + params
f = '_'.join([str(x) for x in to_write]) + '.pt'
torch.save((curves, update_interval, n_examples),
open(os.path.join(curves_path, f), 'wb'))
if train:
if any([x[-1] > best_accuracy for x in curves.values()]):
print(
'New best accuracy! Saving network parameters to disk.'
)
# Save network to disk.
network.save(os.path.join(params_path, model_name + '.pt'))
path = os.path.join(
params_path,
'_'.join(['auxiliary', model_name]) + '.pt')
torch.save((logreg_model.coef_, logreg_model.intercept_),
open(path, 'wb'))
best_accuracy = max([x[-1] for x in curves.values()])
# Refit logistic regression model.
logreg_model = logreg_fit(full_spike_record[:i], labels[:i],
logreg_model)
print()
# Get next input sample.
image = images[i % len(images)]
sample = bernoulli(datum=image, time=time, dt=dt,
max_prob=1).unsqueeze(1).unsqueeze(1)
inpts = {'X': sample}
# Run the network on the input.
network.run(inpts=inpts, time=time)
network.connections['X', 'Y_'].w = network.connections['X', 'Y'].w
# Add to spikes recording.
spike_record[i % update_interval] = spikes['Y_'].get('s').view(
time, -1)
full_spike_record[i] = spikes['Y_'].get('s').view(time, -1).sum(0)
# Optionally plot various simulation information.
if plot and i % plot_update_interval == 0:
_input = inpts['X'].view(time, 400).sum(0).view(20, 20)
w = network.connections['X', 'Y'].w
_spikes = {
'X': spikes['X'].get('s').view(400, time),
'Y': spikes['Y'].get('s').view(n_filters * total_conv_size,
time),
'Y_': spikes['Y_'].get('s').view(n_filters * total_conv_size,
time)
}
inpt_axes, inpt_ims = plot_input(image.view(20, 20),
_input,
label=labels[i % len(labels)],
ims=inpt_ims,
axes=inpt_axes)
spike_ims, spike_axes = plot_spikes(spikes=_spikes,
ims=spike_ims,
axes=spike_axes)
weights_im = plot_conv2d_weights(
w, im=weights_im, wmax=network.connections['X', 'Y'].wmax)
plt.pause(1e-2)
network.reset_() # Reset state variables.
print(f'Progress: {n_examples} / {n_examples} ({t() - start:.4f} seconds)')
i += 1
if i % len(labels) == 0:
current_labels = labels[-update_interval:]
current_record = full_spike_record[-update_interval:]
else:
current_labels = labels[i % len(labels) - update_interval:i %
len(labels)]
current_record = full_spike_record[i % len(labels) -
update_interval:i % len(labels)]
# Update and print accuracy evaluations.
curves, preds = update_curves(curves,
current_labels,
n_classes,
full_spike_record=current_record,
logreg=logreg_model)
print_results(curves)
for scheme in preds:
predictions[scheme] = torch.cat([predictions[scheme], preds[scheme]],
-1)
if train:
if any([x[-1] > best_accuracy for x in curves.values()]):
print('New best accuracy! Saving network parameters to disk.')
# Save network to disk.
network.save(os.path.join(params_path, model_name + '.pt'))
path = os.path.join(params_path,
'_'.join(['auxiliary', model_name]) + '.pt')
torch.save((logreg_model.coef_, logreg_model.intercept_),
open(path, 'wb'))
if train:
print('\nTraining complete.\n')
else:
print('\nTest complete.\n')
print('Average accuracies:\n')
for scheme in curves.keys():
print('\t%s: %.2f' % (scheme, float(np.mean(curves[scheme]))))
# Save accuracy curves to disk.
to_write = ['train'] + params if train else ['test'] + params
to_write = [str(x) for x in to_write]
f = '_'.join(to_write) + '.pt'
torch.save((curves, update_interval, n_examples),
open(os.path.join(curves_path, f), 'wb'))
# Save results to disk.
results = [np.mean(curves['logreg']), np.std(curves['logreg'])]
to_write = params + results if train else test_params + results
to_write = [str(x) for x in to_write]
name = 'train.csv' if train else 'test.csv'
if not os.path.isfile(os.path.join(results_path, name)):
with open(os.path.join(results_path, name), 'w') as f:
if train:
columns = [
'seed', 'n_train', 'kernel_size', 'stride', 'n_filters',
'padding', 'inhib', 'time', 'lr', 'lr_decay', 'dt',
'intensity', 'update_interval', 'mean_logreg', 'std_logreg'
]
header = ','.join(columns) + '\n'
f.write(header)
else:
columns = [
'seed', 'n_train', 'n_test', 'kernel_size', 'stride',
'n_filters', 'padding', 'inhib', 'time', 'lr', 'lr_decay',
'dt', 'intensity', 'update_interval', 'mean_logreg',
'std_logreg'
]
header = ','.join(columns) + '\n'
f.write(header)
with open(os.path.join(results_path, name), 'a') as f:
f.write(','.join(to_write) + '\n')
if labels.numel() > n_examples:
labels = labels[:n_examples]
else:
while labels.numel() < n_examples:
if 2 * labels.numel() > n_examples:
labels = torch.cat(
[labels, labels[:n_examples - labels.numel()]])
else:
labels = torch.cat([labels, labels])
# Compute confusion matrices and save them to disk.
confusions = {}
for scheme in predictions:
confusions[scheme] = confusion_matrix(labels, predictions[scheme])
to_write = ['train'] + params if train else ['test'] + test_params
f = '_'.join([str(x) for x in to_write]) + '.pt'
torch.save(confusions, os.path.join(confusion_path, f))
target=output,
w=0.5 * torch.randn(output.n, output.n))
network.add_connection(C1, source="I", target="O")
network.add_connection(C2, source="O", target="O")
spikes = {}
for l in network.layers:
spikes[l] = Monitor(network.layers[l], ["s"], time=250)
network.add_monitor(spikes[l], name="%s_spikes" % l)
voltages = {"O": Monitor(network.layers["O"], ["v"], time=250)}
network.add_monitor(voltages["O"], name="O_voltages")
# Get MNIST training images and labels.
images, labels = MNIST(path="../../data/MNIST", download=True).get_train()
images *= 0.25
# Create lazily iterating Poisson-distributed data loader.
loader = zip(poisson_loader(images, time=250), iter(labels))
inpt_axes = None
inpt_ims = None
spike_axes = None
spike_ims = None
weights_im = None
weights_im2 = None
voltage_ims = None
voltage_axes = None
# Run training data on reservoir computer and store (spikes per neuron, label) per example.
def main(seed=0,
n_neurons=100,
n_train=60000,
n_test=10000,
inhib=100,
lr=1e-2,
lr_decay=1,
time=350,
dt=1,
theta_plus=0.05,
theta_decay=1e7,
intensity=1,
progress_interval=10,
update_interval=250,
plot=False,
train=True,
gpu=False):
assert n_train % update_interval == 0 and n_test % update_interval == 0, \
'No. examples must be divisible by update_interval'
params = [
seed, n_neurons, n_train, inhib, lr, lr_decay, time, dt, theta_plus,
theta_decay, intensity, progress_interval, update_interval
]
test_params = [
seed, n_neurons, n_train, n_test, inhib, lr, lr_decay, time, dt,
theta_plus, theta_decay, intensity, progress_interval, update_interval
]
model_name = '_'.join([str(x) for x in params])
np.random.seed(seed)
if gpu:
torch.set_default_tensor_type('torch.cuda.FloatTensor')
torch.cuda.manual_seed_all(seed)
else:
torch.manual_seed(seed)
n_examples = n_train if train else n_test
n_sqrt = int(np.ceil(np.sqrt(n_neurons)))
n_classes = 10
# Build network.
if train:
network = DiehlAndCook2015v2(n_inpt=784,
n_neurons=n_neurons,
inh=inhib,
dt=dt,
norm=78.4,
theta_plus=theta_plus,
theta_decay=theta_decay,
nu=[0, lr])
else:
network = load(os.path.join(params_path, model_name + '.pt'))
network.connections['X', 'Y'].update_rule = NoOp(
connection=network.connections['X', 'Y'],
nu=network.connections['X', 'Y'].nu)
network.layers['Y'].theta_decay = 0
network.layers['Y'].theta_plus = 0
# Load MNIST data.
dataset = MNIST(path=data_path, download=True)
if train:
images, labels = dataset.get_train()
else:
images, labels = dataset.get_test()
images = images.view(-1, 784)
images *= intensity
# Record spikes during the simulation.
spike_record = torch.zeros(update_interval, time, n_neurons)
full_spike_record = torch.zeros(n_examples, n_neurons).long()
# Neuron assignments and spike proportions.
if train:
assignments = -torch.ones_like(torch.Tensor(n_neurons))
proportions = torch.zeros_like(torch.Tensor(n_neurons, n_classes))
rates = torch.zeros_like(torch.Tensor(n_neurons, n_classes))
ngram_scores = {}
else:
path = os.path.join(params_path,
'_'.join(['auxiliary', model_name]) + '.pt')
assignments, proportions, rates, ngram_scores = torch.load(
open(path, 'rb'))
# Sequence of accuracy estimates.
curves = {'all': [], 'proportion': [], 'ngram': []}
predictions = {scheme: torch.Tensor().long() for scheme in curves.keys()}
if train:
best_accuracy = 0
spikes = {}
for layer in set(network.layers):
spikes[layer] = Monitor(network.layers[layer],
state_vars=['s'],
time=time)
network.add_monitor(spikes[layer], name='%s_spikes' % layer)
# Train the network.
if train:
print('\nBegin training.\n')
else:
print('\nBegin test.\n')
inpt_axes = None
inpt_ims = None
spike_ims = None
spike_axes = None
weights_im = None
assigns_im = None
perf_ax = None
start = t()
for i in range(n_examples):
if i % progress_interval == 0:
print(f'Progress: {i} / {n_examples} ({t() - start:.4f} seconds)')
start = t()
if i % update_interval == 0 and i > 0:
if train:
network.connections['X', 'Y'].update_rule.nu[1] *= lr_decay
if i % len(labels) == 0:
current_labels = labels[-update_interval:]
else:
current_labels = labels[i % len(images) - update_interval:i %
len(images)]
# Update and print accuracy evaluations.
curves, preds = update_curves(curves,
current_labels,
n_classes,
spike_record=spike_record,
assignments=assignments,
proportions=proportions,
ngram_scores=ngram_scores,
n=2)
print_results(curves)
for scheme in preds:
predictions[scheme] = torch.cat(
[predictions[scheme], preds[scheme]], -1)
# Save accuracy curves to disk.
to_write = ['train'] + params if train else ['test'] + params
f = '_'.join([str(x) for x in to_write]) + '.pt'
torch.save((curves, update_interval, n_examples),
open(os.path.join(curves_path, f), 'wb'))
if train:
if any([x[-1] > best_accuracy for x in curves.values()]):
print(
'New best accuracy! Saving network parameters to disk.'
)
# Save network to disk.
network.save(os.path.join(params_path, model_name + '.pt'))
path = os.path.join(
params_path,
'_'.join(['auxiliary', model_name]) + '.pt')
torch.save((assignments, proportions, rates, ngram_scores),
open(path, 'wb'))
best_accuracy = max([x[-1] for x in curves.values()])
# Assign labels to excitatory layer neurons.
assignments, proportions, rates = assign_labels(
spike_record, current_labels, n_classes, rates)
# Compute ngram scores.
ngram_scores = update_ngram_scores(spike_record,
current_labels, n_classes,
2, ngram_scores)
print()
# Get next input sample.
image = images[i % len(images)]
sample = poisson(datum=image, time=time, dt=dt)
inpts = {'X': sample}
# Run the network on the input.
network.run(inpts=inpts, time=time)
retries = 0
while spikes['Y'].get('s').sum() < 1 and retries < 3:
retries += 1
image *= 2
sample = poisson(datum=image, time=time, dt=dt)
inpts = {'X': sample}
network.run(inpts=inpts, time=time)
# Add to spikes recording.
spike_record[i % update_interval] = spikes['Y'].get('s').t()
full_spike_record[i] = spikes['Y'].get('s').t().sum(0).long()
# Optionally plot various simulation information.
if plot:
# _input = image.view(28, 28)
# reconstruction = inpts['X'].view(time, 784).sum(0).view(28, 28)
_spikes = {layer: spikes[layer].get('s') for layer in spikes}
input_exc_weights = network.connections[('X', 'Y')].w
square_weights = get_square_weights(
input_exc_weights.view(784, n_neurons), n_sqrt, 28)
# square_assignments = get_square_assignments(assignments, n_sqrt)
# inpt_axes, inpt_ims = plot_input(_input, reconstruction, label=labels[i], axes=inpt_axes, ims=inpt_ims)
spike_ims, spike_axes = plot_spikes(_spikes,
ims=spike_ims,
axes=spike_axes)
weights_im = plot_weights(square_weights, im=weights_im)
# assigns_im = plot_assignments(square_assignments, im=assigns_im)
# perf_ax = plot_performance(curves, ax=perf_ax)
plt.pause(1e-8)
network.reset_() # Reset state variables.
print(f'Progress: {n_examples} / {n_examples} ({t() - start:.4f} seconds)')
i += 1
if i % len(labels) == 0:
current_labels = labels[-update_interval:]
else:
current_labels = labels[i % len(images) - update_interval:i %
len(images)]
# Update and print accuracy evaluations.
curves, preds = update_curves(curves,
current_labels,
n_classes,
spike_record=spike_record,
assignments=assignments,
proportions=proportions,
ngram_scores=ngram_scores,
n=2)
print_results(curves)
for scheme in preds:
predictions[scheme] = torch.cat([predictions[scheme], preds[scheme]],
-1)
if train:
if any([x[-1] > best_accuracy for x in curves.values()]):
print('New best accuracy! Saving network parameters to disk.')
# Save network to disk.
if train:
network.save(os.path.join(params_path, model_name + '.pt'))
path = os.path.join(
params_path, '_'.join(['auxiliary', model_name]) + '.pt')
torch.save((assignments, proportions, rates, ngram_scores),
open(path, 'wb'))
if train:
print('\nTraining complete.\n')
else:
print('\nTest complete.\n')
print('Average accuracies:\n')
for scheme in curves.keys():
print('\t%s: %.2f' % (scheme, float(np.mean(curves[scheme]))))
# Save accuracy curves to disk.
to_write = ['train'] + params if train else ['test'] + params
f = '_'.join([str(x) for x in to_write]) + '.pt'
torch.save((curves, update_interval, n_examples),
open(os.path.join(curves_path, f), 'wb'))
# Save results to disk.
results = [
np.mean(curves['all']),
np.mean(curves['proportion']),
np.mean(curves['ngram']),
np.max(curves['all']),
np.max(curves['proportion']),
np.max(curves['ngram'])
]
to_write = params + results if train else test_params + results
to_write = [str(x) for x in to_write]
name = 'train.csv' if train else 'test.csv'
if not os.path.isfile(os.path.join(results_path, name)):
with open(os.path.join(results_path, name), 'w') as f:
if train:
f.write(
'random_seed,n_neurons,n_train,inhib,lr,lr_decay,time,timestep,theta_plus,theta_decay,intensity,'
'progress_interval,update_interval,mean_all_activity,mean_proportion_weighting,'
'mean_ngram,max_all_activity,max_proportion_weighting,max_ngram\n'
)
else:
f.write(
'random_seed,n_neurons,n_train,n_test,inhib,lr,lr_decay,time,timestep,theta_plus,theta_decay,'
'intensity,progress_interval,update_interval,mean_all_activity,mean_proportion_weighting,'
'mean_ngram,max_all_activity,max_proportion_weighting,max_ngram\n'
)
with open(os.path.join(results_path, name), 'a') as f:
f.write(','.join(to_write) + '\n')
if labels.numel() > n_examples:
labels = labels[:n_examples]
else:
while labels.numel() < n_examples:
if 2 * labels.numel() > n_examples:
labels = torch.cat(
[labels, labels[:n_examples - labels.numel()]])
else:
labels = torch.cat([labels, labels])
# Compute confusion matrices and save them to disk.
confusions = {}
for scheme in predictions:
confusions[scheme] = confusion_matrix(labels, predictions[scheme])
to_write = ['train'] + params if train else ['test'] + test_params
f = '_'.join([str(x) for x in to_write]) + '.pt'
torch.save(confusions, os.path.join(confusion_path, f))
# Save full spike record to disk.
torch.save(full_spike_record, os.path.join(spikes_path, f))
# Build Diehl & Cook 2015 network.
network = DiehlAndCook2015(
n_inpt=784,
n_neurons=400,
exc=22.5,
inh=17.5,
dt=1.0,
norm=78.4,
inpt_shape=(1, 1, 28, 28),
)
# Specify dataset
mnist = MNIST(
PoissonEncoder(time=50, dt=1.0),
None,
root="../../data/MNIST",
download=True,
train=True,
transform=transforms.Compose(
[transforms.ToTensor(), transforms.Lambda(lambda x: x * 128.0)]
),
)
# Build pipeline from components.
pipeline = TorchVisionDatasetPipeline(
network, mnist, TensorboardAnalyzer("logs/minimal_mnist"), plot_interval=100
)
pipeline.train()
def main(seed=0, p_remove=0):
model = '0_16_2_250_4_0.01_0.99_60000_250.0_250_1.0_0.05_1e-07_0.5_0.2_10_250.pt'
np.random.seed(seed)
torch.manual_seed(seed)
torch.set_default_tensor_type('torch.cuda.FloatTensor')
torch.cuda.manual_seed_all(seed)
crop = 4
time = 250
n_filters = 250
intensity = 0.5
n_examples = 10000
n_classes = 10
# Load network.
network = load_network(
os.path.join(
ROOT_DIR, 'params', 'mnist', 'crop_locally_connected', model
)
)
network.connections['X', 'Y'].update_rule = NoOp(
connection=network.connections['X', 'Y'], nu=network.connections['X', 'Y'].nu
)
network.layers['Y'].theta_decay = 0
network.layers['Y'].theta_plus = 0
network.connections['X', 'Y'].norm = None
# Remove `p_remove` percentage of neurons (set outgoing synapses to 0).
mask = torch.bernoulli(p_remove * torch.ones(network.layers['Y'].shape)).byte()
network.connections['X', 'Y'].w[:, mask] = 0
conv_size = network.connections['X', 'Y'].conv_size
conv_prod = int(np.prod(conv_size))
n_neurons = n_filters * conv_prod
# Voltage recording for excitatory and inhibitory layers.
voltage_monitor = Monitor(network.layers['Y'], ['v'], time=time)
network.add_monitor(voltage_monitor, name='output_voltage')
# Load MNIST data.
dataset = MNIST(path=data_path, download=True, shuffle=True)
images, labels = dataset.get_test()
images *= intensity
images = images[:, crop:-crop, crop:-crop]
update_interval = 250
# Record spikes during the simulation.
spike_record = torch.zeros(update_interval, time, n_neurons)
# Neuron assignments and spike proportions.
path = os.path.join(
ROOT_DIR, 'params', 'mnist', 'crop_locally_connected', f'auxiliary_{model}'
)
assignments, proportions, rates, ngram_scores = torch.load(open(path, 'rb'))
# Sequence of accuracy estimates.
curves = {'all': [], 'proportion': [], 'ngram': []}
predictions = {
scheme: torch.Tensor().long() for scheme in curves.keys()
}
spikes = {}
for layer in set(network.layers):
spikes[layer] = Monitor(network.layers[layer], state_vars=['s'], time=time)
network.add_monitor(spikes[layer], name=f'{layer}_spikes')
start = t()
for i in range(n_examples):
if i % 10 == 0:
print(f'Progress: {i} / {n_examples} ({t() - start:.4f} seconds)')
start = t()
if i % update_interval == 0 and i > 0:
if i % len(labels) == 0:
current_labels = labels[-update_interval:]
else:
current_labels = labels[i % len(images) - update_interval:i % len(images)]
# Update and print accuracy evaluations.
curves, preds = update_curves(
curves, current_labels, n_classes, spike_record=spike_record, assignments=assignments,
proportions=proportions, ngram_scores=ngram_scores, n=2
)
print_results(curves)
for scheme in preds:
predictions[scheme] = torch.cat([predictions[scheme], preds[scheme]], -1)
# Get next input sample.
image = images[i % len(images)].contiguous().view(-1)
sample = poisson(datum=image, time=time, dt=1)
inpts = {'X': sample}
# Run the network on the input.
network.run(inpts=inpts, time=time)
retries = 0
while spikes['Y'].get('s').sum() < 5 and retries < 3:
retries += 1
image *= 2
sample = poisson(datum=image, time=time, dt=1)
inpts = {'X': sample}
network.run(inpts=inpts, time=time)
# Add to spikes recording.
spike_record[i % update_interval] = spikes['Y'].get('s').t()
network.reset_() # Reset state variables.
print(f'Progress: {n_examples} / {n_examples} ({t() - start:.4f} seconds)')
i += 1
if i % len(labels) == 0:
current_labels = labels[-update_interval:]
else:
current_labels = labels[i % len(images) - update_interval:i % len(images)]
# Update and print accuracy evaluations.
curves, preds = update_curves(
curves, current_labels, n_classes, spike_record=spike_record, assignments=assignments,
proportions=proportions, ngram_scores=ngram_scores, n=2
)
print_results(curves)
for scheme in preds:
predictions[scheme] = torch.cat([predictions[scheme], preds[scheme]], -1)
print('Average accuracies:\n')
for scheme in curves.keys():
print('\t%s: %.2f' % (scheme, float(np.mean(curves[scheme]))))
# Save results to disk.
results = [
np.mean(curves['all']), np.mean(curves['proportion']), np.mean(curves['ngram']),
np.max(curves['all']), np.max(curves['proportion']), np.max(curves['ngram'])
]
to_write = [str(x) for x in [seed, p_remove] + results]
name = 'neuron_robust.csv'
if not os.path.isfile(os.path.join(results_path, name)):
with open(os.path.join(results_path, name), 'w') as f:
f.write(
'random_seed,p_remove\n'
)
with open(os.path.join(results_path, name), 'a') as f:
f.write(','.join(to_write) + '\n')
def main(seed=0, n_examples=100, gpu=False, plot=False):
np.random.seed(seed)
if gpu:
torch.set_default_tensor_type('torch.cuda.FloatTensor')
torch.cuda.manual_seed_all(seed)
else:
torch.manual_seed(seed)
model_name = '0_12_4_150_4_0.01_0.99_60000_250.0_250_1.0_0.05_1e-07_0.5_0.2_10_250'
network = load_network(os.path.join(params_path, f'{model_name}.pt'))
for l in network.layers:
network.layers[l].dt = network.dt
for c in network.connections:
network.connections[c].dt = network.dt
network.layers['Y'].one_spike = True
network.layers['Y'].lbound = None
kernel_size = 12
side_length = 20
n_filters = 150
time = 250
intensity = 0.5
crop = 4
conv_size = network.connections['X', 'Y'].conv_size
locations = network.connections['X', 'Y'].locations
conv_prod = int(np.prod(conv_size))
n_neurons = n_filters * conv_prod
n_classes = 10
# Voltage recording for excitatory and inhibitory layers.
voltage_monitor = Monitor(network.layers['Y'], ['v'], time=time)
network.add_monitor(voltage_monitor, name='output_voltage')
# Load MNIST data.
dataset = MNIST(path=data_path, download=True)
images, labels = dataset.get_test()
images *= intensity
images = images[:, crop:-crop, crop:-crop]
# Neuron assignments and spike proportions.
path = os.path.join(params_path,
'_'.join(['auxiliary', model_name]) + '.pt')
assignments, proportions, rates, ngram_scores = torch.load(open(
path, 'rb'))
spikes = {}
for layer in set(network.layers):
spikes[layer] = Monitor(network.layers[layer],
state_vars=['s'],
time=time)
network.add_monitor(spikes[layer], name=f'{layer}_spikes')
# Train the network.
print('\nBegin black box adversarial attack.\n')
spike_ims = None
spike_axes = None
weights_im = None
inpt_ims = None
inpt_axes = None
max_iters = 25
delta = 0.1
epsilon = 0.1
for i in range(n_examples):
# Get next input sample.
original = images[i % len(images)].contiguous().view(-1)
label = labels[i % len(images)]
# Check if the image is correctly classified.
sample = poisson(datum=original, time=time)
inpts = {'X': sample}
# Run the network on the input.
network.run(inpts=inpts, time=time)
# Check for incorrect classification.
s = spikes['Y'].get('s').view(1, n_neurons, time)
prediction = ngram(spikes=s,
ngram_scores=ngram_scores,
n_labels=10,
n=2).item()
if prediction != label:
continue
# Create adversarial example.
adversarial = False
while not adversarial:
adv_example = 255 * torch.rand(original.size())
sample = poisson(datum=adv_example, time=time)
inpts = {'X': sample}
# Run the network on the input.
network.run(inpts=inpts, time=time)
# Check for incorrect classification.
s = spikes['Y'].get('s').view(1, n_neurons, time)
prediction = ngram(spikes=s,
ngram_scores=ngram_scores,
n_labels=n_classes,
n=2).item()
if prediction == label:
adversarial = True
j = 0
current = original.clone()
while j < max_iters:
# Orthogonal perturbation.
# perturb = orthogonal_perturbation(delta=delta, image=adv_example, target=original)
# temp = adv_example + perturb
# # Forward perturbation.
# temp = temp.clone() + forward_perturbation(epsilon * get_diff(temp, original), temp, adv_example)
# print(temp)
perturbation = torch.randn(original.size())
unnormed_source_direction = original - perturbation
source_norm = torch.norm(unnormed_source_direction)
source_direction = unnormed_source_direction / source_norm
dot = torch.dot(perturbation, source_direction)
perturbation -= dot * source_direction
perturbation *= epsilon * source_norm / torch.norm(perturbation)
D = 1 / np.sqrt(epsilon**2 + 1)
direction = perturbation - unnormed_source_direction
spherical_candidate = current + D * direction
spherical_candidate = torch.clamp(spherical_candidate, 0, 255)
new_source_direction = original - spherical_candidate
new_source_direction_norm = torch.norm(new_source_direction)
# length if spherical_candidate would be exactly on the sphere
length = delta * source_norm
# length including correction for deviation from sphere
deviation = new_source_direction_norm - source_norm
length += deviation
# make sure the step size is positive
length = max(0, length)
# normalize the length
length = length / new_source_direction_norm
candidate = spherical_candidate + length * new_source_direction
candidate = torch.clamp(candidate, 0, 255)
sample = poisson(datum=candidate, time=time)
inpts = {'X': sample}
# Run the network on the input.
network.run(inpts=inpts, time=time)
# Check for incorrect classification.
s = spikes['Y'].get('s').view(1, n_neurons, time)
prediction = ngram(spikes=s,
ngram_scores=ngram_scores,
n_labels=10,
n=2).item()
# Optionally plot various simulation information.
if plot:
_input = original.view(side_length, side_length)
reconstruction = candidate.view(side_length, side_length)
_spikes = {
'X': spikes['X'].get('s').view(side_length**2, time),
'Y': spikes['Y'].get('s').view(n_neurons, time)
}
w = network.connections['X', 'Y'].w
spike_ims, spike_axes = plot_spikes(spikes=_spikes,
ims=spike_ims,
axes=spike_axes)
weights_im = plot_locally_connected_weights(w,
n_filters,
kernel_size,
conv_size,
locations,
side_length,
im=weights_im)
inpt_axes, inpt_ims = plot_input(_input,
reconstruction,
label=labels[i],
ims=inpt_ims,
axes=inpt_axes)
plt.pause(1e-8)
if prediction == label:
print('Attack failed.')
else:
print('Attack succeeded.')
adv_example = candidate
j += 1
network.reset_() # Reset state variables.
print('\nAdversarial attack complete.\n')
def main(seed=0, n_neurons=100, n_train=60000, n_test=10000, inhib=250, time=50, lr=1e-2, lr_decay=0.99, dt=1,
theta_plus=0.05, theta_decay=1e-7, progress_interval=10, update_interval=250, train=True, plot=False,
gpu=False):
assert n_train % update_interval == 0 and n_test % update_interval == 0, \
'No. examples must be divisible by update_interval'
params = [
seed, n_neurons, n_train, inhib, time, lr, lr_decay,
theta_plus, theta_decay, progress_interval, update_interval
]
test_params = [
seed, n_neurons, n_train, n_test, inhib, time, lr, lr_decay,
theta_plus, theta_decay, progress_interval, update_interval
]
model_name = '_'.join([str(x) for x in params])
np.random.seed(seed)
if gpu:
torch.set_default_tensor_type('torch.cuda.FloatTensor')
torch.cuda.manual_seed_all(seed)
else:
torch.manual_seed(seed)
if train:
n_examples = n_train
else:
n_examples = n_test
n_sqrt = int(np.ceil(np.sqrt(n_neurons)))
n_classes = 10
# Build network.
if train:
network = Network(dt=dt)
input_layer = RealInput(n=784, traces=True, trace_tc=5e-2)
network.add_layer(input_layer, name='X')
output_layer = DiehlAndCookNodes(
n=n_neurons, traces=True, rest=0, reset=0, thresh=1, refrac=0,
decay=1e-2, trace_tc=5e-2, theta_plus=theta_plus, theta_decay=theta_decay
)
network.add_layer(output_layer, name='Y')
w = 0.3 * torch.rand(784, n_neurons)
input_connection = Connection(
source=network.layers['X'], target=network.layers['Y'], w=w, update_rule=PostPre,
nu=[0, lr], wmin=0, wmax=1, norm=78.4
)
network.add_connection(input_connection, source='X', target='Y')
w = -inhib * (torch.ones(n_neurons, n_neurons) - torch.diag(torch.ones(n_neurons)))
recurrent_connection = Connection(
source=network.layers['Y'], target=network.layers['Y'], w=w, wmin=-inhib, wmax=0
)
network.add_connection(recurrent_connection, source='Y', target='Y')
else:
path = os.path.join('..', '..', 'params', data, model)
network = load_network(os.path.join(path, model_name + '.pt'))
network.connections['X', 'Y'].update_rule = NoOp(
connection=network.connections['X', 'Y'], nu=network.connections['X', 'Y'].nu
)
network.layers['Y'].theta_decay = 0
network.layers['Y'].theta_plus = 0
# Load MNIST data.
dataset = MNIST(path=os.path.join('..', '..', 'data', 'MNIST'), download=True)
if train:
images, labels = dataset.get_train()
else:
images, labels = dataset.get_test()
images = images.view(-1, 784)
images = images / 255
# if train:
# for i in range(n_neurons):
# network.connections['X', 'Y'].w[:, i] = images[i] + images[i].mean() * torch.randn(784)
# Record spikes during the simulation.
spike_record = torch.zeros(update_interval, time, n_neurons)
# Neuron assignments and spike proportions.
if train:
assignments = -torch.ones_like(torch.Tensor(n_neurons))
proportions = torch.zeros_like(torch.Tensor(n_neurons, n_classes))
rates = torch.zeros_like(torch.Tensor(n_neurons, n_classes))
ngram_scores = {}
else:
path = os.path.join('..', '..', 'params', data, model)
path = os.path.join(path, '_'.join(['auxiliary', model_name]) + '.pt')
assignments, proportions, rates, ngram_scores = torch.load(open(path, 'rb'))
# Sequence of accuracy estimates.
curves = {'all': [], 'proportion': [], 'ngram': []}
if train:
best_accuracy = 0
spikes = {}
for layer in set(network.layers):
spikes[layer] = Monitor(network.layers[layer], state_vars=['s'], time=time)
network.add_monitor(spikes[layer], name='%s_spikes' % layer)
# Train the network.
if train:
print('\nBegin training.\n')
else:
print('\nBegin test.\n')
inpt_axes = None
inpt_ims = None
spike_ims = None
spike_axes = None
weights_im = None
assigns_im = None
perf_ax = None
start = t()
for i in range(n_examples):
if i % progress_interval == 0:
print(f'Progress: {i} / {n_examples} ({t() - start:.4f} seconds)')
start = t()
if i % update_interval == 0 and i > 0:
if train:
network.connections['X', 'Y'].update_rule.nu[1] *= lr_decay
if i % len(labels) == 0:
current_labels = labels[-update_interval:]
else:
current_labels = labels[i % len(images) - update_interval:i % len(images)]
# Update and print accuracy evaluations.
curves, predictions = update_curves(
curves, current_labels, n_classes, spike_record=spike_record, assignments=assignments,
proportions=proportions, ngram_scores=ngram_scores, n=2
)
print_results(curves)
if train:
if any([x[-1] > best_accuracy for x in curves.values()]):
print('New best accuracy! Saving network parameters to disk.')
# Save network to disk.
path = os.path.join('..', '..', 'params', data, model)
if not os.path.isdir(path):
os.makedirs(path)
network.save(os.path.join(path, model_name + '.pt'))
path = os.path.join(path, '_'.join(['auxiliary', model_name]) + '.pt')
torch.save((assignments, proportions, rates, ngram_scores), open(path, 'wb'))
best_accuracy = max([x[-1] for x in curves.values()])
# Assign labels to excitatory layer neurons.
assignments, proportions, rates = assign_labels(spike_record, current_labels, n_classes, rates)
# Compute ngram scores.
ngram_scores = update_ngram_scores(spike_record, current_labels, n_classes, 2, ngram_scores)
print()
# Get next input sample.
image = images[i % n_examples].repeat([time, 1])
inpts = {'X': image}
# Run the network on the input.
network.run(inpts=inpts, time=time)
retries = 0
while spikes['Y'].get('s').sum() < 5 and retries < 3:
retries += 1
image *= 2
inpts = {'X': image}
network.run(inpts=inpts, time=time)
# Add to spikes recording.
spike_record[i % update_interval] = spikes['Y'].get('s').t()
# Optionally plot various simulation information.
if plot and i % update_interval == 0:
_input = images[i % n_examples].view(28, 28)
reconstruction = inpts['X'].view(time, 784).sum(0).view(28, 28)
_spikes = {layer: spikes[layer].get('s') for layer in spikes}
input_exc_weights = network.connections['X', 'Y'].w
square_weights = get_square_weights(input_exc_weights.view(784, n_neurons), n_sqrt, 28)
square_assignments = get_square_assignments(assignments, n_sqrt)
# inpt_axes, inpt_ims = plot_input(_input, reconstruction, label=labels[i], axes=inpt_axes, ims=inpt_ims)
spike_ims, spike_axes = plot_spikes(_spikes, ims=spike_ims, axes=spike_axes)
weights_im = plot_weights(square_weights, im=weights_im)
# assigns_im = plot_assignments(square_assignments, im=assigns_im)
# perf_ax = plot_performance(curves, ax=perf_ax)
plt.pause(1e-1)
network.reset_() # Reset state variables.
print(f'Progress: {n_examples} / {n_examples} ({t() - start:.4f} seconds)')
i += 1
if i % len(labels) == 0:
current_labels = labels[-update_interval:]
else:
current_labels = labels[i % len(images) - update_interval:i % len(images)]
# Update and print accuracy evaluations.
curves, predictions = update_curves(
curves, current_labels, n_classes, spike_record=spike_record, assignments=assignments,
proportions=proportions, ngram_scores=ngram_scores, n=2
)
print_results(curves)
if train:
if any([x[-1] > best_accuracy for x in curves.values()]):
print('New best accuracy! Saving network parameters to disk.')
# Save network to disk.
if train:
path = os.path.join('..', '..', 'params', data, model)
if not os.path.isdir(path):
os.makedirs(path)
network.save(os.path.join(path, model_name + '.pt'))
path = os.path.join(path, '_'.join(['auxiliary', model_name]) + '.pt')
torch.save((assignments, proportions, rates, ngram_scores), open(path, 'wb'))
if train:
print('\nTraining complete.\n')
else:
print('\nTest complete.\n')
print('Average accuracies:\n')
for scheme in curves.keys():
print('\t%s: %.2f' % (scheme, float(np.mean(curves[scheme]))))
# Save accuracy curves to disk.
path = os.path.join('..', '..', 'curves', data, model)
if not os.path.isdir(path):
os.makedirs(path)
if train:
to_write = ['train'] + params
else:
to_write = ['test'] + params
to_write = [str(x) for x in to_write]
f = '_'.join(to_write) + '.pt'
torch.save((curves, update_interval, n_examples), open(os.path.join(path, f), 'wb'))
# Save results to disk.
path = os.path.join('..', '..', 'results', data, model)
if not os.path.isdir(path):
os.makedirs(path)
results = [
np.mean(curves['all']), np.mean(curves['proportion']), np.mean(curves['ngram']),
np.max(curves['all']), np.max(curves['proportion']), np.max(curves['ngram'])
]
if train:
to_write = params + results
else:
to_write = test_params + results
to_write = [str(x) for x in to_write]
if train:
name = 'train.csv'
else:
name = 'test.csv'
if not os.path.isfile(os.path.join(path, name)):
with open(os.path.join(path, name), 'w') as f:
if train:
f.write('random_seed,n_neurons,n_train,inhib,time,lr,lr_decay,theta_plus,theta_decay,'
'progress_interval,update_interval,mean_all_activity,mean_proportion_weighting,'
'mean_ngram,max_all_activity,max_proportion_weighting,max_ngram\n')
else:
f.write('random_seed,n_neurons,n_train,n_test,inhib,time,lr,lr_decay,theta_plus,theta_decay,'
'progress_interval,update_interval,mean_all_activity,mean_proportion_weighting,'
'mean_ngram,max_all_activity,max_proportion_weighting,max_ngram\n')
with open(os.path.join(path, name), 'a') as f:
f.write(','.join(to_write) + '\n')
def main(seed=0,
n_train=60000,
n_test=10000,
inhib=250,
kernel_size=(16, ),
stride=(2, ),
time=100,
n_filters=25,
crop=0,
lr=1e-2,
lr_decay=0.99,
dt=1,
theta_plus=0.05,
theta_decay=1e-7,
intensity=5,
norm=0.2,
progress_interval=10,
update_interval=250,
train=True,
plot=False,
gpu=False):
assert n_train % update_interval == 0 and n_test % update_interval == 0, \
'No. examples must be divisible by update_interval'
params = [
seed, kernel_size, stride, n_filters, crop, lr, lr_decay, n_train,
inhib, time, dt, theta_plus, theta_decay, intensity, norm,
progress_interval, update_interval
]
model_name = '_'.join([str(x) for x in params])
if not train:
test_params = [
seed, kernel_size, stride, n_filters, crop, lr, lr_decay, n_train,
n_test, inhib, time, dt, theta_plus, theta_decay, intensity, norm,
progress_interval, update_interval
]
np.random.seed(seed)
if gpu:
torch.set_default_tensor_type('torch.cuda.FloatTensor')
torch.cuda.manual_seed_all(seed)
else:
torch.manual_seed(seed)
side_length = 28 - crop * 2
n_inpt = side_length**2
n_examples = n_train if train else n_test
n_classes = 10
# Build network.
if train:
network = LocallyConnectedNetwork(
n_inpt=n_inpt,
input_shape=[side_length, side_length],
kernel_size=kernel_size,
stride=stride,
n_filters=n_filters,
inh=inhib,
dt=dt,
nu=[0, lr],
theta_plus=theta_plus,
theta_decay=theta_decay,
wmin=0.0,
wmax=1.0,
norm=norm)
else:
network = load_network(os.path.join(params_path, model_name + '.pt'))
network.connections['X', 'Y'].update_rule = NoOp(
connection=network.connections['X', 'Y'],
nu=network.connections['X', 'Y'].nu)
network.layers['Y'].theta_decay = 0
network.layers['Y'].theta_plus = 0
conv_size = network.connections['X', 'Y'].conv_size
locations = network.connections['X', 'Y'].locations
conv_prod = int(np.prod(conv_size))
n_neurons = n_filters * conv_prod
# Voltage recording for excitatory and inhibitory layers.
voltage_monitor = Monitor(network.layers['Y'], ['v'], time=time)
network.add_monitor(voltage_monitor, name='output_voltage')
# Load MNIST data.
dataset = MNIST(path=data_path, download=True)
if train:
images, labels = dataset.get_train()
else:
images, labels = dataset.get_test()
images *= intensity
images = images[:, crop:-crop, crop:-crop]
# Record spikes during the simulation.
spike_record = torch.zeros(update_interval, time, n_neurons)
full_spike_record = torch.zeros(n_examples, n_neurons)
# Neuron assignments and spike proportions.
if train:
logreg_model = LogisticRegression(warm_start=True,
n_jobs=-1,
solver='lbfgs',
max_iter=1000,
multi_class='multinomial')
logreg_model.coef_ = np.zeros([n_classes, n_neurons])
logreg_model.intercept_ = np.zeros(n_classes)
logreg_model.classes_ = np.arange(n_classes)
else:
path = os.path.join(params_path,
'_'.join(['auxiliary', model_name]) + '.pt')
logreg_coef, logreg_intercept = torch.load(open(path, 'rb'))
logreg_model = LogisticRegression(warm_start=True,
n_jobs=-1,
solver='lbfgs',
max_iter=1000,
multi_class='multinomial')
logreg_model.coef_ = logreg_coef
logreg_model.intercept_ = logreg_intercept
logreg_model.classes_ = np.arange(n_classes)
if train:
best_accuracy = 0
# Sequence of accuracy estimates.
curves = {'logreg': []}
predictions = {scheme: torch.Tensor().long() for scheme in curves.keys()}
spikes = {}
for layer in set(network.layers):
spikes[layer] = Monitor(network.layers[layer],
state_vars=['s'],
time=time)
network.add_monitor(spikes[layer], name=f'{layer}_spikes')
# Train the network.
if train:
print('\nBegin training.\n')
else:
print('\nBegin test.\n')
spike_ims = None
spike_axes = None
weights_im = None
weights2_im = None
start = t()
for i in range(n_examples):
if i % progress_interval == 0:
print(f'Progress: {i} / {n_examples} ({t() - start:.4f} seconds)')
start = t()
if i % update_interval == 0 and i > 0:
if i % len(labels) == 0:
current_labels = labels[-update_interval:]
current_record = full_spike_record[-update_interval:]
else:
current_labels = labels[i % len(labels) - update_interval:i %
len(labels)]
current_record = full_spike_record[i % len(labels) -
update_interval:i %
len(labels)]
# Update and print accuracy evaluations.
curves, preds = update_curves(curves,
current_labels,
n_classes,
full_spike_record=current_record,
logreg=logreg_model)
print_results(curves)
for scheme in preds:
predictions[scheme] = torch.cat(
[predictions[scheme], preds[scheme]], -1)
# Save accuracy curves to disk.
to_write = ['train'] + params if train else ['test'] + params
f = '_'.join([str(x) for x in to_write]) + '.pt'
torch.save((curves, update_interval, n_examples),
open(os.path.join(curves_path, f), 'wb'))
if train:
if any([x[-1] > best_accuracy for x in curves.values()]):
print(
'New best accuracy! Saving network parameters to disk.'
)
# Save network to disk.
network.save(os.path.join(params_path, model_name + '.pt'))
path = os.path.join(
params_path,
'_'.join(['auxiliary', model_name]) + '.pt')
torch.save((logreg_model.coef_, logreg_model.intercept_),
open(path, 'wb'))
best_accuracy = max([x[-1] for x in curves.values()])
# Refit logistic regression model.
logreg_model = logreg_fit(full_spike_record[:i], labels[:i],
logreg_model)
print()
# Get next input sample.
image = images[i % len(images)].contiguous().view(-1)
sample = bernoulli(datum=image, time=time, dt=dt)
inpts = {'X': sample}
# Run the network on the input.
network.run(inpts=inpts, time=time)
retries = 0
while spikes['Y'].get('s').sum() < 5 and retries < 3:
retries += 1
image *= 2
sample = bernoulli(datum=image, time=time, dt=dt)
inpts = {'X': sample}
network.run(inpts=inpts, time=time)
# Add to spikes recording.
spike_record[i % update_interval] = spikes['Y'].get('s').view(time, -1)
full_spike_record[i] = spikes['Y'].get('s').view(time, -1).sum(0)
if plot:
# Optionally plot various simulation information.
_spikes = {
'X': spikes['X'].get('s').view(side_length**2, time),
'Y': spikes['Y'].get('s').view(n_filters * conv_prod, time)
}
spike_ims, spike_axes = plot_spikes(spikes=_spikes,
ims=spike_ims,
axes=spike_axes)
weights_im = plot_locally_connected_weights(
network.connections[('X', 'Y')].w,
n_filters,
kernel_size,
conv_size,
locations,
side_length,
im=weights_im)
weights2_im = plot_weights(logreg_model.coef_, im=weights2_im)
plt.pause(1e-8)
network.reset_() # Reset state variables.
print(f'Progress: {n_examples} / {n_examples} ({t() - start:.4f} seconds)')
i += 1
if i % len(labels) == 0:
current_labels = labels[-update_interval:]
current_record = full_spike_record[-update_interval:]
else:
current_labels = labels[i % len(labels) - update_interval:i %
len(labels)]
current_record = full_spike_record[i % len(labels) -
update_interval:i % len(labels)]
# Update and print accuracy evaluations.
curves, preds = update_curves(curves,
current_labels,
n_classes,
full_spike_record=current_record,
logreg=logreg_model)
print_results(curves)
for scheme in preds:
predictions[scheme] = torch.cat([predictions[scheme], preds[scheme]],
-1)
if train:
if any([x[-1] > best_accuracy for x in curves.values()]):
print('New best accuracy! Saving network parameters to disk.')
# Save network to disk.
network.save(os.path.join(params_path, model_name + '.pt'))
path = os.path.join(params_path,
'_'.join(['auxiliary', model_name]) + '.pt')
torch.save((logreg_model.coef_, logreg_model.intercept_),
open(path, 'wb'))
if train:
print('\nTraining complete.\n')
else:
print('\nTest complete.\n')
print('Average accuracies:\n')
for scheme in curves.keys():
print('\t%s: %.2f' % (scheme, float(np.mean(curves[scheme]))))
# Save accuracy curves to disk.
if train:
to_write = ['train'] + params
f = '_'.join([str(x) for x in to_write]) + '.pt'
torch.save((curves, update_interval, n_examples),
open(os.path.join(curves_path, f), 'wb'))
# Save results to disk.
results = [np.mean(curves['logreg']), np.std(curves['logreg'])]
to_write = params + results if train else test_params + results
to_write = [str(x) for x in to_write]
if train:
name = 'train.csv'
else:
name = 'test.csv'
if not os.path.isfile(os.path.join(results_path, name)):
with open(os.path.join(results_path, name), 'w') as f:
if train:
f.write(
'random_seed,kernel_size,stride,n_filters,crop,lr,lr_decay,n_train,inhib,time,timestep,theta_plus,'
'theta_decay,intensity,norm,progress_interval,update_interval,mean_logreg,std_logreg\n'
)
else:
f.write(
'random_seed,kernel_size,stride,n_filters,crop,lr,lr_decay,n_train,n_test,inhib,time,timestep,'
'theta_plus,theta_decay,intensity,norm,progress_interval,update_interval,mean_logreg,std_logreg\n'
)
with open(os.path.join(results_path, name), 'a') as f:
f.write(','.join(to_write) + '\n')
if labels.numel() > n_examples:
labels = labels[:n_examples]
else:
while labels.numel() < n_examples:
if 2 * labels.numel() > n_examples:
labels = torch.cat(
[labels, labels[:n_examples - labels.numel()]])
else:
labels = torch.cat([labels, labels])
# Compute confusion matrices and save them to disk.
confusions = {}
for scheme in predictions:
confusions[scheme] = confusion_matrix(labels, predictions[scheme])
to_write = ['train'] + params if train else ['test'] + test_params
f = '_'.join([str(x) for x in to_write]) + '.pt'
torch.save(confusions, os.path.join(confusion_path, f))
from bindsnet.datasets import MNIST
from bindsnet.encoding import poisson
from bindsnet.pipeline import Pipeline
from bindsnet.models import DiehlAndCook2015
from bindsnet.environment import DatasetEnvironment
# Build Diehl & Cook 2015 network.
network = DiehlAndCook2015(n_input=784, n_neurons=400, exc=22.5,
inh=17.5, dt=1.0, norm=78.4)
# Specify dataset wrapper environment.
environment = DatasetEnvironment(dataset=MNIST(path='../../data/MNIST', download=True),
train=True, intensity=0.25)
# Build pipeline from components.
pipeline = Pipeline(network=network, environment=environment, plot_type='line',
encoding=poisson, time=350, plot_interval=1)
# Train the network.
for i in range(60000):
pipeline.step()
network.reset_()
output_bias_connection = Connection(source=output_bias,
target=output_layer)
network.add_connection(input_connection, source='X', target='Y')
network.add_connection(hidden_bias_connection, source='Y_b', target='Y')
network.add_connection(hidden_connection, source='Y', target='Z')
network.add_connection(output_bias_connection, source='Z_b', target='Z')
# State variable monitoring.
for l in network.layers:
m = Monitor(network.layers[l], state_vars=['s'], time=time)
network.add_monitor(m, name=l)
else:
network = load_network(os.path.join(params_path, model_name + '.pt'))
# Load MNIST data.
dataset = MNIST(path=data_path, download=True, shuffle=True)
if train:
images, labels = dataset.get_train()
else:
images, labels = dataset.get_test()
images, labels = images[:n_examples], labels[:n_examples]
images, labels = iter(images.view(-1, 784) / 255), iter(labels)
grads = {}
accuracies = []
predictions = []
ground_truth = []
best = -np.inf
spike_ims, spike_axes, weights1_im, weights2_im = None, None, None, None
print(f'Using device = {str(device)}')
network = network.to(device)
if gpu and cuda_avail:
torch.cuda.manual_seed(seed)
else:
torch.manual_seed(seed)
# Get the dataset
print('Fetching the dataset')
kws = ['up', 'down', 'on', 'off', 'yes', 'no', 'go', 'stop']
data_path = './data'
train_data = MNIST(
RankOrderEncoder(time=time, dt=dt),
None,
root=os.path.join('.', 'mnist'),
download=True,
transform=transforms.Compose([
transforms.Resize(size=(22, 22)),
transforms.ToTensor(),
transforms.Lambda(lambda x: x * intensity)
])) if use_mnist else SpeechCommandsDataset(
data_path, download=True, kws=kws, unknown=True, silence=True)
valid_data = MNIST(RankOrderEncoder(time=time, dt=dt),
None,
root=os.path.join('.', 'mnist'),
download=False,
transform=transforms.Compose([
transforms.Resize(size=(22, 22)),
transforms.ToTensor(),
transforms.Lambda(lambda x: x * intensity)
]),
train=False) if use_mnist else SpeechCommandsDataset(
network.add_layer(input_layer, name='X')
network.add_layer(conv_layer, name='Y')
network.add_layer(conv_layer_, name='Y_')
network.add_layer(conv_layer2, name='Z')
network.add_connection(conv_conn, source='X', target='Y')
network.add_connection(conv_conn_, source='X', target='Y_')
network.add_connection(recurrent_conn, source='Y', target='Y')
network.add_connection(conv_conn2, source='Y', target='Z')
network.add_connection(recurrent_conn2, source='Z', target='Z')
# Voltage recording for excitatory and inhibitory layers.
voltage_monitor = Monitor(network.layers['Y'], ['v'], time=time)
network.add_monitor(voltage_monitor, name='output_voltage')
# Load MNIST data.
dataset = MNIST(path=os.path.join('..', '..', 'data', 'MNIST'), download=True)
if train:
images, labels = dataset.get_train()
else:
images, labels = dataset.get_test()
images *= intensity
# Record spikes during the simulation.
spike_record = torch.zeros(update_interval, time, n_neurons)
# Neuron assignments and spike proportions.
if train:
assignments = -torch.ones_like(torch.Tensor(n_neurons))
proportions = torch.zeros_like(torch.Tensor(n_neurons, n_classes))
network = DiehlAndCook2015(n_inpt=784,
n_neurons=n_neurons,
exc=exc,
inh=inh,
dt=dt,
norm=78.4,
theta_plus=1)
# Voltage recording for excitatory and inhibitory layers.
exc_voltage_monitor = Monitor(network.layers["Ae"], ["v"], time=time)
inh_voltage_monitor = Monitor(network.layers["Ai"], ["v"], time=time)
network.add_monitor(exc_voltage_monitor, name="exc_voltage")
network.add_monitor(inh_voltage_monitor, name="inh_voltage")
# Load MNIST data.
images, labels = MNIST(path=os.path.join("..", "..", "data", "MNIST"),
download=True).get_train()
images = images.view(-1, 784)
images *= intensity
# Lazily encode data as Poisson spike trains.
data_loader = poisson_loader(data=images, time=time, dt=dt)
# Record spikes during the simulation.
spike_record = torch.zeros(update_interval, time, n_neurons)
# Neuron assignments and spike proportions.
assignments = -torch.ones_like(torch.Tensor(n_neurons))
proportions = torch.zeros_like(torch.Tensor(n_neurons, 10))
rates = torch.zeros_like(torch.Tensor(n_neurons, 10))
# Sequence of accuracy estimates.
from bindsnet.datasets import MNIST from bindsnet.encoding import poisson from bindsnet.pipeline import Pipeline from bindsnet.models import DiehlAndCook2015 from bindsnet.environment import DatasetEnvironment # Build Diehl & Cook 2015 network. network = DiehlAndCook2015(n_inpt=784, n_neurons=400, exc=22.5, inh=17.5, dt=1.0, norm=78.4) # Specify dataset wrapper environment. environment = DatasetEnvironment(dataset=MNIST(path='../../data/MNIST'), train=True, download=True, intensity=0.25) # Build pipeline from components. pipeline = Pipeline(network=network, environment=environment, encoding=poisson, time=350, plot_interval=1) # Train the network. for i in range(60000): pipeline.step() network._reset()
def main(seed=0, n_epochs=5, batch_size=100, time=50, update_interval=50, plot=False):
np.random.seed(seed)
if torch.cuda.is_available():
torch.set_default_tensor_type('torch.cuda.FloatTensor')
torch.cuda.manual_seed_all(seed)
else:
torch.manual_seed(seed)
print()
print('Creating and training the ANN...')
print()
# Create and train an ANN on the MNIST dataset.
ANN = FullyConnectedNetwork()
# Get the MNIST data.
images, labels = MNIST('../../data/MNIST', download=True).get_train()
images /= images.max() # Standardizing to [0, 1].
images = images.view(-1, 784)
labels = labels.long()
# Specify optimizer and loss function.
optimizer = optim.Adam(params=ANN.parameters(), lr=1e-3)
criterion = nn.CrossEntropyLoss()
# Train the ANN.
batches_per_epoch = int(images.size(0) / batch_size)
for i in range(n_epochs):
losses = []
accuracies = []
for j in range(batches_per_epoch):
batch_idxs = torch.from_numpy(
np.random.choice(np.arange(images.size(0)), size=batch_size, replace=False)
)
im_batch = images[batch_idxs]
label_batch = labels[batch_idxs]
outputs = ANN.forward(im_batch)
loss = criterion(outputs, label_batch)
predictions = torch.max(outputs, 1)[1]
correct = (label_batch == predictions).sum().float() / batch_size
optimizer.zero_grad()
loss.backward()
optimizer.step()
losses.append(loss.item())
accuracies.append(correct.item())
print(f'Epoch: {i+1} / {n_epochs}; Loss: {np.mean(losses):.4f}; Accuracy: {np.mean(accuracies) * 100:.4f}')
print()
print('Converting ANN to SNN...')
# Do ANN to SNN conversion.
SNN = ann_to_snn(ANN, input_shape=(784,), data=images)
for l in SNN.layers:
if l != 'Input':
SNN.add_monitor(
Monitor(SNN.layers[l], state_vars=['s', 'v'], time=time), name=l
)
spike_ims = None
spike_axes = None
correct = []
print()
print('Testing SNN on MNIST data...')
print()
# Test SNN on MNIST data.
start = t()
for i in range(images.size(0)):
if i > 0 and i % update_interval == 0:
print(
f'Progress: {i} / {images.size(0)}; Elapsed: {t() - start:.4f}; Accuracy: {np.mean(correct) * 100:.4f}'
)
start = t()
SNN.run(inpts={'Input': images[i].repeat(time, 1, 1)}, time=time)
spikes = {layer: SNN.monitors[layer].get('s') for layer in SNN.monitors}
voltages = {layer: SNN.monitors[layer].get('v') for layer in SNN.monitors}
prediction = torch.softmax(voltages['5'].sum(1), 0).argmax()
correct.append((prediction == labels[i]).item())
SNN.reset_()
if plot:
spikes = {k: spikes[k].cpu() for k in spikes}
spike_ims, spike_axes = plot_spikes(spikes, ims=spike_ims, axes=spike_axes)
plt.pause(1e-3)
network = DiehlAndCook2015(n_inpt=784,
n_neurons=n_neurons,
exc=exc,
inh=inh,
dt=dt,
nu=[0, 1e-2],
norm=78.4)
# Voltage recording for excitatory and inhibitory layers.
exc_voltage_monitor = Monitor(network.layers['Ae'], ['v'], time=time)
inh_voltage_monitor = Monitor(network.layers['Ai'], ['v'], time=time)
network.add_monitor(exc_voltage_monitor, name='exc_voltage')
network.add_monitor(inh_voltage_monitor, name='inh_voltage')
# Load MNIST data.
images, labels = MNIST(path=os.path.join('..', '..', 'data', 'MNIST'),
download=True).get_train()
images = images.view(-1, 784)
images *= intensity
# Lazily encode data as Poisson spike trains.
data_loader = poisson_loader(data=images, time=time, dt=dt)
# Record spikes during the simulation.
spike_record = torch.zeros(update_interval, time, n_neurons)
# Neuron assignments and spike proportions.
assignments = -torch.ones_like(torch.Tensor(n_neurons))
proportions = torch.zeros_like(torch.Tensor(n_neurons, 10))
rates = torch.zeros_like(torch.Tensor(n_neurons, 10))
# Sequence of accuracy estimates.
norm=78.4,
nu=(1e-4, 1e-2),
theta_plus=theta_plus,
inpt_shape=(1, 28, 28),
)
# Directs network to GPU
if gpu:
network.to("cuda")
# Load MNIST data.
dataset = MNIST(
PoissonEncoder(time=time, dt=dt),
None,
"../../data/MNIST",
download=True,
transform=transforms.Compose(
[transforms.ToTensor(),
transforms.Lambda(lambda x: x * intensity)]),
)
# Neuron assignments and spike proportions.
n_classes = 10
assignments = -torch.ones(n_neurons, device=device)
proportions = torch.zeros((n_neurons, n_classes), device=device)
rates = torch.zeros((n_neurons, n_classes), device=device)
# Sequence of accuracy estimates.
accuracy = {"all": [], "proportion": []}
# Voltage recording for excitatory and inhibitory layers.
def main(args):
if args.gpu:
torch.cuda.manual_seed_all(args.seed)
else:
torch.manual_seed(args.seed)
conv_size = int(
(28 - args.kernel_size + 2 * args.padding) / args.stride) + 1
# Build network.
network = Network()
input_layer = Input(n=784, shape=(1, 28, 28), traces=True)
conv_layer = DiehlAndCookNodes(
n=args.n_filters * conv_size * conv_size,
shape=(args.n_filters, conv_size, conv_size),
traces=True,
)
conv_conn = Conv2dConnection(
input_layer,
conv_layer,
kernel_size=args.kernel_size,
stride=args.stride,
update_rule=PostPre,
norm=0.4 * args.kernel_size**2,
nu=[0, args.lr],
reduction=max_without_indices,
wmax=1.0,
)
w = torch.zeros(args.n_filters, conv_size, conv_size, args.n_filters,
conv_size, conv_size)
for fltr1 in range(args.n_filters):
for fltr2 in range(args.n_filters):
if fltr1 != fltr2:
for i in range(conv_size):
for j in range(conv_size):
w[fltr1, i, j, fltr2, i, j] = -100.0
w = w.view(args.n_filters * conv_size * conv_size,
args.n_filters * conv_size * conv_size)
recurrent_conn = Connection(conv_layer, conv_layer, w=w)
network.add_layer(input_layer, name="X")
network.add_layer(conv_layer, name="Y")
network.add_connection(conv_conn, source="X", target="Y")
network.add_connection(recurrent_conn, source="Y", target="Y")
# Voltage recording for excitatory and inhibitory layers.
voltage_monitor = Monitor(network.layers["Y"], ["v"], time=args.time)
network.add_monitor(voltage_monitor, name="output_voltage")
if args.gpu:
network.to("cuda")
# Load MNIST data.
train_dataset = MNIST(
PoissonEncoder(time=args.time, dt=args.dt),
None,
os.path.join(ROOT_DIR, "data", "MNIST"),
download=True,
train=True,
transform=transforms.Compose([
transforms.ToTensor(),
transforms.Lambda(lambda x: x * args.intensity)
]),
)
spikes = {}
for layer in set(network.layers):
spikes[layer] = Monitor(network.layers[layer],
state_vars=["s"],
time=args.time)
network.add_monitor(spikes[layer], name="%s_spikes" % layer)
voltages = {}
for layer in set(network.layers) - {"X"}:
voltages[layer] = Monitor(network.layers[layer],
state_vars=["v"],
time=args.time)
network.add_monitor(voltages[layer], name="%s_voltages" % layer)
# Train the network.
print("Begin training.\n")
start = time()
weights_im = None
for epoch in range(args.n_epochs):
if epoch % args.progress_interval == 0:
print("Progress: %d / %d (%.4f seconds)" %
(epoch, args.n_epochs, time() - start))
start = time()
train_dataloader = DataLoader(
train_dataset,
batch_size=args.batch_size,
shuffle=True,
num_workers=4,
pin_memory=args.gpu,
)
for step, batch in enumerate(tqdm(train_dataloader)):
# Get next input sample.
inpts = {"X": batch["encoded_image"]}
if args.gpu:
inpts = {k: v.cuda() for k, v in inpts.items()}
# Run the network on the input.
network.run(inpts=inpts, time=args.time, input_time_dim=0)
# Decay learning rate.
network.connections["X", "Y"].nu[1] *= 0.99
# Optionally plot various simulation information.
if args.plot:
weights = conv_conn.w
weights_im = plot_conv2d_weights(weights, im=weights_im)
plt.pause(1e-8)
network.reset_() # Reset state variables.
print("Progress: %d / %d (%.4f seconds)\n" %
(args.n_epochs, args.n_epochs, time() - start))
print("Training complete.\n")
def main(seed=0, n_neurons=100, n_train=60000, n_test=10000, inhib=100, lr=0.01, lr_decay=1, time=350, dt=1,
theta_plus=0.05, theta_decay=1e-7, progress_interval=10, update_interval=250, plot=False,
train=True, gpu=False):
assert n_train % update_interval == 0 and n_test % update_interval == 0, \
'No. examples must be divisible by update_interval'
params = [
seed, n_neurons, n_train, inhib, lr_decay, time, dt,
theta_plus, theta_decay, progress_interval, update_interval
]
model_name = '_'.join([str(x) for x in params])
np.random.seed(seed)
if gpu:
torch.set_default_tensor_type('torch.cuda.FloatTensor')
torch.cuda.manual_seed_all(seed)
else:
torch.manual_seed(seed)
n_examples = n_train if train else n_test
n_classes = 10
# Build network.
if train:
network = Network(dt=dt)
input_layer = Input(n=784, traces=True, trace_tc=5e-2)
network.add_layer(input_layer, name='X')
output_layer = DiehlAndCookNodes(
n=n_classes, rest=0, reset=1, thresh=1, decay=1e-2,
theta_plus=theta_plus, theta_decay=theta_decay, traces=True, trace_tc=5e-2
)
network.add_layer(output_layer, name='Y')
w = torch.rand(784, n_classes)
input_connection = Connection(
source=input_layer, target=output_layer, w=w,
update_rule=MSTDPET, nu=lr, wmin=0, wmax=1,
norm=78.4, tc_e_trace=0.1
)
network.add_connection(input_connection, source='X', target='Y')
else:
network = load(os.path.join(params_path, model_name + '.pt'))
network.connections['X', 'Y'].update_rule = NoOp(
connection=network.connections['X', 'Y'], nu=network.connections['X', 'Y'].nu
)
network.layers['Y'].theta_decay = torch.IntTensor([0])
network.layers['Y'].theta_plus = torch.IntTensor([0])
# Load MNIST data.
environment = MNISTEnvironment(
dataset=MNIST(root=data_path, download=True), train=train, time=time
)
# Create pipeline.
pipeline = Pipeline(
network=network, environment=environment, encoding=repeat,
action_function=select_spiked, output='Y', reward_delay=None
)
spikes = {}
for layer in set(network.layers):
spikes[layer] = Monitor(network.layers[layer], state_vars=('s',), time=time)
network.add_monitor(spikes[layer], name='%s_spikes' % layer)
if train:
network.add_monitor(Monitor(
network.connections['X', 'Y'].update_rule, state_vars=('tc_e_trace',), time=time
), 'X_Y_e_trace')
# Train the network.
if train:
print('\nBegin training.\n')
else:
print('\nBegin test.\n')
spike_ims = None
spike_axes = None
weights_im = None
elig_axes = None
elig_ims = None
start = t()
for i in range(n_examples):
if i % progress_interval == 0:
print(f'Progress: {i} / {n_examples} ({t() - start:.4f} seconds)')
start = t()
if i > 0 and train:
network.connections['X', 'Y'].update_rule.nu[1] *= lr_decay
# Run the network on the input.
# print("Example",i,"Results:")
# for j in range(time):
# result = pipeline.env_step()
# pipeline.step(result,a_plus=1, a_minus=0)
# print(result)
for j in range(time):
pipeline.train()
if not train:
_spikes = {layer: spikes[layer].get('s') for layer in spikes}
if plot:
_spikes = {layer: spikes[layer].get('s') for layer in spikes}
w = network.connections['X', 'Y'].w
square_weights = get_square_weights(w.view(784, n_classes), 4, 28)
spike_ims, spike_axes = plot_spikes(_spikes, ims=spike_ims, axes=spike_axes)
weights_im = plot_weights(square_weights, im=weights_im)
elig_ims, elig_axes = plot_voltages(
{'Y': network.monitors['X_Y_e_trace'].get('e_trace').view(-1, time)[1500:2000]},
plot_type='line', ims=elig_ims, axes=elig_axes
)
plt.pause(1e-8)
pipeline.reset_state_variables() # Reset state variables.
network.connections['X', 'Y'].update_rule.tc_e_trace = torch.zeros(784, n_classes)
print(f'Progress: {n_examples} / {n_examples} ({t() - start:.4f} seconds)')
if train:
network.save(os.path.join(params_path, model_name + '.pt'))
print('\nTraining complete.\n')
else:
print('\nTest complete.\n')
# End of network creation
# Monitors:
spikes = {}
for l in network.layers:
spikes[l] = Monitor(network.layers[l], ["s"], time=num_timesteps)
network.add_monitor(spikes[l], name="%s_spikes" % l)
# Data and initial encoding:
dataset = MNIST(
RampNoLeakTNNEncoder(time=num_timesteps, dt=1),
None,
root=os.path.join("..", "..", "data", "MNIST"),
download=True,
transform=transforms.Compose(
[transforms.ToTensor(), transforms.Lambda(lambda x: x * intensity)]
),
)
# Create a dataloader to iterate and batch data
dataloader = torch.utils.data.DataLoader(
dataset, batch_size=1, shuffle=True, num_workers=0, pin_memory=False
)
inpt_axes = None
inpt_ims = None
spike_axes = None
def main(args):
update_interval = args.update_steps * args.batch_size
# Sets up GPU use
torch.backends.cudnn.benchmark = False
if args.gpu and torch.cuda.is_available():
torch.cuda.manual_seed_all(args.seed)
else:
torch.manual_seed(args.seed)
# Determines number of workers to use
if args.n_workers == -1:
args.n_workers = args.gpu * 4 * torch.cuda.device_count()
n_sqrt = int(np.ceil(np.sqrt(args.n_neurons)))
if args.reduction == "sum":
reduction = torch.sum
elif args.reduction == "mean":
reduction = torch.mean
elif args.reduction == "max":
reduction = max_without_indices
else:
raise NotImplementedError
# Build network.
network = DiehlAndCook2015v2(
n_inpt=784,
n_neurons=args.n_neurons,
inh=args.inh,
dt=args.dt,
norm=78.4,
nu=(0.0, 1e-2),
reduction=reduction,
theta_plus=args.theta_plus,
inpt_shape=(1, 28, 28),
)
# Directs network to GPU.
if args.gpu:
network.to("cuda")
# Load MNIST data.
dataset = MNIST(
PoissonEncoder(time=args.time, dt=args.dt),
None,
root=os.path.join(ROOT_DIR, "data", "MNIST"),
download=True,
train=True,
transform=transforms.Compose([
transforms.ToTensor(),
transforms.Lambda(lambda x: x * args.intensity)
]),
)
dataset, valid_dataset = torch.utils.data.random_split(
dataset, [59000, 1000])
test_dataset = MNIST(
PoissonEncoder(time=args.time, dt=args.dt),
None,
root=os.path.join(ROOT_DIR, "data", "MNIST"),
download=True,
train=False,
transform=transforms.Compose([
transforms.ToTensor(),
transforms.Lambda(lambda x: x * args.intensity)
]),
)
# Neuron assignments and spike proportions.
n_classes = 10
assignments = -torch.ones(args.n_neurons)
proportions = torch.zeros(args.n_neurons, n_classes)
rates = torch.zeros(args.n_neurons, n_classes)
# Set up monitors for spikes and voltages
spikes = {}
for layer in set(network.layers):
spikes[layer] = Monitor(network.layers[layer],
state_vars=["s"],
time=args.time)
network.add_monitor(spikes[layer], name="%s_spikes" % layer)
weights_im = None
spike_ims, spike_axes = None, None
# Record spikes for length of update interval.
spike_record = torch.zeros(update_interval, args.time, args.n_neurons)
if os.path.isdir(args.log_dir):
shutil.rmtree(args.log_dir)
# Summary writer.
writer = SummaryWriter(log_dir=args.log_dir, flush_secs=60)
for epoch in range(args.n_epochs):
print(f"\nEpoch: {epoch}\n")
labels = []
# Get training data loader.
dataloader = DataLoader(
dataset=dataset,
batch_size=args.batch_size,
shuffle=True,
num_workers=args.n_workers,
pin_memory=args.gpu,
)
for step, batch in enumerate(dataloader):
print(f"Step: {step} / {len(dataloader)}")
global_step = 60000 * epoch + args.batch_size * step
if step % args.update_steps == 0 and step > 0:
# Disable learning.
network.train(False)
# Get test data loader.
valid_dataloader = DataLoader(
dataset=valid_dataset,
batch_size=args.test_batch_size,
shuffle=True,
num_workers=args.n_workers,
pin_memory=args.gpu,
)
test_labels = []
test_spike_record = torch.zeros(len(valid_dataset), args.time,
args.n_neurons)
t0 = time()
for test_step, test_batch in enumerate(valid_dataloader):
# Prep next input batch.
inpts = {"X": test_batch["encoded_image"]}
if args.gpu:
inpts = {k: v.cuda() for k, v in inpts.items()}
# Run the network on the input (inference mode).
network.run(inpts=inpts,
time=args.time,
one_step=args.one_step)
# Add to spikes recording.
s = spikes["Y"].get("s").permute((1, 0, 2))
test_spike_record[(test_step * args.test_batch_size
):(test_step * args.test_batch_size) +
s.size(0)] = s
# Plot simulation data.
if args.valid_plot:
input_exc_weights = network.connections["X", "Y"].w
square_weights = get_square_weights(
input_exc_weights.view(784, args.n_neurons),
n_sqrt, 28)
spikes_ = {
layer: spikes[layer].get("s")[:, 0]
for layer in spikes
}
spike_ims, spike_axes = plot_spikes(spikes_,
ims=spike_ims,
axes=spike_axes)
weights_im = plot_weights(square_weights,
im=weights_im)
plt.pause(1e-8)
# Reset state variables.
network.reset_()
test_labels.extend(test_batch["label"].tolist())
t1 = time() - t0
writer.add_scalar(tag="time/test",
scalar_value=t1,
global_step=global_step)
# Convert the list of labels into a tensor.
test_label_tensor = torch.tensor(test_labels)
# Get network predictions.
all_activity_pred = all_activity(
spikes=test_spike_record,
assignments=assignments,
n_labels=n_classes,
)
proportion_pred = proportion_weighting(
spikes=test_spike_record,
assignments=assignments,
proportions=proportions,
n_labels=n_classes,
)
writer.add_scalar(
tag="accuracy/valid/all vote",
scalar_value=100 * torch.mean(
(test_label_tensor.long()
== all_activity_pred).float()),
global_step=global_step,
)
writer.add_scalar(
tag="accuracy/valid/proportion weighting",
scalar_value=100 * torch.mean(
(test_label_tensor.long() == proportion_pred).float()),
global_step=global_step,
)
square_weights = get_square_weights(
network.connections["X", "Y"].w.view(784, args.n_neurons),
n_sqrt,
28,
)
img_tensor = colorize(square_weights, cmap="hot_r")
writer.add_image(
tag="weights",
img_tensor=img_tensor,
global_step=global_step,
dataformats="HWC",
)
# Convert the array of labels into a tensor
label_tensor = torch.tensor(labels)
# Get network predictions.
all_activity_pred = all_activity(spikes=spike_record,
assignments=assignments,
n_labels=n_classes)
proportion_pred = proportion_weighting(
spikes=spike_record,
assignments=assignments,
proportions=proportions,
n_labels=n_classes,
)
writer.add_scalar(
tag="accuracy/train/all vote",
scalar_value=100 * torch.mean(
(label_tensor.long() == all_activity_pred).float()),
global_step=global_step,
)
writer.add_scalar(
tag="accuracy/train/proportion weighting",
scalar_value=100 * torch.mean(
(label_tensor.long() == proportion_pred).float()),
global_step=global_step,
)
# Assign labels to excitatory layer neurons.
assignments, proportions, rates = assign_labels(
spikes=spike_record,
labels=label_tensor,
n_labels=n_classes,
rates=rates,
)
# Re-enable learning.
network.train(True)
labels = []
labels.extend(batch["label"].tolist())
# Prep next input batch.
inpts = {"X": batch["encoded_image"]}
if args.gpu:
inpts = {k: v.cuda() for k, v in inpts.items()}
# Run the network on the input (training mode).
t0 = time()
network.run(inpts=inpts, time=args.time, one_step=args.one_step)
t1 = time() - t0
writer.add_scalar(tag="time/train/step",
scalar_value=t1,
global_step=global_step)
# Add to spikes recording.
s = spikes["Y"].get("s").permute((1, 0, 2))
spike_record[(step * args.batch_size) %
update_interval:(step * args.batch_size %
update_interval) + s.size(0)] = s
# Plot simulation data.
if args.plot:
input_exc_weights = network.connections["X", "Y"].w
square_weights = get_square_weights(
input_exc_weights.view(784, args.n_neurons), n_sqrt, 28)
spikes_ = {
layer: spikes[layer].get("s")[:, 0]
for layer in spikes
}
spike_ims, spike_axes = plot_spikes(spikes_,
ims=spike_ims,
axes=spike_axes)
weights_im = plot_weights(square_weights, im=weights_im)
plt.pause(1e-8)
# Reset state variables.
network.reset_()
def main(args):
if args.update_steps is None:
args.update_steps = max(
250 // args.batch_size, 1
) #Its value is 16 # why is it always multiplied with step? #update_steps is how many batch to classify before updating the graphs
update_interval = args.update_steps * args.batch_size # Value is 240 #update_interval is how many pictures to classify before updating the graphs
# Sets up GPU use
torch.backends.cudnn.benchmark = False
if args.gpu and torch.cuda.is_available():
torch.cuda.manual_seed_all(
args.seed
) #to enable reproducability of the code to get the same result
else:
torch.manual_seed(args.seed)
# Determines number of workers to use
if args.n_workers == -1:
args.n_workers = args.gpu * 4 * torch.cuda.device_count()
n_sqrt = int(np.ceil(np.sqrt(args.n_neurons)))
if args.reduction == "sum": #could have used switch to improve performance
reduction = torch.sum #weight updates for the batch
elif args.reduction == "mean":
reduction = torch.mean
elif args.reduction == "max":
reduction = max_without_indices
else:
raise NotImplementedError
# Build network.
network = DiehlAndCook2015v2( #Changed here
n_inpt=784, # input dimensions are 28x28=784
n_neurons=args.n_neurons,
inh=args.inh,
dt=args.dt,
norm=78.4,
nu=(1e-4, 1e-2),
reduction=reduction,
theta_plus=args.theta_plus,
inpt_shape=(1, 28, 28),
)
# Directs network to GPU
if args.gpu:
network.to("cuda")
# Load MNIST data.
dataset = MNIST(
PoissonEncoder(time=args.time, dt=args.dt),
None,
root=os.path.join(ROOT_DIR, "data", "MNIST"),
download=True,
train=True,
transform=transforms.Compose( #Composes several transforms together
[
transforms.ToTensor(),
transforms.Lambda(lambda x: x * args.intensity)
]),
)
test_dataset = MNIST(
PoissonEncoder(time=args.time, dt=args.dt),
None,
root=os.path.join(ROOT_DIR, "data", "MNIST"),
download=True,
train=False,
transform=transforms.Compose([
transforms.ToTensor(),
transforms.Lambda(lambda x: x * args.intensity)
]),
)
# Neuron assignments and spike proportions.
n_classes = 10 #changed
assignments = -torch.ones(args.n_neurons) #assignments is set to -1
proportions = torch.zeros(args.n_neurons,
n_classes) #matrix of 100x10 filled with zeros
rates = torch.zeros(args.n_neurons,
n_classes) #matrix of 100x10 filled with zeros
# Set up monitors for spikes and voltages
spikes = {}
for layer in set(network.layers):
spikes[layer] = Monitor(
network.layers[layer], state_vars=["s"], time=args.time
) # Monitors: Records state variables of interest. obj:An object to record state variables from during network simulation.
network.add_monitor(
spikes[layer], name="%s_spikes" % layer
) #state_vars: Iterable of strings indicating names of state variables to record.
#param time: If not ``None``, pre-allocate memory for state variable recording.
weights_im = None
spike_ims, spike_axes = None, None
# Record spikes for length of update interval.
spike_record = torch.zeros(update_interval, args.time, args.n_neurons)
if os.path.isdir(
args.log_dir): #checks if the path is a existing directory
shutil.rmtree(
args.log_dir) # is used to delete an entire directory tree
# Summary writer.
writer = SummaryWriter(
log_dir=args.log_dir, flush_secs=60
) #SummaryWriter: these utilities let you log PyTorch models and metrics into a directory for visualization
#flush_secs: in seconds, to flush the pending events and summaries to disk.
for epoch in range(args.n_epochs): #default is 1
print("\nEpoch: {epoch}\n")
labels = []
# Create a dataloader to iterate and batch data
dataloader = DataLoader( #It represents a Python iterable over a dataset
dataset,
batch_size=args.batch_size, #how many samples per batch to load
shuffle=
True, #set to True to have the data reshuffled at every epoch
num_workers=args.n_workers,
pin_memory=args.
gpu, #If True, the data loader will copy Tensors into CUDA pinned memory before returning them.
)
for step, batch in enumerate(
dataloader
): #Enumerate() method adds a counter to an iterable and returns it in a form of enumerate object
print("Step:", step)
global_step = 60000 * epoch + args.batch_size * step
if step % args.update_steps == 0 and step > 0:
# Convert the array of labels into a tensor
label_tensor = torch.tensor(labels)
# Get network predictions.
all_activity_pred = all_activity(spikes=spike_record,
assignments=assignments,
n_labels=n_classes)
proportion_pred = proportion_weighting(
spikes=spike_record,
assignments=assignments,
proportions=proportions,
n_labels=n_classes,
)
writer.add_scalar(
tag="accuracy/all vote",
scalar_value=torch.mean(
(label_tensor.long() == all_activity_pred).float()),
global_step=global_step,
)
#Vennila: Records the accuracies in each step
value = torch.mean(
(label_tensor.long() == all_activity_pred).float())
value = value.item()
accuracy.append(value)
print("ACCURACY:", value)
writer.add_scalar(
tag="accuracy/proportion weighting",
scalar_value=torch.mean(
(label_tensor.long() == proportion_pred).float()),
global_step=global_step,
)
writer.add_scalar(
tag="spikes/mean",
scalar_value=torch.mean(torch.sum(spike_record, dim=1)),
global_step=global_step,
)
square_weights = get_square_weights(
network.connections["X", "Y"].w.view(784, args.n_neurons),
n_sqrt,
28,
)
img_tensor = colorize(square_weights, cmap="hot_r")
writer.add_image(
tag="weights",
img_tensor=img_tensor,
global_step=global_step,
dataformats="HWC",
)
# Assign labels to excitatory layer neurons.
assignments, proportions, rates = assign_labels(
spikes=spike_record,
labels=label_tensor,
n_labels=n_classes,
rates=rates,
)
labels = []
labels.extend(
batch["label"].tolist()
) #for each batch or 16 pictures the labels of it is added to this list
# Prep next input batch.
inpts = {"X": batch["encoded_image"]}
if args.gpu:
inpts = {
k: v.cuda()
for k, v in inpts.items()
} #.cuda() is used to set up and run CUDA operations in the selected GPU
# Run the network on the input.
t0 = time()
network.run(inputs=inpts, time=args.time, one_step=args.one_step
) # Simulate network for given inputs and time.
t1 = time() - t0
# Add to spikes recording.
s = spikes["Y"].get("s").permute((1, 0, 2))
spike_record[(step * args.batch_size) %
update_interval:(step * args.batch_size %
update_interval) + s.size(0)] = s
writer.add_scalar(tag="time/simulation",
scalar_value=t1,
global_step=global_step)
# if(step==1):
# input_exc_weights = network.connections["X", "Y"].w
# an_array = input_exc_weights.detach().cpu().clone().numpy()
# #print(np.shape(an_array))
# data = asarray(an_array)
# savetxt('data.csv',data)
# print("Beginning weights saved")
# if(step==3749):
# input_exc_weights = network.connections["X", "Y"].w
# an_array = input_exc_weights.detach().cpu().clone().numpy()
# #print(np.shape(an_array))
# data2 = asarray(an_array)
# savetxt('data2.csv',data2)
# print("Ending weights saved")
# Plot simulation data.
if args.plot:
input_exc_weights = network.connections["X", "Y"].w
# print("Weights:",input_exc_weights)
square_weights = get_square_weights(
input_exc_weights.view(784, args.n_neurons), n_sqrt, 28)
spikes_ = {
layer: spikes[layer].get("s")[:, 0]
for layer in spikes
}
spike_ims, spike_axes = plot_spikes(spikes_,
ims=spike_ims,
axes=spike_axes)
weights_im = plot_weights(square_weights, im=weights_im)
plt.pause(1e-8)
# Reset state variables.
network.reset_state_variables()
print(end_accuracy()) #Vennila
def main(seed=0,
n_neurons=100,
n_train=60000,
n_test=10000,
inhib=250,
lr=1e-2,
lr_decay=1,
time=100,
dt=1,
theta_plus=0.05,
theta_decay=1e-7,
intensity=1,
progress_interval=10,
update_interval=100,
plot=False,
train=True,
gpu=False,
no_inhib=False,
no_theta=False):
assert n_train % update_interval == 0 and n_test % update_interval == 0, \
'No. examples must be divisible by update_interval'
params = [
seed, n_neurons, n_train, inhib, lr, lr_decay, time, dt, theta_plus,
theta_decay, intensity, progress_interval, update_interval
]
test_params = [
seed, n_neurons, n_train, n_test, inhib, lr, lr_decay, time, dt,
theta_plus, theta_decay, intensity, progress_interval, update_interval
]
model_name = '_'.join([str(x) for x in params])
np.random.seed(seed)
if gpu:
torch.set_default_tensor_type('torch.cuda.FloatTensor')
torch.cuda.manual_seed_all(seed)
else:
torch.manual_seed(seed)
n_examples = n_train if train else n_test
n_sqrt = int(np.ceil(np.sqrt(n_neurons)))
n_classes = 10
# Build network.
if train:
network = Network()
input_layer = Input(n=784, traces=True, trace_tc=5e-2)
network.add_layer(input_layer, name='X')
output_layer = DiehlAndCookNodes(n=n_neurons,
traces=True,
rest=0,
reset=0,
thresh=5,
refrac=0,
decay=1e-2,
trace_tc=5e-2,
theta_plus=theta_plus,
theta_decay=theta_decay)
network.add_layer(output_layer, name='Y')
w = 0.3 * torch.rand(784, n_neurons)
input_connection = Connection(source=network.layers['X'],
target=network.layers['Y'],
w=w,
update_rule=WeightDependentPostPre,
nu=[0, lr],
wmin=0,
wmax=1,
norm=78.4)
network.add_connection(input_connection, source='X', target='Y')
w = -inhib * (torch.ones(n_neurons, n_neurons) -
torch.diag(torch.ones(n_neurons)))
recurrent_connection = Connection(source=network.layers['Y'],
target=network.layers['Y'],
w=w,
wmin=-inhib,
wmax=0,
update_rule=WeightDependentPostPre,
nu=[0, -100 * lr],
norm=inhib / 2 * n_neurons)
network.add_connection(recurrent_connection, source='Y', target='Y')
mask = network.connections['Y', 'Y'].w == 0
masks = {('Y', 'Y'): mask}
else:
network = load_network(os.path.join(params_path, model_name + '.pt'))
network.connections['X', 'Y'].update_rule = NoOp(
connection=network.connections['X', 'Y'],
nu=network.connections['X', 'Y'].nu)
network.connections['Y', 'Y'].update_rule = NoOp(
connection=network.connections['X', 'Y'],
nu=network.connections['X', 'Y'].nu)
network.layers['Y'].theta_decay = 0
network.layers['Y'].theta_plus = 0
if no_inhib:
del network.connections['Y', 'Y']
if no_theta:
network.layers['Y'].theta = 0
# Load MNIST data.
dataset = MNIST(path=data_path, download=True)
if train:
images, labels = dataset.get_train()
else:
images, labels = dataset.get_test()
images = images.view(-1, 784)
images *= intensity
labels = labels.long()
monitors = {}
for layer in set(network.layers):
if 'v' in network.layers[layer].__dict__:
monitors[layer] = Monitor(network.layers[layer],
state_vars=['s', 'v'],
time=time)
else:
monitors[layer] = Monitor(network.layers[layer],
state_vars=['s'],
time=time)
network.add_monitor(monitors[layer], name=layer)
# Train the network.
if train:
print('\nBegin training.\n')
else:
print('\nBegin test.\n')
inpt_axes = None
inpt_ims = None
spike_ims = None
spike_axes = None
voltage_ims = None
voltage_axes = None
weights_im = None
weights2_im = None
unclamps = {}
per_class = int(n_neurons / n_classes)
for label in range(n_classes):
unclamp = torch.ones(n_neurons).byte()
unclamp[label * per_class:(label + 1) * per_class] = 0
unclamps[label] = unclamp
predictions = torch.zeros(n_examples)
corrects = torch.zeros(n_examples)
start = t()
for i in range(n_examples):
if i % progress_interval == 0:
print(f'Progress: {i} / {n_examples} ({t() - start:.4f} seconds)')
start = t()
if i % update_interval == 0 and i > 0 and train:
network.save(os.path.join(params_path, model_name + '.pt'))
network.connections['X', 'Y'].update_rule.nu[1] *= lr_decay
# Get next input sample.
image = images[i % len(images)]
label = labels[i % len(images)].item()
sample = poisson(datum=image, time=time, dt=dt)
inpts = {'X': sample}
# Run the network on the input.
if train:
network.run(inpts=inpts,
time=time,
unclamp={'Y': unclamps[label]},
masks=masks)
else:
network.run(inpts=inpts, time=time)
if not train:
retries = 0
while monitors['Y'].get('s').sum() == 0 and retries < 3:
retries += 1
image *= 1.5
sample = poisson(datum=image, time=time, dt=dt)
inpts = {'X': sample}
if train:
network.run(inpts=inpts,
time=time,
unclamp={'Y': unclamps[label]},
masks=masks)
else:
network.run(inpts=inpts, time=time)
output = monitors['Y'].get('s')
summed_neurons = output.sum(dim=1).view(n_classes, per_class)
summed_classes = summed_neurons.sum(dim=1)
prediction = torch.argmax(summed_classes).item()
correct = prediction == label
predictions[i] = prediction
corrects[i] = int(correct)
# Optionally plot various simulation information.
if plot:
# _input = image.view(28, 28)
# reconstruction = inpts['X'].view(time, 784).sum(0).view(28, 28)
# v = {'Y': monitors['Y'].get('v')}
s = {layer: monitors[layer].get('s') for layer in monitors}
input_exc_weights = network.connections['X', 'Y'].w
square_weights = get_square_weights(
input_exc_weights.view(784, n_neurons), n_sqrt, 28)
recurrent_weights = network.connections['Y', 'Y'].w
# inpt_axes, inpt_ims = plot_input(_input, reconstruction, label=labels[i], axes=inpt_axes, ims=inpt_ims)
# voltage_ims, voltage_axes = plot_voltages(v, ims=voltage_ims, axes=voltage_axes)
spike_ims, spike_axes = plot_spikes(s,
ims=spike_ims,
axes=spike_axes)
weights_im = plot_weights(square_weights, im=weights_im)
weights2_im = plot_weights(recurrent_weights,
im=weights2_im,
wmin=-inhib,
wmax=0)
plt.pause(1e-8)
network.reset_() # Reset state variables.
print(f'Progress: {n_examples} / {n_examples} ({t() - start:.4f} seconds)')
if train:
network.save(os.path.join(params_path, model_name + '.pt'))
if train:
print('\nTraining complete.\n')
else:
print('\nTest complete.\n')
accuracy = torch.mean(corrects).item() * 100
print(f'\nAccuracy: {accuracy}\n')
to_write = params + [accuracy] if train else test_params + [accuracy]
to_write = [str(x) for x in to_write]
name = 'train.csv' if train else 'test.csv'
if not os.path.isfile(os.path.join(results_path, name)):
with open(os.path.join(results_path, name), 'w') as f:
if train:
f.write(
'random_seed,n_neurons,n_train,inhib,lr,lr_decay,time,timestep,theta_plus,'
'theta_decay,intensity,progress_interval,update_interval,accuracy\n'
)
else:
f.write(
'random_seed,n_neurons,n_train,n_test,inhib,lr,lr_decay,time,timestep,'
'theta_plus,theta_decay,intensity,progress_interval,update_interval,accuracy\n'
)
with open(os.path.join(results_path, name), 'a') as f:
f.write(','.join(to_write) + '\n')
if labels.numel() > n_examples:
labels = labels[:n_examples]
else:
while labels.numel() < n_examples:
if 2 * labels.numel() > n_examples:
labels = torch.cat(
[labels, labels[:n_examples - labels.numel()]])
else:
labels = torch.cat([labels, labels])
# Compute confusion matrices and save them to disk.
confusion = confusion_matrix(labels, predictions)
to_write = ['train'] + params if train else ['test'] + test_params
f = '_'.join([str(x) for x in to_write]) + '.pt'
torch.save(confusion, os.path.join(confusion_path, f))
norm=78.4,
nu=[0.0, 1e-2],
inpt_shape=(1, 28, 28))
# Voltage recording for excitatory and inhibitory layers.
exc_voltage_monitor = Monitor(network.layers["Ae"], ["v"], time=time)
inh_voltage_monitor = Monitor(network.layers["Ai"], ["v"], time=time)
network.add_monitor(exc_voltage_monitor, name="exc_voltage")
network.add_monitor(inh_voltage_monitor, name="inh_voltage")
# Load MNIST data.
dataset = MNIST(
PoissonEncoder(time=time, dt=dt),
None,
root=os.path.join("..", "..", "data", "MNIST"),
download=True,
transform=transforms.Compose(
[transforms.ToTensor(),
transforms.Lambda(lambda x: x * intensity)]),
)
# Create a dataloader to iterate and batch data
dataloader = torch.utils.data.DataLoader(dataset,
batch_size=1,
shuffle=True,
num_workers=0,
pin_memory=gpu)
# Record spikes during the simulation.
spike_record = torch.zeros(update_interval, time, n_neurons)
def main(seed=0,
n_epochs=5,
batch_size=100,
time=50,
update_interval=50,
plot=False,
save=True):
np.random.seed(seed)
if torch.cuda.is_available():
torch.set_default_tensor_type('torch.cuda.FloatTensor')
torch.cuda.manual_seed_all(seed)
else:
torch.manual_seed(seed)
print()
print('Loading MNIST data...')
print()
# Get the CIFAR-10 data.
images, labels = MNIST('../../data/MNIST', download=True).get_train()
images /= images.max() # Standardizing to [0, 1].
images = images.view(-1, 784)
labels = labels.long()
test_images, test_labels = MNIST('../../data/MNIST',
download=True).get_test()
test_images /= test_images.max() # Standardizing to [0, 1].
test_images = test_images.view(-1, 784)
test_labels = test_labels.long()
if torch.cuda.is_available():
images = images.cuda()
labels = labels.cuda()
test_images = test_images.cuda()
test_labels = test_labels.cuda()
ANN = FullyConnectedNetwork()
model_name = '_'.join(
[str(x) for x in [seed, n_epochs, batch_size, time, update_interval]])
# Specify loss function.
criterion = nn.CrossEntropyLoss()
if save and os.path.isfile(os.path.join(params_path, model_name + '.pt')):
print()
print('Loading trained ANN from disk...')
ANN.load_state_dict(
torch.load(os.path.join(params_path, model_name + '.pt')))
if torch.cuda.is_available():
ANN = ANN.cuda()
else:
print()
print('Creating and training the ANN...')
print()
# Specify optimizer.
optimizer = optim.Adam(params=ANN.parameters(),
lr=1e-3,
weight_decay=1e-4)
batches_per_epoch = int(images.size(0) / batch_size)
# Train the ANN.
for i in range(n_epochs):
losses = []
accuracies = []
for j in range(batches_per_epoch):
batch_idxs = torch.from_numpy(
np.random.choice(np.arange(images.size(0)),
size=batch_size,
replace=False))
im_batch = images[batch_idxs]
label_batch = labels[batch_idxs]
outputs = ANN.forward(im_batch)
loss = criterion(outputs, label_batch)
predictions = torch.max(outputs, 1)[1]
correct = (label_batch
== predictions).sum().float() / batch_size
optimizer.zero_grad()
loss.backward()
optimizer.step()
losses.append(loss.item())
accuracies.append(correct.item() * 100)
outputs = ANN.forward(test_images)
loss = criterion(outputs, test_labels).item()
predictions = torch.max(outputs, 1)[1]
test_accuracy = ((test_labels == predictions).sum().float() /
test_labels.numel()).item() * 100
avg_loss = np.mean(losses)
avg_acc = np.mean(accuracies)
print(
f'Epoch: {i+1} / {n_epochs}; Train Loss: {avg_loss:.4f}; Train Accuracy: {avg_acc:.4f}'
)
print(
f'\tTest Loss: {loss:.4f}; Test Accuracy: {test_accuracy:.4f}')
if save:
torch.save(ANN.state_dict(),
os.path.join(params_path, model_name + '.pt'))
outputs = ANN.forward(test_images)
loss = criterion(outputs, test_labels)
predictions = torch.max(outputs, 1)[1]
accuracy = ((test_labels == predictions).sum().float() /
test_labels.numel()).item() * 100
print()
print(
f'(Post training) Test Loss: {loss:.4f}; Test Accuracy: {accuracy:.4f}'
)
print()
print('Evaluating ANN on adversarial examples from FSGM method...')
# Convert pytorch model to a tf_model and wrap it in cleverhans.
tf_model_fn = convert_pytorch_model_to_tf(ANN)
cleverhans_model = CallableModelWrapper(tf_model_fn, output_layer='logits')
sess = tf.Session()
x_op = tf.placeholder(tf.float32, shape=(
None,
784,
))
# Create an FGSM attack.
fgsm_op = FastGradientMethod(cleverhans_model, sess=sess)
fgsm_params = {'eps': 0.2, 'clip_min': 0.0, 'clip_max': 1.0}
adv_x_op = fgsm_op.generate(x_op, **fgsm_params)
adv_preds_op = tf_model_fn(adv_x_op)
# Run an evaluation of our model against FGSM white-box attack.
total = 0
correct = 0
adv_preds = sess.run(adv_preds_op, feed_dict={x_op: test_images})
correct += (np.argmax(adv_preds, axis=1) == test_labels).sum()
total += len(test_images)
accuracy = float(correct) / total
print()
print('Adversarial accuracy: {:.3f}'.format(accuracy * 100))
print()
print('Converting ANN to SNN...')
with sess.as_default():
test_images = adv_x_op.eval(feed_dict={x_op: test_images})
test_images = torch.tensor(test_images)
# Do ANN to SNN conversion.
SNN = ann_to_snn(ANN,
input_shape=(784, ),
data=test_images,
percentile=100)
for l in SNN.layers:
if l != 'Input':
SNN.add_monitor(Monitor(SNN.layers[l],
state_vars=['s', 'v'],
time=time),
name=l)
print()
print('Testing SNN on FGSM-modified MNIST data...')
print()
# Test SNN on MNIST data.
spike_ims = None
spike_axes = None
correct = []
n_images = test_images.size(0)
start = t()
for i in range(n_images):
if i > 0 and i % update_interval == 0:
accuracy = np.mean(correct) * 100
print(
f'Progress: {i} / {n_images}; Elapsed: {t() - start:.4f}; Accuracy: {accuracy:.4f}'
)
start = t()
SNN.run(inpts={'Input': test_images[i].repeat(time, 1, 1)}, time=time)
spikes = {
layer: SNN.monitors[layer].get('s')
for layer in SNN.monitors
}
voltages = {
layer: SNN.monitors[layer].get('v')
for layer in SNN.monitors
}
prediction = torch.softmax(voltages['fc3'].sum(1), 0).argmax()
correct.append((prediction == test_labels[i]).item())
SNN.reset_()
if plot:
spikes = {k: spikes[k].cpu() for k in spikes}
spike_ims, spike_axes = plot_spikes(spikes,
ims=spike_ims,
axes=spike_axes)
plt.pause(1e-3)
