def global_model_fn(self):
"""
The function of the global model used in the graph dependent part of the network
:return: Sequential sonnet model
"""
return snt.Sequential([
sonnet_nets.ActivatedLinear(16, self.global_layer_activation),
sonnet_nets.ActivatedLinear(32, self.global_layer_activation),
sonnet_nets.ActivatedLinear(16, self.global_layer_activation),
snt.LayerNorm()
])
def edge_model_fn(self):
"""
The function of the edge model used in the graph dependent part of the network
:return: Sequential sonnet model
"""
return snt.Sequential([
sonnet_nets.ActivatedLinear(HIDDEN_SIZE,
self.edge_layer_activation),
sonnet_nets.ActivatedLinear(int(HIDDEN_SIZE / 2),
self.edge_layer_activation),
sonnet_nets.ActivatedLinear(int(HIDDEN_SIZE / 4),
self.edge_layer_activation)
])
def node_model_fn(self):
"""
The function of the node model used in the graph dependent part of the network
:return: Sequential sonnet model
"""
return snt.Sequential([
#sonnet_nets.ActivatedLSTM(hidden_size=HIDDEN_SIZE, activation_funcion=self.node_layer_activation),
sonnet_nets.ActivatedLinear(HIDDEN_SIZE,
self.node_layer_activation),
sonnet_nets.ActivatedLinear(int(HIDDEN_SIZE / 2),
self.node_layer_activation),
sonnet_nets.ActivatedLinear(int(HIDDEN_SIZE / 4),
self.node_layer_activation),
sonnet_nets.ActivatedLinear(2, self.node_layer_activation)
])
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)