def analyze_model_from_file(results, lesion=False, simulation=True):
#results = pickle.load(open(save_fn,'rb'))
# generate batch of batch_train_size - to run analyses on this batch
if lesion:
trial_type = 'spatialDMS'
updates = {'trial_type': trial_type}
update_parameters(updates)
stim = stimulus.Stimulus()
trial_info = stim.generate_trial(trial_type)
#trial_info=results['trial_info'] #from the saved file - not useful when running models with more than one task
#syn_x_stacked = np.stack(results['syn_x'], axis=1)
#syn_u_stacked = np.stack(results['syn_u'], axis=1)
h_stacked = np.stack(results['h'], axis=1)
trial_time = np.arange(0, h_stacked.shape[1] * par['dt'], par['dt'])
weights = results['weights']
updates = {}
for k, v in results['parameters'].items():
updates[k] = v
update_parameters(updates)
save_fn = par['save_dir'] + par['save_fn'] + '_DMS'
#lesion_results = lesion_weights_spatialcat(trial_info, h_stacked, syn_x_stacked, syn_u_stacked, weights)
lesion_results = lesion_weights_spatialcat(trial_info, h_stacked,
weights)
for key, val in lesion_results.items():
results[key] = val
pickle.dump(results, open(save_fn, 'wb'))
print('Analysis results saved in ', save_fn)
"""
Simulate the network for a new set of trials
"""
if simulation:
print('simulating network...')
h_stacked = np.stack(results['h'], axis=1)
trial_time = np.arange(0, h_stacked.shape[1] * par['dt'], par['dt'])
weights = results['weights']
updates = {}
for k, v in results['parameters'].items():
updates[k] = v
update_parameters(updates)
trial_type = 'spatial_cat'
updates = {'color_flag': False, 'trial_type': trial_type}
update_parameters(updates)
stim = stimulus.Stimulus()
trial_info_whiteonly = stim.generate_trial(trial_type)
save_fn = par['save_dir'] + 'color_flag_' + par['save_fn']
simulation_results = simulate_network_spatialcat(
trial_info_whiteonly, h_stacked, weights)
for key, val in simulation_results.items():
results[key] = val
pickle.dump(results, open(save_fn, 'wb'))
print('Analysis results saved in ', save_fn)
def main(gpu_id=None):
if gpu_id is not None:
os.environ['CUDA_VISIBLE_DEVICES'] = gpu_id
print_key_params()
tf.reset_default_graph()
x = tf.placeholder(tf.float32, [
par['num_time_steps'] * par['trials_per_seq'], par['batch_size'],
par['n_input']
], 'stim')
r = tf.placeholder(tf.float32, [
par['num_time_steps'] * par['trials_per_seq'], par['batch_size'],
par['n_output']
], 'reward')
m = tf.placeholder(
tf.float32,
[par['num_time_steps'] * par['trials_per_seq'], par['batch_size']],
'mask')
stim = stimulus.Stimulus()
gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=0.9
) # if gpu_id == '0' else tf.GPUOptions()
results_dict = {'reward_list': [], 'novel_reward_list': []}
with tf.Session(config=tf.ConfigProto(gpu_options=gpu_options)) as sess:
device = '/cpu:0' if gpu_id is None else '/gpu:0'
with tf.device(device):
model = Model(x, r, m)
sess.run(tf.global_variables_initializer())
reward_list = []
for i in range(par['n_iters']):
name, trial_info = stim.generate_trial(test_tasks=False)
_, reward, pol_loss, action, h, mask = \
sess.run([model.train, model.reward_full, model.pol_loss, model.action, model.h, model.mask], \
feed_dict={x:trial_info['neural_input'], r:trial_info['reward_data'],m:trial_info['train_mask']})
results_dict['reward_list'].append(reshape_reward(reward))
if i % 100 == 0:
name, trial_info = stim.generate_trial(test_tasks=True)
reward_novel = sess.run(model.reward_full, feed_dict={x:trial_info['neural_input'], \
r:trial_info['reward_data'], m:trial_info['train_mask']})
results_dict['novel_reward_list'].append(
reshape_reward(reward_novel))
print('Iter {:>4} | Reward: {:6.3f} | Reward novel: {:6.3f} | Pol. Loss: {:6.3f}'.format(\
i, np.mean(np.sum(reward, axis=0)), np.mean(np.sum(reward_novel, axis=0)),pol_loss))
weights = sess.run(model.var_dict)
results = {'weights': weights, 'results_dict': results_dict}
pickle.dump(results, open(par['save_fn'], 'wb'))
def main(gpu_id=None):
""" Run training """
# Isolate requested GPU
if gpu_id is not None:
os.environ['CUDA_VISIBLE_DEVICES'] = gpu_id
# Reset TensorFlow graph
tf.reset_default_graph()
# Define placeholders
x = tf.placeholder(
tf.float32, [par['num_time_steps'], par['batch_size'], par['n_input']],
'stim')
y = tf.placeholder(
tf.float32,
[par['num_time_steps'], par['batch_size'], par['n_output']], 'out')
m = tf.placeholder(tf.float32, [par['num_time_steps'], par['batch_size']],
'mask')
# Set up stimulus and recording
stim = stimulus.Stimulus()
# Start TensorFlow session
gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=0.8
) if gpu_id == '0' else tf.GPUOptions()
with tf.Session(config=tf.ConfigProto(gpu_options=gpu_options)) as sess:
# Select CPU or GPU
device = '/cpu:0' if gpu_id is None else '/gpu:0'
with tf.device(device):
model = Model(x, y, m)
# Initialize variables and start timer
sess.run(tf.global_variables_initializer())
t_start = time.time()
# Begin training loop, iterating over tasks
for i in range(par['n_train_batches']):
# Generate a batch of stimulus data for training
trial_info = stim.make_batch()
# Put together the feed dictionary
feed_dict = {
x: trial_info['neural_input'],
y: trial_info['desired_output'],
m: trial_info['train_mask']
}
# Run the model
_, task_loss, output, state, spike = \
sess.run([model.train, model.task_loss, \
model.output, model.h, model.h_out], feed_dict=feed_dict)
def main():
# Reset TensorFlow graph
tf.reset_default_graph()
# Create placeholders for the model
x = tf.placeholder(tf.float32, [par['batch_size'], 32, 32, 3], 'stim')
print('Batch size:', par['batch_size'])
print('Input size:', 32 * 32 * 3, '\n')
with tf.Session() as sess:
with tf.device('/gpu:0'):
model = Autoencoder(x)
sess.run(tf.global_variables_initializer())
t_start = time.time()
s = stimulus.Stimulus()
for i in range(2000):
task_id = np.random.randint(10)
y_hat, _, _ = s.make_batch(task_id, test=False)
_, loss, spike_loss, y = sess.run(
[model.train_op, model.loss, model.spike_loss, model.y],
{x: y_hat})
if i % 100 == 0:
print('\n Iter. | Task | Loss')
print('----------------------------')
print('',
str(i).ljust(5), '|',
str(task_id).ljust(4), '|', loss, '|', spike_loss)
if i % 100 == 0:
print(y_hat.shape, y.shape)
f, axarr = plt.subplots(2, 2)
axarr[0, 0].imshow(y_hat[0],
aspect='auto',
interpolation='none')
axarr[0, 0].set_title('Original 0')
axarr[0, 1].imshow(y[0], aspect='auto', interpolation='none')
axarr[0, 1].set_title('Reconstructed 0')
axarr[1, 0].imshow(y_hat[1],
aspect='auto',
interpolation='none')
axarr[1, 0].set_title('Original 1')
axarr[1, 1].imshow(y[1], aspect='auto', interpolation='none')
axarr[1, 1].set_title('Reconstructed 1')
plt.savefig('./encoder_testing/' + str(i))
W = {}
for var in tf.trainable_variables():
W[var.op.name] = var.eval()
def analyze(x, filename):
update_parameters(x['parameters'])
v_neuronal = np.zeros(
[par['num_pulses'], par['num_motion_dirs'], par['n_hidden']])
v_synaptic = np.zeros(
[par['num_pulses'], par['num_motion_dirs'], par['n_hidden']])
stim = stimulus.Stimulus()
trial_info = stim.generate_trial()
input_data = np.squeeze(
np.split(trial_info['neural_input'],
x['parameters']['num_time_steps'],
axis=1))
y_hat, h, syn_x, syn_u = run_model(input_data, x['parameters']['h_init'], \
x['parameters']['syn_x_init'], x['parameters']['syn_u_init'], x['weights'])
h = np.squeeze(np.split(h, x['parameters']['num_time_steps'], axis=1))
syn_x = np.squeeze(
np.split(syn_x, x['parameters']['num_time_steps'], axis=1))
syn_u = np.squeeze(
np.split(syn_u, x['parameters']['num_time_steps'], axis=1))
for i in range(par['num_pulses']):
time = x['timeline'][2 * i + 1] + 10
dir = trial_info['sample'][:, i]
for d in range(par['num_motion_dirs']):
ind = np.where(dir == d)[0]
v_neuronal[i, d, :] = np.mean(h[time][:, ind], axis=1)
v_synaptic[i, d, :] = np.mean(syn_x[time][:, ind] *
syn_u[time][:, ind],
axis=1)
print(np.mean(v_neuronal[i], axis=0).shape)
v_neuronal[i] = v_neuronal[i] - np.mean(v_neuronal[i], axis=0)
v_synaptic[i] = v_synaptic[i] - np.mean(v_synaptic[i], axis=0)
dot_neuronal = np.zeros(
(par['num_pulses'], par['num_pulses'], par['num_motion_dirs']))
for j in range(par['num_pulses'] - 1):
for i in range(j + 1, par['num_pulses']):
dot_neuronal[j, i, :] = np.diag(
np.dot(v_neuronal[j], np.transpose(v_neuronal[i])))
print(dot_neuronal)
def __init__(self):
# Reset TensorFlow graph
tf.reset_default_graph()
# Train on CIFAR-10 task
task_id = 0
# Create placeholders for the model
input_data = tf.placeholder(tf.float32, [par['batch_size'], 32, 32, 3],
'stim')
target_data = tf.placeholder(tf.float32, [par['batch_size'], 100],
'target')
mask = tf.placeholder(tf.float32,
[par['batch_size'], par['layer_dims'][-1]],
'mask')
print('Batch size:', par['batch_size'])
print('Input size:', 32 * 32 * 3, '\n')
with tf.Session() as sess:
cifar_model = self.model(input_data, target_data, mask)
sess.run(tf.global_variables_initializer())
t_start = time.time()
s = stimulus.Stimulus(include_cifar10=True)
for i in range(par['n_batches_top_down']):
x, y, m = s.make_batch(task_id, test=False)
_, loss = sess.run([self.train_op, self.loss],
feed_dict={
input_data: x,
target_data: y,
mask: m
})
if i % 1000 == 0:
print('Iteration ', i, ' Loss ', loss)
#print('', str(i).ljust(5), '|', str(task_id).ljust(4), '|', loss)
W = {}
for var in tf.trainable_variables():
W[var.op.name] = var.eval()
fn = './encoder_testing/conv_weights.pkl'
pickle.dump(W, open(fn, 'wb'))
print('Convolutional Weights saved in ', fn)
def run_simulation(weights):
update_parameters({'batch_train_size': 256})
update_parameters({'noise_in_sd': 1e-9})
update_parameters({'noise_rnn_sd': 1e-9})
update_parameters({'n_recurrent': 10})
stim = stimulus.Stimulus()
trial_info = stim.generate_trial(0)
trial_type = np.zeros((par['batch_train_size']))
for i in range(par['batch_train_size']):
trial_type[i] = 2*(trial_info['sample'][i]//3) \
+ trial_info['test'][i]//3
trial_length = trial_info['neural_input'].shape[1]
x = np.split(trial_info['neural_input'], trial_length, axis=1)
y_hat, hidden_state_hist = run_model(x, par['h_init'], weights)
return y_hat, hidden_state_hist, trial_type
def analyze_model_from_file(filename, savefile = None):
x = pickle.load(open(filename, 'rb'))
if savefile is None:
x['parameters']['save_fn'] = 'test.pkl'
else:
x['parameters']['save_fn'] = savefile
update_parameters(x['parameters'])
stim = stimulus.Stimulus()
trial_info = stim.generate_trial()
input_data = np.squeeze(np.split(trial_info['neural_input'], x['parameters']['num_time_steps'], axis=1))
print('input_data', len(input_data), input_data[0].shape)
y_hat, h, syn_x, syn_u = run_model(input_data, x['parameters']['h_init'], \
x['parameters']['syn_x_init'], x['parameters']['syn_u_init'], x['weights'])
h = np.squeeze(np.split(h, x['parameters']['num_time_steps'], axis=1))
syn_x = np.squeeze(np.split(syn_x, x['parameters']['num_time_steps'], axis=1))
syn_u = np.squeeze(np.split(syn_u, x['parameters']['num_time_steps'], axis=1))
analyze_model(trial_info, y_hat, h, syn_x, syn_u, None, x['weights'], simulation = False, \
lesion = False, tuning = True, decoding = False, load_previous_file = False, save_raw_data = False)
def main(save_fn, gpu_id=None):
""" Run supervised learning training """
# Update all dependencies in parameters
update_dependencies()
# Isolate requested GPU
if gpu_id is not None:
os.environ["CUDA_VISIBLE_DEVICES"] = gpu_id
# If desired, train the convolutional layers with the CIFAR datasets
# Otherwise, the network will load convolutional weights from the saved file
if (par['task'] == 'cifar' or par['task']
== 'imagenet') and par['train_convolutional_layers']:
convolutional_layers.ConvolutionalLayers()
print('\nRunning model.\n')
# Reset TensorFlow graph
tf.reset_default_graph()
# Create placeholders for the model
if par['task'] == 'mnist':
x = tf.placeholder(tf.float32,
[par['batch_size'], par['layer_dims'][0]], 'stim')
elif par['task'] == 'cifar' or par['task'] == 'imagenet':
x = tf.placeholder(tf.float32, [par['batch_size'], 32, 32, 3], 'stim')
y = tf.placeholder(tf.float32, [par['batch_size'], par['layer_dims'][-1]],
'out')
mask = tf.placeholder(tf.float32,
[par['batch_size'], par['layer_dims'][-1]], 'mask')
rule = tf.placeholder(tf.float32, [par['batch_size'], par['n_tasks']],
'rulecue')
gating = [
tf.placeholder(tf.float32, [par['layer_dims'][n + 1]], 'gating')
for n in range(par['n_layers'] - 1)
]
droput_keep_pct = tf.placeholder(tf.float32, [], 'dropout')
input_droput_keep_pct = tf.placeholder(tf.float32, [], 'input_dropout')
# Set up stimulus
stim = stimulus.Stimulus(labels_per_task=par['labels_per_task'])
# Initialize accuracy records
accuracy_full = []
accuracy_grid = np.zeros((par['n_tasks'], par['n_tasks']))
# Enter TensorFlow session
with tf.Session() as sess:
# sess = tf_debug.LocalCLIDebugWrapperSession(sess)
# Select CPU or GPU
device = '/cpu:0' if gpu_id is None else '/gpu:0'
with tf.device(device):
model = Model(x, y, gating, mask, droput_keep_pct,
input_droput_keep_pct, rule)
# Initialize variables
sess.run(tf.global_variables_initializer())
sess.run(model.reset_prev_vars)
# test if importance vals change
prev_imp = tf.Variable(tf.zeros(par['layer_dims'][1]), trainable=False)
# Begin training loop, iterating over tasks
for task in range(par['n_tasks']):
if par['gating_type'] is 'iXdG':
# test if imp vals change
sess.run(tf.assign(prev_imp, model.importance[1]))
# Update the importance of each unit
sess.run(model.update_importance)
# Create gates by importance for each task
sess.run(model.update_gates)
curr_task_gate = {}
for layer in range(1, len(par['layer_dims']) - 1):
copy = tf.Variable(tf.zeros(
tf.shape(model.curr_gate[layer])),
trainable=False)
sess.run(tf.assign(copy, model.curr_gate[layer]))
curr_task_gate[layer] = sess.run(copy)
model.gates[task] = curr_task_gate
# testing gates
if True:
if task == 0:
for var in model.variables_stabilization:
print(var.op.name)
print(var.get_shape())
assert sess.run(tf.count_nonzero(
model.curr_gate[1])) == 400
assert sess.run(
tf.equal(model.curr_gate[1],
model.gates[task][1])).all()
if task > 0:
layer_test = 1
expected_active_units = round(
(1 - par['gate_pct']) *
par['layer_dims'][layer_test])
print('Gates equal?')
print(
sess.run(
tf.equal(model.gates[task][layer_test],
model.gates[task -
1][layer_test])).all())
print('Importance vals equal?')
print(
sess.run(tf.equal(prev_imp,
model.importance[1])).all())
# vals_imp, idxs_imp = tf.math.top_k(model.importance[layer_test], k=2000)
# print(sess.run(idxs_imp)[-400:])
# print(sess.run(vals_imp)[-400:])
# print(sess.run(model.curr_gate[layer_test]))
# _, idxs = tf.math.top_k(model.curr_gate[layer_test], k=expected_active_units)
# print(sess.run(idxs))
# assert sess.run(tf.equal(idxs, sess.run(idxs_imp)[-400:])).all()
# _, idxs_copy = tf.math.top_k(model.gates[task][layer_test], k=expected_active_units)
# _, idxs_prev = tf.math.top_k(model.gates[task-1][layer_test], k=expected_active_units)
# print(sess.run(tf.shape(model.curr_gate[layer_test])))
# print(sess.run(idxs_prev))
# print(sess.run(idxs_copy))
# Create dictionary of gating signals applied to each hidden layer for this task
gating_dict = {
k: v
for k, v in zip(gating, list(model.gates[task].values()))
}
else:
gating_dict = {
k: v
for k, v in zip(gating, par['gating'][task])
}
# Create rule cue vector for this task
rule_cue = np.zeros([par['batch_size'], par['n_tasks']])
rule_cue[:, task] = 1
# Iterate over batches
for i in range(par['n_train_batches']):
# Make batch of training data
stim_in, y_hat, mk = stim.make_batch(task, test=False)
# Run the model using one of the available stabilization methods
if par['stabilization'] == 'pathint':
_, _, loss, AL = sess.run([model.train_op, model.update_small_omega, model.task_loss, model.aux_loss], \
feed_dict={x:stim_in, y:y_hat, **gating_dict, mask:mk, droput_keep_pct:par['drop_keep_pct'], \
input_droput_keep_pct:par['input_drop_keep_pct'], rule:rule_cue})
elif par['stabilization'] == 'EWC':
_, loss, AL = sess.run([model.train_op, model.task_loss, model.aux_loss], \
feed_dict={x:stim_in, y:y_hat, **gating_dict, mask:mk, droput_keep_pct:par['drop_keep_pct'], \
input_droput_keep_pct:par['input_drop_keep_pct'], rule:rule_cue})
# Display network performance
# if i%500 == 0:
if i % 9 == 0:
print('Iter: ', i, 'Loss: ', loss, 'Aux Loss: ', AL)
# Update big omegaes, and reset other values before starting new task
if par['stabilization'] == 'pathint':
sess.run(model.update_big_omega)
elif par['stabilization'] == 'EWC':
for _ in range(par['EWC_batch_divisor'] *
par['EWC_fisher_num_batches']):
stim_in, _, mk = stim.make_batch(task, test=False)
sess.run([model.update_big_omega], feed_dict = \
{x:stim_in, **gating_dict, mask:mk, droput_keep_pct:par['drop_keep_pct'], \
input_droput_keep_pct:par['input_drop_keep_pct'], rule:rule_cue})
# Reset the Adam Optimizer, and set the prev_weight values to their current values
sess.run(model.reset_adam_op)
sess.run(model.reset_prev_vars)
if par['stabilization'] == 'pathint':
sess.run(model.reset_small_omega)
"""
# update unit importance values
if par['gating_type'] is 'iXdG':
sess.run(model.update_importance)
"""
# Test the networks on all trained tasks
num_test_reps = 10
accuracy = np.zeros((task + 1))
for test_task in range(task + 1):
# Use appropriate gating and rule cues
if par['gating_type'] is 'iXdG':
gating_dict = {
k: v
for k, v in zip(gating,
list(model.gates[test_task].values()))
}
else:
gating_dict = {
k: v
for k, v in zip(gating, par['gating'][test_task])
}
test_rule_cue = np.zeros([par['batch_size'], par['n_tasks']])
test_rule_cue[:, test_task] = 1
# Repeat the test as desired
for r in range(num_test_reps):
stim_in, y_hat, mk = stim.make_batch(test_task, test=True)
acc = sess.run(model.accuracy, feed_dict={x:stim_in, y:y_hat, \
**gating_dict, mask:mk, droput_keep_pct:1.0, input_droput_keep_pct:1.0, rule:test_rule_cue})/num_test_reps
accuracy_grid[task, test_task] += acc
accuracy[test_task] += acc
# Display network performance after testing is complete
print('Task ', task, ' Mean ', np.mean(accuracy), ' First ',
accuracy[0], ' Last ', accuracy[-1])
accuracy_full.append(np.mean(accuracy))
# Reset weights between tasks if called upon
if par['reset_weights']:
sess.run(model.reset_weights)
# Save model performance and parameters if desired
if par['save_analysis']:
save_results = {'task': task, 'accuracy': accuracy, 'accuracy_full': accuracy_full, \
'accuracy_grid': accuracy_grid, 'par': par}
pickle.dump(save_results, open(par['save_dir'] + save_fn, 'wb'))
print('\nModel execution complete.')
def main(task):
"""
This function will take the task object as input and
creates a model object to learn the task
It would run the iterations
At each iteration, a new batch of trials will be created
"""
# Reset TensorFlow before running anything
tf.reset_default_graph()
# Tensorflow finds the supported CPU and GPU devices you can use
config = tf.ConfigProto()
trial_length = par['num_time_steps']
# Calculate shape of the stimulus for this task
# Define a placeholder for the stimulus the agent sees
#stimulus = tf.placeholder(tf.float64, shape=[task.total_dur, task.num_inputs, par['batch_train_size']])
stimulus = tf.placeholder(
tf.float64,
shape=[trial_length, task.num_inputs, par['batch_train_size']])
# Define a placeholder for the truth or correct answer about each trial
truth = tf.placeholder(tf.float64, shape=par['batch_train_size'])
# A TEMPORARY placeholder for target
#target = tf.placeholder(tf.float64, shape=[task.total_dur, 3, par['batch_train_size']])
target = tf.placeholder(tf.float64,
shape=[trial_length, 3, par['batch_train_size']])
# Create a model for the given task object
M = Model()
# Build the tf structure that runs trials
M.run_model(task, stimulus, truth)
M.optimize(task, target)
# Create a model from Nick's code
stim = stm.Stimulus()
n_input = task.num_inputs
'''
mask = tf.placeholder(tf.float64, shape=[task.total_dur, par['batch_train_size']])
x = tf.placeholder(tf.float64, shape=[n_input, task.total_dur, par['batch_train_size']]) # input data
target2 = tf.placeholder(tf.float64, shape=[3, task.total_dur, par['batch_train_size']]) # input data
actual_reward = tf.placeholder(tf.float64, shape=[task.total_dur,par['batch_train_size']])
pred_reward = tf.placeholder(tf.float64, shape=[task.total_dur, par['batch_train_size']])
actual_action = tf.placeholder(tf.float64, shape=[task.total_dur, 3, par['batch_train_size']])
'''
mask = tf.placeholder(tf.float64,
shape=[trial_length, par['batch_train_size']])
x = tf.placeholder(tf.float64,
shape=[n_input, trial_length,
par['batch_train_size']]) # input data
target2 = tf.placeholder(tf.float64,
shape=[3, trial_length,
par['batch_train_size']]) # input data
actual_reward = tf.placeholder(
tf.float64, shape=[trial_length, par['batch_train_size']])
pred_reward = tf.placeholder(tf.float64,
shape=[trial_length, par['batch_train_size']])
actual_action = tf.placeholder(
tf.float64, shape=[trial_length, 3, par['batch_train_size']])
M_Nick = model_RL.Model(x, target2, actual_reward, pred_reward,
actual_action, mask)
#M_Nick.run_model(task, stimulus, truth)
#M_Nick.optimize(task, target)
with tf.Session(config=config) as sess:
init = tf.global_variables_initializer()
sess.run(init)
t_start = time.time()
vloss = np.zeros((1, par['num_iterations']))
ploss = np.zeros((1, par['num_iterations']))
perf = np.zeros((1, par['num_iterations']))
for it in range(par['num_iterations']):
# Create a batch of stimuli, stores in attribute stimulus for the task
#task.create_stimulus()
# generate batch of batch_train_size with Nick's code
trial_info = stim.generate_trial()
"""
Run the model
"""
my_truth = np.zeros(par['batch_train_size'])
my_truth[trial_info['desired_output'][1, -1, :] ==
1] = 1 # Match trials
my_truth[trial_info['desired_output'][2, -1, :] ==
1] = 2 # Non_match trials
_, _, vloss[0, it], ploss[0, it], pol_grads, pol_out, pol_out0, actions, logpi, my_reward, action_array, time_mask, cumsum_logpi, pol_r, temp1, temp2, ideal, my_baseline, entropy = \
sess.run([M.pol_train_op, M.val_train_op, M.Loss_val, M.Loss_pol, M.pol_capped_gvs,task.pol_out_history, task.pol_out_history0, task.actions, task.logpi, M.reward, task.action_array, task.time_mask, task.cumsum_logpi, \
task.pol_r_history, task.temp1, task.temp2, task.ideal, M.baseline, \
M.entropy], {stimulus: np.swapaxes(trial_info['neural_input'], 1, 0), truth: my_truth, target: np.swapaxes(trial_info['desired_output'], 1, 0)})
# Run Nick's model
pol_out, val_out, pol_rnn, action, stacked_mask, reward = sess.run([M_Nick.pol_out, M_Nick.val_out, M_Nick.h_pol, M_Nick.action, \
M_Nick.stacked_mask,M_Nick.reward], {x: trial_info['neural_input'], target2: trial_info['desired_output'], mask: trial_info['train_mask']})
trial_reward = np.squeeze(np.stack(reward))
trial_action = np.stack(action)
_, _, pol_loss, val_loss = sess.run([M_Nick.train_pol, M_Nick.train_val, M_Nick.pol_loss, M_Nick.val_loss], \
{x: trial_info['neural_input'], target2: trial_info['desired_output'], mask: trial_info['train_mask'], \
actual_reward: trial_reward, pred_reward: np.squeeze(val_out), actual_action:trial_action })
pol_out = np.array(pol_out)
pol_out0 = np.array(pol_out0)
temp1 = np.array(temp1)
temp2 = np.array(temp2)
if it % 100 == 0:
fig = plt.plot(pol_out[:, :, 0])
plt.legend(['Fixate', 'match', 'Non-match'])
plt.title(str(my_truth[0]))
plt.savefig('Iteration_' + str(it) +
'.png') # save the figure to file
plt.close()
print('%6d, %6.1f, %6.1f, %6.1f, %6.1f, %6.2f' %
(it, my_reward.sum(), my_baseline.sum(), ploss[0, it],
vloss[0, it], entropy))
print('%6d, %6.1f, %6.1f, %6.1f, %6.1f' %
(it, np.array(trial_reward).sum(),
np.array(val_out).sum(), pol_loss, val_loss))
#pdb.set_trace()
#plt.plot(pol_out[:,:,0]); plt.show()
#if np.isnan(ploss[0, it]):
# pdb.set_trace()
#if it>=1000:
# pdb.set_trace()
pdb.set_trace()
a = 5
def main(gpu_id=None):
if gpu_id is not None:
os.environ['CUDA_VISIBLE_DEVICES'] = gpu_id
tf.reset_default_graph()
x = tf.placeholder(
tf.float32, [par['num_time_steps'], par['batch_size'], par['n_input']],
'stim')
y = tf.placeholder(
tf.float32,
[par['num_time_steps'], par['batch_size'], par['n_output']], 'out')
m = tf.placeholder(tf.float32, [par['num_time_steps'], par['batch_size']],
'mask')
stim = stimulus.Stimulus()
gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=0.8
) if gpu_id == '0' else tf.GPUOptions()
with tf.Session(config=tf.ConfigProto(gpu_options=gpu_options)) as sess:
device = '/cpu:0' if gpu_id is None else '/gpu:0'
with tf.device(device):
cortex = Cortex(x, y, m)
sess.run(tf.global_variables_initializer())
event_data = {'stimuli': [], 'actions': [], 'rewards': []}
print('\nTraining cortex for {} task:'.format(par['task']))
for i in range(par['n_batches']):
t0 = time.time()
name, trial_info = stim.generate_trial()
feed_dict = {
x: trial_info['neural_input'],
y: trial_info['reward_data'],
m: trial_info['train_mask']
}
_, pol_loss, val_loss, spike_loss, ent_loss, h, reward, action = \
sess.run([cortex.train_op, cortex.pol_loss, cortex.val_loss, cortex.spike_loss, \
cortex.ent_loss, cortex.h, cortex.reward, cortex.action], feed_dict=feed_dict)
# Select appropriate events
inds = list(np.where(reward != 0.))
rew_zero_times = np.array(
[np.random.randint(ts) for ts in inds[0]])
inds[0] = np.array([
np.random.choice([rt, zt], p=[0.75, 0.25])
for rt, zt in zip(inds[0], rew_zero_times)
])
event_stimuli = trial_info['neural_input'][inds[0], inds[1], :]
event_actions = action[inds[0], inds[1], :]
event_rewards = reward[inds[0], inds[1], :]
# Sample from events and save the event data
event_data['stimuli'].append(
event_stimuli[::par['sample_step'], :])
event_data['actions'].append(
event_actions[::par['sample_step'], :])
event_data['rewards'].append(
event_rewards[::par['sample_step'], :])
if i % 100 == 0:
print('Iter: {:>4} | Rew: {:6.3f} | Pol. Loss: {:8.5f} | Val. Loss: {:8.5f} | Ent. Loss: {:8.5f} | Spiking: {:8.5f} |'.format(\
i, np.mean(np.sum(reward, axis=0)), pol_loss, val_loss, ent_loss, np.mean(h)))
for (key, val) in event_data.items():
event_data[key] = np.concatenate(val, axis=0)
pickle.dump(
event_data,
open('./datadir/{}task_cortex_event_data.pkl'.format(par['task']),
'wb'))
print('Event samples saved. Model complete. \n')
def train_encoder(gpu_id=None):
if gpu_id is not None:
os.environ['CUDA_VISIBLE_DEVICES'] = gpu_id
tf.reset_default_graph()
x = tf.placeholder(tf.float32,
[par['num_time_steps'], None, par['n_input']], 'input')
update_parameters({'noise_in': 0.05})
stim = stimulus.Stimulus()
with tf.Session() as sess:
device = '/cpu:0' if gpu_id is None else '/gpu:0'
with tf.device(device):
encoder = Encoder(x)
sess.run(tf.global_variables_initializer())
print('\nTraining encoder for {} task:'.format(par['task']))
for i in range(10001):
name, trial_info = stim.generate_trial()
feed_dict = {x: trial_info['neural_input']}
_, rec_loss, act_loss, wei_loss, enc, rec, loss_plot = \
sess.run([encoder.train, encoder.rec_loss, \
encoder.act_loss, encoder.wei_loss, encoder.E, \
encoder.R, encoder.loss_plot], feed_dict=feed_dict)
if i % 500 == 0:
fig, ax = plt.subplots(3, 5, figsize=(14, 8))
for t in range(5):
ax[0, t].imshow(trial_info['neural_input'][:, t, :].T,
aspect='auto')
ax[1, t].imshow(rec[:, t, :].T, aspect='auto')
ax[2, t].imshow(enc[:, t, :].T, aspect='auto')
ax[0, t].set_title('Trial {}'.format(t))
for p in range(3):
ax[p, t].set_xticks([])
ax[p, t].set_yticks([])
ax[0, 0].set_ylabel('Stimulus Input')
ax[1, 0].set_ylabel('Reconstruction')
ax[2, 0].set_ylabel('Encoding')
ax[2, 0].set_xlabel('Time')
fig.suptitle('Input, Reconstruction, and Encoding')
plt.savefig('./savedir/input_rec_and_enc.png',
bbox_inches='tight')
plt.clf()
plt.close()
print(
'Iter: {:>6} | Rec. Loss: {:7.5f} | Act. Loss: {:7.5f} | Wei. Loss: {:7.5f}'
.format(i, rec_loss, act_loss, wei_loss))
print('\nEncoder training complete.')
# Record weights to check reconstruction
W, U = sess.run([encoder.W, encoder.U])
weights = {'W': W, 'U': U}
if par['internal_sampling']:
print('Sampling stimulus:')
# Sample stimuli (enough for 200 batches, or ~100k trials, which should sample most of the distribution)
for j in range(4):
trial_info_list = []
for i in range(50):
name, trial_info = stim.generate_trial()
feed_dict = {x: trial_info['neural_input']}
enc, = sess.run([encoder.E], feed_dict=feed_dict)
trial_info['encoded_input'] = enc
trial_info_list.append(trial_info)
aggregated_trial_info = {}
for key in (list(trial_info.keys()) + ['encoded_input']):
aggregated_trial_info[key] = np.concatenate(
[item[key] for item in trial_info_list], axis=1)
# Put together save data and save
save_data = {
'task': name,
'trial_info': aggregated_trial_info,
'weights': weights,
'rec_loss': rec_loss,
'act_loss': act_loss
}
pickle.dump(
save_data,
open(
'./datadir/{}task_{}unit_input_encoding_part{}.pkl'.
format(name, par['n_latent'], j), 'wb'))
print('Encoded stimulus samples saved. Model complete. \n')
else:
print('Saving weights...')
save_data = {
'task': name,
'weights': weights,
'rec_loss': rec_loss,
'act_loss': act_loss
}
pickle.dump(
save_data,
open(
'./datadir/{}task_{}unit_input_encoder_weights.pkl'.format(
name, par['n_latent']), 'wb'))
print('Encoder weights saved. Model complete. \n')
h *= suppress_activity
if par['synapse_config'] is None:
syn_x_new = np.ones_like(h)
syn_u_new = np.ones_like(h)
return h, syn_x_new, syn_u_new
data_dir = './savedir/'
filename = data_dir + 'WMnew.pkl' # 保存的模型文件
results = pickle.load(open(filename, 'rb'))
update_parameters(results['parameters'])
stim = stimulus.Stimulus()
# generate trials with match probability at 50%
trial_info = stim.generate_trial(test_mode = True)
input_data = np.squeeze(np.split(trial_info['neural_input'], par['num_time_steps'], axis=0))
h_init = results['weights']['h']
y, h, syn_x, syn_u = run_model(input_data, h_init, \
results['parameters']['syn_x_init'], results['parameters']['syn_u_init'], results['weights'])
syn_efficacy = syn_x*syn_u
'''
# generate trials with random sample and test stimuli, used for decoding
trial_info_decode = stim.generate_trial(test_mode = True)
def make_matrix(self, filename, method, N):
x = pickle.load(open(filename, 'rb'))
beh_threshold = 0.1
val_th = 0.1
ind_accurate = np.where(np.array(x['accuracy_hist']) > 0.98)[0]
#N = np.argmax(ind_accurate)
#N = 6
print('N = ', N)
if method == 'elim_lesion' or method == 'elim':
parameters.update_parameters(x['par'])
s = stimulus.Stimulus()
trial_info = s.generate_trial()
if method == 'lesion':
significant_weights_rnn = x['model_performance']['accuracy'][
-1] - x['lesion_accuracy_rnn'][0, :, :] > beh_threshold
significant_weights_out = x['model_performance']['accuracy'][
-1] - x['lesion_accuracy_out'][0, :, :] > beh_threshold
v = np.array([0]*x['parameters']['num_exc_units'] + [1]*x['parameters']['num_inh_units'] \
+ [2]*x['parameters']['n_output'])
W = np.vstack((significant_weights_rnn, significant_weights_out))
d = W.shape[0] - W.shape[1]
W = np.hstack((W, np.zeros((W.shape[0], d))))
elif method == 'elim':
num_units = 50 - N
w1 = np.zeros((num_units, num_units))
w2 = np.zeros((3, num_units))
ind = np.where(x['gate_hist'][N] > 0)[0]
for i in range(num_units):
for j in range(num_units):
w1[i, j] = x['weights_hist'][N]['w_rnn'][ind[i],
ind[j]] > val_th
for j in range(3):
w2[j,
i] = x['weights_hist'][N]['w_out'][j, ind[i]] > val_th
n_exc = int(np.sum(x['gate_hist'][N][:x['par']['num_exc']]))
n_inh = int(np.sum(x['gate_hist'][N][x['par']['num_exc']:]))
v = np.array([0] * n_exc + [1] * n_inh +
[2] * x['par']['n_output'])
W = np.vstack((w1, w2))
d = W.shape[0] - W.shape[1]
W = np.hstack((W, np.zeros((W.shape[0], d))))
elif method == 'elim_lesion':
num_units = 50 - N
r = analysis.lesion_weights(trial_info, x['par']['h_init'], x['par']['syn_x_init'], x['par']['syn_u_init'], \
x['weights_hist'][N], x['gate_hist'][N])
#plt.imshow(np.squeeze(r['lesion_accuracy_rnn']), aspect='auto', interpolation = 'none')
#plt.colorbar()
#plt.show()
w1_full = np.tile(
x['accuracy_hist'][N],
(x['par']['n_hidden'], x['par']['n_hidden'])) - np.squeeze(
r['lesion_accuracy_rnn']) > beh_threshold
w2_full = np.tile(
x['accuracy_hist'][N],
(x['par']['n_output'], x['par']['n_hidden'])) - np.squeeze(
r['lesion_accuracy_out']) > beh_threshold
w1 = np.zeros((num_units, num_units))
w2 = np.zeros((3, num_units))
ind = np.where(x['gate_hist'][N] > 0)[0]
for i in range(num_units):
for j in range(num_units):
w1[i, j] = w1_full[ind[i], ind[j]]
for j in range(3):
w2[j, i] = w2_full[j, ind[i]]
#plt.imshow(w1, aspect='auto', interpolation = 'none')
#plt.colorbar()
#plt.show()
print('accuracy ', x['accuracy_hist'][N])
n_exc = int(np.sum(x['gate_hist'][N][:x['par']['num_exc']]))
n_inh = int(np.sum(x['gate_hist'][N][x['par']['num_exc']:]))
v = np.array([0] * n_exc + [1] * n_inh +
[2] * x['par']['n_output'])
W = np.vstack((w1, w2))
d = W.shape[0] - W.shape[1]
W = np.hstack((W, np.zeros((W.shape[0], d))))
plt.imshow(W, aspect='auto', interpolation='none')
plt.colorbar()
plt.show()
print(v)
elif method == 'stacked':
W = []
for i in range(x['W_rnn'].shape[0]):
w1 = np.reshape(x['W_rnn'][i, :], (50, 50)) > 0.2
w2 = np.reshape(x['W_out'][i, :], (3, 50)) > 0.2
v = np.array([0] * 40 + [1] * 10 + [2] * 3)
W1 = np.vstack((w1, w2))
d = W1.shape[0] - W1.shape[1]
W1 = np.hstack((W1, np.zeros((W1.shape[0], d))))
W.append(W1)
return W, v
def main(gpu_id=None):
if gpu_id is not None:
os.environ["CUDA_VISIBLE_DEVICES"] = gpu_id
"""
Reset TensorFlow before running anything
"""
tf.reset_default_graph()
"""
Create the stimulus class to generate trial paramaters and input activity
"""
stim = stimulus.Stimulus()
n_input, n_hidden, n_output = par['shape']
N = par[
'batch_train_size'] # trials per iteration, calculate gradients after batch_train_size
"""
Define all placeholder
"""
mask = tf.placeholder(
tf.float64, shape=[par['num_time_steps'], par['batch_train_size']])
x = tf.placeholder(
tf.float64,
shape=[n_input, par['num_time_steps'],
par['batch_train_size']]) # input data
target = tf.placeholder(tf.float64,
shape=[
par['n_output'], par['num_time_steps'],
par['batch_train_size']
]) # input data
actual_reward = tf.placeholder(
tf.float64, shape=[par['num_time_steps'], par['batch_train_size']])
pred_reward = tf.placeholder(
tf.float64, shape=[par['num_time_steps'], par['batch_train_size']])
actual_action = tf.placeholder(tf.float64,
shape=[
par['num_time_steps'], par['n_output'],
par['batch_train_size']
])
config = tf.ConfigProto()
#config.gpu_options.allow_growth=True
# enter "config=tf.ConfigProto(log_device_placement=True)" inside Session to check whether CPU/GPU in use
with tf.Session(config=config) as sess:
if gpu_id is not None:
model = Model(x, target, actual_reward, pred_reward, actual_action,
mask)
else:
#with tf.device("/gpu:0"):
model = Model(x, target, actual_reward, pred_reward, actual_action,
mask)
init = tf.global_variables_initializer()
sess.run(init)
# keep track of the model performance across training
model_performance = {
'accuracy': [],
'loss': [],
'perf_loss': [],
'spike_loss': [],
'trial': []
}
for i in range(par['num_iterations']):
# generate batch of batch_train_size
trial_info = stim.generate_trial()
"""
Run the model
"""
pol_out, val_out, pol_rnn, action, stacked_mask, reward = sess.run([model.pol_out, model.val_out, model.h_pol, model.action, \
model.stacked_mask,model.reward], {x: trial_info['neural_input'], target: trial_info['desired_output'], mask: trial_info['train_mask']})
trial_reward = np.squeeze(np.stack(reward))
trial_action = np.stack(action)
#plt.imshow(np.squeeze(trial_reward))
#plt.colorbar()
#plt.show()
_, _, pol_loss, val_loss = sess.run([model.train_pol, model.train_val, model.pol_loss, model.val_loss], \
{x: trial_info['neural_input'], target: trial_info['desired_output'], mask: trial_info['train_mask'], \
actual_reward: trial_reward, pred_reward: np.squeeze(val_out), actual_action:trial_action })
accuracy, _, _ = analysis.get_perf(trial_info['desired_output'],
action,
trial_info['train_mask'])
#model_performance = append_model_performance(model_performance, accuracy, val_loss, pol_loss, spike_loss, (i+1)*N)
"""
Save the network model and output model performance to screen
"""
if i % par['iters_between_outputs'] == 0 and i > 0:
print_results(i, N, pol_loss, 0., pol_rnn, accuracy)
r = np.squeeze(np.sum(np.stack(trial_reward), axis=0))
print('Mean mask', np.mean(stacked_mask), ' val loss ',
val_loss, ' reward ', np.mean(r), np.max(r))
#plt.imshow(np.squeeze(stacked_mask[:,:]))
#plt.colorbar()
#plt.show()
#plt.imshow(np.squeeze(trial_reward))
#plt.colorbar()
#plt.show()
"""
Save model, analyze the network model and save the results
"""
#save_path = saver.save(sess, par['save_dir'] + par['ckpt_save_fn'])
if par['analyze_model']:
weights = eval_weights()
analysis.analyze_model(trial_info, y_hat, state_hist, syn_x_hist, syn_u_hist, model_performance, weights, \
simulation = True, lesion = False, tuning = False, decoding = False, load_previous_file = False, save_raw_data = False)
# Generate another batch of trials with test_mode = True (sample and test stimuli
# are independently drawn), and then perform tuning and decoding analysis
trial_info = stim.generate_trial(test_mode=True)
y_hat, state_hist, syn_x_hist, syn_u_hist = \
sess.run([model.y_hat, model.hidden_state_hist, model.syn_x_hist, model.syn_u_hist], \
{x: trial_info['neural_input'], y: trial_info['desired_output'], mask: trial_info['train_mask']})
analysis.analyze_model(trial_info, y_hat, state_hist, syn_x_hist, syn_u_hist, model_performance, weights, \
simulation = False, lesion = False, tuning = par['analyze_tuning'], decoding = True, load_previous_file = True, save_raw_data = False)
def main(gpu_id = None):
if gpu_id is not None:
os.environ["CUDA_VISIBLE_DEVICES"] = gpu_id
print(os.environ["CUDA_VISIBLE_DEVICES"])
# Print key parameters
print_important_params()
# Reset TensorFlow before running anything
tf.reset_default_graph()
# Create the stimulus class to generate trial paramaters and input activity
stim = stimulus.Stimulus()
# Define all placeholder
m = tf.placeholder(tf.float32, [par['num_time_steps'], par['batch_size']], 'mask')
x = tf.placeholder(tf.float32, [par['num_time_steps'], par['batch_size'], par['n_input']], 'input')
t = tf.placeholder(tf.float32, [par['num_time_steps'], par['batch_size'], par['n_output']], 'target')
# Make sure savedir exists
if not os.path.exists(f'savedir/{gpu_id}/'):
os.makedirs(f'savedir/{gpu_id}/')
save_increment = 0
# enter "config=tf.ConfigProto(log_device_placement=True)" inside Session to check whether CPU/GPU in use
with tf.Session(config=tf.ConfigProto()) as sess:
device = '/cpu:0' if gpu_id is None else '/gpu:0'
t0 = time.time()
with tf.device(device):
model = Model(x, t, m)
print(f"Model initialized. Time elapsed: {str(time.time() - t0)}")
print(par['shape'])
sess.run(tf.global_variables_initializer())
# keep track of the model performance across training
t0 = time.time()
# Set up records for memory bank performance storage
memory_bank_performance = dict()
for k in range(par['n_memory_banks']):
memory_bank_performance[k] = []
memory_bank_id = 0
model_performance = {'accuracy': [], 'loss': [], 'perf_loss': [], 'spike_loss': [], \
'weight_loss': [], 'iteration': []}
i = 0
while i < par['num_iterations']:
# Generate batch of batch_train_size
trial_info = stim.generate_trial(set_rule = None)
# Run the model
_, loss, perf_loss, spike_loss, weight_loss, y, h, syn_x, syn_u = \
sess.run([model.train_op, model.loss, model.perf_loss, model.spike_loss, \
model.weight_loss, model.y, model.h, model.syn_x, model.syn_u], \
{x: trial_info['neural_input'], t: trial_info['desired_output'], m: trial_info['train_mask']})
accuracies = analysis.get_perf_sr(trial_info['desired_output'], y, trial_info['train_mask'])
model_performance = append_model_performance(model_performance, accuracies, loss, perf_loss, spike_loss, weight_loss, i)
if (np.mean(model_performance['accuracy'][-50:])) > 0.90 or (i == par['num_iterations'] - 1):
memory_bank_id, i = update_memory_bank(memory_bank_id, memory_bank_performance)
old_w = sess.run(model.var_dict)['w_rnn']
old_w[old_w <= 1e-5] = 1e-3
par['new_init'] = old_w
_ = sess.run(model.resaturation_op)
# Save the network model and output model performance to screen
elif (i+1)%par['iters_between_outputs']==0:
t1 = time.time()
#accuracies = [analysis.get_perf_sr(trial_info['desired_output'], y[n,:,:,:], trial_info['train_mask'], True) for n in range(par['n_networks'])]
print_results(i, perf_loss, spike_loss, weight_loss, h, accuracies)
print(f"Elapsed time: {str(t1 - t0)}")
t0 = time.time()
#memory_bank_performance[memory_bank_id].append(accuracies)
# Also: compute accuracies on each of the previous memory banks
# (IDEA HERE: SHOULD WE CONSTRAIN S.T. NO INDIVIDUAL MEMORIES OVERLAP?)
if memory_bank_id > 0:
print("Accuracy on previous memory banks:")
for k in range(par['n_memory_banks']):
if k <= memory_bank_id:
par['memory_bank_id'] = k
update_dependencies()
trial_info = stim.generate_trial(set_rule = None, memory_bank_id = k)
h, y, syn_x, syn_u = sess.run([model.h, model.y_output, model.syn_x, model.syn_u],
{x: trial_info['neural_input'], t: trial_info['desired_output'], m: trial_info['train_mask']})
accuracies = analysis.get_perf_sr(trial_info['desired_output'], y, trial_info['train_mask'])
print(f"\tbank {k}: {accuracies}")
memory_bank_performance[k].append(accuracies)
else:
memory_bank_performance[k].append(0)
i += 1
# Save out activities
trial_info = stim.generate_trial(set_rule = None, memory_bank_id = memory_bank_id)
h, y, syn_x, syn_u = sess.run([model.h, model.y_output, model.syn_x, model.syn_u],
{x: trial_info['neural_input'], t: trial_info['desired_output'], m: trial_info['train_mask']})
# Save model and results
weights = sess.run(model.var_dict)
save_results(model_performance, weights, save_fn="sequence_reproduction_proof_of_concept.pkl")
def main(gpu_id=None):
print('\nRunning model.\n')
##################
### Setting Up ###
##################
tf.reset_default_graph()
""" Set up GPU """
if gpu_id is not None:
os.environ["CUDA_VISIBLE_DEVICES"] = gpu_id
""" Reset TensorFlow before running anything """
tf.reset_default_graph()
""" Set up performance recording """
model_performance = {'accuracy': [], 'par': [], 'task_list': []}
stim = stimulus.Stimulus()
mask = tf.placeholder(
tf.float32, shape=[p.par['num_time_steps'], p.par['batch_train_size']])
x = tf.placeholder(tf.float32,
shape=[
p.par['n_input'], p.par['num_time_steps'],
p.par['batch_train_size']
])
y = tf.placeholder(tf.float32,
shape=[
p.par['n_output'], p.par['num_time_steps'],
p.par['batch_train_size']
])
z = tf.placeholder(
tf.float32, shape=[p.par['num_networks'], p.par['generator_dims'][0]])
""" Start TensorFlow session """
with tf.Session() as sess:
if gpu_id is None:
model = Model(x, y, mask, z)
else:
with tf.device("/gpu:0"):
model = Model(x, y, mask, z)
# Initialize session variables
init = tf.global_variables_initializer()
sess.run(init)
t_start = time.time()
# Restore variables from previous model if desired
saver = tf.train.Saver()
if p.par['load_previous_model']:
saver.restore(sess, p.par['save_dir'] + p.par['ckpt_load_fn'])
print('Model ' + p.par['ckpt_load_fn'] + ' restored.')
generator_var_hist = []
for k in range(p.par['num_network_iters']):
print('NETWORK ITERATION ', k)
generator_var = np.random.normal(0,
1,
size=(p.par['num_networks'],
p.par['generator_dims'][0]))
generator_var_hist.append(generator_var)
for i in range(p.par['num_iterations']):
# generate batch of batch_train_size
trial_info = stim.generate_trial(k)
"""
Run the model
"""
_,_, total_loss, perf_loss, spike_loss, aux_loss, network_output = sess.run([model.train_op, \
model.update_small_omega, model.total_loss, model.perf_loss, \
model.spike_loss, model.aux_loss, model.networks_output], {x: trial_info['neural_input'], \
y: trial_info['desired_output'], mask: trial_info['train_mask'], z: generator_var})
if (i + 1
) % p.par['iters_between_outputs'] == 0: # and i != 0:
accuracy = get_perf(trial_info['desired_output'],
network_output,
trial_info['train_mask'])
iteration_time = time.time() - t_start
iterstr = 'Iter. {:>4}'.format(i)
timestr = 'Time. {:>7.4}'.format(iteration_time)
lossstr = 'Total Loss: {:>7.4}'.format(total_loss)
auxstr = 'Aux Loss: {:>7.4}'.format(aux_loss)
#perfstr = 'Perf. Loss: {:>7.4} +/- {:<7.4}'.format(np.mean(perf_loss), np.std(perf_loss))
#spikstr = 'Spike Loss: {:>7.4} +/- {:<7.4}'.format(np.mean(spike_loss), np.std(spike_loss))
#wirestr = 'Wiring Loss: {:>7.4} +/- {:<7.4}'.format(np.mean(wiring_loss), np.std(wiring_loss))
perfstr = 'Perf. Loss: {:>7.4}'.format(np.mean(perf_loss))
spikstr = 'Spike Loss: {:>7.4}'.format(np.mean(spike_loss))
#wirestr = 'Wiring Loss: {:>7.4}'.format(np.mean(wiring_loss))
accuracystr = 'Accuracy: {:>7.4} +/- {:<7.4}'.format(
np.mean(accuracy), np.std(accuracy))
print(' | '.join([
str(x) for x in [
iterstr, timestr, perfstr, spikstr, accuracystr,
auxstr
]
]))
if (i + 1) % 1000 == 0:
num_reps = 50
accuracy = get_perf(trial_info['desired_output'],
network_output,
trial_info['train_mask'])
novel_accuracy = []
for r in range(num_reps):
novel_var = np.random.normal(
0, 1, size=(1, p.par['generator_dims'][0]))
tt = np.mean(novel_var) > 0
trial_info = stim.generate_trial()
network_output = sess.run(model.networks_output, {x: trial_info['neural_input'], \
y: trial_info['desired_output'], mask: trial_info['train_mask'], z: novel_var})
novel_accuracy.append(get_perf(trial_info['desired_output'], network_output, \
trial_info['train_mask']))
previous_accuracy = []
for r in range(len(generator_var_hist)):
network_output = sess.run(model.networks_output, {x: trial_info['neural_input'], \
y: trial_info['desired_output'], mask: trial_info['train_mask'], z: generator_var_hist[r]})
previous_accuracy.append(get_perf(trial_info['desired_output'], network_output, \
trial_info['train_mask']))
accuracystr = 'Accuracy: {:>7.4} +/- {:<7.4}'.format(
np.mean(novel_accuracy), np.std(novel_accuracy))
prev_accuracystr = 'Accuracy: {:>7.4} +/- {:<7.4}'.format(
np.mean(previous_accuracy), np.std(previous_accuracy))
print('Novel ', accuracystr)
print('Previous ', prev_accuracystr)
print('Saving data...')
save_data(accuracy, novel_accuracy, previous_accuracy)
# Update big omegaes, and reset other values before starting new task
big_omegas = sess.run(
[model.update_big_omega, model.big_omega_var])
# Reset the Adam Optimizer, and set the previous parater values to their current values
sess.run(model.reset_adam_op)
sess.run(model.reset_prev_vars)
sess.run(model.reset_small_omega)
def main(gpu_id=None):
# Print out context
print_important_params()
# Select GPU
if gpu_id is not None:
os.environ['CUDA_VISIBLE_DEVICES'] = gpu_id
# Reduce memory consumption for GPU 0
gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=0.8) \
if gpu_id == '3' else tf.GPUOptions()
# Initialize stimulus environment
environment = stimulus.Stimulus()
# Reset graph and designate placeholders
tf.reset_default_graph()
stim_pl = tf.placeholder(tf.float32, [par['batch_size'], par['n_input']],
'stim')
prev_latent_pl = tf.placeholder(tf.float32,
[par['batch_size'], par['n_latent']],
'prev_latent')
latent_pl = tf.placeholder(tf.float32,
[par['batch_size'], par['n_latent']], 'latent')
hop_latent_pl = tf.placeholder(tf.float32,
[par['batch_size'], par['n_latent']],
'hp_latent')
prev_action_pl = tf.placeholder(tf.float32,
[par['batch_size'], par['n_pol']],
'prev_action')
action_pl = tf.placeholder(tf.float32, [par['batch_size'], par['n_pol']],
'action')
reward_pl = tf.placeholder(tf.float32, [par['batch_size'], par['n_val']],
'reward')
prev_reward_pl = tf.placeholder(tf.float32,
[par['batch_size'], par['n_unique_vals']],
'prev_reward')
prev_val_pl = tf.placeholder(tf.float32, [par['batch_size'], par['n_val']],
'prev_val')
future_val_pl = tf.placeholder(tf.float32,
[par['batch_size'], par['n_val']],
'future_val')
# Start TensorFlow session
with tf.Session(config=tf.ConfigProto(gpu_options=gpu_options)) as sess:
# Set up and initialize model on desired device
device = '/cpu:0' if gpu_id is None else '/gpu:0'
with tf.device(device):
model = Model(stim_pl, prev_latent_pl, latent_pl, hop_latent_pl, prev_reward_pl, \
reward_pl, prev_action_pl, action_pl, prev_val_pl, future_val_pl)
sess.run(tf.global_variables_initializer())
# Start training loop
for i in range(par['num_batches']):
# Reset environment at the start of each iteration
environment.reset_agents()
environment.reset_rewards()
x_read_list = []
# Pre-allocate prev_val and total_reward
prev_latent = np.zeros((par['batch_size'], par['n_latent']),
dtype=np.float32)
prev_action = np.zeros((par['batch_size'], par['n_pol']),
dtype=np.float32)
prev_reward_scalar = np.zeros((par['batch_size'], par['n_val']),
dtype=np.float32)
prev_reward_matrix = np.zeros(
(par['batch_size'], par['n_unique_vals']), dtype=np.float32)
prev_val = np.zeros((par['batch_size'], par['n_val']),
dtype=np.float32)
total_reward = np.zeros((par['batch_size'], 1), dtype=np.float32)
action_record = []
H_stim_record = []
H_act_f_record = []
# Iterate through time
for t in range(par['num_time_steps']):
# Make inputs
stim_in = environment.make_inputs()
# Generate the latent, train encoder weights
_, reconstruction_loss, latent_loss, latent, latent_mu, latent_log_var, sparsity_loss = \
sess.run([model.train_encoder, model.reconstruction_loss, model.latent_loss, \
model.latent, model.latent_mu, model.latent_log_var, model.sparsity_loss], \
feed_dict = {stim_pl: stim_in})
agent_location = environment.get_agent_locs()
# Generate the policy and value functions
pol, val, hopfield_read = sess.run([model.pol_out, model.val_out, model.hopfield_read], \
feed_dict = {latent_pl: latent})
# Choose action, calculate reward and determine next state
action = np.array([
np.random.multinomial(1, pol[t, :] - 1e-6)
for t in range(par['batch_size'])
])
agent_locs = environment.get_agent_locs()
reward = environment.agent_action(action)
reward_one_hot = np.zeros(
(par['batch_size'], par['n_unique_vals']),
dtype=np.float32)
reward_one_hot[np.arange(par['batch_size']),
np.int8(np.squeeze(reward))] = 1.
if t < 0:
latent /= (
1e-9 +
np.sqrt(np.sum(latent**2, axis=1, keepdims=True)))
prev_latent /= (
1e-9 +
np.sqrt(np.sum(prev_latent**2, axis=1, keepdims=True)))
z = np.sum(latent * prev_latent, axis=1)
print(t, agent_location[0, :], action[0, :], ' Dot prod ',
z[0])
plt.imshow(np.reshape(hopfield_read[0, :], (4, 2)),
aspect='auto')
plt.colorbar()
plt.title('Location ' + str(agent_location[0, :]))
plt.show()
#action_record.append(action)
# Update total reward and prev_val
total_reward += reward
# Update the Hopfield network
sess.run([model.update_hopfield_with_dep, model.train_RL, model.hopfield_read], \
feed_dict = {latent_pl: prev_latent, action_pl: prev_action, reward_pl: prev_reward_scalar, \
future_val_pl: val, prev_latent_pl : prev_latent, hop_latent_pl : latent, \
prev_action_pl: prev_action, prev_reward_pl : prev_reward_matrix})
prev_latent = latent + 0.
prev_reward_matrix = reward_one_hot + 0.
prev_reward_scalar = reward + 0.
prev_action = action + 0.
prev_val = val + 0.
if t == -1:
z_read = np.stack(x_read_list, axis=0)
fig, ax = plt.subplots(2, 2, figsize=[8, 8])
#ax[0,0].imshow(z_read[:, 0, :], aspect = 'auto')
#ax[0,1].imshow(z_read[:, 1, :], aspect = 'auto')
#ax[1,0].imshow(z_read[:, 2, :], aspect = 'auto')
#ax[1,1].imshow(z_read[:, 3, :], aspect = 'auto')
ax[0, 0].imshow(H_stim_record[0], aspect='auto')
ax[0, 1].imshow(H_stim_record[100], aspect='auto')
ax[1, 0].imshow(H_stim_record[300], aspect='auto')
ax[1, 1].imshow(H_stim_record[-1], aspect='auto')
fig, ax = plt.subplots(2, 2, figsize=[8, 8])
ax[0, 0].imshow(H_act_f_record[0], aspect='auto')
ax[0, 1].imshow(H_act_f_record[100], aspect='auto')
ax[1, 0].imshow(H_act_f_record[300], aspect='auto')
ax[1, 1].imshow(H_act_f_record[-1], aspect='auto')
plt.show()
# Reset agents that have obtained a reward
environment.reset_agents(reward != 0.)
# Analyze actions
#action_record = np.concatenate(action_record, axis=0)
#action_record = np.round(np.mean(action_record, axis=0), 2).tolist()
# Output network performance
if i % 10 == 0:
H_sas, H_sar = sess.run([model.H_sas, model.H_sar])
#print('Iter {:>4} | Mean Reward: {:6.3f} | Recon Loss: {:8.6f} | Sparsity Loss: {:8.6f} | Sim act.: {:8.6f} | Action Dist: {}'.format(\
# i, np.mean(total_reward), rec_loss, sparsity_loss, sim_active, action_record))
print('Iter {:>4} | Mean Reward: {:6.3f} | Recon Loss: {:8.6f} | Latent Loss: {:8.6f} |'.format(\
i, np.mean(total_reward), reconstruction_loss, latent_loss))
print(' | Latent Mu: {:8.6f} | Latent Var: {:8.6f} | Sparisty Loss: {:8.6f}'.format(\
np.mean(latent_mu), np.mean(latent_log_var), sparsity_loss))
print(' | H_sas', np.mean(H_sas**2), '| H_sar',
np.mean(H_sar**2))
if i % 200 == 0 and par['train_encoder'] and i > 0:
weights = sess.run(model.var_dict)
results = {'weights': weights, 'latent': latent}
pickle.dump(results,
open('./savedir/VAE_8x8_model_weights3.pkl', 'wb'))
# Update model weights
sess.run(model.update_weights)
sess.run(model.reset_hopfield)
def main(gpu_id=None):
# Print out context
# Select GPU
if gpu_id is not None:
os.environ['CUDA_VISIBLE_DEVICES'] = gpu_id
# Reduce memory consumption for GPU 0
gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=0.8) \
if gpu_id == '3' else tf.GPUOptions()
# Initialize stimulus environment
stim = stimulus.Stimulus()
# Reset graph and designate placeholders
tf.reset_default_graph()
x = tf.placeholder(tf.float32, [par['batch_size'], 28, 28, 1], 'input')
y = tf.placeholder(tf.int32, [par['batch_size']], 'output')
k = tf.placeholder(tf.float32, [], 'keep_prob')
# Start TensorFlow session
with tf.Session(config=tf.ConfigProto(gpu_options=gpu_options)) as sess:
# Set up and initialize model on desired device
device = '/cpu:0' if gpu_id is None else '/gpu:0'
with tf.device(device):
model = Model(x, y, k)
sess.run(tf.global_variables_initializer())
# Start training loop
print('-' * 20 + '\nStarting training.')
for i in range(par['iterations']):
images, labels = stim.make_batch()
train_get_list = [model.train]
_, = sess.run(train_get_list,
feed_dict={
x: images,
y: labels,
k: par['dropout_keep_prob']
})
if i % 100 == 0:
images, labels = stim.make_batch(test=True)
test_get_list = [
model.output, model.recon, model.task_loss,
model.recon_loss, model.latent_loss
]
output, recon, task_loss, recon_loss, latent_loss = sess.run(
test_get_list, feed_dict={
x: images,
y: labels,
k: 1.
})
acc = np.mean(np.argmax(output, axis=-1) == labels)
print('Iter {:>5} | Accuracy: {:5.3f} | Task Loss: {:5.3f} | Recon Loss: {:5.3f} | Latent Loss: {:5.3f}'.format(\
i, acc, task_loss, recon_loss, latent_loss))
print('\nTraining complete. Recording results.')
images, labels = stim.make_batch(test=True)
test_get_list = [
model.output, model.recon, model.task_loss, model.recon_loss,
model.latent_loss
]
output, recon, task_loss, recon_loss, latent_loss = sess.run(
test_get_list, feed_dict={
x: images,
y: labels,
k: 1.
})
acc = np.mean(np.argmax(output, axis=-1) == labels)
data = {
'parameters': par,
'images': images,
'labels': labels,
'recons': recon,
'weights': sess.run([model.var_dict])[0],
'accuracy': acc
}
pickle.dump(data, open('./savedir/' + par['savefn'] + '.pkl', 'wb'))
print('Results saved. Model complete.')
def main():
"""
Reset TensorFlow before running anything
"""
tf.reset_default_graph()
"""
Create the stimulus class to generate trial paramaters and input activity
"""
stim = stimulus.Stimulus()
n_input, n_hidden, n_output = par['shape']
N = par['batch_train_size'] * par[
'num_batches'] # trials per iteration, calculate gradients after batch_train_size
"""
Define all placeholder
"""
mask = tf.placeholder(
tf.float32, shape=[par['num_time_steps'], par['batch_train_size']])
x = tf.placeholder(
tf.float32,
shape=[n_input, par['num_time_steps'],
par['batch_train_size']]) # input data
y = tf.placeholder(
tf.float32,
shape=[n_output, par['num_time_steps'],
par['batch_train_size']]) # target data
# enter "config=tf.ConfigProto(log_device_placement=True)" inside Session to check whether CPU/GPU in use
with tf.Session() as sess:
#with tf.device("/gpu:0"):
model = Model(x, y, mask)
init = tf.global_variables_initializer()
sess.run(init)
t_start = time.time()
saver = tf.train.Saver()
# Restore variables from previous model if desired
if par['load_previous_model']:
saver.restore(sess, par['save_dir'] + par['ckpt_load_fn'])
print('Model ' + par['ckpt_load_fn'] + ' restored.')
# keep track of the model performance across training
model_performance = {
'accuracy': [],
'loss': [],
'perf_loss': [],
'spike_loss': [],
'trial': [],
'time': []
}
for i in range(par['num_iterations']):
# generate batch of N (batch_train_size X num_batches) trials
trial_info = stim.generate_trial()
# keep track of the model performance for this batch
loss = np.zeros((par['num_batches']))
perf_loss = np.zeros((par['num_batches']))
spike_loss = np.zeros((par['num_batches']))
accuracy = np.zeros((par['num_batches']))
for j in range(par['num_batches']):
"""
Select batches of size batch_train_size
"""
ind = range(j * par['batch_train_size'],
(j + 1) * par['batch_train_size'])
target_data = trial_info['desired_output'][:, :, ind]
input_data = trial_info['neural_input'][:, :, ind]
train_mask = trial_info['train_mask'][:, ind]
"""
Run the model
If learning rate > 0, then also run the optimizer;
if learning rate = 0, then skip optimizer
"""
if par['learning_rate'] > 0:
_, loss[j], perf_loss[j], spike_loss[j], y_hat, state_hist, W_rnn, = \
sess.run([model.train_op, model.loss, model.perf_loss, model.spike_loss, model.y_hat, \
model.hidden_state_hist, model.WY], {x: input_data, y: target_data, mask: train_mask})
else:
loss[j], perf_loss[j], spike_loss[j], y_hat, state_hist, W_rnn = \
sess.run([model.loss, model.perf_loss, model.spike_loss, model.y_hat, model.hidden_state_hist, \
model.WY], {x: input_data, y: target_data, mask: train_mask})
accuracy[j] = analysis.get_perf(target_data, y_hat, train_mask)
iteration_time = time.time() - t_start
model_performance = append_model_performance(
model_performance, accuracy, loss, perf_loss, spike_loss,
(i + 1) * N, iteration_time)
"""
Save the network model and output model performance to screen
"""
if (i + 1) % par['iters_between_outputs'] == 0 or i + 1 == par[
'num_iterations']:
print_results(i, N, iteration_time, perf_loss, spike_loss,
state_hist, accuracy)
save_path = saver.save(sess,
par['save_dir'] + par['ckpt_save_fn'])
"""
Analyze the network model and save the results
"""
if par['analyze_model']:
weights = eval_weights(W_rnn)
analysis.analyze_model(trial_info, y_hat, state_hist,
model_performance, weights)
def main(gpu_id):
os.environ["CUDA_VISIBLE_DEVICES"] = gpu_id
"""
Reset TensorFlow before running anything
"""
tf.reset_default_graph()
"""
Create the stimulus class to generate trial paramaters and input activity
"""
stim = stimulus.Stimulus()
"""
Define all placeholder
"""
mask = tf.placeholder(
tf.float32, shape=[par['num_time_steps'], par['batch_train_size']])
x = tf.placeholder(tf.float32,
shape=[
par['num_motion_dirs'], par['num_time_steps'],
par['batch_train_size']
]) # input data
y = tf.placeholder(tf.float32,
shape=[
par['n_output'], par['num_time_steps'],
par['batch_train_size']
]) # target data
with tf.Session() as sess:
#with tf.device("/gpu:0"):
model = Model(x, y, mask, True)
init = tf.global_variables_initializer()
sess.run(init)
t_start = time.time()
for i in range(par['num_iterations']):
# generate batch of batch_train_size
trial_info = stim.generate_trial(i // par['iters_per_group'])
"""
plt.imshow(trial_info['desired_output'][:,:, 0], aspect = 'auto')
plt.show()
plt.imshow(trial_info['neural_input'][:,:, 0], aspect = 'auto')
plt.show()
"""
_, loss, perf_loss, spike_loss, weight_loss, y_hat, state_hist = \
sess.run([model.train_op, model.loss, model.perf_loss, model.spike_loss, \
model.weight_loss, model.y_hat, \
model.hidden_state_hist], {x: trial_info['neural_input'], \
y: trial_info['desired_output'], mask: trial_info['train_mask']})
if i == par['iters_per_group']:
sess.run(model.reset_weights)
iteration_time = time.time() - t_start
if (i + 1) % par['iters_between_outputs'] == 0 or i + 1 == par[
'num_iterations']:
print_results(i, iteration_time, perf_loss, spike_loss,
weight_loss, trial_info['desired_output'])
y1 = np.stack(y_hat, axis=1)
"""
ind = np.argsort(np.var(np.mean(trial_info['desired_output'],axis=2), axis=1))
plt.subplot(2,2,1)
plt.plot(trial_info['desired_output'][ind[-1], :, 0])
plt.plot(y1[ind[-1], :, 0],'r')
plt.subplot(2,2,2)
plt.plot(trial_info['desired_output'][ind[-2], :, 0])
plt.plot(y1[ind[-2], :, 0],'r')
plt.subplot(2,2,3)
plt.plot(trial_info['desired_output'][ind[-3], :, 0])
plt.plot(y1[ind[-3], :, 0],'r')
plt.subplot(2,2,4)
plt.plot(trial_info['desired_output'][ind[-4], :, 0])
plt.plot(y1[ind[-4], :, 0],'r')
plt.show()
"""
weights = eval_weights()
save_results = {'weights': weights, 'par': par, 'y_out': y1, 'y_target': trial_info['desired_output'], \
'h': state_hist, 'input': trial_info['neural_input']}
pickle.dump(
save_results,
open(par['save_dir'] + 'saved_results_wc1e6_v0.pkl', 'wb'))
def main(gpu_id=None, code_state=historian.record_code_state()):
""" Run supervised learning training """
# Isolate requested GPU
if gpu_id is not None:
os.environ["CUDA_VISIBLE_DEVICES"] = gpu_id
# Reset TensorFlow graph before running anything
tf.reset_default_graph()
# Define all placeholders
x, y, m, l, ci, cj, h, sx, su = get_placeholders()
# Set up stimulus and model performance recording
stim = stimulus.Stimulus()
model_performance = {
'accuracy': [],
'pulse_accuracy': [],
'loss': [],
'perf_loss': [],
'spike_loss': [],
'trial': []
}
# Start TensorFlow session
with tf.Session() as sess:
# Select CPU or GPU
device = '/cpu:0' if gpu_id is None else '/gpu:0'
with tf.device(device):
model = Model(x, y, m, l, ci, cj, h, sx, su)
# Initialize variables and start the timer
sess.run(tf.global_variables_initializer())
t_start = time.time()
# Begin training loop
print('\nStarting training...\n')
acc_count = int(0)
accuracy_threshold = \
np.array([0.0, 0.6, 0.7, 0.8, 0.9, 0.95, 0.96, \
0.97, 0.98])
save_fn = par['save_dir'] + par['save_fn']
save_fn_ind = save_fn[1:].find('.') - 1
for i in range(par['num_iterations']):
# Generate a batch of stimulus for training
trial_info = shuffle_trials(stim)
# Put together the feed dictionary
feed_dict = {
x: trial_info['neural_input'],
y: trial_info['desired_output'],
m: trial_info['train_mask']
}
# Run the model
_, loss, perf_loss, spike_loss, y_hat, state_hist, syn_x_hist, syn_u_hist = \
sess.run([model.train_op, model.loss, model.perf_loss, model.spike_loss, model.y_hat, \
model.hidden_hist, model.syn_x_hist, model.syn_u_hist], feed_dict=feed_dict)
# Calculate accuracy from the model's output
if par['output_type'] == 'directional':
accuracy, pulse_accuracy = analysis.get_coord_perf(
trial_info['desired_output'], y_hat,
trial_info['train_mask'], trial_info['pulse_id'])
elif par['output_type'] == 'one_hot':
accuracy, pulse_accuracy = analysis.get_perf(
trial_info['desired_output'], y_hat,
trial_info['train_mask'], trial_info['pulse_id'])
# Record the model's performance
model_performance = append_model_performance(
model_performance, accuracy, pulse_accuracy, loss, perf_loss,
spike_loss, (i + 1) * par['batch_train_size'])
# Save and show the model's performance
if i % par[
'iters_between_outputs'] == 0: #in list(range(len(par['trial_type']))):
print_results(i, par['trial_type'], perf_loss, spike_loss,
state_hist, accuracy, pulse_accuracy)
# if i%200 in list(range(len(par['trial_type']))):
# weights = sess.run(model.var_dict)
# results = {
# 'model_performance': model_performance,
# 'parameters': par,
# 'weights': weights}
# pickle.dump(results, open(par['save_dir'] + par['save_fn'], 'wb') )
# if i>=5 and all(np.array(model_performance['accuracy'][-5:]) > accuracy_threshold[acc_count]):
# break
if i > 5 and all(
np.array(model_performance['accuracy'][-5:]) >
accuracy_threshold[acc_count]):
print('SAVING')
weights = sess.run(model.var_dict)
results = {
'model_performance': model_performance,
'parameters': par,
'weights': weights,
'code_state': code_state
}
acc_str = str(int(accuracy_threshold[acc_count] * 100))
sf = save_fn[:-4] + '_acc' + acc_str + save_fn[-4:]
print(sf)
pickle.dump(results, open(sf, 'wb'))
acc_count += 1
if acc_count >= len(accuracy_threshold):
break
# If required, save the model, analyze it, and save the results
if par['analyze_model']:
weights = sess.run(model.var_dict)
syn_x_stacked = np.stack(syn_x_hist, axis=1)
syn_u_stacked = np.stack(syn_u_hist, axis=1)
h_stacked = np.stack(state_hist, axis=1)
trial_time = np.arange(0, h_stacked.shape[1] * par['dt'],
par['dt'])
mean_h = np.mean(np.mean(h_stacked, axis=2), axis=1)
results = {
'model_performance': model_performance,
'parameters': par,
'weights': weights,
'trial_time': trial_time,
'mean_h': mean_h
}
pickle.dump(results, open(par['save_dir'] + par['save_fn'], 'wb'))
def main(gpu_id=None):
# Print out context
print_important_params()
# Select GPU
if gpu_id is not None:
os.environ['CUDA_VISIBLE_DEVICES'] = gpu_id
# Reduce memory consumption for GPU 0
gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=0.8) \
if gpu_id == '0' else tf.GPUOptions()
# Initialize stimulus environment
environment = stimulus.Stimulus()
# Reset graph and designate placeholders
tf.reset_default_graph()
stim_pl = tf.placeholder(tf.float32, [par['batch_size'], par['n_input']],
'stim')
action_pl = tf.placeholder(tf.float32, [par['batch_size'], par['n_pol']],
'action')
reward_pl = tf.placeholder(tf.float32, [par['batch_size'], par['n_val']],
'reward')
future_val_pl = tf.placeholder(tf.float32,
[par['batch_size'], par['n_val']],
'prev_val')
time_step_pl = tf.placeholder(tf.int32, [], 'time_step')
# Start TensorFlow session
with tf.Session(config=tf.ConfigProto(gpu_options=gpu_options)) as sess:
# Set up and initialize model on desired device
device = '/cpu:0' if gpu_id is None else '/gpu:0'
with tf.device(device):
model = Model(stim_pl, reward_pl, action_pl, future_val_pl,
time_step_pl)
sess.run(tf.global_variables_initializer())
# Start training loop
for i in range(par['num_batches']):
# Reset environment at the start of each iteration
environment.reset_agents()
environment.reset_rewards()
# Pre-allocate prev_val and total_reward
prev_stim = np.zeros((par['batch_size'], par['n_input']),
dtype=np.float32)
future_val = np.zeros((par['batch_size'], par['n_val']),
dtype=np.float32)
prev_action = np.zeros((par['batch_size'], par['n_pol']),
dtype=np.float32)
prev_reward = np.zeros((par['batch_size'], par['n_val']),
dtype=np.float32)
prev_t = 0
total_reward = np.zeros((par['batch_size'], 1), dtype=np.float32)
action_record = []
# Iterate through time
for t in range(par['num_time_steps']):
# Make inputs
stim_in = environment.make_inputs()
# Train encoder weights, output the policy and value functions
_, pol, val, rec_loss, sparsity_loss, x_read = sess.run([model.train_encoder, model.pol_out, model.val_out, \
model.reconstruction_loss, model.sparsity_loss, model.x_read], feed_dict = {stim_pl: stim_in})
W = sess.run(model.var_dict)
# Choose action, calculate reward and determine next state
action = np.array([
np.random.multinomial(1, pol[t, :] - 1e-6)
for t in range(par['batch_size'])
])
reward = environment.agent_action(action)
action_record.append(action)
# Update total reward and prev_val
total_reward += reward
if i > 100:
# Update the Hopfield network
sess.run([model.update_hopfield, model.train_RL], \
feed_dict = {stim_pl: prev_stim, action_pl: prev_action, reward_pl: prev_reward, \
future_val_pl: val, time_step_pl:prev_t})
prev_stim = stim_in
prev_reward = reward
prev_action = action
prev_t = t
# Reset agents that have obtained a reward
environment.reset_agents(reward != 0.)
# Update model weights
sess.run(model.update_weights)
sess.run(model.reset_hopfield)
# Analyze actions
action_record = np.concatenate(action_record, axis=0)
action_record = np.round(np.mean(action_record, axis=0),
2).tolist()
# Output network performance
print('Iter {:>4} | Mean Reward: {:6.3f} | Recon Loss: {:8.6f} | Sparsity Loss: {:8.6f} | Action Dist: {}'.format(\
i, np.mean(total_reward), rec_loss, sparsity_loss, action_record))
def main(save_fn, gpu_id=None):
if gpu_id is not None:
os.environ["CUDA_VISIBLE_DEVICES"] = gpu_id
# train the convolutional layers with the CIFAR-10 dataset
# otherwise, it will load the convolutional weights from the saved file
if (par['task'] == 'cifar' or par['task']
== 'imagenet') and par['train_convolutional_layers']:
convolutional_layers.ConvolutionalLayers()
print('\nRunning model.\n')
# Reset TensorFlow graph
tf.reset_default_graph()
# Create placeholders for the model
# input_data, target_data, gating, mask, dropout keep pct hidden layers, dropout keep pct input layers
if par['task'] == 'mnist':
x = tf.placeholder(tf.float32,
[par['batch_size'], par['layer_dims'][0]], 'stim')
elif par['task'] == 'cifar' or par['task'] == 'imagenet':
x = tf.placeholder(tf.float32, [par['batch_size'], 32, 32, 3], 'stim')
y = tf.placeholder(tf.float32, [par['batch_size'], par['layer_dims'][-1]],
'out')
mask = tf.placeholder(tf.float32,
[par['batch_size'], par['layer_dims'][-1]], 'mask')
droput_keep_pct = tf.placeholder(tf.float32, [], 'dropout')
input_droput_keep_pct = tf.placeholder(tf.float32, [], 'input_dropout')
gating = [
tf.placeholder(tf.float32, [par['layer_dims'][n + 1]], 'gating')
for n in range(par['n_layers'] - 1)
]
context_vector = tf.placeholder(tf.float32, [1, par['n_tasks']],
'context_vector')
stim = stimulus.Stimulus(labels_per_task=par['labels_per_task'])
accuracy_full = []
accuracy_grid = np.zeros((par['n_tasks'], par['n_tasks']))
with tf.Session() as sess:
if gpu_id is None:
model = Model(x, y, gating, mask, droput_keep_pct,
input_droput_keep_pct, context_vector)
else:
with tf.device("/gpu:0"):
model = Model(x, y, gating, mask, droput_keep_pct,
input_droput_keep_pct, context_vector)
init = tf.global_variables_initializer()
sess.run(init)
t_start = time.time()
sess.run(model.reset_prev_vars)
for task in range(par['n_tasks']):
cont_vect = np.zeros((1, par['n_tasks']), dtype=np.float32)
cont_vect[0, task] = 1.
# create dictionary of gating signals applied to each hidden layer for this task
gating_dict = {k: v for k, v in zip(gating, par['gating'][task])}
for i in range(par['n_train_batches']):
# make batch of training data
stim_in, y_hat, mk = stim.make_batch(task, test=False)
if par['stabilization'] == 'pathint':
_, _, loss, AL, gl = sess.run([model.train_op, model.update_small_omega, model.task_loss, model.aux_loss, model.gate_loss], \
feed_dict = {x:stim_in, y:y_hat, **gating_dict, mask:mk, droput_keep_pct:par['drop_keep_pct'], \
input_droput_keep_pct:par['input_drop_keep_pct'], context_vector:cont_vect})
elif par['stabilization'] == 'EWC':
_,loss, AL, gl, weight_grads, h, entropy_loss = sess.run([model.train_op, model.task_loss, model.aux_loss, model.gate_loss,\
model.weight_grads, model.h, model.entropy_loss], feed_dict = \
{x:stim_in, y:y_hat, **gating_dict, mask:mk, droput_keep_pct:par['drop_keep_pct'], \
input_droput_keep_pct:par['input_drop_keep_pct'], context_vector:cont_vect})
if i // 500 == i / 500:
print('Iter: ', i, 'Loss: ', loss, 'Aux Loss: ', AL,
'gate loss ', gl, 'entropy loss', entropy_loss)
# Update big omegaes, and reset other values before starting new task
if par['stabilization'] == 'pathint':
big_omegas = sess.run(
[model.update_big_omega, model.big_omega_var])
elif par['stabilization'] == 'EWC':
for n in range(par['EWC_fisher_num_batches']):
stim_in, y_hat, mk = stim.make_batch(task, test=False)
_, _ = sess.run([model.update_big_omega,model.big_omega_var], feed_dict = \
{x:stim_in, y:y_hat, **gating_dict, mask:mk, droput_keep_pct:1.0, \
input_droput_keep_pct:1.0, context_vector:cont_vect})
big_omegas = sess.run([model.big_omega_var])
sess.run([model.reset_shunted_weights])
# Reset the Adam Optimizer, and set the previous parater values to their current values
sess.run(model.reset_adam_op)
sess.run(model.reset_prev_vars)
if par['stabilization'] == 'pathint':
sess.run(model.reset_small_omega)
# Test the netwroks on all trained tasks
num_test_reps = 10
accuracy = np.zeros((task + 1))
for test_task in range(task + 1):
cont_vect = np.zeros((1, par['n_tasks']), dtype=np.float32)
cont_vect[0, test_task] = 1
gating_dict = {
k: v
for k, v in zip(gating, par['gating'][test_task])
}
for r in range(num_test_reps):
stim_in, y_hat, mk = stim.make_batch(test_task, test=True)
acc = sess.run(model.accuracy, feed_dict={x:stim_in, y:y_hat, \
**gating_dict, mask:mk, droput_keep_pct:1.0, input_droput_keep_pct:1.0,\
context_vector:cont_vect})/num_test_reps
accuracy_grid[task, test_task] += acc
accuracy[test_task] += acc
print('Task ', task, ' Mean ', np.mean(accuracy), ' First ',
accuracy[0], ' Last ', accuracy[-1])
accuracy_full.append(np.mean(accuracy))
# reset weights between tasks if called upon
if par['reset_weights']:
sess.run(model.reset_weights)
above_zeros = []
for i in range(len(h)):
above_zeros.append(
np.float32(np.sum(h[i], axis=0, keepdims=True) > 1e-16))
print('mean h above zero ', np.mean(above_zeros[i]))
"""
for k in big_omegas[0].keys():
plt.imshow(big_omegas[0][k], aspect = 'auto')
plt.colorbar()
plt.show()
print(k, big_omegas[0][k].shape)
"""
if par['save_analysis']:
save_results = {'task': task, 'accuracy': accuracy, 'accuracy_full': accuracy_full, \
'accuracy_grid': accuracy_grid, 'big_omegas': big_omegas, 'par': par}
pickle.dump(save_results, open(par['save_dir'] + save_fn, 'wb'))
print('\nModel execution complete.')
def main(gpu_id=None):
""" Run training """
# Isolate requested GPU
if gpu_id is not None:
os.environ['CUDA_VISIBLE_DEVICES'] = gpu_id
# Reset TensorFlow graph
tf.reset_default_graph()
# Define placeholders
x = tf.placeholder(
tf.float32, [par['num_time_steps'], par['batch_size'], par['n_input']],
'stim')
y = tf.placeholder(
tf.float32,
[par['num_time_steps'], par['batch_size'], par['n_output']], 'out')
m = tf.placeholder(tf.float32, [par['num_time_steps'], par['batch_size']],
'mask')
# Set up stimulus and recording
stim = stimulus.Stimulus()
data_record = {
n: []
for n in
['iter', 'acc', 'task_loss', 'spike_loss', 'entropy_loss', 'spiking']
}
# Start TensorFlow session
gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=0.8
) if gpu_id == '0' else tf.GPUOptions()
with tf.Session(config=tf.ConfigProto(gpu_options=gpu_options)) as sess:
# Select CPU or GPU
device = '/cpu:0' if gpu_id is None else '/gpu:0'
with tf.device(device):
model = Model(x, y, m)
# Initialize variables and start timer
sess.run(tf.global_variables_initializer())
t_start = time.time()
# Begin training loop, iterating over tasks
for i in range(par['n_train_batches']):
# Generate a batch of stimulus data for training
trial_info = stim.make_batch()
# Put together the feed dictionary
feed_dict = {
x: trial_info['neural_input'],
y: trial_info['desired_output'],
m: trial_info['train_mask']
}
# Run the model
_, task_loss, output, state, spike, syn_x, syn_u = \
sess.run([model.train, model.task_loss, \
model.output, model.h, model.h_out, model.syn_x, model.syn_u], feed_dict=feed_dict)
# Display network performance
if i % 25 == 0:
spiking = (par['num_time_steps'] / par['dt']) * np.mean(spike)
acc = get_perf(trial_info['desired_output'], output,
trial_info['train_mask'])
data_record['iter'].append(i)
data_record['acc'].append(acc)
data_record['task_loss'].append(task_loss)
data_record['spiking'].append(spiking)
trials = 4
if i % 600 == -1 and i > 0:
for k in range(200):
fig, ax = plt.subplots(5, trials, figsize=[12, 8])
for b in range(trials):
ax[0, b].plot(state[:, b, k])
ax[1, b].plot(spike[:, b, k])
ax[2, b].plot(syn_x[:, b, k])
ax[3, b].plot(syn_u[:, b, k])
ax[4, b].plot(syn_x[:, b, k] * syn_u[:, b, k])
plt.savefig(
'./savedir/mt128_neuron{}_outputs.png'.format(k))
plt.clf()
plt.close()
trials = 4
fig, ax = plt.subplots(6, trials, figsize=[12, 8])
for b in range(trials):
ax[0, b].imshow(trial_info['neural_input'][:, b, :].T,
aspect='auto')
ax[1, b].imshow(trial_info['desired_output'][:, b, :].T,
aspect='auto')
ax[2, b].imshow(output[:, b, :].T, aspect='auto')
ax[3, b].imshow(state[:, b, :].T, aspect='auto')
ax[4, b].imshow(spike[:, b, :].T, aspect='auto')
ax[5, b].imshow((syn_u[:, b, :] * syn_x[:, b, :]).T,
aspect='auto')
ax[0, 0].set_ylabel('Network input')
ax[1, 0].set_ylabel('Expected Output')
ax[2, 0].set_ylabel('Network Output')
ax[3, 0].set_ylabel('Membrane Voltage')
ax[4, 0].set_ylabel('Spike Output')
ax[5, 0].set_ylabel('Synaptic Eff.')
ax[4, 0].set_xlabel('Time')
plt.savefig('./savedir/iter{}_outputs.png'.format(i))
plt.clf()
plt.close()
pickle.dump(
data_record,
open(par['savedir'] + par['save_fn'] + '.pkl', 'wb'))
print('Iter: {:>6} | Accuracy: {:5.3f} | Task Loss: {:5.3f} | Spike Rate: {:6.2f} Hz'.format(\
i, acc, task_loss, spiking))
if par['save_analysis']:
save_results = {
'task': task,
'accuracy_curve': accuracy_curve,
'loss_curve': loss_curve,
'par': par
}
pickle.dump(save_results,
open(par['savedir'] + save_fn + '.pkl', 'wb'))
print('\nModel execution complete.')
def main(gpu_id=None):
if gpu_id is not None:
os.environ["CUDA_VISIBLE_DEVICES"] = gpu_id
"""
Reset TensorFlow before running anything
"""
tf.reset_default_graph()
"""
Create the stimulus class to generate trial paramaters and input activity
"""
stim = stimulus.Stimulus()
f = pickle.load(open('./savedir/var_pulses_8_cue_off.pkl', 'rb'))
par['weights_trained'] = f['weights']
update_parameters(f['parameters'])
n_input, n_hidden, n_output = par['shape']
N = par[
'batch_train_size'] # trials per iteration, calculate gradients after batch_train_size
"""
Define all placeholder
"""
mask = tf.placeholder(
tf.float32, shape=[par['num_time_steps'], par['batch_train_size']])
x = tf.placeholder(
tf.float32,
shape=[n_input, par['num_time_steps'],
par['batch_train_size']]) # input data
y = tf.placeholder(
tf.float32,
shape=[n_output, par['num_time_steps'],
par['batch_train_size']]) # target data
config = tf.ConfigProto()
#config.gpu_options.allow_growth=True
# enter "config=tf.ConfigProto(log_device_placement=True)" inside Session to check whether CPU/GPU in use
with tf.Session(config=config) as sess:
device = '/cpu:0' if gpu_id is None else '/gpu:0'
with tf.device(device):
model = Model(x, y, mask)
init = tf.global_variables_initializer()
sess.run(init)
# keep track of the model performance across training
model_performance = {
'accuracy': [],
'pulse_accuracy': [],
'loss': [],
'perf_loss': [],
'spike_loss': [],
'trial': []
}
for i in range(par['num_iterations']):
# generate batch of batch_train_size
trial_info = stim.generate_trial(
analysis=False,
num_fixed=0,
var_delay=par['var_delay'],
var_resp_delay=par['var_resp_delay'],
var_num_pulses=par['var_num_pulses'])
if not par['var_num_pulses']:
onset = np.array([
np.unique(np.array(trial_info['timeline']))[-2 * p - 2]
for p in range(par['num_pulses'])
][::-1])
pulse_masks = np.array([
np.zeros((par['num_time_steps'], par['batch_train_size']),
dtype=np.float32)
] * par['num_pulses'])
for p in range(par['num_pulses']):
pulse_masks[p, onset[p] +
par['mask_duration'] // par['dt']:onset[p] +
par['sample_time'] // par['dt'], :] = 1
"""
Run the model
"""
_, loss, perf_loss, spike_loss, y_hat, state_hist, syn_x_hist, syn_u_hist = \
sess.run([model.train_op, model.loss, model.perf_loss, model.spike_loss, model.y_hat, \
model.hidden_state_hist, model.syn_x_hist, model.syn_u_hist], {x: trial_info['neural_input'], \
y: trial_info['desired_output'], mask: trial_info['train_mask']})
accuracy = analysis.get_perf(trial_info['desired_output'], y_hat,
trial_info['train_mask'])
pulse_accuracy = []
if not par['var_num_pulses']:
for p in range(par['num_pulses']):
pulse_accuracy.append(
analysis.get_perf(trial_info['desired_output'], y_hat,
pulse_masks[p]))
model_performance = append_model_performance(
model_performance, accuracy, pulse_accuracy, loss, perf_loss,
spike_loss, (i + 1) * N)
"""
Save the network model and output model performance to screen
"""
if i % par['iters_between_outputs'] == 0 and i > 0:
print_results(i, N, perf_loss, spike_loss, state_hist,
accuracy)
if i % 5000 == 0:
weights = eval_weights()
syn_x_stacked = np.stack(syn_x_hist, axis=1)
syn_u_stacked = np.stack(syn_u_hist, axis=1)
h_stacked = np.stack(state_hist, axis=1)
trial_time = np.arange(0, h_stacked.shape[1] * par['dt'],
par['dt'])
mean_h = np.mean(np.mean(h_stacked, axis=2), axis=1)
results = {
'model_performance': model_performance,
'parameters': par,
'weights': weights,
'trial_time': trial_time,
'mean_h': mean_h,
'timeline': trial_info['timeline']
}
pickle.dump(results,
open(par['save_dir'] + par['save_fn'], 'wb'))
if accuracy > 0.995:
weights = eval_weights()
syn_x_stacked = np.stack(syn_x_hist, axis=1)
syn_u_stacked = np.stack(syn_u_hist, axis=1)
h_stacked = np.stack(state_hist, axis=1)
trial_time = np.arange(0, h_stacked.shape[1] * par['dt'],
par['dt'])
mean_h = np.mean(np.mean(h_stacked, axis=2), axis=1)
results = {
'model_performance': model_performance,
'parameters': par,
'weights': weights,
'trial_time': trial_time,
'mean_h': mean_h,
'timeline': trial_info['timeline']
}
pickle.dump(results,
open(par['save_dir'] + par['save_fn'], 'wb'))
for b in range(10):
plot_list = [
trial_info['desired_output'][:, :, b],
softmax(
np.array(y_hat)[:, :, b].T -
np.max(np.array(y_hat)[:, :, b].T))
]
fig, axes = plt.subplots(nrows=2, ncols=1, figsize=(7, 7))
j = 0
for ax in axes.flat:
im = ax.imshow(plot_list[j], aspect='auto')
j += 1
cax, kw = mpl.colorbar.make_axes([ax for ax in axes.flat])
plt.colorbar(im, cax=cax, **kw)
plt.savefig("./savedir/output_" + str(par['num_pulses']) +
"pulses_iter_" + str(i) + "_" + str(b) +
".png")
plt.close()
plt.imshow(trial_info['neural_input'][:, :, b])
plt.savefig("./savedir/input_" + str(par['num_pulses']) +
"pulses_iter_" + str(i) + "_" + str(b) +
".png")
plt.close()
break
"""
Save model, analyze the network model and save the results
"""
#save_path = saver.save(sess, par['save_dir'] + par['ckpt_save_fn'])
if par['analyze_model']:
weights = eval_weights()
syn_x_stacked = np.stack(syn_x_hist, axis=1)
syn_u_stacked = np.stack(syn_u_hist, axis=1)
h_stacked = np.stack(state_hist, axis=1)
trial_time = np.arange(0, h_stacked.shape[1] * par['dt'],
par['dt'])
mean_h = np.mean(np.mean(h_stacked, axis=2), axis=1)
results = {
'model_performance': model_performance,
'parameters': par,
'weights': weights,
'trial_time': trial_time,
'mean_h': mean_h,
'timeline': trial_info['timeline']
}
pickle.dump(results, open(par['save_dir'] + par['save_fn'], 'wb'))
def main(task, gpu_id):
# If a gpu id has been provided, use it
if gpu_id is not None:
# Specify the gpu id to be used
os.environ["CUDA_VISIBLE_DEVICES"] = gpu_id
# Reset TensorFlow before running anything
tf.reset_default_graph()
# Tensorflow finds the supported CPU and GPU devices you can use
config = tf.ConfigProto()
trial_length = par['num_time_steps']
# Calculate shape of the stimulus for this task
# Define a placeholder for the stimulus the agent sees
stimulus = tf.placeholder(
tf.float64,
shape=[trial_length, task.num_inputs, par['batch_train_size']])
# Define a placeholder for the truth or correct answer about each trial
truth = tf.placeholder(tf.float64, shape=par['batch_train_size'])
# A TEMPORARY placeholder for target
target = tf.placeholder(tf.float64,
shape=[trial_length, 3, par['batch_train_size']])
# Create a model for the given task object
M = Model()
# Build the tf structure that runs trials
M.run_model(task, stimulus, target)
M.optimize(task, target)
par['num_receptive_fields'] = 1
par['num_fix_tuned'] = 0
par['num_rule_tuned'] = 0
par['num_rules'] = 1
par['variable_delay_max'] = 100
par['mask_duration'] = 0
par['n_output'] = 3
par['input_mean'] = 0
par['n_input'] = 36
par['catch_trial_pct'] = 0
par['rotation_match'] = 0
# Create a model from Nick's code
stim = stm.Stimulus()
with tf.Session(config=config) as sess:
init = tf.global_variables_initializer()
sess.run(init)
t_start = time.time()
for it in range(par['num_iterations']):
# Create a batch of stimuli, stores in attribute stimulus for the task
trial_info = stim.generate_trial()
"""
Run the model
"""
this_truth = np.zeros(par['batch_train_size'])
this_truth[trial_info['desired_output'][1, -1, :] ==
1] = 1 # Match trials
this_truth[trial_info['desired_output'][2, -1, :] ==
1] = 2 # Non_match trials
_, _, vloss, ploss, pol_out, reward, this_action_array, this_time_mask, baseline, advantage = \
sess.run([M.pol_train_op, M.val_train_op, M.Loss_val, M.Loss_pol, task.pol_out_history, task.reward, task.action_array, \
task.time_mask, task.val_out_history, M.advantage], \
{stimulus: np.swapaxes(trial_info['neural_input'], 1, 0), truth: this_truth, target: np.swapaxes(trial_info['desired_output'], 1, 0)})
reward = reward * this_time_mask
baseline = np.array(baseline) * this_time_mask
pol_out = np.array(pol_out)
if it % 100 == 0:
plt.subplot(2, 2, 1)
plt.plot(pol_out[:, :, 0])
plt.legend(['Fixate', 'match', 'Non-match'])
if this_truth[0] == 1:
plt.title('match' + '__' + str(reward.sum()))
elif this_truth[0] == 2:
plt.title('non-match' + '__' + str(reward.sum()))
plt.subplot(2, 2, 2)
temp = np.tile(this_time_mask[:, 0].transpose(),
[3, 1]).transpose() * this_action_array[:, :, 0]
plt.plot(temp)
#plt.plot(advantage[:,0])
#plt.plot(Pan_results['ploss'].transpose())
#plt.plot(Pan_results['vloss'].transpose())
#plt.legend(['-J', 'E'])
plt.subplot(2, 2, 3)
plt.plot(np.array(baseline)[:, 0])
#plt.plot(my_reward_cum[:,0])
#pdb.set_trace()
plt.plot(reward[:, 0])
plt.title('baseline')
plt.plot(advantage[:, 0])
plt.subplot(2, 2, 4)
#plt.plot(np.squeeze(np.array(val_out))[:,0])
#plt.title('his baseline')
#plt.plot(trial_info['neural_input'][:,:,0].transpose())
plt.plot(
np.swapaxes(trial_info['desired_output'], 1, 0)[:, :, 0])
plt.title('Neural input')
plt.savefig(par['save_path'] + 'Iteration_' + str(it) +
'.png') # save the figure to file
plt.close()
print('%5d | Vloss: %6.6f | Reward: %4d' %
(it, vloss, reward.sum()))
def main(gpu_id=None):
if gpu_id is not None:
os.environ["CUDA_VISIBLE_DEVICES"] = gpu_id
"""
Print key parameters
"""
print_important_params()
"""
Reset TensorFlow before running anything
"""
tf.reset_default_graph()
"""
Create the stimulus class to generate trial paramaters and input activity
"""
stim = stimulus.Stimulus()
n_input, n_hidden, n_output = par['shape']
N = par[
'batch_train_size'] # trials per iteration, calculate gradients after batch_train_size
"""
Define all placeholder
"""
mask = tf.placeholder(
tf.float32, shape=[par['num_time_steps'], par['batch_train_size']])
x = tf.placeholder(
tf.float32,
shape=[n_input, par['num_time_steps'],
par['batch_train_size']]) # input data
y = tf.placeholder(
tf.float32,
shape=[n_output, par['num_time_steps'],
par['batch_train_size']]) # target data
config = tf.ConfigProto()
# enter "config=tf.ConfigProto(log_device_placement=True)" inside Session to check whether CPU/GPU in use
with tf.Session(config=config) as sess:
device = '/cpu:0' if gpu_id is None else '/gpu:0'
with tf.device(device):
model = Model(x, y, mask)
sess.run(tf.global_variables_initializer())
# keep track of the model performance across training
model_performance = {
'accuracy': [],
'loss': [],
'perf_loss': [],
'spike_loss': [],
'weight_loss': [],
'trial': []
}
for i in range(par['num_iterations']):
# generate batch of batch_train_size
trial_info = stim.generate_trial(set_rule=None)
"""
Run the model
"""
_, loss, perf_loss, spike_loss, weight_loss, y_hat, state_hist, syn_x_hist, syn_u_hist = \
sess.run([model.train_op, model.loss, model.perf_loss, model.spike_loss, model.weight_loss, model.y_hat, \
model.hidden_state_hist, model.syn_x_hist, model.syn_u_hist], {x: trial_info['neural_input'], \
y: trial_info['desired_output'], mask: trial_info['train_mask']})
accuracy, _, _ = analysis.get_perf(trial_info['desired_output'],
y_hat, trial_info['train_mask'])
model_performance = append_model_performance(
model_performance, accuracy, loss, perf_loss, spike_loss,
weight_loss, (i + 1) * N)
"""
Save the network model and output model performance to screen
"""
if i % par['iters_between_outputs'] == 0:
print_results(i, N, perf_loss, spike_loss, weight_loss,
state_hist, accuracy)
"""
Save model, analyze the network model and save the results
"""
save_results(model_performance)
def main(save_fn, gpu_id=None):
""" Run supervised learning training """
# Update all dependencies in parameters
update_dependencies()
# Isolate requested GPU
if gpu_id is not None:
os.environ["CUDA_VISIBLE_DEVICES"] = gpu_id
# If desired, train the convolutional layers with the CIFAR datasets
# Otherwise, the network will load convolutional weights from the saved file
if (par['task'] in ['cifar', 'imagenet', 'colored_mnist'
]) and par['train_convolutional_layers']:
convolutional_layers.ConvolutionalLayers()
print('\nRunning model.\n')
# Reset TensorFlow graph
tf.reset_default_graph()
# Create placeholders for the model
if par['task'] == 'mnist':
x = tf.placeholder(tf.float32,
[par['batch_size'], par['layer_dims'][0]], 'stim')
elif par['task'] == 'colored_mnist':
x = tf.placeholder(tf.float32, [par['batch_size'], 32, 32, 3], 'stim')
elif par['task'] == 'cifar' or par['task'] == 'imagenet':
x = tf.placeholder(tf.float32, [par['batch_size'], 32, 32, 3], 'stim')
y = tf.placeholder(tf.float32, [par['batch_size'], par['layer_dims'][-1]],
'out')
mask = tf.placeholder(tf.float32,
[par['batch_size'], par['layer_dims'][-1]], 'mask')
rule = tf.placeholder(tf.float32, [par['batch_size'], par['n_tasks']],
'rulecue')
gating = [
tf.placeholder(tf.float32, [par['layer_dims'][n + 1]], 'gating')
for n in range(par['n_layers'] - 1)
]
droput_keep_pct = tf.placeholder(tf.float32, [], 'dropout')
input_droput_keep_pct = tf.placeholder(tf.float32, [], 'input_dropout')
# Set up stimulus
if par['task'] == 'colored_mnist':
stim = stimulus.Stimulus(labels_per_task=par['labels_per_task'],
sep=par['separability'])
else:
stim = stimulus.Stimulus(labels_per_task=par['labels_per_task'])
# Initialize accuracy records
accuracy_full = []
accuracy_grid = np.zeros((par['n_tasks'], par['n_tasks']))
# Enter TensorFlow session
with tf.Session() as sess:
# Select CPU or GPU
device = '/cpu:0' if gpu_id is None else '/gpu:0'
with tf.device(device):
model = Model(x, y, gating, mask, droput_keep_pct,
input_droput_keep_pct, rule)
# Initialize variables
sess.run(tf.global_variables_initializer())
sess.run(model.reset_prev_vars)
# Begin training loop, iterating over tasks
for task in range(par['n_tasks']):
# Create dictionary of gating signals applied to each hidden layer for this task
gating_dict = {k: v for k, v in zip(gating, par['gating'][task])}
# Create rule cue vector for this task
rule_cue = np.zeros([par['batch_size'], par['n_tasks']])
rule_cue[:, task] = 1
# Iterate over batches
for i in range(par['n_train_batches']):
# Make batch of training data
stim_in, y_hat, mk = stim.make_batch(task, test=False)
# Run the model using one of the available stabilization methods
if par['stabilization'] == 'pathint':
_, _, loss, AL = sess.run([model.train_op, model.update_small_omega, model.task_loss, model.aux_loss], \
feed_dict={x:stim_in, y:y_hat, **gating_dict, mask:mk, droput_keep_pct:par['drop_keep_pct'], \
input_droput_keep_pct:par['input_drop_keep_pct'], rule:rule_cue})
elif par['stabilization'] == 'EWC':
_, loss, AL = sess.run([model.train_op, model.task_loss, model.aux_loss], \
feed_dict={x:stim_in, y:y_hat, **gating_dict, mask:mk, droput_keep_pct:par['drop_keep_pct'], \
input_droput_keep_pct:par['input_drop_keep_pct'], rule:rule_cue})
# Display network performance
if i % 500 == 0:
print('Iter: ', i, 'Loss: ', loss, 'Aux Loss: ', AL)
# Update big omegaes, and reset other values before starting new task
if par['stabilization'] == 'pathint':
sess.run(model.update_big_omega)
elif par['stabilization'] == 'EWC':
for _ in range(par['EWC_batch_divisor'] *
par['EWC_fisher_num_batches']):
stim_in, _, mk = stim.make_batch(task, test=False)
sess.run([model.update_big_omega], feed_dict = \
{x:stim_in, **gating_dict, mask:mk, droput_keep_pct:par['drop_keep_pct'], \
input_droput_keep_pct:par['input_drop_keep_pct'], rule:rule_cue})
# Reset the Adam Optimizer, and set the prev_weight values to their current values
sess.run(model.reset_adam_op)
sess.run(model.reset_prev_vars)
if par['stabilization'] == 'pathint':
sess.run(model.reset_small_omega)
# Test the networks on all trained tasks
num_test_reps = 10
accuracy = np.zeros((task + 1))
for test_task in range(task + 1):
# Use appropriate gating and rule cues
gating_dict = {
k: v
for k, v in zip(gating, par['gating'][test_task])
}
test_rule_cue = np.zeros([par['batch_size'], par['n_tasks']])
test_rule_cue[:, test_task] = 1
# Repeat the test as desired
for r in range(num_test_reps):
stim_in, y_hat, mk = stim.make_batch(test_task, test=True)
acc = sess.run(model.accuracy, feed_dict={x:stim_in, y:y_hat, \
**gating_dict, mask:mk, droput_keep_pct:1.0, input_droput_keep_pct:1.0, rule:test_rule_cue})/num_test_reps
accuracy_grid[task, test_task] += acc
accuracy[test_task] += acc
# Display network performance after testing is complete
print('Task ', task, ' Mean ', np.mean(accuracy), ' First ',
accuracy[0], ' Last ', accuracy[-1])
accuracy_full.append(np.mean(accuracy))
# Reset weights between tasks if called upon
if par['reset_weights']:
sess.run(model.reset_weights)
# Save model performance and parameters if desired
if par['save_analysis']:
save_results = {'task': task, 'accuracy': accuracy, 'accuracy_full': accuracy_full, \
'accuracy_grid': accuracy_grid, 'par': par}
pickle.dump(save_results, open(par['save_dir'] + save_fn, 'wb'))
print('\nModel execution complete.')
# write accuracy full to text
with open('accs_200_squared.txt', 'a') as f:
f.write(str(accuracy_full[0]) + "\n")
def analyze_model_from_file(filename, savefile = None, update_params = {}):
""" The first section loads the model weights and simulates the network on
several different task conditions, saving the network activity and output """
results = pickle.load(open(filename, 'rb'))
if savefile is None:
results['parameters']['save_fn'] = 'test.pkl'
else:
results['parameters']['save_fn'] = savefile
update_parameters(results['parameters'])
update_parameters(update_params)
stim = stimulus.Stimulus()
# generate trials with match probability at 50%
trial_info = stim.generate_trial(test_mode = True)
input_data = np.squeeze(np.split(trial_info['neural_input'], par['num_time_steps'], axis=0))
h_init = results['weights']['h']
y, h, syn_x, syn_u = run_model(input_data, h_init, \
results['parameters']['syn_x_init'], results['parameters']['syn_u_init'], results['weights'])
# generate trials with random sample and test stimuli, used for decoding
trial_info_decode = stim.generate_trial(test_mode = True)
input_data = np.squeeze(np.split(trial_info_decode['neural_input'], par['num_time_steps'], axis=0))
_, h_decode, syn_x_decode, syn_u_decode = run_model(input_data, h_init, \
results['parameters']['syn_x_init'], results['parameters']['syn_u_init'], results['weights'])
# generate trials using DMS task, only used for measuring how neuronal and synaptic representations evolve in
# a standardized way, used for figure correlating persistent activity and manipulation
update_parameters({'trial_type': 'DMS'})
trial_info_dms = stim.generate_trial(test_mode = True)
input_data = np.squeeze(np.split(trial_info_dms['neural_input'], trial_info_dms['neural_input'].shape[0], axis=0))
_, h_dms, syn_x_dms, syn_u_dms = run_model(input_data, h_init, \
results['parameters']['syn_x_init'], results['parameters']['syn_u_init'], results['weights'])
update_parameters(results['parameters']) # reset trial type to original value
update_parameters(update_params)
""" The next section performs various analysis """
# calculate task accuracy
results['task_accuracy'],_,_ = get_perf(trial_info['desired_output'], y, trial_info['train_mask'])
results['task_accuracy_per_rule'] = []
for r in np.unique(trial_info['rule']):
ind = np.where(trial_info['rule'] == r)[0]
acc, _, _ = get_perf(trial_info['desired_output'][:,ind,:], y[:,ind,:], trial_info['train_mask'][:, ind])
results['task_accuracy_per_rule'].append(acc)
print('Task accuracy',results['task_accuracy'])
if par['calculate_resp_matrix']:
print('calculate response matrix...')
resp_matrix_results = calculate_response_matrix(trial_info_decode, results['weights'])
for key, val in resp_matrix_results.items():
if np.var(val) > 0:
results[key] = val
# Decode the sample direction from neuronal activity and synaptic efficacies using support vector machines
trial_time = np.arange(0,h.shape[0]*par['dt'], par['dt'])
trial_time_dms = np.arange(0,h_dms.shape[0]*par['dt'], par['dt'])
if par['decoding_reps'] > 0:
print('decoding activity...')
decoding_results = calculate_svms(h_decode, syn_x_decode, syn_u_decode, trial_info_decode, trial_time, \
num_reps = par['decoding_reps'], num_reps_stability = 10, decode_test = par['decode_test'], \
decode_rule = par['decode_rule'], decode_match = par['decode_match'])
for key, val in decoding_results.items():
if np.var(val) > 0:
results[key] = val
if par['trial_type'] in ['DMS', 'DMC', 'DMRS90', 'DMRS90ccw', 'DMRS45', 'DMRS180', 'location_DMS']:
for key, val in decoding_results.items():
if np.var(val) > 0:
results[key + '_dms'] = val
else:
# Calculate decoding for a DMS trial, used to correlate persistent activity and manipulation
update_parameters({'trial_type': 'DMS'})
decoding_results = calculate_svms(h_dms, syn_x_dms, syn_u_dms, trial_info_dms, trial_time_dms, \
num_reps = par['decoding_reps'], num_reps_stability = 0, decode_test = par['decode_test'], decode_rule = par['decode_rule'])
for key, val in decoding_results.items():
if np.var(val) > 0:
results[key + '_dms'] = val
update_parameters(results['parameters'])
update_parameters(update_params)
# Calculate neuronal and synaptic sample motion tuning
if par['analyze_tuning']:
print('calculate tuning...')
tuning_results = calculate_tuning(h_decode, syn_x_decode, syn_u_decode, \
trial_info_decode, trial_time, results['weights'], calculate_test = par['decode_test'])
for key, val in tuning_results.items():
if np.var(val) > 0:
results[key] = val
# Calculate tuning for a DMS trial, used to correlate persistent activity and manipulation
if par['trial_type'] in ['DMS', 'DMC', 'DMRS90', 'DMRS90ccw','DMRS45', 'DMRS180', 'location_DMS']:
for key, val in tuning_results.items():
if np.var(val) > 0:
results[key + '_dms'] = val
else:
update_parameters({'trial_type': 'DMS'})
tuning_results = calculate_tuning(h_dms, syn_x_dms, syn_u_dms, \
trial_info_dms, trial_time_dms, results['weights'], calculate_test = False)
for key, val in tuning_results.items():
if np.var(val) > 0:
results[key + '_dms'] = val
update_parameters(results['parameters'])
update_parameters(update_params)
# Calculate mean sample traces
results['h_sample_mean'] = np.zeros((results['parameters']['num_time_steps'], results['parameters']['n_hidden'], \
results['parameters']['num_motion_dirs']), dtype = np.float32)
for i in range(results['parameters']['num_motion_dirs']):
ind = np.where(trial_info_decode['sample'] == i)[0]
results['h_sample_mean'][:,:,i] = np.mean(h_decode[:,ind,:], axis = 1)
# Calculate the neuronal and synaptic contributions towards solving the task
if par['simulation_reps'] > 0:
print('simulating network...')
simulation_results = simulate_network(trial_info, h, syn_x, \
syn_u, results['weights'], num_reps = par['simulation_reps'])
for key, val in simulation_results.items():
if np.var(val) > 0:
results[key] = val
# Save the analysis results
pickle.dump(results, open(savefile, 'wb') )
print('Analysis results saved in ', savefile)
