如何在Unet中翻转/旋转张量（Tensorflow 2.0）中的特定特征贴图？功能问题张量散射

Question

我在TensorFlow2.0（和克拉斯特遣部队)进行语义切分。在网络中的某些点上，我只想使用张量并翻转/旋转一些特征图。具体来说，初始张量有一个形状=[？，256256,5]（NHWC）。目前，批量大小为N=8。你知道吗

For clarification: Note that I am not aiming to do augmentation but to analyze the network behavior when the intermediate feature maps suffer a slight orientation change. Similar ideas have been studied, either flipping/rotating kernels or the feature maps (Gao et al. 2017, Efficient and Invariant CNNs for Dense Prediction, arXiv:1711.09064).

我创建了一个简单的层来执行此操作（因为尝试使用函数来执行此操作会产生问题；请参见类似的问题here）。你知道吗

函数/层应执行以下操作：

• 输入参数：张量（NHWC）、C值（代码中的增长率）和要翻转/旋转的特征映射的百分比（rotperc）。完成！
• 我随机选择所需数量的特征图。完成！
• 我翻转/旋转这些特征图。完成！
• 我用翻转/旋转的特征映射更新初始张量。要解决！

代码如下：

class MyPartialRotationFlipLayer(layers.Layer):
    def __init__(self, growth_rate, rotperc=0.2):
      super(MyPartialRotationFlipLayer, self).__init__()
      self.nCh = int(tf.multiply(rotperc, growth_rate))
      if self.nCh < 1:
        self.nCh = 1

    def call(self, input):
      ixCh = tf.random.shuffle(tf.range(0, input.shape[-1]))
      ixCh = ixCh[0:self.nCh]           # Randomly select some channel indices
      selChn = tf.gather(input, ixCh, axis=-1) # Get the tensors for those channels (4D tensor)
      selChn = tf.image.random_flip_left_right(selChn) 
      selChn = tf.image.rot90(selChn, k=tf.random.uniform(shape=[], minval=1, 
                                                          maxval=4, 
                                                          dtype=tf.int32))

      # From this point onwards, it is wrong:
      indGrid = tf.stack(tf.meshgrid(tf.range(input.shape[0]), 
                                     tf.range(input.shape[1]), 
                                     tf.range(input.shape[2]), 
                                     ixCh,
                                     indexing='ij'))
      return tf.tensor_scatter_nd_update(input, indGrid, selChn)

例如，使用默认的rotperc=0.2，我只翻转/旋转一个通道。假设ixCh=3（要提取和修改的通道）。因此，selChn是形状的张量=[？，256256,1]和input的形状=[？,256,256,5]. 现在我需要用selChn更新input中的通道ixCh=3。你知道吗

我的问题来自函数tf.tensor\u scatter\u nd\u更新。我的第一个想法是创建一个meshgrid，可以在中用作索引tf.tensor\u scatter\u nd\u更新，但显然此时批处理的数量未知，因此我无法创建meshgrid。你知道吗

如何更新/覆盖初始张量？ 提前谢谢！你知道吗

更新：

我提供了一个可行的解决方案，尽管它对所有选定的通道应用相同的修改（我还希望对每个选定的通道执行不同的翻转/旋转。。。欢迎对此提出任何建议！）你知道吗

因此，解决方案如下：

def build_channel_mask(tensor, channels):
    # Given a 4D tensor (NHWC) and an array of channels (of C), it builds    
    # boolean a 4D mask with True in the indicated channels.
    channels = tf.convert_to_tensor(channels)
    shape = tf.shape(tensor)
    mask = tf.equal(tf.range(shape[3], dtype=channels.dtype), 
                    tf.expand_dims(channels, 1))
    mask = tf.reduce_any(mask, axis=0)
    mask = tf.expand_dims(mask, 0)
    mask = tf.tile(mask, (shape[2], 1))
    mask = tf.expand_dims(mask, 0)
    mask = tf.tile(mask, (shape[1], 1, 1))
    mask = tf.expand_dims(mask, 0)
    mask = tf.tile(mask, (shape[0], 1, 1, 1))
    return mask


class MyPartialRotationFlipLayer(layers.Layer):
    def __init__(self, growth_rate, rotperc=0.2):
      super(MyPartialRotationFlipLayer, self).__init__()
      self.nCh = int(tf.multiply(rotperc, growth_rate))
      if self.nCh < 1:
        self.nCh = 1

    def call(self, x):
      # Select randomly the channels to flip/rotate
      ixCh = tf.random.shuffle(tf.range(0, x.shape[-1]))
      ixCh = ixCh[0:self.nCh]

      # Make a new tensor with the flip/rotation modification
      xNew = tf.image.random_flip_left_right(x)
      xNew = tf.image.rot90(xNew, k=tf.random.uniform(shape=[], minval=1, 
                                                      maxval=4, 
                                                      dtype=tf.int32))
      # Build a mask to use it in tf.where function
      mask = build_channel_mask(x, ixCh)
      return tf.where(mask, xNew, x)

那么，作为最后一个问题，您知道如何执行此操作，但每个选定通道都有独立的fip/旋转吗？