def test_output(self, none_fields):
"""Tests that this function produces the identity."""
graph = self._graph.map(lambda _: None, none_fields)
with tf.name_scope("test"):
graph_id = utils_tf.identity(graph)
graph = utils_tf.make_runnable_in_session(graph)
graph_id = utils_tf.make_runnable_in_session(graph_id)
with self.test_session() as sess:
expected_out, actual_out = sess.run([graph, graph_id])
for field in [
"nodes", "edges", "globals", "receivers", "senders", "n_node",
"n_edge"
]:
if field in none_fields:
self.assertEqual(None, getattr(actual_out, field))
else:
self.assertNDArrayNear(getattr(expected_out, field),
getattr(actual_out, field),
err=1e-4)
def test_getitem_one(self):
index = 2
expected = self.graphs_dicts_out[index]
graphs_tuple = utils_tf.data_dicts_to_graphs_tuple(
self.graphs_dicts_in)
graph_op = utils_tf.get_graph(graphs_tuple, index)
graph_op = utils_tf.make_runnable_in_session(graph_op)
with self.test_session() as sess:
graph = sess.run(graph_op)
actual, = utils_np.graphs_tuple_to_data_dicts(graph)
for k, v in expected.items():
self.assertAllClose(v, actual[k])
self.assertEqual(expected["nodes"].shape[0], actual["n_node"])
self.assertEqual(expected["edges"].shape[0], actual["n_edge"])
def test_getitem(self):
index = slice(1, 3)
expected = self.graphs_dicts_out[index]
graphs_tuple = utils_tf.data_dicts_to_graphs_tuple(
self.graphs_dicts_in)
graphs2_op = utils_tf.get_graph(graphs_tuple, index)
graphs2_op = utils_tf.make_runnable_in_session(graphs2_op)
with self.test_session() as sess:
graphs2 = sess.run(graphs2_op)
actual = utils_np.graphs_tuple_to_data_dicts(graphs2)
for ex, ac in zip(expected, actual):
for k, v in ex.items():
self.assertAllClose(v, ac[k])
self.assertEqual(ex["nodes"].shape[0], ac["n_node"])
self.assertEqual(ex["edges"].shape[0], ac["n_edge"])
def make_all_runnable_in_session(*args):
"""Apply make_runnable_in_session to an iterable of graphs."""
return [utils_tf.make_runnable_in_session(a) for a in args]
def _make_all_runnable_in_session(self, *args): """Lets an iterable of TF graphs be output from a session as NP graphs.""" return [utils_tf.make_runnable_in_session(a) for a in args]
def __init__(self,
obs_dim,
act_dim,
batch_size,
n_objects,
state_embedding_dim,
reward_lr=0.001,
transition_lr=0.001,
hidden_dim_forward_model=32,
hidden_dim_reward_model=32,
folder='results/',
random_samples=True,
model_name='model'):
self.act_dim = act_dim # dimension of the action (one-hot encoded)
self.obs_dim = obs_dim # dimension of the observation (state) vector
self.batch_size = batch_size # batch size to use for the replay buffer
self.state_dim_embedding = state_embedding_dim # dimension of the state embedding per object
self.learning_rate_transition = transition_lr # learning rate for the transition model
self.learning_rate_reward = reward_lr # learning rate for the reward model
self.n_objects = n_objects # one for the agent, one for the exit, object
self.hidden_dim_transition_model = hidden_dim_forward_model # hidden dimension forward model
self.hidden_dim_reward_model = hidden_dim_reward_model # hidden dimension reward model
self.random_samples = random_samples # randomly sample in replay buffer
self.l2_reg = 0.0 # weight normalization norm
# normalization constants for the contrastive loss
self.sigma = 0.5
self.hinge = 1 # hinge
self.model_folder = os.path.join(folder, 'trained_models', model_name)
self.model_path = os.path.join(folder, 'trained_models', model_name,
model_name +
'.ckpt') # path where model is stored
self.model_name = model_name # name of the model
self.graph = tf.Graph()
with self.graph.as_default():
# DEFINE PLACEHOLDERS
self.obs_var = tf.placeholder(tf.float32,
shape=[None, self.obs_dim],
name="obs_var") # state
self.next_obs_var = tf.placeholder(
tf.float32, shape=[None, self.obs_dim],
name="next_obs_var") # next state
self.neg_obs_var = tf.placeholder(
tf.float32, shape=[None, self.obs_dim],
name="neg_obs_var") # negative sample
self.action_var = tf.placeholder(
tf.float32, shape=[None, self.act_dim],
name="action_var") # action (one-hot)
self.reward_var = tf.placeholder(tf.float32,
shape=[None],
name='reward_var') # reward
self.is_training = tf.placeholder(
tf.bool, shape=[],
name="train_cond") # training bool (Training -> True)
# NN FORWARD (transition) MODEL AND REWARD MODEL
self.forward_model_nn = snt.Module(self.forward_model_nn_network,
name='Forward_nn_Network')
self.reward_model_nn = snt.Module(self.reward_model_nn_network,
name='Reward_nn_Network')
self.delta_state_nn_pred = self.forward_model_nn(
self.obs_var, self.action_var)
self.reward_nn_pred = self.reward_model_nn(self.obs_var,
self.action_var)
# GRAPH RELATED PLACEHOLDERS
# self.n_nodes_ph =tf.placeholder(tf.int32, shape=[None], name="n_nodes")
# self.n_edges_ph = tf.placeholder(tf.int32, shape=[None], name="n_edges")
# self.node_attributes_ph = tf.placeholder(tf.float32, shape=[None, self.n_objects, self.state_dim_embedding], name="node_attributes")
# put the object encodings in a graph, nodes are the 2D position hopefully and the global can be the action
self.obs_per_object = self.reshape_observation_vec_to_object_encoding(
self.obs_var)
self.obs_graph = self.object_encoder_graph(self.obs_per_object)
self.next_obs_graph = self.object_encoder_graph(
self.reshape_observation_vec_to_object_encoding(
self.next_obs_var))
# GNN FORWARD (TRANSITION) MODEL AND REWARD MODEL
self.forward_model_gnn = snt.Module(self.forward_model_gnn_network,
name='Forward_gnn_Network')
self.reward_model_gnn = snt.Module(self.reward_model_gnn_network,
name='Reward_gnn_Network')
self.delta_state_gnn_pred = utils_tf.make_runnable_in_session(
self.forward_model_gnn(self.obs_graph, self.action_var))
self.reward_gnn_pred = self.reward_model_gnn(
self.obs_graph, self.action_var)
# MSE LOSS NN
self.reward_loss_nn_mse = tf.reduce_mean(
tf.squared_difference(self.reward_var,
self.reward_nn_pred)) # MSE(r, r^)
self.state_transition_loss_nn_mse = tf.reduce_mean(
tf.squared_difference(
(self.delta_state_nn_pred + self.obs_var),
self.next_obs_var)) # MSE(delta_s + s, s')
# MSE LOSS GNN
self.predicted_next_obs_graph = self.obs_graph # should we make a copy or something?
self.predicted_next_obs_graph = self.predicted_next_obs_graph.replace(
nodes=self.obs_graph.nodes + self.delta_state_gnn_pred.nodes
) # sum node features to get the predicted next state
self.state_transition_loss_gnn_mse = tf.reduce_mean(
tf.squared_difference(self.predicted_next_obs_graph.nodes,
self.next_obs_graph.nodes))
self.reward_loss_gnn_mse = tf.reduce_mean(
tf.squared_difference(self.reward_var, self.reward_gnn_pred))
# # CONTRASTIVE LOSSES
# # self.state_transition_nn_loss_hinge = tf.reduce_mean(
# # tf.maximum(0.0, self.hinge - 0.5 * tf.squared_difference(self.encoded_state + self.latent_action_prediction, self.encoded_state_2)))
#
# # SOME CONSTANTS:
# self.norm = tf.constant((0.5 / self.sigma**2))
# self.H = tf.reduce_mean(self.norm * tf.squared_difference((self.delta_state_nn_pred + self.encoded_state_vec), self.encoded_next_state_vec)) # batch_sum( d(s_t + delta s_t, s_t+1) ) eqn 5, first part
# self.H_tilde = tf.reduce_mean(self.norm * tf.squared_difference(self.encoded_neg_state_vec, self.encoded_next_state_vec)) # batch_sum ( d(s_tilde, s_t+1) ) eqn 5, second part
# self.state_transition_nn_loss_hinge = self.H + tf.maximum(0.0, (self.hinge - self.H_tilde)) # loss function as in eqn 6
# optimizer = tf.train.AdamOptimizer(learning_rate=self.learning_rate)
# # self.train_nn_model = optimizer_reward.minimize(sum(self.reward_loss_nn_mse, self.state_transition_nn_loss))
# self.train_fwd_model_op = optimizer.minimize(self.state_transition_nn_loss_hinge)
# TRAINING - FORWARD MODEL NN
optimizer_transition_nn = tf.train.AdamOptimizer(
self.learning_rate_transition)
self.train_op_fwd_model_nn = optimizer_transition_nn.minimize(
self.state_transition_loss_nn_mse)
# TRAINING - REWARD FUNCTION NN
optimizer_reward_nn = tf.train.AdamOptimizer(
self.learning_rate_reward)
self.train_op_rwrd_model_nn = optimizer_reward_nn.minimize(
self.reward_loss_nn_mse)
# TRAINING - FORWARD MODEL GNN
optimizer_transition_gnn = tf.train.AdamOptimizer(
self.learning_rate_transition)
self.train_op_fwd_model_gnn = optimizer_transition_gnn.minimize(
self.state_transition_loss_gnn_mse)
# TRAINING - REWARD FUNCTION NN
optimizer_reward_gnn = tf.train.AdamOptimizer(
self.learning_rate_reward)
self.train_op_rwrd_model_gnn = optimizer_reward_gnn.minimize(
self.reward_loss_gnn_mse)
# SET THE REPLAY MEMORY FOR THE NETWORK --------------------------------------------------------------
self.replay_buffer = ReplayBuffer(
self.obs_dim,
self.act_dim,
30000,
random_samples=self.random_samples)
# Init session --------------------------------------------------------------------------------------------
self.saver = tf.train.Saver()
self.init = tf.global_variables_initializer()
# Initialize the session
self.sess = tf.Session(graph=self.graph)
self.sess.run(self.init)
def make_all_runnable_in_session(*args):
from graph_nets import utils_tf
"""Lets an iterable of TF graphs be output from a session as NP graphs."""
return [utils_tf.make_runnable_in_session(a) for a in args]
def main():
# A bunch of configuration stuff to clean up...
parser = argparse.ArgumentParser(
description='Train nx-graph with configurations')
add_arg = parser.add_argument
add_arg('name', nargs='?', default='unnamed')
args = parser.parse_args()
results_dir = 'results/{}'.format(args.name)
os.makedirs(results_dir, exist_ok=True)
config = load_config('configs/nxgraph_default.yaml')
base_dir = config['data']['input_dir']
config_tr = config['train']
log_every_seconds = config_tr['time_lapse']
batch_size = config_tr['batch_size'] # need optimization
num_training_iterations = config_tr['iterations']
iter_per_job = config_tr['iter_per_job']
num_processing_steps_tr = config_tr['n_iters'] ## level of message-passing
prod_name = config['prod_name']
learning_rate = config_tr['learning_rate']
output_dir = os.path.join(config['output_dir'], prod_name)
# Start to build tensorflow sessions
tf.reset_default_graph()
# Creates a placeholder for training examples. The placeholders define a
# slot for training examples given in feed dict to be assigned. We create
# graphs.GraphsTuple placeholders using the graph_nets utility functions.
# They are automatically generated from the first graph in the first batch.
# By assigning force_dynamic_num_graphs=True, we ensure the the placeholders
# accepts graphs of any size.
_, _, input_graphs, truth_values = batch_iterator(base_dir,
batch_size).__next__()
input_ph = utils_tf.placeholders_from_data_dicts(
input_graphs[0:1], force_dynamic_num_graphs=True)
truth_ph = tf.placeholder(tf.float64, shape=[None])
# Here, we define our computational graphs.
# - First, we compute the model output using the graph_nets library.
# - Then, we compute our loss function only on edge features, where we utilize a log_loss
# function between the truth values and the model output. There is also some factor
# 'num_processing_steps_tr' that describes the level of message passing that somehow
# plays into this. I need to figure out the details.
# - Finally, we will minimize training loss using the Adam Optimizer algorithm.
model_outputs = SegmentClassifier()(input_ph, num_processing_steps_tr)
triplet_output = triplets_ph[1]
edge_losses = tf.losses.log_loss(truth_ph,
tf.transpose(model_outputs[-1].edges)[0])
training_loss = edge_losses
training_optimizer = tf.train.AdamOptimizer(learning_rate).minimize(
training_loss)
# Allows a graph containing `None` fields to be run in a Tensorflow
# session. This is currently not needed since we have data for all
# elements in the graph, including useless data for the global variable.
input_ph = utils_tf.make_runnable_in_session(input_ph)
# According to documentation, represent a connection between the client
# program and a C++ runtime. See the following link for more information.
# https://www.tensorflow.org/guide/graphs
sess = tf.Session()
# Create session saver
saver = tf.train.Saver()
# Our computation graph uses global variables, so we are required to
# initialize them for the first pass. See the following link for more
# information on Tensorflow variables
# https://www.tensorflow.org/guide/variables
sess.run(tf.global_variables_initializer())
output_index = 0
last_output = time.time()
# We will iterate through our dataset many times to train.
for iteration in range(0, num_training_iterations):
# Iterate through all of the batches and retrieve batch data accordingly.
for batch_index, batch_count, input_batch, truth_batch in batch_iterator(
base_dir, batch_size):
# Turn our data dictionary into a proper graphs.GraphsTuple
# object for use with graph_nets library.
input_graphs = utils_np.data_dicts_to_graphs_tuple(input_batch)
# The utility function make_runnable_in_session to fix problems resulting from
# None fields in graph.
input_graphs = utils_tf.make_runnable_in_session(input_graphs)
# Create a feed dictionary that properly maps graph properties.
# Documentation states that this is only necessary in the case of
# missing properties, but we will do it anyway just to be safe.
feed_dict = utils_tf.get_feed_dict(input_ph, input_graphs)
# We must pass both the input and target graphs into our computation
# graph, so we update our feed dictionary with new properties using
# the same method described above.
feed_dict.update({truth_ph: truth_batch})
# Run our computation graph using the feed_dictionary created above.
# Currently, we appear to be computing multiple values... I need
# to figure out what each of them means.
train_values = sess.run(
{
"step": training_optimizer,
"loss": training_loss,
"outputs": model_outputs
},
feed_dict=feed_dict)
# Compute the time lapse from last save-evaluate-visualize action
current_time = time.time()
output_time_lapse = current_time - last_output
if output_time_lapse > 120:
last_output = current_time
# Create a feed dict with 10 training events. These events have not been
# used during testing, so
_, _, input_batch, truth_batch = batch_iterator(
base_dir, 10, test=True).__next__()
input_graphs = utils_np.data_dicts_to_graphs_tuple(input_batch)
input_graphs = utils_tf.make_runnable_in_session(input_graphs)
feed_dict = utils_tf.get_feed_dict(input_ph, input_graphs)
feed_dict.update({truth_ph: truth_batch})
train_values = sess.run(
{
"loss": training_loss,
"target": truth_ph,
"outputs": model_outputs
},
feed_dict=feed_dict)
cutoff_list = []
purity_list = []
efficiency_list = []
# Compute purity and efficiency for every cutoff from 0 to 1 in steps of 0.01
for filter_cutoff in np.linspace(0, 1, 100):
result = np.transpose(
np.where(
train_values['outputs'][-1].edges > filter_cutoff,
1, 0))[0]
correct = np.sum(
np.where(
np.logical_and(result == truth_batch,
result == np.ones(result.shape)), 1,
0))
purity = correct / np.sum(result) if np.sum(
result) != 0 else 1.0
purity_list.append(purity)
efficiency = correct / np.sum(truth_batch)
efficiency_list.append(efficiency)
cutoff_list.append(filter_cutoff)
# Create purity-efficiency plot and save to folder
plt.figure()
plt.plot(purity_list, efficiency_list)
plt.axis([0, 1, 0, 1])
plt.xlabel('Purity')
plt.ylabel('Efficiency')
os.makedirs(os.path.join(results_dir, 'figures'),
exist_ok=True)
plt.savefig(
os.path.join(
results_dir,
'figures/purity_vs_efficiency{:02d}.png'.format(
output_index)))
plt.close()
# Write the purity-efficiency
csv_path = os.path.join(
results_dir,
'figures/purity_vs_efficiency{:02d}.csv'.format(
output_index))
with open(csv_path, 'w') as csvfile:
csv_writer = csv.writer(csvfile)
csv_writer.writerow(['cutoff', 'purity', 'efficiency'])
for (cutoff, purity,
efficiency) in zip(cutoff_list, purity_list,
efficiency_list):
csv_writer.writerow([cutoff, purity, efficiency])
os.makedirs(os.path.join(results_dir, 'models'), exist_ok=True)
saver.save(
sess,
os.path.join(results_dir,
'models/model{}.ckpt'.format(output_index)))
visualize_hitgraph(
os.path.join(results_dir, 'images'), output_index, {
'nodes': input_batch[0]['nodes'],
'edges': truth_batch,
'senders': input_batch[0]['senders'],
'receivers': input_batch[0]['receivers']
})
print('\repoch: {} progress: {:.4f} loss: {:.4f}'.format(
iteration, batch_index / batch_count,
train_values['loss']))
output_index += 1
sess.close()
graphs_tuple_np = sess.run(graphs_tuple_tf)
# plot_graphs_tuple_np(graphs_tuple_np)
# plt.show()
# If the GraphsTuple has None's we need to make use of `utils_tf.make_runnable_in_session`.
tf.reset_default_graph()
graphs_tuple_tf = utils_tf.data_dicts_to_graphs_tuple(graph_dicts)
# Removing the edges from a graph.
graph_with_nones = graphs_tuple_tf.replace(edges=None,
senders=None,
receivers=None,
n_edge=graphs_tuple_tf.n_edge * 0)
runnable_in_session_graph = utils_tf.make_runnable_in_session(graph_with_nones)
with tf.Session() as sess:
graphs_tuple_np = sess.run(runnable_in_session_graph)
# plot_graphs_tuple_np(graphs_tuple_np)
# plt.show()
tf.reset_default_graph()
OUTPUT_EDGE_SIZE = 10
OUTPUT_NODE_SIZE = 11
OUTPUT_GLOBAL_SIZE = 12
graph_network = modules.GraphNetwork(
edge_model_fn=lambda: snt.Linear(output_size=OUTPUT_EDGE_SIZE),
node_model_fn=lambda: snt.Linear(output_size=OUTPUT_NODE_SIZE),
global_model_fn=lambda: snt.Linear(output_size=OUTPUT_GLOBAL_SIZE))
