class PropagationTransformation( object ):
"""
Transforms a graph state by propagating info across the graph
"""
def __init__(self, transfer_size, graph_spec, transfer_activation=identity, dropout_keep=1):
"""
Params:
transfer_size: Integer, how much to transfer
graph_spec: Instance of GraphStateSpec giving graph spec
transfer_activation: Activation function to use during transfer
"""
self._transfer_size = transfer_size
self._transfer_activation = transfer_activation
self._graph_spec = graph_spec
self._process_input_size = graph_spec.num_node_ids + graph_spec.node_state_size
self._transfer_stack = LayerStack(self._process_input_size, 2 * graph_spec.num_edge_types * transfer_size, activation=self._transfer_activation, name="propagation_transfer", dropout_keep=dropout_keep, dropout_input=False, dropout_output=True)
self._propagation_gru = BaseGRULayer(graph_spec.num_node_ids + self._transfer_size, graph_spec.node_state_size, name="propagation", dropout_keep=dropout_keep, dropout_input=False, dropout_output=True)
@property
def params(self):
return self._propagation_gru.params + self._transfer_stack.params
def dropout_masks(self, srng, state_mask=None):
return self._transfer_stack.dropout_masks(srng) + self._propagation_gru.dropout_masks(srng, use_output=state_mask)
def split_dropout_masks(self, dropout_masks):
transfer_used, dropout_masks = self._transfer_stack.split_dropout_masks(dropout_masks)
gru_used, dropout_masks = self._propagation_gru.split_dropout_masks(dropout_masks)
return (transfer_used+gru_used), dropout_masks
def process(self, gstate, dropout_masks=Ellipsis):
"""
Process a graph state.
1. Data is transfered from each node to each other node along both forward and backward edges.
This data is processed with a Wx+b style update, and an optional transformation is applied
2. Nodes sum the transfered data, weighted by the existence of the other node and the edge.
3. Nodes perform a GRU update with this input
Params:
gstate: A GraphState giving the current state
"""
if dropout_masks is Ellipsis:
dropout_masks = None
append_masks = False
else:
append_masks = True
node_obs = T.concatenate([gstate.node_ids, gstate.node_states],2)
flat_node_obs = node_obs.reshape([-1, self._process_input_size])
transformed, dropout_masks = self._transfer_stack.process(flat_node_obs,dropout_masks)
transformed = transformed.reshape([gstate.n_batch, gstate.n_nodes, 2*self._graph_spec.num_edge_types, self._transfer_size])
scaled_transformed = transformed * T.shape_padright(T.shape_padright(gstate.node_strengths))
# scaled_transformed is of shape (n_batch, n_nodes, 2*num_edge_types, transfer_size)
# We want to multiply through by edge strengths, which are of shape
# (n_batch, n_nodes, n_nodes, num_edge_types), both fwd and backward
edge_strength_scale = T.concatenate([gstate.edge_strengths, gstate.edge_strengths.swapaxes(1,2)], 3)
# edge_strength_scale is of (n_batch, n_nodes, n_nodes, 2*num_edge_types)
intermed = T.shape_padaxis(scaled_transformed, 2) * T.shape_padright(edge_strength_scale)
# intermed is of shape (n_batch, n_nodes "source", n_nodes "dest", 2*num_edge_types, transfer_size)
# now reduce along the "source" and "edge_types" dimensions to get dest activations
# of shape (n_batch, n_nodes, transfer_size)
reduced_result = T.sum(T.sum(intermed, 3), 1)
# now add information fom current node id
full_input = T.concatenate([gstate.node_ids, reduced_result], 2)
# we flatten to apply GRU
flat_input = full_input.reshape([-1, self._graph_spec.num_node_ids + self._transfer_size])
flat_state = gstate.node_states.reshape([-1, self._graph_spec.node_state_size])
new_flat_state, dropout_masks = self._propagation_gru.step(flat_input, flat_state, dropout_masks)
new_node_states = new_flat_state.reshape(gstate.node_states.shape)
new_gstate = gstate.with_updates(node_states=new_node_states)
if append_masks:
return new_gstate, dropout_masks
else:
return new_gstate
def process_multiple(self, gstate, iterations, dropout_masks=Ellipsis):
"""
Run multiple propagagtion steps.
Params:
gstate: A GraphState giving the current state
iterations: An integer. How many steps to propagate
"""
if dropout_masks is Ellipsis:
dropout_masks = None
append_masks = False
else:
append_masks = True
def _scan_step(cur_node_states, node_strengths, node_ids, edge_strengths, *dmasks):
curstate = GraphState(node_strengths, node_ids, cur_node_states, edge_strengths)
newstate, _ = self.process(curstate, dmasks if dropout_masks is not None else None)
return newstate.node_states
outputs_info = [gstate.node_states]
used_dropout_masks, dropout_masks = self.split_dropout_masks(dropout_masks)
all_node_states, _ = theano.scan(_scan_step, n_steps=iterations, non_sequences=[gstate.node_strengths, gstate.node_ids, gstate.edge_strengths] + used_dropout_masks, outputs_info=outputs_info)
final_gstate = gstate.with_updates(node_states=all_node_states[-1,:,:,:])
if append_masks:
return final_gstate, dropout_masks
else:
return final_gstate
class NewNodesInformTransformation( object ):
"""
Transforms a graph state by adding nodes, conditioned on an input vector
"""
def __init__(self, input_width, inform_width, proposal_width, graph_spec, use_old_aggregate=False, dropout_keep=1):
"""
Params:
input_width: Integer giving size of input
inform_width: Size of internal aggregate
proposal_width: Size of internal proposal
graph_spec: Instance of GraphStateSpec giving graph spec
use_old_aggregate: Use the old aggregation mode
"""
self._input_width = input_width
self._graph_spec = graph_spec
self._proposal_width = proposal_width
self._inform_width = inform_width
aggregate_type = AggregateRepresentationTransformationSoftmax \
if use_old_aggregate \
else AggregateRepresentationTransformation
self._inform_aggregate = aggregate_type(inform_width, graph_spec, dropout_keep, dropout_output=True)
self._proposer_gru = BaseGRULayer(input_width+inform_width, proposal_width, name="newnodes_proposer", dropout_keep=dropout_keep, dropout_input=False, dropout_output=True)
self._proposer_stack = LayerStack(proposal_width, 1+graph_spec.num_node_ids, [proposal_width], bias_shift=3.0, name="newnodes_proposer_post", dropout_keep=dropout_keep, dropout_input=False)
@property
def params(self):
return self._proposer_gru.params + self._proposer_stack.params + self._inform_aggregate.params
def dropout_masks(self, srng):
return self._inform_aggregate.dropout_masks(srng) + self._proposer_gru.dropout_masks(srng) + self._proposer_stack.dropout_masks(srng)
def get_candidates(self, gstate, input_vector, max_candidates, dropout_masks=Ellipsis):
"""
Get the current candidate new nodes. This is accomplished as follows:
1. Using the aggregate transformation, we gather information from nodes (who should have performed
a state update already)
1. The proposer network, conditioned on the input and info, proposes multiple candidate nodes,
along with a confidence
3. A new node is created for each candidate node, with an existence strength given by
confidence, and an initial id as proposed
This method directly returns these new nodes for comparision
Params:
gstate: A GraphState giving the current state
input_vector: A tensor of the form (n_batch, input_width)
max_candidates: Integer, limit on the number of candidates to produce
Returns:
new_strengths: A tensor of the form (n_batch, new_node_idx)
new_ids: A tensor of the form (n_batch, new_node_idx, num_node_ids)
"""
if dropout_masks is Ellipsis:
dropout_masks = None
append_masks = False
else:
append_masks = True
n_batch = gstate.n_batch
n_nodes = gstate.n_nodes
aggregated_repr, dropout_masks = self._inform_aggregate.process(gstate, dropout_masks)
# aggregated_repr is of shape (n_batch, inform_width)
full_input = T.concatenate([input_vector, aggregated_repr],1)
outputs_info = [self._proposer_gru.initial_state(n_batch)]
gru_dropout_masks, dropout_masks = self._proposer_gru.split_dropout_masks(dropout_masks)
proposer_step = lambda st,ipt,*dm: self._proposer_gru.step(ipt, st, dm if dropout_masks is not None else None)[0]
raw_proposal_acts, _ = theano.scan(proposer_step, n_steps=max_candidates, non_sequences=[full_input]+gru_dropout_masks, outputs_info=outputs_info)
# raw_proposal_acts is of shape (candidate, n_batch, blah)
flat_raw_acts = raw_proposal_acts.reshape([-1, self._proposal_width])
flat_processed_acts, dropout_masks = self._proposer_stack.process(flat_raw_acts, dropout_masks)
candidate_strengths = T.nnet.sigmoid(flat_processed_acts[:,0]).reshape([max_candidates, n_batch])
candidate_ids = T.nnet.softmax(flat_processed_acts[:,1:]).reshape([max_candidates, n_batch, self._graph_spec.num_node_ids])
new_strengths = candidate_strengths.dimshuffle([1,0])
new_ids = candidate_ids.dimshuffle([1,0,2])
if append_masks:
return new_strengths, new_ids, dropout_masks
else:
return new_strengths, new_ids
def process(self, gstate, input_vector, max_candidates, dropout_masks=Ellipsis):
"""
Process an input vector and update the state accordingly.
"""
if dropout_masks is Ellipsis:
dropout_masks = None
append_masks = False
else:
append_masks = True
new_strengths, new_ids, dropout_masks = self.get_candidates(gstate, input_vector, max_candidates, dropout_masks)
new_gstate = gstate.with_additional_nodes(new_strengths, new_ids)
if append_masks:
return new_gstate, dropout_masks
else:
return new_gstate