def plot_conv_weights(): for f_idx in range(N_SIZE_FEATURES): for g_size in G_SIZES: conv_monitor = net.monitors[s2_name(f_idx, g_size)] w = conv_monitor.get("w") plt.ioff() print(f"showing {s2_name(f_idx, g_size)}") plot_conv2d_weights(w[0]) plot_conv2d_weights(w[-1]) plt.show()
if plot: image = batch["image"].view(28, 28) inpt = inputs["X"].view(time, 784).sum(0).view(28, 28) weights1 = conv_conn.w _spikes = { "X": spikes["X"].get("s").view(time, -1), "Y": spikes["Y"].get("s").view(time, -1), } _voltages = {"Y": voltages["Y"].get("v").view(time, -1)} inpt_axes, inpt_ims = plot_input(image, inpt, label=label, axes=inpt_axes, ims=inpt_ims) spike_ims, spike_axes = plot_spikes(_spikes, ims=spike_ims, axes=spike_axes) weights1_im = plot_conv2d_weights(weights1, im=weights1_im) voltage_ims, voltage_axes = plot_voltages(_voltages, ims=voltage_ims, axes=voltage_axes) plt.pause(1) network.reset_state_variables() # Reset state variables. print("Progress: %d / %d (%.4f seconds)\n" % (n_epochs, n_epochs, t() - start)) print("Training complete.\n")
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))
train_dataloader = DataLoader( train_dataset, batch_size=batch_size, shuffle=True, num_workers=4, pin_memory=gpu, ) for step, batch in enumerate(tqdm(train_dataloader)): # Get next input sample. inpts = {"X": batch["encoded_image"]} if gpu: inpts = {k: v.cuda() for k, v in inpts.items()} label = batch["label"] # Run the network on the input. network.run(inpts=inpts, time=time, input_time_dim=1) # Optionally plot various simulation information. if 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" % (n_epochs, n_epochs, t() - start)) print("Training complete.\n")
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")
inpt = inpts['X'].view(time, 784).sum(0).view(28, 28) weights1 = conv_conn.w _spikes = { 'X': spikes['X'].get('s').view(28**2, time), 'Y': spikes['Y'].get('s').view(n_filters * conv_size**2, time) } _voltages = { 'Y': voltages['Y'].get('v').view(n_filters * conv_size**2, time) } if i == 0: inpt_axes, inpt_ims = plot_input(images[i].view(28, 28), inpt, label=labels[i]) spike_ims, spike_axes = plot_spikes(_spikes) weights1_im = plot_conv2d_weights(weights1, wmax=conv_conn.wmax) voltage_ims, voltage_axes = plot_voltages(_voltages) else: inpt_axes, inpt_ims = plot_input(images[i].view(28, 28), inpt, label=labels[i], axes=inpt_axes, ims=inpt_ims) spike_ims, spike_axes = plot_spikes(_spikes, ims=spike_ims, axes=spike_axes) weights1_im = plot_conv2d_weights(weights1, im=weights1_im) voltage_ims, voltage_axes = plot_voltages(_voltages, ims=voltage_ims, axes=voltage_axes)
def main(model='diehl_and_cook_2015', data='mnist', param_string=None, cmap='hot_r', p_destroy=0, p_delete=0): assert param_string is not None, 'Pass "--param_string" argument on command line or main method.' f = os.path.join(ROOT_DIR, 'params', data, model, f'{param_string}.pt') if not os.path.isfile(f): print('File not found locally. Attempting download from swarm2 cluster.') download_params.main(model=model, data=data, param_string=param_string) network = torch.load(open(f, 'rb')) if data in ['mnist', 'letters', 'fashion_mnist']: if model in ['diehl_and_cook_2015', 'two_level_inhibition', 'real_dac', 'unclamp_dac']: params = param_string.split('_') n_sqrt = int(np.ceil(np.sqrt(int(params[1])))) side = int(np.sqrt(network.layers['X'].n)) w = network.connections['X', 'Y'].w w = get_square_weights(w, n_sqrt, side) plot_weights(w, cmap=cmap) elif model in ['conv']: w = network.connections['X', 'Y'].w plot_conv2d_weights(w, wmax=network.connections['X', 'Y'].wmax, cmap=cmap) elif model in ['fully_conv', 'locally_connected', 'crop_locally_connected', 'bern_crop_locally_connected']: params = param_string.split('_') kernel_size = int(params[1]) stride = int(params[2]) n_filters = int(params[3]) if model in ['crop_locally_connected', 'bern_crop_locally_connected']: crop = int(params[4]) side_length = 28 - crop * 2 else: side_length = 28 if kernel_size == side_length: conv_size = 1 else: conv_size = int((side_length - kernel_size) / stride) + 1 locations = torch.zeros(kernel_size, kernel_size, conv_size, conv_size).long() for c1 in range(conv_size): for c2 in range(conv_size): for k1 in range(kernel_size): for k2 in range(kernel_size): location = c1 * stride * side_length + c2 * stride + k1 * side_length + k2 locations[k1, k2, c1, c2] = location locations = locations.view(kernel_size ** 2, conv_size ** 2) w = network.connections[('X', 'Y')].w mask = torch.bernoulli(p_destroy * torch.ones(w.size())).byte() w[mask] = 0 mask = torch.bernoulli(p_delete * torch.ones(w.size(1))).byte() w[:, mask] = 0 plot_locally_connected_weights( w, n_filters, kernel_size, conv_size, locations, side_length, wmin=w.min(), wmax=w.max(), cmap=cmap ) elif model in ['backprop']: w = network.connections['X', 'Y'].w weights = [ w[:, i].view(28, 28) for i in range(10) ] w = torch.zeros(5 * 28, 2 * 28) for i in range(5): for j in range(2): w[i * 28: (i + 1) * 28, j * 28: (j + 1) * 28] = weights[i + j * 5] plot_weights(w, wmin=-1, wmax=1, cmap=cmap) elif model in ['two_layer_backprop']: params = param_string.split('_') sqrt = int(np.ceil(np.sqrt(int(params[1])))) w = network.connections['Y', 'Z'].w weights = [ w[:, i].view(sqrt, sqrt) for i in range(10) ] w = torch.zeros(5 * sqrt, 2 * sqrt) for i in range(5): for j in range(2): w[i * sqrt: (i + 1) * sqrt, j * sqrt: (j + 1) * sqrt] = weights[i + j * 5] plot_weights(w, wmin=-1, wmax=1, cmap=cmap) w = network.connections['X', 'Y'].w square_weights = get_square_weights(w, sqrt, 28) plot_weights(square_weights, wmin=-1, wmax=1, cmap=cmap) else: raise NotImplementedError('Weight plotting not implemented for this data, model combination.') elif data in ['breakout']: if model in ['crop', 'rebalance', 'two_level']: params = param_string.split('_') n_sqrt = int(np.ceil(np.sqrt(int(params[1])))) side = (50, 72) if model in ['crop', 'rebalance']: w = network.connections[('X', 'Ae')].w else: w = network.connections[('X', 'Y')].w w = get_square_weights(w, n_sqrt, side) plot_weights(w, cmap=cmap) elif model in ['backprop']: w = network.connections['X', 'Y'].w weights = [ w[:, i].view(50, 72) for i in range(4) ] w = torch.zeros(2 * 50, 2 * 72) for j in range(2): for k in range(2): w[j * 50: (j + 1) * 50, k * 72: (k + 1) * 72] = weights[j + k * 2] plot_weights(w, cmap=cmap) else: raise NotImplementedError('Weight plotting not implemented for this data, model combination.') if data in ['cifar10']: if model in ['backprop']: wmin = network.connections['X', 'Y'].wmin wmax = network.connections['X', 'Y'].wmax w = network.connections['X', 'Y'].w weights = [w[:, i].view(32, 32, 3).mean(2) for i in range(10)] w = torch.zeros(5 * 32, 2 * 32) for i in range(5): for j in range(2): w[i * 32: (i + 1) * 32, j * 32: (j + 1) * 32] = weights[i + j * 5] plot_weights(w, wmin=wmin, wmax=wmax, cmap=cmap) else: raise NotImplementedError('Weight plotting not implemented for this data, model combination.') path = os.path.join(ROOT_DIR, 'plots', data, model, 'weights') if not os.path.isdir(path): os.makedirs(path) plt.savefig(os.path.join(path, f'{param_string}.png')) plt.ioff() plt.show()
w2 = conv_conn2.w _spikes = { 'X': spikes['X'].get('s').view(28**2, 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(images[i].view(28, 28), _input, label=labels[i], ims=inpt_ims, axes=inpt_axes) spike_ims, spike_axes = plot_spikes(spikes=_spikes, ims=spike_ims, axes=spike_axes) weights1_im = plot_conv2d_weights(w1, im=weights1_im) weights2_im = plot_conv2d_weights(w2, 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:] else: current_labels = labels[i % len(images) - update_interval:i % len(images)]
for e in range(EPOCHS - 1): train(network, train_data) print("Add decision layers") add_decision_layers(network) train(network, train_data) network.save(TRAINED_NETWORK_PATH) else: network = load(TRAINED_NETWORK_PATH) print("Trained network loaded from file") for feature in FEATURES: for size in FILTER_SIZES: weights = network.monitors["conv%d%d" % (feature, size)].get("w") plot_conv2d_weights(weights[0], cmap='Greys') plt.show() # # for feature in FEATURES: # for size in FILTER_SIZES: # # voltages = network.monitors[get_s2_name(size, feature)].get("v") # # spikes = network.monitors[get_s2_name(size, feature)].get("s") # plot_voltages({"C2": voltages[-300: ]}) # plot_spikes({"C2": spikes[-300: ]}) voltages = network.monitors["OUT"].get("v") spikes = network.monitors["OUT"].get("s") plot_voltages({"Output": voltages}) plot_spikes({"output": spikes}) plt.show()
w = conv_conn.w _spikes = { 'X': spikes['X'].get('s').view(32 * 32 * 3, 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(images[i].view(32, 32, 3), reconstruction, label=labels[i], 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=0.1) 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.
def main(args): # Random seed. 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. if args.n_workers == -1: args.n_workers = args.gpu * 4 * torch.cuda.device_count() # Build network. network = bindsnet.network.Network(dt=args.dt, batch_size=args.batch_size) # Layers. input_layer = Input(shape=(1, 28, 28), traces=True) conv1_layer = LIFNodes(shape=(20, 24, 24), traces=True) pool1_layer = PassThroughNodes(shape=(20, 12, 12), traces=True) conv2_layer = LIFNodes(shape=(50, 8, 8), traces=True) pool2_layer = PassThroughNodes(shape=(50, 4, 4), traces=True) dense_layer = LIFNodes(shape=(200, ), traces=True) output_layer = LIFNodes(shape=(10, ), traces=True) network.add_layer(input_layer, name="I") network.add_layer(conv1_layer, name="C1") network.add_layer(pool1_layer, name="P1") network.add_layer(conv2_layer, name="C2") network.add_layer(pool2_layer, name="P2") network.add_layer(dense_layer, name="D") network.add_layer(output_layer, name="O") # Connections. conv1_connection = Conv2dConnection( source=input_layer, target=conv1_layer, update_rule=WeightDependentPost, nu=(0.0, args.nu), kernel_size=5, stride=1, wmin=-1.0, wmax=1.0, ) pool1_connection = SpatialPooling2dConnection(source=conv1_layer, target=pool1_layer, kernel_size=2, stride=2) conv2_connection = Conv2dConnection( source=pool1_layer, target=conv2_layer, update_rule=WeightDependentPost, nu=(0.0, args.nu), kernel_size=5, stride=1, wmin=-1.0, wmax=1.0, ) pool2_connection = SpatialPooling2dConnection(source=conv2_layer, target=pool2_layer, kernel_size=2, stride=2) dense_connection = Connection( source=pool2_layer, target=dense_layer, update_rule=WeightDependentPost, nu=(0.0, args.nu), wmin=-1.0, wmax=1.0, ) output_connection = Connection( source=dense_layer, target=output_layer, update_rule=WeightDependentPost, nu=(0.0, args.nu), wmin=-1.0, wmax=1.0, ) network.add_connection(connection=conv1_connection, source="I", target="C1") network.add_connection(connection=pool1_connection, source="C1", target="P1") network.add_connection(connection=conv2_connection, source="P1", target="C2") network.add_connection(connection=pool2_connection, source="C2", target="P2") network.add_connection(connection=dense_connection, source="P2", target="D") network.add_connection(connection=output_connection, source="D", target="O") # Monitors. for name, layer in network.layers.items(): monitor = Monitor(obj=layer, state_vars=("s", ), time=args.time) network.add_monitor(monitor=monitor, name=name) # 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(bindsnet.ROOT_DIR, "data", "MNIST"), download=True, transform=transforms.Compose([ transforms.ToTensor(), transforms.Lambda(lambda x: x * args.intensity) ]), ) spike_ims = None spike_axes = None conv1_weights_im = None conv2_weights_im = None dense_weights_im = None output_weights_im = None for epoch in range(args.n_epochs): # Create a dataloader to iterate over dataset. dataloader = DataLoader( dataset, batch_size=args.batch_size, shuffle=True, num_workers=args.n_workers, pin_memory=args.gpu, ) for step, batch in enumerate(tqdm(dataloader)): # Prep next input batch. inpts = {"I": 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) # Plot simulation data. if args.plot: spikes = {} for name, monitor in network.monitors.items(): spikes[name] = monitor.get("s")[:, 0].view(args.time, -1) spike_ims, spike_axes = plot_spikes(spikes, ims=spike_ims, axes=spike_axes) conv1_weights_im = plot_conv2d_weights(conv1_connection.w, im=conv1_weights_im, wmin=-1.0, wmax=1.0) conv2_weights_im = plot_conv2d_weights(conv2_connection.w, im=conv2_weights_im, wmin=-1.0, wmax=1.0) dense_weights_im = plot_weights(dense_connection.w, im=dense_weights_im, wmin=-1.0, wmax=1.0) output_weights_im = plot_weights(output_connection.w, im=output_weights_im, wmin=-1.0, wmax=1.0) plt.pause(1e-8) # Reset state variables. network.reset_()
def main(seed=0, n_train=60000, n_test=10000, kernel_size=16, stride=4, n_filters=25, padding=0, inhib=500, lr=0.01, lr_decay=0.99, time=50, dt=1, intensity=1, progress_interval=10, update_interval=250, train=True, plot=False, gpu=False): if gpu: torch.set_default_tensor_type('torch.cuda.FloatTensor') torch.cuda.manual_seed_all(seed) else: torch.manual_seed(seed) if not train: update_interval = n_test if kernel_size == 32: conv_size = 1 else: conv_size = int((32 - kernel_size + 2 * padding) / stride) + 1 per_class = int((n_filters * conv_size * conv_size) / 10) # Build network. network = Network() input_layer = Input(n=1024, shape=(1, 1, 32, 32), traces=True) conv_layer = DiehlAndCookNodes(n=n_filters * conv_size * conv_size, shape=(1, n_filters, conv_size, conv_size), traces=True) conv_conn = Conv2dConnection(input_layer, conv_layer, kernel_size=kernel_size, stride=stride, update_rule=PostPre, norm=0.4 * kernel_size**2, nu=[0, lr], wmin=0, wmax=1) w = -inhib * torch.ones(n_filters, conv_size, conv_size, n_filters, conv_size, conv_size) for f in range(n_filters): for i in range(conv_size): for j in range(conv_size): w[f, i, j, f, i, j] = 0 w = w.view(n_filters * conv_size**2, n_filters * conv_size**2) 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=time) network.add_monitor(voltage_monitor, name='output_voltage') # Load CIFAR-10 data. dataset = CIFAR10(path=os.path.join('..', '..', 'data', 'CIFAR10'), download=True) if train: images, labels = dataset.get_train() else: images, labels = dataset.get_test() images *= intensity images = images.mean(-1) # Lazily encode data as Poisson spike trains. data_loader = poisson_loader(data=images, time=time, dt=dt) 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) network.add_monitor(voltages[layer], name='%s_voltages' % layer) inpt_axes = None inpt_ims = None spike_ims = None spike_axes = None weights_im = None voltage_ims = None voltage_axes = None # Train the network. print('Begin training.\n') start = t() for i in range(n_train): if i % progress_interval == 0: print('Progress: %d / %d (%.4f seconds)' % (i, n_train, t() - start)) start = t() if train and i > 0: network.connections['X', 'Y'].nu[1] *= lr_decay # Get next input sample. sample = next(data_loader).unsqueeze(1).unsqueeze(1) inpts = {'X': sample} # Run the network on the input. network.run(inpts=inpts, time=time) # Optionally plot various simulation information. if plot: # inpt = inpts['X'].view(time, 1024).sum(0).view(32, 32) weights1 = conv_conn.w _spikes = { 'X': spikes['X'].get('s').view(32**2, time), 'Y': spikes['Y'].get('s').view(n_filters * conv_size**2, time) } _voltages = { 'Y': voltages['Y'].get('v').view(n_filters * conv_size**2, time) } # inpt_axes, inpt_ims = plot_input( # images[i].view(32, 32), inpt, label=labels[i], axes=inpt_axes, ims=inpt_ims # ) # voltage_ims, voltage_axes = plot_voltages(_voltages, ims=voltage_ims, axes=voltage_axes) spike_ims, spike_axes = plot_spikes(_spikes, ims=spike_ims, axes=spike_axes) weights_im = plot_conv2d_weights(weights1, im=weights_im) plt.pause(1e-8) network.reset_() # Reset state variables. print('Progress: %d / %d (%.4f seconds)\n' % (n_train, n_train, t() - start)) print('Training complete.\n')