def call(self, inputs): if self.data_mode == 'graph': X = inputs[0] I = inputs[1] else: X = inputs if self.data_mode == 'graph': return tf.segment_sum(X, I) else: return K.sum(X, axis=-2, keepdims=(self.data_mode == 'single'))
def top_k(scores, I, ratio, top_k_var): """ Returns indices to get the top K values in `scores` segment-wise, with segments defined by I. K is not fixed, but it is defined as a ratio of the number of elements in each segment. :param scores: a rank 1 tensor with scores; :param I: a rank 1 tensor with segment IDs; :param ratio: float, ratio of elements to keep for each segment; :param top_k_var: a tf.Variable without shape validation (e.g., `tf.Variable(0.0, validate_shape=False)`); :return: a rank 1 tensor containing the indices to get the top K values of each segment in `scores`. """ num_nodes = tf.segment_sum(tf.ones_like(I), I) # Number of nodes in each graph cumsum = tf.cumsum(num_nodes) # Cumulative number of nodes (A, A+B, A+B+C) cumsum_start = cumsum - num_nodes # Start index of each graph n_graphs = tf.shape(num_nodes)[0] # Number of graphs in batch max_n_nodes = tf.reduce_max(num_nodes) # Order of biggest graph in batch batch_n_nodes = tf.shape(I)[0] # Number of overall nodes in batch to_keep = tf.ceil(ratio * tf.cast(num_nodes, tf.float32)) to_keep = tf.cast(to_keep, tf.int32) # Nodes to keep in each graph index = tf.range(batch_n_nodes) index = (index - tf.gather(cumsum_start, I)) + (I * max_n_nodes) y_min = tf.reduce_min(scores) dense_y = tf.ones((n_graphs * max_n_nodes, )) dense_y = dense_y * tf.cast( y_min - 1, tf.float32 ) # subtract 1 to ensure that filler values do not get picked dense_y = tf.assign( top_k_var, dense_y, validate_shape=False ) # top_k_var is a variable with unknown shape defined in the elsewhere dense_y = tf.scatter_update(dense_y, index, scores) dense_y = tf.reshape(dense_y, (n_graphs, max_n_nodes)) perm = tf.argsort(dense_y, direction='DESCENDING') perm = perm + cumsum_start[:, None] perm = tf.reshape(perm, (-1, )) to_rep = tf.tile(tf.constant([1., 0.]), (n_graphs, )) rep_times = tf.reshape( tf.concat((to_keep[:, None], (max_n_nodes - to_keep)[:, None]), -1), (-1, )) mask = tf_repeat_1d(to_rep, rep_times) perm = tf.boolean_mask(perm, mask) return perm
def call(self, inputs): if self.data_mode == 'graph': X, I = inputs else: X = inputs attn_coeff = K.dot(X, self.attn_kernel) attn_coeff = K.squeeze(attn_coeff, -1) attn_coeff = K.softmax(attn_coeff) if self.data_mode == 'single': output = K.dot(attn_coeff[None, ...], X) elif self.data_mode == 'batch': output = K.batch_dot(attn_coeff, X) else: output = attn_coeff[:, None] * X output = tf.segment_sum(output, I) return output
def call(self, inputs): if self.data_mode == 'graph': X, I = inputs else: X = inputs inputs_linear = K.dot(X, self.lg_kernel) + self.lg_bias attn_map = K.dot(X, self.attn_kernel) + self.attn_bias attn_map = K.sigmoid(attn_map) masked_inputs = inputs_linear * attn_map if self.data_mode in {'single', 'batch'}: output = K.sum(masked_inputs, axis=-2, keepdims=self.data_mode == 'single') else: output = tf.segment_sum(masked_inputs, I) return output