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, 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():
"""
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=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(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(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()
# 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')
# 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'
with tf.device(device):
model = Model(x, t, m)
sess.run(tf.global_variables_initializer())
# keep track of the model performance across training
model_performance = {'accuracy': [], 'loss': [], 'perf_loss': [], 'spike_loss': [], \
'weight_loss': [], 'iteration': []}
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, 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']})
accuracy, _, _ = analysis.get_perf(trial_info['desired_output'], y,
trial_info['train_mask'])
model_performance = append_model_performance(
model_performance, accuracy, loss, perf_loss, spike_loss,
weight_loss, i)
# Save the network model and output model performance to screen
if i % par['iters_between_outputs'] == 0:
print_results(i, perf_loss, spike_loss, weight_loss, h,
accuracy)
# Save model and results
weights = sess.run(model.var_dict)
save_results(model_performance, weights)
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()
# 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}/')
# Find the highest ID saved for this task;
save_increment = 0
ids = glob.glob(f'savedir/{gpu_id}/{par["trial_type"]}*.pkl')
if len(ids) > 0:
nums = [
int(i[i.find(par['trial_type']) + len(par['trial_type']):-4])
for i in ids
]
save_increment = max(nums) + 1
# 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)}")
sess.run(tf.global_variables_initializer())
# keep track of the model performance across training
model_performance = {'accuracy': [], 'loss': [], 'perf_loss': [], 'spike_loss': [], \
'weight_loss': [], 'iteration': []}
t0 = time.time()
for j in range(save_increment, par['num_network_sets_per_gpu']):
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, 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(trial_info['desired_output'],
y[n, :, :, :],
trial_info['train_mask'])[0]
for n in range(par['n_networks'])
]
model_performance = append_model_performance(
model_performance, accuracies, loss, perf_loss, spike_loss,
weight_loss, i)
# Save the network model and output model performance to screen
if (i + 1) % par['iters_between_outputs'] == 0:
t1 = time.time()
print_results(i, perf_loss, spike_loss, weight_loss, h,
accuracies)
print(f"Elapsed time: {str(t1 - t0)}")
t0 = time.time()
# Save model and results
weights = sess.run(model.var_dict)
save_results(model_performance,
weights,
save_fn=f'{str(gpu_id)}/{par["trial_type"]}{j}.pkl')
# After each bunch: clear history, reset all weights, run again
update_trial_params()
update_dependencies()
sess.run([model.reset_vars_ops, model.reset_opt])
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()
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:
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 batch_train_size
trial_info = stim.generate_trial()
"""
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'])
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 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, tuning = False, decoding = False, load_previous_file = False, save_raw_data = False)
# Generate another batch of trials with decoding_test_mode = True (sample and test stimuli
# are independently drawn), and then perform tuning and decoding analysis
update = {'decoding_test_mode': True}
update_parameters(update)
trial_info = stim.generate_trial()
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, tuning = par['analyze_tuning'], decoding = True, load_previous_file = True, save_raw_data = False)
if par['trial_type'] == 'dualDMS':
# run an additional session with probe stimuli
save_fn = 'probe_' + par['save_fn']
update = {'probe_trial_pct': 1, 'save_fn': save_fn}
update_parameters(update)
trial_info = stim.generate_trial()
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, tuning = False, decoding = True, \
load_previous_file = False, save_raw_data = False)
