def build_graph(parameters): """Build the segment_sum op testing graph.""" data = tf.compat.v1.placeholder(dtype=parameters["data_dtype"], name="data", shape=parameters["data_shape"]) segment_ids = tf.constant(parameters["segment_ids"], dtype=tf.int32) out = tf.segment_sum(data, segment_ids) return [data], [out]
def forward(self, self_vecs, neigh_vecs, segment_ids=None): """Calculates attention coefficients. The following code is used to implement equation (1) (2) (3) and (4) in the paper, which is called **self-attention** process. \alpha_{ij}^{l} & = \mathrm{softmax_i} (e_{ij}^{l}) e_{ij}^{l} = \mathrm{LeakyReLU}(\vec{a}^T [W h_{i} \vert W h_{j}]) Args: self_vecs: Tensor, batch nodes' embedding vector, shape [B, D] neigh_vecs: Tensor, corresponding neighbor nodes' embedding vector, shape [total_nbrs, D] segment_ids: Tensor, segment ids indicates neighbor nodes' belonging, shape [total_nbrs] Returns: updated neighbor nodes' embedding vector. """ self_vecs = self._fc(self_vecs) neigh_vecs = self._fc(neigh_vecs) if segment_ids is None: # sampled GAT num_neibors = neigh_vecs.shape[1] self_vecs_extend = tf.tile(tf.expand_dims(self_vecs, 1), [1, num_neibors, 1]) coefficients = tf.nn.softmax(tf.nn.leaky_relu(self._attn_fc( tf.concat([self_vecs_extend, neigh_vecs], axis=-1)))) coefficients = tf.nn.dropout(coefficients, 1 - self._attn_drop) neigh_vecs = tf.multiply(coefficients, neigh_vecs) neigh_vecs = tf.reduce_sum(neigh_vecs, axis=1) else: # full neighbor GAT self_vecs_extend = tf.gather(self_vecs, segment_ids) coefficients = tf.math.exp(tf.nn.leaky_relu(self._attn_fc( tf.concat([self_vecs_extend, neigh_vecs], axis=-1)))) seg_sum = tf.gather(tf.segment_sum(coefficients, segment_ids), segment_ids) coefficients = coefficients / seg_sum coefficients = tf.nn.dropout(coefficients, 1 - self._attn_drop) neigh_vecs = tf.multiply(coefficients, neigh_vecs) neigh_vecs = tf.segment_sum(neigh_vecs, segment_ids) if self._activation is not None: neigh_vecs = self._activation(neigh_vecs) return neigh_vecs
def _random_segmentation(num_items, num_segments): """Partition a sequence of items randomly into non-empty segments. Args: num_items: an integer scalar > 0 num_segments: an integer scalar in [1, num_items] Returns: a Tensor with shape [num_segments] containing positive integers that add up to num_items """ first_in_segment = tf.pad( tf.random.shuffle( to_int(tf.range(num_items - 1) < num_segments - 1), seed=123), [[1, 0]]) segment_id = tf.cumsum(first_in_segment) segment_length = tf.segment_sum(tf.ones_like(segment_id), segment_id) return segment_length
def forward(self, self_vecs, neigh_vecs, segment_ids=None): """ Update node's embedding based on its neighbors. Args: self_vecs: batch nodes' embeddings with shape [B, D] neigh_vecs: neighbor nodes' embeddings with shape [total_nbrs, D] segment_ids: segment ids that indicates neighbor nodes' belonging, shape [total_nbrs] Returns: updated batch nodes' embedding vector [B, H] """ if segment_ids is None: # sampled GCN neigh_vecs = tf.reduce_sum(neigh_vecs, axis=1) else: # full neighbor GCN neigh_vecs = tf.segment_sum(data=neigh_vecs, segment_ids=segment_ids) updated_vecs = tf.reduce_sum([self_vecs, neigh_vecs], axis=0) return self._fc(updated_vecs)