def __init__(self, model_id, name="MLPGraphNetwork"): super(MLPGraphNetwork, self).__init__(name=name) with self._enter_variable_scope(): self._network = modules.GraphNetwork( edge_model_fn=lambda: get_model_from_config(model_id, model_type="mlp") (n_neurons=EncodeProcessDecode_v5_no_skip_batch_norm. n_neurons_edges, n_layers=EncodeProcessDecode_v5_no_skip_batch_norm. n_layers_edges, output_size=None, typ="mlp_layer_norm", name="mlp_core_edge"), node_model_fn=lambda: get_model_from_config(model_id, model_type="mlp") (n_neurons=EncodeProcessDecode_v5_no_skip_batch_norm. n_neurons_nodes, n_layers=EncodeProcessDecode_v5_no_skip_batch_norm. n_layers_nodes, output_size=None, typ="mlp_layer_norm", activation_final=False, name="mlp_core_node"), global_model_fn=lambda: get_model_from_config(model_id, model_type="mlp") (n_neurons=EncodeProcessDecode_v5_no_skip_batch_norm. n_neurons_globals, n_layers=EncodeProcessDecode_v5_no_skip_batch_norm. n_layers_globals, output_size=None, typ="mlp_layer_norm", name="mlp_core_global"))
def __init__(self, name="MLPGraphNetwork"): super(MLPGraphNetwork, self).__init__(name=name) with self._enter_variable_scope(): self._network = modules.GraphNetwork( EncodeProcessDecode_v2.make_mlp_model_edges, EncodeProcessDecode_v2.make_mlp_model, EncodeProcessDecode_v2.make_mlp_model)
def __init__(self, name="MLPGraphNetwork"): super(MLPGraphNetwork, self).__init__(name=name) with self._enter_variable_scope(): self._network = modules.GraphNetwork( edge_model_fn=make_mlp_model_small, node_model_fn=make_mlp_model, global_model_fn=make_mlp_model_small)
def __init__(self, model_id, name="MLPGraphNetwork"): super(MLPGraphNetwork, self).__init__(name=name) with self._enter_variable_scope(): self._network = modules.GraphNetwork( edge_model_fn=lambda: get_model_from_config(model_id, model_type="mlp") (n_neurons=EncodeProcessDecode_v3_172_latent_dim2. n_neurons_edges, n_layers=EncodeProcessDecode_v3_172_latent_dim2. n_layers_edges, output_size=None, typ="mlp_layer_norm", name="mlp_core_edge"), node_model_fn=lambda: get_model_from_config(model_id, model_type="mlp") (n_neurons=EncodeProcessDecode_v3_172_latent_dim2. n_neurons_nodes_total_dim, n_layers=EncodeProcessDecode_v3_172_latent_dim2. n_layers_nodes, output_size=None, typ="mlp_layer_norm", name="mlp_core_node"), global_model_fn=lambda: get_model_from_config(model_id, model_type="mlp") (n_neurons=EncodeProcessDecode_v3_172_latent_dim2. n_neurons_globals, n_layers=EncodeProcessDecode_v3_172_latent_dim2. n_layers_globals, output_size=None, typ="mlp_layer_norm", name="mlp_core_global"))
def __init__(self, model_id, name="MLPGraphNetwork"): super(MLPGraphNetwork, self).__init__(name=name) with self._enter_variable_scope(): self._network = modules.GraphNetwork( edge_model_fn=lambda: get_model_from_config(model_id, model_type="mlp") (n_neurons=EncodeProcessDecode_v5_512_improve_shapes_exp3. n_neurons_edges, n_layers=EncodeProcessDecode_v5_512_improve_shapes_exp3. n_layers_edges, output_size=None, typ="mlp_layer_norm", name="mlp_core_edge"), node_model_fn=lambda: get_model_from_config(model_id, model_type="mlp") ( n_neurons=EncodeProcessDecode_v5_512_improve_shapes_exp3. n_neurons_nodes, n_layers=EncodeProcessDecode_v5_512_improve_shapes_exp3. n_layers_nodes, output_size=EncodeProcessDecode_v5_512_improve_shapes_exp3. n_neurons_nodes_total_dim, typ="mlp_transform", # todo: was earlier "mlp_layer_norm" activation_final=False, name="mlp_core_node"), global_model_fn=lambda: get_model_from_config(model_id, model_type="mlp") (n_neurons=EncodeProcessDecode_v5_512_improve_shapes_exp3. n_neurons_globals, n_layers=EncodeProcessDecode_v5_512_improve_shapes_exp3. n_layers_globals, output_size=None, typ="mlp_layer_norm", name="mlp_core_global"))
def __init__(self, latent_sizes_edge, latent_sizes_node, latent_sizes_global, name="MLPGraphNetwork"): super(MLPGraphNetwork, self).__init__(name=name) self.edge_fun = lambda: snt.Sequential([ snt.nets.MLP(latent_sizes_edge, activate_final=True), snt.LayerNorm() ]) self.node_fun = lambda: snt.Sequential([ snt.nets.MLP(latent_sizes_node, activate_final=True), snt.LayerNorm() ]) self.global_fun = lambda: snt.Sequential( [snt.nets.MLP(latent_sizes_global, activate_final=False)]) with self._enter_variable_scope(): self._network = modules.GraphNetwork( self.edge_fun, self.node_fun, self.global_fun, edge_block_opt=EDGE_BLOCK_OPT, node_block_opt=NODE_BLOCK_OPT, global_block_opt=GLOBAL_BLOCK_OPT)
def __init__(self, name="MLPGraphNetwork"): super(MLPGraphNetwork, self).__init__(name=name) with self._enter_variable_scope(): self._network = \ modules.GraphNetwork(make_mlp_model, make_mlp_model, make_mlp_model, global_block_opt={"use_edges":False,"use_nodes":False})
def test_incompatible_higher_rank_partial_outputs_raises(self): """A error should be raised if partial outputs have incompatible shapes.""" input_graph = self._get_shaped_input_graph() edge_model_fn, node_model_fn, global_model_fn = self._get_shaped_model_fns() edge_model_fn_2 = functools.partial( snt.Conv2D, output_channels=10, kernel_shape=[3, 3], stride=[1, 2]) graph_network = modules.GraphNetwork( edge_model_fn_2, node_model_fn, global_model_fn) with self.assertRaisesRegexp(ValueError, "in both shapes must be equal"): graph_network(input_graph) node_model_fn_2 = functools.partial( snt.Conv2D, output_channels=10, kernel_shape=[3, 3], stride=[1, 2]) graph_network = modules.GraphNetwork( edge_model_fn, node_model_fn_2, global_model_fn) with self.assertRaisesRegexp(ValueError, "in both shapes must be equal"): graph_network(input_graph)
def __init__(self, latent_size, num_layers, name="MLPGraphNetwork"): super(MLPGraphNetwork, self).__init__(name=name) with self._enter_variable_scope(): self._network = modules.GraphNetwork( make_mlp_model(latent_size, num_layers), make_mlp_model(latent_size, num_layers), make_mlp_model(latent_size, num_layers))
def __init__(self, name="MLPGraphNetwork"): super(MLPGraphNetwork, self).__init__(name=name) self._network = modules.GraphNetwork( edge_model_fn=make_mlp_model, node_model_fn=make_mlp_model, global_model_fn=make_mlp_model )
def __init__(self, node_layer_s, edge_layer_s, globals_layer_s, name="MLPGraphNetwork"): super(MLPGraphNetwork, self).__init__(name=name) with self._enter_variable_scope(): self._network = modules.GraphNetwork(node_model_fn=make_mlp(node_layer_s), edge_model_fn=make_mlp(edge_layer_s), global_model_fn=make_mlp(globals_layer_s), reducer = tf.math.unsorted_segment_mean)
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 test_edge_block_options(self, use_edges, use_receiver_nodes, use_sender_nodes, use_globals): """Test for configuring the EdgeBlock options.""" reducer = tf.math.unsorted_segment_sum input_graph = self._get_input_graph() edge_model_fn = functools.partial(snt.Linear, output_size=10) edge_block_opt = {"use_edges": use_edges, "use_receiver_nodes": use_receiver_nodes, "use_sender_nodes": use_sender_nodes, "use_globals": use_globals} # 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} # Identity global model global_model_fn = lambda: tf.identity global_block_opt = {"use_globals": True, "use_nodes": False, "use_edges": False} 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, reducer=reducer) output_graph = graph_network(input_graph) edge_block = blocks.EdgeBlock( edge_model_fn=lambda: graph_network._edge_block._edge_model, use_edges=use_edges, use_receiver_nodes=use_receiver_nodes, use_sender_nodes=use_sender_nodes, use_globals=use_globals) expected_output_edge_block = edge_block(input_graph) expected_output_node_block = expected_output_edge_block expected_output_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.compat.v1.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.assertAllEqual(expected_edges_out, output_graph_out.edges) self.assertAllEqual(expected_nodes_out, output_graph_out.nodes) self.assertAllEqual(expected_globals_out, output_graph_out.globals)
def _get_model(self): 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) return modules.GraphNetwork( edge_model_fn=edge_model_fn, node_model_fn=node_model_fn, global_model_fn=global_model_fn)
def test_incompatible_higher_rank_inputs_raises(self, field_to_reshape): """A exception should be raised if the inputs have incompatible shapes.""" input_graph = self._get_shaped_input_graph() edge_model_fn, node_model_fn, global_model_fn = self._get_shaped_model_fns() input_graph = input_graph.map( lambda v: tf.transpose(v, [0, 2, 1, 3]), [field_to_reshape]) graph_network = modules.GraphNetwork( edge_model_fn, node_model_fn, global_model_fn) with self.assertRaisesRegexp(ValueError, "in both shapes must be equal"): graph_network(input_graph)
def __init__(self, name="MLPGraphNetwork"): super(MLPGraphNetwork, self).__init__(name=name) #with self._enter_variable_scope(): if 2 > 1: self._network = modules.GraphNetwork( make_mlp_model('edge'), make_mlp_model('node'), make_mlp_model('global'), #name=name )
def __init__(self, latent_size=16, num_layers=2, global_block=True, last_round=False, name="MLPGraphNetwork"): super(MLPGraphNetwork, self).__init__(name=name) partial_make_mlp_model = partial(make_mlp_model, latent_size=latent_size, num_layers=num_layers, last_round_edges=False) if last_round: partial_make_mlp_model_edges = partial(make_mlp_model, latent_size=latent_size, num_layers=num_layers, last_round_edges=True) else: partial_make_mlp_model_edges = partial_make_mlp_model with self.name_scope: if global_block: self._network = modules.GraphNetwork( partial_make_mlp_model_edges, partial_make_mlp_model, partial_make_mlp_model, edge_block_opt={"use_globals": True}, node_block_opt={"use_globals": True}, global_block_opt={ "use_globals": True, "edges_reducer": tf.unsorted_segment_mean, "nodes_reducer": tf.unsorted_segment_mean }) else: self._network = modules.GraphNetwork( partial_make_mlp_model_edges, partial_make_mlp_model, make_identity_model, edge_block_opt={"use_globals": False}, node_block_opt={"use_globals": False}, global_block_opt={ "use_globals": False, })
def __init__(self, n_recurrences, mlp_sizes, mlp_kwargs = None, name='Graph'): """Creates a new GraphBasedModel object. Args: n_recurrences: the number of message passing steps in the graph network. mlp_sizes: the number of neurons in each layer of the MLP. mlp_kwargs: additional keyword aguments passed to the MLP. name: the name of the Sonnet module. """ super(GraphBasedModel, self).__init__(name=name) self._n_recurrences = n_recurrences if mlp_kwargs is None: mlp_kwargs = {} model_fn = functools.partial( snt.nets.MLP, output_sizes=mlp_sizes, activate_final=True, **mlp_kwargs) final_model_fn = functools.partial( snt.nets.MLP, output_sizes=mlp_sizes + (1,), activate_final=False, **mlp_kwargs) with self._enter_variable_scope(): self._encoder = gn_modules.GraphIndependent( node_model_fn=model_fn, edge_model_fn=model_fn) if self._n_recurrences > 0: self._propagation_network = gn_modules.GraphNetwork( node_model_fn=model_fn, edge_model_fn=model_fn, # We do not use globals, hence we just pass the identity function. global_model_fn=lambda: lambda x: x, reducer=tf.unsorted_segment_sum, edge_block_opt=dict(use_globals=False), node_block_opt=dict(use_globals=False), global_block_opt=dict(use_globals=False)) self._decoder = gn_modules.GraphIndependent( node_model_fn=final_model_fn, edge_model_fn=model_fn)
def _build(self, graphs): with self._enter_variable_scope(): self.graph_transformer = modules.GraphNetwork( edge_model_fn=self._build_linear(self.edge_dimension, self.edge_activation, "edge-dense"), node_model_fn=self._build_linear(self.node_dimension, self.node_activation, "node-dense"), global_model_fn=self._build_linear(self.global_dimension, self.global_activation, "global-dense"), node_block_opt={"use_sent_edges": True}, reducer=self.reducer, name="graph_transformer") return self.graph_transformer(graphs)
def __init__( self, # for simplicity, all layers have the same size and the edge, # node and global models use the same structure latent_size, num_layers, name="MLPGraphNetwork", ): super(MLPGraphNetwork, self).__init__(name=name) model_fn = partial(make_mlp_model, latent_size=latent_size, num_layers=num_layers) with self._enter_variable_scope(): self._network = modules.GraphNetwork( edge_model_fn=model_fn, node_model_fn=model_fn, global_model_fn=model_fn, )
def test_created_variables(self, name=None): """Verifies variable names and shapes created by a GraphNetwork.""" name = name if name is not None else "graph_network" expected_var_shapes_dict = { name + "/edge_block/mlp/linear_0/b:0": [5], name + "/edge_block/mlp/linear_0/w:0": [4 + 4 + 3, 5], name + "/node_block/mlp/linear_0/b:0": [10], name + "/node_block/mlp/linear_0/w:0": [5 + 2 + 3, 10], name + "/global_block/mlp/linear_0/b:0": [15], name + "/global_block/mlp/linear_0/w:0": [10 + 5 + 3, 15], } input_graph = self._get_input_graph() extra_kwargs = {"name": name} if name else {} model = modules.GraphNetwork( edge_model_fn=functools.partial(snt.nets.MLP, output_sizes=[5]), node_model_fn=functools.partial(snt.nets.MLP, output_sizes=[10]), global_model_fn=functools.partial(snt.nets.MLP, output_sizes=[15]), **extra_kwargs) model(input_graph) variables = model.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 __init__(self, name="MLPGraphNetwork"): super(MLPGraphNetwork, self).__init__(name=name) self._network = modules.GraphNetwork(make_mlp_model, make_mlp_model, make_mlp_model)
def __init__(self, name="MLPGraphNetwork"): super(MLPGraphNetwork, self).__init__(name=name) with self._enter_variable_scope(): self._network = modules.GraphNetwork(make_mlp_model, make_mlp_model, make_mlp_model)
def __init__(self, edge_output_size=None, node_output_size=None, global_output_size=None, edge_layer_activation=tf.nn.relu, node_layer_activation=tf.nn.relu, global_layer_activation=tf.nn.relu, last_edge_layer_activation=tf.nn.softmax, last_node_layer_activation=tf.nn.softmax, last_global_layer_activation=tf.keras.activations.linear, edge_vocab_size=20, edge_embed_dim=100, node_vocab_size=1000, node_embed_dim=100, name="GraphAttention"): """ This network structure is supposed to handle NLP problems. :param edge_output_size: The size of the output vector corresponding to each edge :param node_output_size: The size of the output vector corresponding to each node :param global_output_size: The size of the output vector corresponding to the global feature :param edge_layer_activation: The activation used in each layer considering the edges. ReLU by default. :param node_layer_activation: The activation used in each layer considering the nodes. ReLU by default. :param global_layer_activation: The activation used in each layer considering the global feature. ReLU by default. :param last_edge_layer_activation: The activation function of the output layer corresponding to the edges. SoftMax by default. :param last_node_layer_activation: The activation function of the output layer corresponding to the nodes. SoftMax by default. :param last_global_layer_activation: The activation function of the output layer corresponding to the global features. Linear by default. :param edge_vocab_size: The size of the vocabulary containing the edges, if we use a non-pretrained embedding. :param edge_embed_dim: The dimension of the edge embedding, if we use a non-pretrained embedding. :param node_vocab_size: The size of the vocabulary containing the nodes, if we use a non-pretrained embedding. :param node_embed_dim: The dimension of the node embedding, if we use a non-pretrained embedding. :param name: The name of the network """ super(SimpleGraphAttention, self).__init__(name=name) self.edge_layer_activation = edge_layer_activation self.node_layer_activation = node_layer_activation self.global_layer_activation = global_layer_activation self.edge_vocab_size = edge_vocab_size self.edge_embed_dim = edge_embed_dim self.node_vocab_size = node_vocab_size self.node_embed_dim = node_embed_dim self._encoder = Encoder() self._network = graph_net_modules.GraphNetwork( edge_model_fn=self.edge_model_fn, node_model_fn=self.node_model_fn, global_model_fn=self.global_model_fn, reducer=tf.unsorted_segment_sum) # Transforms the outputs into the appropriate shapes. edge_fn = None if edge_output_size is None else \ lambda: sonnet_nets.ActivatedLinear(edge_output_size, last_edge_layer_activation) node_fn = None if node_output_size is None else \ lambda: sonnet_nets.ActivatedLinear(node_output_size, last_node_layer_activation) global_fn = None if global_output_size is None else \ lambda: sonnet_nets.ActivatedLinear(global_output_size, last_global_layer_activation) with self._enter_variable_scope(): self._output_transform = graph_net_modules.GraphIndependent( edge_fn, node_fn, global_fn)
print_graphs_tuple(graphs_tuple) recovered_data_dict_list = utils_np.graphs_tuple_to_data_dicts(graphs_tuple) #print(recovered_data_dict_list) tf.reset_default_graph() graph_dicts = data_dict_list 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)) input_graphs = utils_tf.data_dicts_to_graphs_tuple(graph_dicts) print(len(graph_dicts)) print(graph_dicts) print(input_graphs) output_graphs = graph_network(input_graphs) print(output_graphs.globals) #print_graphs_tuple(output_graphs) def zeros_graph(sample_graph, edge_size, node_size, global_size): zeros_graphs = sample_graph.replace(nodes=None, edges=None, globals=None)
def __init__(self, num_processing_steps=None, latent_size=None, n_layers=None, edge_output_size=None, node_output_size=None, global_output_size=None, reducer=None, out_init_scale=5.0, name="AggregationNet"): super(AggregationDiffNet, self).__init__(name=name) if num_processing_steps is None: self._proc_hops = [1, 1, 1, 1, 1, 1, 1, 1, 1, 1] else: self._proc_hops = num_processing_steps if reducer is None or reducer == 'max': reducer = unsorted_segment_max_or_zero elif reducer == 'logsumexp': reducer = segment_logsumexp elif reducer == 'softmax': reducer = segment_transformer elif reducer == 'sum': reducer = tf.math.unsorted_segment_sum else: raise ValueError('Unkown reducer!') if latent_size is None: latent_size = 16 if n_layers is None: n_layers = 2 self._num_processing_steps = len(self._proc_hops) self._n_stacked = latent_size * self._num_processing_steps def make_mlp(): return snt.nets.MLP([latent_size] * n_layers, activate_final=True) # def make_linear(): # return snt.nets.MLP([latent_size], activate_final=False) self._core = modules.GraphNetwork( edge_model_fn=make_mlp, node_model_fn=make_mlp, global_model_fn=make_mlp, edge_block_opt={'use_globals': False}, node_block_opt={ 'use_globals': False, 'use_sent_edges': False }, name="graph_net", reducer=reducer) self._encoder = modules.GraphIndependent(make_mlp, make_mlp, make_mlp, name="encoder") self._decoder = modules.GraphIndependent(make_mlp, make_mlp, make_mlp, name="decoder") inits = { 'w': ortho_init(out_init_scale), 'b': tf.constant_initializer(0.0) } # Transforms the outputs into the appropriate shapes. edge_fn = None if edge_output_size is None else lambda: snt.Linear( edge_output_size, initializers=inits, name="edge_output") node_fn = None if node_output_size is None else lambda: snt.Linear( node_output_size, initializers=inits, name="node_output") global_fn = None if global_output_size is None else lambda: snt.Linear( global_output_size, initializers=inits, name="global_output") with self._enter_variable_scope(): self._output_transform = modules.GraphIndependent(edge_fn, node_fn, global_fn, name="output")
graph_dicts = [ graph_3_nodes_4_edges, graph_5_nodes_8_edges, graph_7_nodes_13_edges, graph_9_nodes_25_edges ] ########## # Connecting a GraphNetwork recurrently input_graphs = utils_tf.data_dicts_to_graphs_tuple(graph_dicts) # graph_network = modules.GraphNetwork( # edge_model_fn=lambda: snt.Linear(output_size=EDGE_SIZE), # node_model_fn=lambda: snt.Linear(output_size=NODE_SIZE), # global_model_fn=lambda: snt.Linear(output_size=GLOBAL_SIZE)) graph_network = modules.GraphNetwork( edge_model_fn=snt.Linear(output_size=EDGE_SIZE), node_model_fn=snt.Linear(output_size=NODE_SIZE), global_model_fn=snt.Linear(output_size=GLOBAL_SIZE)) num_recurrent_passes = 3 previous_graphs = input_graphs for unused_pass in range(num_recurrent_passes): previous_graphs = graph_network(previous_graphs) print(previous_graphs.nodes[0]) output_graphs = previous_graphs tvars = graph_network.trainable_variables print('') ############### # broadcast
def __init__(self, num_processing_steps=None, latent_size=None, n_layers=None, edge_output_size=None, node_output_size=None, global_output_size=None, reducer=None, out_init_scale=5.0, name="AggregationNet"): super(NonLinearGraphNet, self).__init__(name=name) if num_processing_steps is None: self._num_processing_steps = 5 else: self._num_processing_steps = num_processing_steps if reducer is None or reducer == 'max': reducer = unsorted_segment_max_or_zero elif reducer == 'mean': reducer = tf.math.unsorted_segment_mean elif reducer == 'sum': reducer = tf.math.unsorted_segment_sum else: raise ValueError('Unknown reducer!') if latent_size is None: latent_size = 16 if n_layers is None: n_layers = 2 def make_mlp(): return snt.nets.MLP([latent_size] * n_layers, activate_final=False) if self._num_processing_steps > 0: # Edge model f^e(v_sender, v_receiver, e) - in the linear linear model, f^e = v_sender # Average over all the received edge features to get e' # Node model f^v(v, e'), but in the linear model, it was just f^v = e' self._core = modules.GraphNetwork( edge_model_fn=make_mlp, node_model_fn=make_mlp, global_model_fn=make_mlp, edge_block_opt={'use_globals': False}, node_block_opt={ 'use_globals': False, 'use_sent_edges': False }, name="graph_net", reducer=reducer) self._encoder = modules.GraphIndependent(make_mlp, make_mlp, make_mlp, name="encoder") self._decoder = modules.GraphIndependent(make_mlp, make_mlp, make_mlp, name="decoder") inits = { 'w': ortho_init(out_init_scale), 'b': tf.constant_initializer(0.0) } # Transforms the outputs into the appropriate shapes. edge_fn = None if edge_output_size is None else lambda: snt.Linear( edge_output_size, initializers=inits, name="edge_output") node_fn = None if node_output_size is None else lambda: snt.Linear( node_output_size, initializers=inits, name="node_output") global_fn = None if global_output_size is None else lambda: snt.Linear( global_output_size, initializers=inits, name="global_output") with self._enter_variable_scope(): self._output_transform = modules.GraphIndependent(edge_fn, node_fn, global_fn, name="output")
def test_higher_rank_outputs(self): """Tests that a graph net can be build with higher rank inputs/outputs.""" input_graph = self._get_shaped_input_graph() network = modules.GraphNetwork(*self._get_shaped_model_fns()) self._assert_build_and_run(network, input_graph)
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)