How to explicitly broadcast a tensor to match another's shape in tensorflow?

前端 未结 4 1830
爱一瞬间的悲伤
爱一瞬间的悲伤 2021-02-13 09:17

I have three tensors, A, B and C in tensorflow, A and B are both of shape (m, n, r), C is a binary tensor of sha

4条回答
  •  再見小時候
    2021-02-13 09:42

    import tensorflow as tf
    
    def broadcast(tensor, shape):
         """Broadcasts ``x`` to have shape ``shape``.
                                                                       |
         Uses ``tf.Assert`` statements to ensure that the broadcast is
         valid.
    
         First calculates the number of missing dimensions in 
         ``tf.shape(x)`` and left-pads the shape of ``x`` with that many 
         ones. Then identifies the dimensions of ``x`` that require
         tiling and tiles those dimensions appropriately.
    
         Args:
             x (tf.Tensor): The tensor to broadcast.
             shape (Union[tf.TensorShape, tf.Tensor, Sequence[int]]): 
                 The shape to broadcast to.
    
         Returns:
             tf.Tensor: ``x``, reshaped and tiled to have shape ``shape``.
    
         """
         with tf.name_scope('broadcast') as scope:
             shape_x = tf.shape(x)
             rank_x = tf.shape(shape0)[0]
             shape_t = tf.convert_to_tensor(shape, preferred_dtype=tf.int32)
             rank_t = tf.shape(shape1)[0]
    
             with tf.control_dependencies([tf.Assert(
                 rank_t >= rank_x,
                 ['len(shape) must be >= tf.rank(x)', shape_x, shape_t],
                 summarize=255
             )]):
                 missing_dims = tf.ones(tf.stack([rank_t - rank_x], 0), tf.int32)
    
             shape_x_ = tf.concat([missing_dims, shape_x], 0)
             should_tile = tf.equal(shape_x_, 1)
    
             with tf.control_dependencies([tf.Assert(
                 tf.reduce_all(tf.logical_or(tf.equal(shape_x_, shape_t), should_tile),
                 ['cannot broadcast shapes', shape_x, shape_t],
                 summarize=255
             )]):
                 multiples = tf.where(should_tile, shape_t, tf.ones_like(shape_t))
                 out = tf.tile(tf.reshape(x, shape_x_), multiples, name=scope)
    
             try:
                 out.set_shape(shape)
             except:
                 pass
    
             return out
    
    A = tf.random_normal([20, 100, 10])
    B = tf.random_normal([20, 100, 10])
    C = tf.random_normal([20, 100, 1])
    
    C = broadcast(C, A.shape)
    D = tf.select(C, A, B)
    

提交回复
热议问题