def __init__(self, name="DecaySimulator"):
super(DecaySimulator, self).__init__(name=name)
self._node_linear = make_mlp_model()
self._node_rnn = snt.GRU(hidden_size=LATENT_SIZE, name='node_rnn')
self._node_proper = snt.nets.MLP([4], activate_final=False)
self._edge_block = blocks.EdgeBlock(edge_model_fn=make_mlp_model,
use_edges=False,
use_receiver_nodes=True,
use_sender_nodes=True,
use_globals=False,
name='edge_encoder_block')
self._node_encoder_block = blocks.NodeBlock(
node_model_fn=make_mlp_model,
use_received_edges=False,
use_sent_edges=False,
use_nodes=True,
use_globals=False,
name='node_encoder_block')
self._global_encoder_block = blocks.GlobalBlock(
global_model_fn=make_mlp_model,
use_edges=True,
use_nodes=True,
use_globals=False,
nodes_reducer=tf.math.unsorted_segment_sum,
edges_reducer=tf.math.unsorted_segment_sum,
name='global_encoder_block')
self._core = MLPGraphNetwork()
# self._core = InteractionNetwork(
# edge_model_fn=make_mlp_model,
# node_model_fn=make_mlp_model,
# reducer=tf.math.unsorted_segment_sum
# )
# # Transforms the outputs into appropriate shapes.
node_output_size = 64
node_fn = lambda: snt.Sequential([
snt.nets.MLP(
[node_output_size],
activation=tf.nn.relu, # default is relu
name='node_output')
])
global_output_size = 1
global_fn = lambda: snt.Sequential([
snt.nets.MLP(
[global_output_size],
activation=tf.nn.relu, # default is relu
name='global_output'),
tf.sigmoid
])
self._output_transform = modules.GraphIndependent(
edge_model_fn=None,
node_model_fn=node_fn,
global_model_fn=global_fn)
def __init__(self, node_model_fn, global_model_fn, name="rwrd_gnn"):
# aggregator_fn = tf.math.unsorted_segment_sum,
"""Initializes the reward model"""
super(GraphNeuralNetwork_reward, self).__init__(name=name)
with self._enter_variable_scope():
# self._edge_block = blocks.EdgeBlock(
# edge_model_fn=edge_model_fn,
# use_edges=False,
# use_receiver_nodes=True,
# use_sender_nodes=True,
# use_globals=True,
# name='edge_block')
self._node_block = blocks.NodeBlock(
node_model_fn=node_model_fn,
use_received_edges=False,
use_sent_edges=False,
use_nodes=True,
use_globals=True,
received_edges_reducer=tf.math.unsorted_segment_sum,
sent_edges_reducer=tf.math.unsorted_segment_sum,
name="node_block")
self._global_block = blocks.GlobalBlock(
global_model_fn=global_model_fn,
use_edges=False,
use_nodes=True,
use_globals=True,
nodes_reducer=tf.math.unsorted_segment_sum,
edges_reducer=tf.math.unsorted_segment_sum,
name="global_block")
def __init__(self, name, edge_model_fn, node_model_fn, global_model_fn):
"""
description: initializes the model
:param edge_model_fn: Function passed to the edge block, in this paper, it is an MLP
:param node_model_fn: Function passed to the node block, in this paper for the task of bAbI, it is an LSTM over timesteps
:param edge_model_fn: Function passed to the global block, in this paper, it is an MLP
"""
super(RRN, self).__init__(name=name)
self._edge_block = blocks.EdgeBlock(edge_model_fn=edge_model_fn,
use_edges=False,
use_receiver_nodes=True,
use_sender_nodes=True,
use_globals=False)
self._global_block = blocks.GlobalBlock(
global_model_fn=global_model_fn,
use_edges=False,
use_nodes=True,
use_globals=False,
nodes_reducer=tf.unsorted_segment_sum)
self._node_model_fn = node_model_fn
def __init__(self, name="FourTopPredictor"):
super(FourTopPredictor, self).__init__(name=name)
self._edge_block = blocks.EdgeBlock(edge_model_fn=make_mlp_model,
use_edges=False,
use_receiver_nodes=True,
use_sender_nodes=True,
use_globals=False,
name='edge_encoder_block')
self._node_encoder_block = blocks.NodeBlock(
node_model_fn=make_mlp_model,
use_received_edges=False,
use_sent_edges=False,
use_nodes=True,
use_globals=False,
name='node_encoder_block')
self._global_block = blocks.GlobalBlock(
global_model_fn=make_mlp_model,
use_edges=True,
use_nodes=True,
use_globals=False,
)
self._core = MLPGraphNetwork()
# Transforms the outputs into appropriate shapes.
global_output_size = n_target_node_features * n_max_tops
self._global_nn = snt.nets.MLP(
[128, 128, global_output_size],
activation=tf.nn.leaky_relu, # default is relu, tanh
dropout_rate=0.30,
name='global_output')
def test_compatible_higher_rank_no_raise(self):
"""No exception should occur with higher ranks tensors."""
input_graph = self._get_shaped_input_graph()
input_graph = input_graph.map(lambda v: tf.transpose(v, [0, 2, 1, 3]))
network = blocks.GlobalBlock(
functools.partial(snt.Conv2D, output_channels=10, kernel_shape=[3, 3]))
self._assert_build_and_run(network, input_graph)
def test_same_as_subblocks(self, reducer):
"""Compares the output to explicit subblocks output.
Args:
reducer: The reducer used in the `NodeBlock` and `GlobalBlock`.
"""
input_graph = self._get_input_graph()
edge_model_fn = functools.partial(snt.Linear, output_size=5)
node_model_fn = functools.partial(snt.Linear, output_size=10)
global_model_fn = functools.partial(snt.Linear, output_size=15)
graph_network = modules.GraphNetwork(
edge_model_fn=edge_model_fn,
node_model_fn=node_model_fn,
global_model_fn=global_model_fn,
reducer=reducer)
output_graph = graph_network(input_graph)
edge_block = blocks.EdgeBlock(
edge_model_fn=lambda: graph_network._edge_block._edge_model,
use_sender_nodes=True,
use_edges=True,
use_receiver_nodes=True,
use_globals=True)
node_block = blocks.NodeBlock(
node_model_fn=lambda: graph_network._node_block._node_model,
use_nodes=True,
use_sent_edges=False,
use_received_edges=True,
use_globals=True,
received_edges_reducer=reducer)
global_block = blocks.GlobalBlock(
global_model_fn=lambda: graph_network._global_block._global_model,
use_nodes=True,
use_edges=True,
use_globals=True,
edges_reducer=reducer,
nodes_reducer=reducer)
expected_output_edge_block = edge_block(input_graph)
expected_output_node_block = node_block(expected_output_edge_block)
expected_output_global_block = global_block(expected_output_node_block)
expected_edges = expected_output_edge_block.edges
expected_nodes = expected_output_node_block.nodes
expected_globals = expected_output_global_block.globals
with self.test_session() as sess:
sess.run(tf.global_variables_initializer())
(output_graph_out,
expected_edges_out, expected_nodes_out, expected_globals_out) = sess.run(
(output_graph, expected_edges, expected_nodes, expected_globals))
self._assert_all_none_or_all_close(expected_edges_out,
output_graph_out.edges)
self._assert_all_none_or_all_close(expected_nodes_out,
output_graph_out.nodes)
self._assert_all_none_or_all_close(expected_globals_out,
output_graph_out.globals)
def __init__(self, name="GlobalClassifierNoEdgeInfo"):
super(GlobalClassifierNoEdgeInfo, self).__init__(name=name)
self._edge_block = blocks.EdgeBlock(edge_model_fn=make_mlp_model,
use_edges=False,
use_receiver_nodes=True,
use_sender_nodes=True,
use_globals=False,
name='edge_encoder_block')
self._node_encoder_block = blocks.NodeBlock(
node_model_fn=make_mlp_model,
use_received_edges=False,
use_sent_edges=False,
use_nodes=True,
use_globals=False,
name='node_encoder_block')
self._global_block = blocks.GlobalBlock(
global_model_fn=make_mlp_model,
use_edges=True,
use_nodes=True,
use_globals=False,
)
self._core = MLPGraphNetwork()
# Transforms the outputs into appropriate shapes.
global_output_size = 1
global_fn = lambda: snt.Sequential([
snt.nets.MLP([LATENT_SIZE, global_output_size],
name='global_output'), tf.sigmoid
])
self._output_transform = modules.GraphIndependent(
None, None, global_fn)
def test_output_values(
self, use_edges, use_nodes, use_globals, edges_reducer, nodes_reducer):
"""Compares the output of a GlobalBlock to an explicit computation."""
input_graph = self._get_input_graph()
global_block = blocks.GlobalBlock(
global_model_fn=self._global_model_fn,
use_edges=use_edges,
use_nodes=use_nodes,
use_globals=use_globals,
edges_reducer=edges_reducer,
nodes_reducer=nodes_reducer)
output_graph = global_block(input_graph)
model_inputs = []
if use_edges:
model_inputs.append(
blocks.EdgesToGlobalsAggregator(edges_reducer)(input_graph))
if use_nodes:
model_inputs.append(
blocks.NodesToGlobalsAggregator(nodes_reducer)(input_graph))
if use_globals:
model_inputs.append(input_graph.globals)
model_inputs = tf.concat(model_inputs, axis=-1)
self.assertEqual(input_graph.edges, output_graph.edges)
self.assertEqual(input_graph.nodes, output_graph.nodes)
with self.test_session() as sess:
output_graph_out, model_inputs_out = sess.run(
(output_graph, model_inputs))
expected_output_globals = model_inputs_out * self._scale
self.assertNDArrayNear(
expected_output_globals, output_graph_out.globals, err=1e-4)
def __init__(self,
node_model_fn,
global_model_fn,
reducer=tf.unsorted_segment_sum,
name="deep_sets"):
"""Initializes the DeepSets module.
Args:
node_model_fn: A callable to be passed to NodeBlock. The callable must
return a Sonnet module (or equivalent; see NodeBlock for details). The
shape of this module's output must equal the shape of the input graph's
global features, but for the first and last axis.
global_model_fn: A callable to be passed to GlobalBlock. The callable must
return a Sonnet module (or equivalent; see GlobalBlock for details).
reducer: Reduction to be used when aggregating the nodes in the globals.
This should be a callable whose signature matches
tf.unsorted_segment_sum.
name: The module name.
"""
super(DeepSets, self).__init__(name=name)
with self._enter_variable_scope():
self._node_block = blocks.NodeBlock(node_model_fn=node_model_fn,
use_received_edges=False,
use_sent_edges=False,
use_nodes=True,
use_globals=True)
self._global_block = blocks.GlobalBlock(
global_model_fn=global_model_fn,
use_edges=False,
use_nodes=True,
use_globals=False,
nodes_reducer=reducer)
def test_unused_field_can_be_none(
self, use_edges, use_nodes, use_globals, none_field):
"""Checks that computation can handle non-necessary fields left None."""
input_graph = self._get_input_graph([none_field])
global_block = blocks.GlobalBlock(
global_model_fn=self._global_model_fn,
use_edges=use_edges,
use_nodes=use_nodes,
use_globals=use_globals)
output_graph = global_block(input_graph)
model_inputs = []
if use_edges:
model_inputs.append(
blocks.EdgesToGlobalsAggregator(tf.unsorted_segment_sum)(input_graph))
if use_nodes:
model_inputs.append(
blocks.NodesToGlobalsAggregator(tf.unsorted_segment_sum)(input_graph))
if use_globals:
model_inputs.append(input_graph.globals)
model_inputs = tf.concat(model_inputs, axis=-1)
self.assertEqual(input_graph.edges, output_graph.edges)
self.assertEqual(input_graph.nodes, output_graph.nodes)
with self.test_session() as sess:
actual_globals, model_inputs_out = sess.run(
(output_graph.globals, model_inputs))
expected_output_globals = model_inputs_out * self._scale
self.assertNDArrayNear(expected_output_globals, actual_globals, err=1e-4)
def test_created_variables(self, use_edges, use_nodes, use_globals,
expected_first_dim_w):
"""Verifies the variable names and shapes created by a GlobalBlock."""
output_size = 10
expected_var_shapes_dict = {
"global_block/mlp/linear_0/b:0": [output_size],
"global_block/mlp/linear_0/w:0":
[expected_first_dim_w, output_size]
}
input_graph = self._get_input_graph()
global_block = blocks.GlobalBlock(global_model_fn=functools.partial(
snt.nets.MLP, output_sizes=[output_size]),
use_edges=use_edges,
use_nodes=use_nodes,
use_globals=use_globals)
global_block(input_graph)
variables = global_block.get_variables()
var_shapes_dict = {
var.name: var.get_shape().as_list()
for var in variables
}
self.assertDictEqual(expected_var_shapes_dict, var_shapes_dict)
def test_no_input_raises_exception(self):
"""Checks that receiving no input raises an exception."""
with self.assertRaisesRegexp(ValueError, "At least one of "):
blocks.GlobalBlock(global_model_fn=self._global_model_fn,
use_edges=False,
use_nodes=False,
use_globals=False)
def __init__(self,
edge_model_fn,
global_model_fn,
reducer=tf.unsorted_segment_sum,
name="relation_network"):
"""Initializes the RelationNetwork module.
Args:
edge_model_fn: A callable that will be passed to EdgeBlock to perform
per-edge computations. The callable must return a Sonnet module (or
equivalent; see EdgeBlock for details).
global_model_fn: A callable that will be passed to GlobalBlock to perform
per-global computations. The callable must return a Sonnet module (or
equivalent; see GlobalBlock for details).
reducer: Reducer to be used by GlobalBlock to aggregate edges. Defaults to
tf.unsorted_segment_sum.
name: The module name.
"""
super(RelationNetwork, self).__init__(name=name)
with self._enter_variable_scope():
self._edge_block = blocks.EdgeBlock(edge_model_fn=edge_model_fn,
use_edges=False,
use_receiver_nodes=True,
use_sender_nodes=True,
use_globals=False)
self._global_block = blocks.GlobalBlock(
global_model_fn=global_model_fn,
use_edges=True,
use_nodes=False,
use_globals=False,
edges_reducer=reducer)
def __init__(self,
attention_node_projection_model,
attention_edge_projection_model,
query_key_product_model,
node_model_fn,
edge_model_fn,
global_model_fn,
num_heads,
key_size,
value_size,
edge_block_opt=None,
global_block_opt=None,
name="GAT"):
"""
Args:
attention_node_projection_model: Model used for projection to get
query, key and values
Final layer dim should be key_size * num_heads
attention_edge_projection_model: Model used for projection to get
query, key and values
Final layer dim should be (key_size + value_size) * num_heads
query_key_product_model: Model used to find "dot product" between
queries and keys.
Final layer dim should be 1.
node_model_fn: Model applied to node embeddings finally.
edge_model_fn: Model applied to node embeddings finally.
num_heads: Number of attention heads
key_size: Key dimension
value_size: value dimension
edge_block_opt: Additional options to be passed to the EdgeBlock. Can
contain keys `use_edges`, `use_receiver_nodes`, `use_sender_nodes`,
`use_globals`. By default, these are all True.
global_block_opt: Additional options to be passed to the GlobalBlock. Can
contain the keys `use_edges`, `use_nodes`, `use_globals` (all set to
True by default), and `edges_reducer`, `nodes_reducer` (defaults to
`reducer`).
name: The module name.
"""
super().__init__(name=name)
self._attention_node_projection_model = attention_node_projection_model
self._attention_edge_projection_model = attention_edge_projection_model
self._query_key_product_model = query_key_product_model
self.num_heads = num_heads
self.key_size = key_size
self.value_size = value_size
edge_block_opt = _make_default_edge_block_opt(edge_block_opt)
global_block_opt = _make_default_global_block_opt(
global_block_opt, tf.unsorted_segment_sum)
# does not make sense without using sender nodes.
assert edge_block_opt['use_sender_nodes']
with self._enter_variable_scope():
self._node_model = node_model_fn()
self._edge_block = blocks.EdgeBlock(edge_model_fn=edge_model_fn,
**edge_block_opt)
self._global_block = blocks.GlobalBlock(
global_model_fn=global_model_fn, **global_block_opt)
def tinet(input_graph):
embedding = blocks.NodeBlock(
node_model_fn=lambda: tf.keras.layers.Embedding(800, 32),
use_received_edges=False,
use_sent_edges=False,
use_nodes=True,
use_globals=False)
graph_network_layer1 = blocks.NodeBlock(
# edge_model_fn = lambda: tf.layers.Dense(16, activation=tf.nn.relu),
node_model_fn=lambda: tf.layers.Dense(32, activation=tf.nn.relu))
# global_model_fn = lambda: tf.layers.Dense(8, activation=tf.nn.relu))
graph_network_layer2 = blocks.NodeBlock(
# edge_model_fn = lambda: tf.layers.Dense(16, activation=tf.nn.relu),
node_model_fn=lambda: tf.layers.Dense(32, activation=tf.nn.relu))
# global_model_fn = lambda: tf.layers.Dense(8, activation=tf.nn.relu))
graph_network_layer3 = blocks.NodeBlock(
# edge_model_fn = lambda: tf.layers.Dense(16, activation=tf.nn.relu),
node_model_fn=lambda: tf.layers.Dense(32, activation=tf.nn.relu))
# global_model_fn = lambda: tf.layers.Dense(8, activation=tf.nn.relu))
graph_network_layer4 = blocks.NodeBlock(
# edge_model_fn = lambda: tf.layers.Dense(16, activation=tf.nn.relu),
node_model_fn=lambda: tf.layers.Dense(32, activation=tf.nn.relu))
# global_model_fn = lambda: tf.layers.Dense(8, activation=tf.nn.relu))
graph_network_layer5 = blocks.GlobalBlock(
# edge_model_fn = lambda: tf.layers.Dense(16, activation=tf.nn.relu),
# node_model_fn = lambda: tf.layers.Dense(16, activation=tf.nn.relu),
global_model_fn=lambda: tf.layers.Dense(40, activation=tf.nn.relu))
graph_network_layer6 = blocks.GlobalBlock(
# edge_model_fn = lambda: tf.layers.Dense(16, activation=tf.nn.relu),
# node_model_fn = lambda: tf.layers.Dense(16, activation=tf.nn.relu),
global_model_fn=lambda: tf.layers.Dense(40, activation=tf.nn.relu))
h0 = embedding(input_graph)
h1 = graph_network_layer1(h0)
h2 = graph_network_layer2(h1)
h3 = graph_network_layer3(h2)
h4 = graph_network_layer4(h3)
h5 = graph_network_layer5(h4)
h6 = graph_network_layer6(h5)
out = h6.globals
return tf.layers.dense(out, 4, activation=None)
def test_missing_field_raises_exception(
self, use_edges, use_nodes, use_globals, none_field):
"""Checks that missing a required field raises an exception."""
input_graph = self._get_input_graph([none_field])
global_block = blocks.GlobalBlock(
global_model_fn=self._global_model_fn,
use_edges=use_edges,
use_nodes=use_nodes,
use_globals=use_globals)
with self.assertRaisesRegexp(ValueError, "field cannot be None"):
global_block(input_graph)
def test_missing_aggregation_raises_exception(self, use_edges, use_nodes,
edges_reducer,
nodes_reducer):
"""Checks that missing a required aggregation argument raises an error."""
with self.assertRaisesRegexp(ValueError, "should not be None"):
blocks.GlobalBlock(global_model_fn=self._global_model_fn,
use_edges=use_edges,
use_nodes=use_nodes,
use_globals=False,
edges_reducer=edges_reducer,
nodes_reducer=nodes_reducer)
def test_optional_arguments(self, scale, offset):
"""Assesses the correctness of the GlobalBlock using arguments."""
input_graph = self._get_input_graph()
global_block = blocks.GlobalBlock(
global_model_fn=self._global_model_args_fn)
output_graph_out = global_block(
input_graph, global_model_kwargs=dict(scale=scale, offset=offset))
fixed_scale = scale
fixed_offset = offset
model_fn = lambda: lambda features: features * fixed_scale + fixed_offset
hardcoded_global_block = blocks.GlobalBlock(global_model_fn=model_fn)
expected_graph_out = hardcoded_global_block(input_graph)
self.assertIs(expected_graph_out.edges, output_graph_out.edges)
self.assertIs(expected_graph_out.nodes, output_graph_out.nodes)
self.assertNDArrayNear(
expected_graph_out.globals.numpy(),
output_graph_out.globals.numpy(),
err=1e-4)
def __init__(self,
edge_model_fn,
node_model_fn,
global_model_fn,
reducer=tf.math.unsorted_segment_sum,
edge_block_opt=None,
node_block_opt=None,
global_block_opt=None,
name="graph_network"):
"""Initializes the GraphNetwork module.
Args:
edge_model_fn: A callable that will be passed to EdgeBlock to perform
per-edge computations. The callable must return a Sonnet module (or
equivalent; see EdgeBlock for details).
node_model_fn: A callable that will be passed to NodeBlock to perform
per-node computations. The callable must return a Sonnet module (or
equivalent; see NodeBlock for details).
global_model_fn: A callable that will be passed to GlobalBlock to perform
per-global computations. The callable must return a Sonnet module (or
equivalent; see GlobalBlock for details).
reducer: Reducer to be used by NodeBlock and GlobalBlock to aggregate
nodes and edges. Defaults to tf.unsorted_segment_sum. This will be
overridden by the reducers specified in `node_block_opt` and
`global_block_opt`, if any.
edge_block_opt: Additional options to be passed to the EdgeBlock. Can
contain keys `use_edges`, `use_receiver_nodes`, `use_sender_nodes`,
`use_globals`. By default, these are all True.
node_block_opt: Additional options to be passed to the NodeBlock. Can
contain the keys `use_received_edges`, `use_sent_edges`, `use_nodes`,
`use_globals` (all set to True by default), and
`received_edges_reducer`, `sent_edges_reducer` (default to `reducer`).
global_block_opt: Additional options to be passed to the GlobalBlock. Can
contain the keys `use_edges`, `use_nodes`, `use_globals` (all set to
True by default), and `edges_reducer`, `nodes_reducer` (defaults to
`reducer`).
name: The module name.
"""
super(GraphNetwork, self).__init__(name=name)
edge_block_opt = _make_default_edge_block_opt(edge_block_opt)
node_block_opt = _make_default_node_block_opt(node_block_opt, reducer)
global_block_opt = _make_default_global_block_opt(
global_block_opt, reducer)
#with self._enter_variable_scope():
if 2 > 1:
self._edge_block = blocks.EdgeBlock(edge_model_fn=edge_model_fn,
**edge_block_opt)
self._node_block = blocks.NodeBlock(node_model_fn=node_model_fn,
**node_block_opt)
self._global_block = blocks.GlobalBlock(
global_model_fn=global_model_fn, **global_block_opt)
def test_incompatible_higher_rank_inputs_no_raise(self, use_edges,
use_nodes, use_globals,
field):
"""No exception should occur if a differently shapped field is not used."""
input_graph = self._get_shaped_input_graph()
input_graph = input_graph.replace(
**{field: tf.transpose(getattr(input_graph, field), [0, 2, 1, 3])})
network = blocks.GlobalBlock(functools.partial(snt.Conv2D,
output_channels=10,
kernel_shape=[3, 3]),
use_edges=use_edges,
use_nodes=use_nodes,
use_globals=use_globals)
self._assert_build_and_run(network, input_graph)
def test_incompatible_higher_rank_inputs_raises(self, use_edges, use_nodes,
use_globals, field):
"""A exception should be raised if the inputs have incompatible shapes."""
input_graph = self._get_shaped_input_graph()
input_graph = input_graph.replace(
**{field: tf.transpose(getattr(input_graph, field), [0, 2, 1, 3])})
network = blocks.GlobalBlock(functools.partial(snt.Conv2D,
output_channels=10,
kernel_shape=[3, 3]),
use_edges=use_edges,
use_nodes=use_nodes,
use_globals=use_globals)
with self.assertRaisesRegexp(ValueError,
"in both shapes must be equal"):
network(input_graph)
def test_same_as_subblocks(self, reducer, none_fields):
"""Compares the output to explicit subblocks output.
Args:
reducer: The reducer used in the NodeBlock.
none_fields: (list of strings) The corresponding fields are removed from
the input graph.
"""
input_graph = self._get_input_graph()
input_graph = input_graph.map(lambda _: None, none_fields)
deep_sets = self._get_model(reducer)
output_graph = deep_sets(input_graph)
output_nodes = output_graph.nodes
output_globals = output_graph.globals
node_block = blocks.NodeBlock(
node_model_fn=lambda: deep_sets._node_block._node_model,
use_received_edges=False,
use_sent_edges=False,
use_nodes=True,
use_globals=True)
global_block = blocks.GlobalBlock(
global_model_fn=lambda: deep_sets._global_block._global_model,
use_edges=False,
use_nodes=True,
use_globals=False,
nodes_reducer=reducer)
node_block_out = node_block(input_graph)
expected_nodes = node_block_out.nodes
expected_globals = global_block(node_block_out).globals
self.assertAllEqual(input_graph.edges, output_graph.edges)
self.assertAllEqual(input_graph.receivers, output_graph.receivers)
self.assertAllEqual(input_graph.senders, output_graph.senders)
with self.test_session() as sess:
sess.run(tf.global_variables_initializer())
(output_nodes_, output_globals_, expected_nodes_,
expected_globals_) = sess.run(
[output_nodes, output_globals, expected_nodes, expected_globals])
self._assert_all_none_or_all_close(expected_nodes_, output_nodes_)
self._assert_all_none_or_all_close(expected_globals_, output_globals_)
def __init__(self, model_id, name="MLPGraphNetwork"):
super(MLPGraphNetwork, self).__init__(name=name)
node_block_options = {
"use_received_edges": False,
"use_sent_edges": False,
"use_nodes": True,
"use_globals": True,
}
global_block_options = {
"use_edges": False,
"use_nodes": True,
"use_globals": True,
}
node_block_opt = modules._make_default_node_block_opt(
node_block_options, tf.unsorted_segment_sum)
global_block_opt = modules._make_default_global_block_opt(
global_block_options, tf.unsorted_segment_sum)
with self._enter_variable_scope():
node_model_fn = lambda: get_model_from_config(
model_id, model_type="mlp"
)(n_neurons=EncodeProcessDecode_v4_172_no_edges.n_neurons_nodes,
n_layers=EncodeProcessDecode_v4_172_no_edges.n_layers_nodes,
output_size=EncodeProcessDecode_v4_172_no_edges.
n_neurons_nodes_total_dim,
typ="mlp_transform",
activation_final=False,
name="mlp_core_node")
global_model_fn = lambda: get_model_from_config(
model_id, model_type="mlp"
)(n_neurons=EncodeProcessDecode_v4_172_no_edges.n_neurons_globals,
n_layers=EncodeProcessDecode_v4_172_no_edges.n_layers_globals,
output_size=None,
typ="mlp_layer_norm",
name="mlp_core_global")
self._node_block = blocks.NodeBlock(node_model_fn=node_model_fn,
**node_block_opt)
self._global_block = blocks.GlobalBlock(
global_model_fn=global_model_fn, **global_block_opt)
def __init__(self,
num_components,
reducer=tf.math.unsorted_segment_mean,
properties_size=11,
name=None):
super(DecoderNetwork, self).__init__(name=name)
self.num_components = num_components
self.components_dim = properties_size * 10
edge_enc_size = 32
node_model_fn = lambda: snt.nets.MLP([self.components_dim],
activate_final=True,
activation=tf.nn.sigmoid)
edge_model_fn = lambda: snt.nets.MLP(
[edge_enc_size], activate_final=True, activation=tf.nn.sigmoid)
global_model_fn = lambda: snt.nets.MLP([self.components_dim],
activate_final=True,
activation=tf.nn.sigmoid)
self.node_block = blocks.NodeBlock(node_model_fn,
use_received_edges=False,
use_sent_edges=False,
use_nodes=True,
use_globals=False)
self.edge_block = blocks.EdgeBlock(edge_model_fn,
use_edges=False,
use_receiver_nodes=True,
use_sender_nodes=True,
use_globals=False)
self.global_block = blocks.GlobalBlock(global_model_fn,
use_edges=True,
use_nodes=False,
use_globals=False,
edges_reducer=reducer)
def test_same_as_subblocks(self, reducer, none_field=None):
"""Compares the output to explicit subblocks output.
Args:
reducer: The reducer used in the `GlobalBlock`.
none_field: (string, default=None) If not None, the corresponding field
is removed from the input graph.
"""
input_graph = self._get_input_graph(none_field)
relation_network = self._get_model(reducer)
output_graph = relation_network(input_graph)
edge_block = blocks.EdgeBlock(
edge_model_fn=lambda: relation_network._edge_block._edge_model,
use_edges=False,
use_receiver_nodes=True,
use_sender_nodes=True,
use_globals=False)
global_block = blocks.GlobalBlock(
global_model_fn=lambda: relation_network._global_block._global_model,
use_edges=True,
use_nodes=False,
use_globals=False,
edges_reducer=reducer,
nodes_reducer=reducer)
expected_output_edge_block = edge_block(input_graph)
expected_output_global_block = global_block(expected_output_edge_block)
self.assertEqual(input_graph.edges, output_graph.edges)
self.assertEqual(input_graph.nodes, output_graph.nodes)
with self.test_session() as sess:
sess.run(tf.global_variables_initializer())
(actual_globals_out, expected_globals_out) = sess.run(
(output_graph.globals, expected_output_global_block.globals))
self._assert_all_none_or_all_close(expected_globals_out, actual_globals_out)
def test_global_block_options(self,
use_edges,
use_nodes,
use_globals,
edges_reducer,
nodes_reducer):
"""Test for configuring the NodeBlock options."""
input_graph = self._get_input_graph()
if use_edges:
edges_reducer = edges_reducer or tf.unsorted_segment_sum
if use_nodes:
nodes_reducer = nodes_reducer or tf.unsorted_segment_sum
# Identity edge model.
edge_model_fn = lambda: tf.identity
edge_block_opt = {"use_edges": True,
"use_receiver_nodes": False,
"use_sender_nodes": False,
"use_globals": False}
# Identity node model
node_model_fn = lambda: tf.identity
node_block_opt = {"use_received_edges": False,
"use_sent_edges": False,
"use_nodes": True,
"use_globals": False}
global_model_fn = functools.partial(snt.Linear, output_size=10)
global_block_opt = {"use_globals": use_globals,
"use_nodes": use_nodes,
"use_edges": use_edges,
"edges_reducer": edges_reducer,
"nodes_reducer": nodes_reducer}
graph_network = modules.GraphNetwork(
edge_model_fn=edge_model_fn,
edge_block_opt=edge_block_opt,
node_model_fn=node_model_fn,
node_block_opt=node_block_opt,
global_model_fn=global_model_fn,
global_block_opt=global_block_opt)
output_graph = graph_network(input_graph)
global_block = blocks.GlobalBlock(
global_model_fn=lambda: graph_network._global_block._global_model,
use_edges=use_edges,
use_nodes=use_nodes,
use_globals=use_globals,
edges_reducer=edges_reducer,
nodes_reducer=nodes_reducer)
expected_output_edge_block = input_graph
expected_output_node_block = expected_output_edge_block
expected_output_global_block = global_block(expected_output_node_block)
expected_edges = expected_output_edge_block.edges
expected_nodes = expected_output_node_block.nodes
expected_globals = expected_output_global_block.globals
with self.test_session() as sess:
sess.run(tf.global_variables_initializer())
(output_graph_out,
expected_edges_out, expected_nodes_out, expected_globals_out) = sess.run(
(output_graph, expected_edges, expected_nodes, expected_globals))
self._assert_all_none_or_all_close(expected_edges_out,
output_graph_out.edges)
self._assert_all_none_or_all_close(expected_nodes_out,
output_graph_out.nodes)
self._assert_all_none_or_all_close(expected_globals_out,
output_graph_out.globals)
def __init__(self,
num_output: int,
output_size: int,
node_size: int = 4,
edge_size: int = 4,
starting_global_size: int = 10,
inter_graph_connect_prob: float = 0.01,
crossing_steps: int = 4,
reducer=tf.math.unsorted_segment_mean,
properties_size=10,
name=None):
super(GraphMappingNetwork, self).__init__(name=name)
self.num_output = num_output
self.output_size = output_size
self.crossing_steps = crossing_steps
self.empty_node_variable = tf.Variable(initial_value=tf.random.truncated_normal((node_size,)),
name='empty_token_node')
# values for different kinds of edges in the graph, which will be learned
self.intra_graph_edge_variable = tf.Variable(initial_value=tf.random.truncated_normal((edge_size,)),
name='intra_graph_edge_var')
self.intra_token_graph_edge_variable = tf.Variable(initial_value=tf.random.truncated_normal((edge_size,)),
name='intra_token_graph_edge_var')
self.inter_graph_edge_variable = tf.Variable(initial_value=tf.random.truncated_normal((edge_size,)),
name='inter_graph_edge_var')
self.starting_global_variable = tf.Variable(initial_value=tf.random.truncated_normal((starting_global_size,)),
name='starting_global_var')
self.inter_graph_connect_prob = inter_graph_connect_prob
self.projection_node_block = blocks.NodeBlock(lambda: snt.Linear(node_size, name='project'),
use_received_edges=False,
use_sent_edges=False,
use_nodes=True,
use_globals=False)
node_model_fn = lambda: snt.nets.MLP([node_size, node_size], activate_final=True, activation=tf.nn.leaky_relu)
edge_model_fn = lambda: snt.nets.MLP([edge_size, edge_size], activate_final=True, activation=tf.nn.leaky_relu)
global_model_fn = lambda: snt.nets.MLP([starting_global_size, starting_global_size], activate_final=True,
activation=tf.nn.leaky_relu)
self.edge_block = blocks.EdgeBlock(edge_model_fn,
use_edges=True,
use_receiver_nodes=True,
use_sender_nodes=True,
use_globals=True)
self.node_block = blocks.NodeBlock(node_model_fn,
use_received_edges=True,
use_sent_edges=True,
use_nodes=True,
use_globals=True,
received_edges_reducer=reducer,
sent_edges_reducer=reducer)
self.global_block = blocks.GlobalBlock(global_model_fn,
use_edges=True,
use_nodes=True,
use_globals=True,
edges_reducer=reducer,
nodes_reducer=reducer)
self.output_projection_node_block = blocks.NodeBlock(lambda: snt.Linear(self.output_size, name='project'),
use_received_edges=False,
use_sent_edges=False,
use_nodes=True,
use_globals=False)
def __init__(self, scope: str, model: GraphModel, reg_param):
# Process parameters
self.scope = scope
self.model = model
self.n_out = self.model.get_global_output_size()
# Configure regularization
self.regularizer = tf.contrib.layers.l2_regularizer(scale=reg_param)
self.reg_linear = {"w": self.regularizer, "b": self.regularizer}
# self.reg_embed = {}
self.reg_embed = {"embeddings": self.regularizer}
# Set up input tensors
with tf.variable_scope(self.scope + "/state_input"):
self.input_graphs = utils_tf.placeholders_from_data_dicts(
[self.model.placeholder_graph()],
force_dynamic_num_graphs=True,
name="local_state")
with tf.variable_scope(self.scope + "/ground_truth"):
# Reinforcement learning inputs
self.true_action = tf.placeholder(tf.int32,
shape=(None, ),
name="action")
self.n_objects = tf.placeholder(tf.int32,
shape=(None, ),
name="n_objects")
self.target_q = tf.placeholder(tf.float32,
shape=(None, ),
name="target_q")
self.target_value = tf.placeholder(tf.float32,
shape=(None, ),
name="target_value")
with tf.variable_scope(self.scope):
# Separately embed different categorical variables as dense vectors
self.encoder_module = GraphEncoder(model,
name="encoder",
regularizers=self.reg_embed)
self.embedded_graphs = self.encoder_module(self.input_graphs)
# Apply an intermediate transformation to pass information between neighboring nodes
self.intermediate_graphs = DenseGraphTransform(
model.hidden_edge_dimension,
model.hidden_node_dimension,
model.hidden_global_dimension,
name="intermediate",
regularizer=self.regularizer)(self.embedded_graphs)
# Then apply a final transformation to produce a global output and node-level evaluations
self.output_graphs = DenseGraphTransform(
model.action_dimension,
1,
1,
node_activation=None,
global_activation=None,
name="output",
regularizer=self.regularizer)(self.intermediate_graphs)
with tf.variable_scope(self.scope + "/outputs"):
# If given a true action, get the corresponding output
self.graph_indices = tf.math.cumsum(self.output_graphs.n_node,
exclusive=True,
name="starting_node_index")
self.true_indices = self.graph_indices + self.true_action
self.chosen_node_outputs = tf.reshape(
tf.gather(self.output_graphs.nodes,
self.true_indices,
name="chosen_action_outputs"), [-1])
# In case we need a policy output, build the following tensors:
# 1) a learned stochastic policy for all possible actions,
# 2) the individual probability of the chosen action
# 3) the log of that individual probability."""
# First, get each node's index
node_indices = tf.range(tf.shape(self.output_graphs.nodes)[0])
# Then, get the index of each graphs' first action
first_action_indices = self.graph_indices + self.n_objects
# broadcast action indices to nodes and compare to node indices
first_action_broadcast = blocks.broadcast_globals_to_nodes(
self.output_graphs.replace(
globals=tf.reshape(first_action_indices, [-1, 1])))
action_mask = tf.greater_equal(
node_indices, tf.reshape(first_action_broadcast, [-1]))
# Zero out the objects and apply softmax to the actions (treat action-nodes as logits)
exp_or_zero = self.output_graphs.replace(nodes=tf.where(
action_mask, tf.math.exp(self.output_graphs.nodes),
tf.zeros_like(self.output_graphs.nodes)))
# Sum the node values so that the global for each graph is the softmax denominator
sum_nodes = blocks.GlobalBlock(lambda: tf.identity,
use_edges=False,
use_globals=False)
softmax_graph = sum_nodes(exp_or_zero)
# Then divide each node's value by that denominator, or set to 1 where denominator is 0
def node_value_to_prob(node_inputs):
p = tf.div_no_nan(node_inputs[:, 0], node_inputs[:, 1])
return tf.where(p > 0, p, tf.ones_like(p))
policy_graph = blocks.NodeBlock(
lambda: node_value_to_prob,
use_received_edges=False,
use_sent_edges=False)(softmax_graph)
self.policy = policy_graph.nodes
self.p_chosen = tf.gather(self.policy,
self.true_indices,
name="p_true_action")
self.log_p_chosen = tf.log(self.p_chosen, name="logp_true_action")
# Configure metrics for training and display
self.TRAIN_METRIC_OPS = self.scope + "/TRAIN_METRIC_OPS"
self.VAL_METRIC_OPS = self.scope + "/VAL_METRIC_OPS"
self.reg_term = tf.reduce_sum(
tf.get_collection(tf.GraphKeys.REGULARIZATION_LOSSES))
tf.summary.scalar('reg_loss', self.reg_term)
