Python tensorflow.segment_sum方法代码示例

# 需要导入模块: import tensorflow [as 别名]
# 或者: from tensorflow import segment_sum [as 别名]
def create_tensor(self, in_layers=None, set_tensors=True, **kwargs):
    """
    parent layers: atom_features, atom_membership
    """
    if in_layers is None:
      in_layers = self.in_layers
    in_layers = convert_to_layers(in_layers)

    self.build()
    output = in_layers[0].out_tensor
    atom_membership = in_layers[1].out_tensor
    for i, W in enumerate(self.W_list[:-1]):
      output = tf.matmul(output, W) + self.b_list[i]
      output = self.activation(output)
    output = tf.matmul(output, self.W_list[-1]) + self.b_list[-1]
    if self.output_activation:
      output = self.activation(output)
    output = tf.segment_sum(output, atom_membership)
    out_tensor = output
    if set_tensors:
      self.variables = self.trainable_weights
      self.out_tensor = out_tensor
    return out_tensor 

# 需要导入模块: import tensorflow [as 别名]
# 或者: from tensorflow import segment_sum [as 别名]
def call(self, x):
    """Execute this layer on input tensors.

    Parameters
    ----------
    x: list of Tensor
      should be [embedding tensor of molecules, of shape (batch_size*max_n_atoms*n_embedding),
                 mask tensor of molecules, of shape (batch_size*max_n_atoms)]

    Returns
    -------
    list of tf.Tensor
      Of shape (batch_size)
    """
    self.build()
    output = x[0]
    atom_membership = x[1]
    for i, W in enumerate(self.W_list[:-1]):
      output = tf.matmul(output, W) + self.b_list[i]
      output = self.activation(output)
    output = tf.matmul(output, self.W_list[-1]) + self.b_list[-1]
    if self.output_activation:
      output = self.activation(output)
    output = tf.segment_sum(output, atom_membership)
    return output 

# 需要导入模块: import tensorflow [as 别名]
# 或者: from tensorflow import segment_sum [as 别名]
def test_SegmentSum(self):
        t = tf.segment_sum(self.random(4, 2, 3), np.array([0, 1, 1, 2]))
        self.check(t) 

# 需要导入模块: import tensorflow [as 别名]
# 或者: from tensorflow import segment_sum [as 别名]
def testValues(self):
    dtypes = [tf.float32,
              tf.float64,
              tf.int64,
              tf.int32,
              tf.complex64,
              tf.complex128]

    # Each item is np_op1, np_op2, tf_op
    ops_list = [(np.add, None, tf.segment_sum),
                (self._mean_cum_op, self._mean_reduce_op,
                 tf.segment_mean),
                (np.ndarray.__mul__, None, tf.segment_prod),
                (np.minimum, None, tf.segment_min),
                (np.maximum, None, tf.segment_max)]

    # A subset of ops has been enabled for complex numbers
    complex_ops_list = [(np.add, None, tf.segment_sum),
                        (np.ndarray.__mul__, None, tf.segment_prod)]

    n = 10
    shape = [n, 2]
    indices = [i // 3 for i in range(n)]
    for dtype in dtypes:
      if dtype in (tf.complex64, tf.complex128):
        curr_ops_list = complex_ops_list
      else:
        curr_ops_list = ops_list

      with self.test_session(use_gpu=False):
        tf_x, np_x = self._input(shape, dtype=dtype)
        for np_op1, np_op2, tf_op in curr_ops_list:
          np_ans = self._segmentReduce(indices, np_x, np_op1, np_op2)
          s = tf_op(data=tf_x, segment_ids=indices)
          tf_ans = s.eval()
          self._assertAllClose(indices, np_ans, tf_ans)
          # NOTE(mrry): The static shape inference that computes
          # `tf_ans.shape` can only infer that sizes from dimension 1
          # onwards, because the size of dimension 0 is data-dependent
          # and may therefore vary dynamically.
          self.assertAllEqual(np_ans.shape[1:], tf_ans.shape[1:]) 

# 需要导入模块: import tensorflow [as 别名]
# 或者: from tensorflow import segment_sum [as 别名]
def testSegmentIdsShape(self):
    shape = [4, 4]
    tf_x, _ = self._input(shape)
    indices = tf.constant([0, 1, 2, 2], shape=[2, 2])
    with self.assertRaises(ValueError):
      tf.segment_sum(data=tf_x, segment_ids=indices) 

# 需要导入模块: import tensorflow [as 别名]
# 或者: from tensorflow import segment_sum [as 别名]
def testSegmentIdsSize(self):
    shape = [4, 4]
    with self.test_session():
      tf_x, _ = self._input(shape)
      indices = [0, 1]
      s = tf.segment_sum(data=tf_x, segment_ids=indices)
      with self.assertRaisesOpError("segment_ids should be the same size"):
        s.eval() 

# 需要导入模块: import tensorflow [as 别名]
# 或者: from tensorflow import segment_sum [as 别名]
def testSegmentIdsValid(self):
    # This is a baseline for the following SegmentIdsInvalid* tests.
    shape = [4, 4]
    with self.test_session():
      tf_x, _ = self._input(shape)
      indices = [0, 0, 0, 1]
      result = tf.segment_sum(data=tf_x, segment_ids=indices).eval()
      self.assertAllEqual([[15, 18, 21, 24], [13, 14, 15, 16]], result) 

# 需要导入模块: import tensorflow [as 别名]
# 或者: from tensorflow import segment_sum [as 别名]
def testSegmentIdsInvalid1(self):
    shape = [4, 4]
    with self.test_session():
      tf_x, _ = self._input(shape)
      indices = [-1, -1, 0, 0]
      s = tf.segment_sum(data=tf_x, segment_ids=indices)
      with self.assertRaisesOpError("segment ids do not start at 0"):
        s.eval() 

# 需要导入模块: import tensorflow [as 别名]
# 或者: from tensorflow import segment_sum [as 别名]
def testSegmentIdsInvalid3(self):
    shape = [4, 4]
    with self.test_session():
      tf_x, _ = self._input(shape)
      indices = [0, 0, 2, 2]
      s = tf.segment_sum(data=tf_x, segment_ids=indices)
      with self.assertRaisesOpError("segment ids are not increasing by 1"):
        s.eval() 

# 需要导入模块: import tensorflow [as 别名]
# 或者: from tensorflow import segment_sum [as 别名]
def testSegmentIdsInvalid4(self):
    shape = [4, 4]
    with self.test_session():
      tf_x, _ = self._input(shape)
      indices = [0, 1, 0, 1]
      s = tf.segment_sum(data=tf_x, segment_ids=indices)
      with self.assertRaisesOpError("segment ids are not increasing by 1"):
        s.eval() 

# 需要导入模块: import tensorflow [as 别名]
# 或者: from tensorflow import segment_sum [as 别名]
def testSegmentIdsInvalid5(self):
    shape = [4, 4]
    with self.test_session():
      tf_x, _ = self._input(shape)
      indices = [0, 1, 2, 0]
      s = tf.segment_sum(data=tf_x, segment_ids=indices)
      with self.assertRaisesOpError(
          r"Segment id 1 out of range \[0, 1\), probably "
          "because 'segment_ids' input is not sorted."):
        s.eval() 

# 需要导入模块: import tensorflow [as 别名]
# 或者: from tensorflow import segment_sum [as 别名]
def testSegmentIdsInvalid6(self):
    shape = [4, 4]
    with self.test_session():
      tf_x, _ = self._input(shape)
      indices = [0, 0, 0, -1]
      s = tf.segment_sum(data=tf_x, segment_ids=indices)
      with self.assertRaisesOpError("segment ids must be >= 0"):
        s.eval() 

# 需要导入模块: import tensorflow [as 别名]
# 或者: from tensorflow import segment_sum [as 别名]
def testSegmentIdsInvalid7(self):
    shape = [4, 4]
    with self.test_session():
      tf_x, _ = self._input(shape)
      indices = [0, 0, 0, -2]
      s = tf.segment_sum(data=tf_x, segment_ids=indices)
      with self.assertRaisesOpError("segment ids must be >= 0"):
        s.eval() 

# 需要导入模块: import tensorflow [as 别名]
# 或者: from tensorflow import segment_sum [as 别名]
def testGradientMatchesSegmentSum(self):
    # Strategy: compute the gradient for UnsortedSegmentSum and SegmentSum
    # and compare the outputs, which should be identical.
    # NB: for this test to work, indices must be valid for SegmentSum, namely
    # it must be sorted, the indices must be contiguous, and num_segments
    # must be max(indices) + 1.
    indices = [0, 0, 1, 1, 1, 2, 3, 4, 5]
    n = len(indices)
    num_cols = 2
    shape = [n, num_cols]
    num_segments = max(indices) + 1
    with self.test_session(use_gpu=self.use_gpu):
      tf_x, np_x = self._input(shape, dtype=tf.float64)
      # Results from UnsortedSegmentSum
      unsorted_s = tf.unsorted_segment_sum(data=tf_x,
                                                 segment_ids=indices,
                                                 num_segments=num_segments)
      (unsorted_jacob_t, unsorted_jacob_n) = tf.test.compute_gradient(
          tf_x,
          shape,
          unsorted_s,
          [num_segments, num_cols],
          x_init_value=np_x.astype(np.double),
          delta=1)
      # Results from SegmentSum
      sorted_s = tf.segment_sum(data=tf_x, segment_ids=indices)
      sorted_jacob_t, sorted_jacob_n = tf.test.compute_gradient(
          tf_x,
          shape,
          sorted_s,
          [num_segments, num_cols],
          x_init_value=np_x.astype(np.double),
          delta=1)
    self.assertAllClose(unsorted_jacob_t, sorted_jacob_t, rtol=1e-3, atol=1e-3)
    self.assertAllClose(unsorted_jacob_n, sorted_jacob_n, rtol=1e-3, atol=1e-3) 

# 需要导入模块: import tensorflow [as 别名]
# 或者: from tensorflow import segment_sum [as 别名]
def _compute_vert_context_soft(self, edge_factor, vert_factor, reuse=False):
        """
        attention-based vertex(node) message pooling
        """

        out_edge = utils.pad_and_gather(edge_factor, self.edge_pair_mask_inds[:,0])
        in_edge = utils.pad_and_gather(edge_factor, self.edge_pair_mask_inds[:,1])
        # gather correspounding vert factors
        vert_factor_gathered = tf.gather(vert_factor, self.edge_pair_segment_inds)

        # concat outgoing edges and ingoing edges with gathered vert_factors
        out_edge_w_input = tf.concat(concat_dim=1, values=[out_edge, vert_factor_gathered])
        in_edge_w_input = tf.concat(concat_dim=1, values=[in_edge, vert_factor_gathered])

        # compute compatibility scores
        (self.feed(out_edge_w_input)
             .fc(1, relu=False, reuse=reuse, name='out_edge_w_fc')
             .sigmoid(name='out_edge_score'))
        (self.feed(in_edge_w_input)
             .fc(1, relu=False, reuse=reuse, name='in_edge_w_fc')
             .sigmoid(name='in_edge_score'))

        out_edge_w = self.get_output('out_edge_score')
        in_edge_w = self.get_output('in_edge_score')

        # weight the edge factors with computed weigths
        out_edge_weighted = tf.mul(out_edge, out_edge_w)
        in_edge_weighted = tf.mul(in_edge, in_edge_w)


        edge_sum = out_edge_weighted + in_edge_weighted
        vert_ctx = tf.segment_sum(edge_sum, self.edge_pair_segment_inds)
        return vert_ctx