通过嵌套的tf.map_fn反向传播梯度

Question

我希望能够在每个像素的深度通道上映射一个TensorFlow函数，该函数对应于具有尺寸[batch_size, H, W, n_channels]的矩阵中的每个向量。换句话说，对于批次中的每个大小为H x W的图像：

1. 我提取一些功能图F_k（其数量为n_channels），具有相同的大小H x W（因此，所有功能图共同形成一个形状的张量[H，W，n_channels];
2. 然后，我希望将自定义函数应用于与每个特征图F_k的第i行和第j列相关联的向量v_ij，但是完全探索了深度通道（例如，v的维数为[1 x 1 x n_channels]）。理想情况下，所有这些都会同时发生。

下面是解释过程的图片。唯一的区别是输入和输出“接受域”均为1x1（独立地将函数应用于每个像素）。

这类似于将1x1卷积应用于矩阵; 但是，我需要在深度通道上应用更一般的函数，而不是简单的求和操作。

我认为tf.map_fn()可能是一个选项，我尝试了以下解决方案，其中我递归使用tf.map_fn()来访问与每个像素相关联的特征。然而，这似乎有些次优，最重要的是，在尝试反向传播梯度时会引发错误。

您有任何想法，为什么会出现这种情况，以及我应该如何构造我的代码以避免错误？

这是我当前的函数实现：

import tensorflow as tf
from tensorflow import layers


def apply_function_on_pixel_features(incoming):
    # at first the input is [None, W, H, n_channels]
    if len(incoming.get_shape()) > 1:
        return tf.map_fn(lambda x: apply_function_on_pixel_features(x), incoming)
    else:
        # here the input is [n_channels]
        # apply some function that applies a transfomration and returns a vetor of the same size
        output = my_custom_fun(incoming) # my_custom_fun() doesn't change the shape
        return output

和我的代码主体：

H = 128
W = 132
n_channels = 8

x1 = tf.placeholder(tf.float32, [None, H, W, 1])
x2 = layers.conv2d(x1, filters=n_channels, kernel_size=3, padding='same')

# now apply a function to the features vector associated to each pixel
x3 = apply_function_on_pixel_features(x2)  
x4 = tf.nn.softmax(x3)

loss = cross_entropy(x4, labels)
optimizer = tf.train.AdamOptimizer(lr)
train_op = optimizer.minimize(loss)  # <--- ERROR HERE!

特别地，错误如下：

File "/home/venvs/tensorflowGPU/lib/python3.6/site-packages/tensorflow/python/ops/control_flow_ops.py", line 2481, in AddOp
    self._AddOpInternal(op)

File "/home/venvs/tensorflowGPU/lib/python3.6/site-packages/tensorflow/python/ops/control_flow_ops.py", line 2509, in _AddOpInternal
    self._MaybeAddControlDependency(op)
File "/home/venvs/tensorflowGPU/lib/python3.6/site-packages/tensorflow/python/ops/control_flow_ops.py", line 2547, in _MaybeAddControlDependency
    op._add_control_input(self.GetControlPivot().op)

AttributeError: 'NoneType' object has no attribute 'op'

整个错误堆栈和代码可以在这里找到。 感谢帮助， G.

---

更新：

根据@thushv89的建议，我添加了一个可能的解决方案来解决问题。我仍然不知道为什么我的先前代码不起作用。任何关于此的见解仍将非常感激。

- gab

---

请参见 https://dev59.com/Q6rka4cB1Zd3GeqPjtSg。 - geometrikal

@geometrikal 谢谢你的回答。恐怕我没有解释清楚问题。我更新了问题，也许更清楚了。如果您仍然认为广播是最好的选择，能否请您更好地解释一下如何在我的情况下使用它？（我没明白） - gab

@gabriele，从图片上看，您似乎正在尝试在特征图中的每个像素上应用一些自定义函数？如果是这样，为什么需要递归？只需进行重塑，执行map_fn，然后再次重塑回原始形状即可。 - thushv89

@gabriele 或许加上你想要做什么的原因会提供更多背景信息，帮助读者找到解决方案。从外观上看，这需要进行超级数学计算。 - Chaitanya Bapat

@thushv89 我按照你的建议操作后，成功地实现了梯度传播。虽然我仍然不明白我的实现哪里出了问题，但现在看起来一切都正常了。非常感谢 :) 我将当前的解决方案添加到答案中。 - gab

Answer 1

按照@thushv89的建议，我重塑了数组，应用了功能，然后再次重塑（以避免tf.map_fn递归）。 我仍然不知道之前的代码为什么不起作用，但当前的实现允许将梯度传播回前一层。 对于可能感兴趣的人，我将其保留在下面：

def apply_function_on_pixel_features(incoming, batch_size):

    # get input shape:
    _, W, H, C = incoming.get_shape().as_list()
    incoming_flat = tf.reshape(incoming, shape=[batch_size * W * H, C])

    # apply function on every vector of shape [1, C]
    out_matrix = my_custom_fun(incoming_flat)  # dimension remains unchanged

    # go back to the input shape shape [None, W, H, C]
    out_shape = tf.convert_to_tensor([batch_size, W, H, C])
    out_matrix = tf.reshape(out_matrix, shape=out_shape)

    return out_matrix

请注意，现在我需要给出批量大小以正确地重塑张量，因为如果我给出无限制或-1作为维度，TensorFlow会抱怨。

对上述代码的任何评论和见解仍将非常感激。

- gab

---

嘿，谢谢你的问题，很有趣。有一件事：你不应该写成incoming_flat = tf.reshape(incoming, shape=[-1, C])和out_shape = tf.convert_to_tensor([-1, W, H, C])吗？ - MPKenning

Answer 2

@gabriele，关于需要依赖batch_size的问题，你尝试过以下方法吗？这个函数不依赖于batch_size。您可以将map_fn替换为任何您喜欢的内容。

def apply_function_on_pixel_features(incoming):

    # get input shape:
    _, W, H, C = incoming.get_shape().as_list()
    incoming_flat = tf.reshape(incoming, shape=[-1, C])

    # apply function on every vector of shape [1, C]
    out_matrix = tf.map_fn(lambda x: x+1, incoming_flat)  # dimension remains unchanged

    # go back to the input shape shape [None, W, H, C]
    out_matrix = tf.reshape(out_matrix, shape=[-1, W, H, C])

    return out_matrix

我测试的完整代码如下。

import numpy as np
import tensorflow as tf
from tensorflow.keras.losses import categorical_crossentropy

def apply_function_on_pixel_features(incoming):

    # get input shape:
    _, W, H, C = incoming.get_shape().as_list()
    incoming_flat = tf.reshape(incoming, shape=[-1])

    # apply function on every vector of shape [1, C]
    out_matrix = tf.map_fn(lambda x: x+1, incoming_flat)  # dimension remains unchanged

    # go back to the input shape shape [None, W, H, C]
    out_matrix = tf.reshape(out_matrix, shape=[-1, W, H, C])

    return out_matrix

H = 32
W = 32
x1 = tf.placeholder(tf.float32, [None, H, W, 1])
labels = tf.placeholder(tf.float32, [None, 10])
x2 = tf.layers.conv2d(x1, filters=1, kernel_size=3, padding='same')

# now apply a function to the features vector associated to each pixel
x3 = apply_function_on_pixel_features(x2)  
x4 = tf.layers.flatten(x3)
x4 = tf.layers.dense(x4, units=10, activation='softmax')

loss = categorical_crossentropy(labels, x4)
optimizer = tf.train.AdamOptimizer(0.001)
train_op = optimizer.minimize(loss)


x = np.zeros(shape=(10, H, W, 1))
y = np.random.choice([0,1], size=(10, 10))


with tf.Session() as sess:
  tf.global_variables_initializer().run()
  sess.run(train_op, feed_dict={x1: x, labels:y})

- thushv89

---

嗨@thushv89，感谢您的建议。然而，根据您的建议，我将重塑张量的形状为[-1]而不是[batch_size * W * H，C]（这是我需要将函数一致应用于每个像素的所有特征所需的形状）。此外，我认为重塑为[-1，C]然后[-1，W，H，C]会导致错误。 TensorFlow似乎抱怨无法将具有未知形状的对象转换为张量。 - gab

@gabriele，对我来说运行得很好。你有错误吗？ - thushv89

嗨@thushv89，抱歉回复晚了，但我之前无法测试这个。我再次尝试用-1替换batch_size，现在似乎可以工作了。可能我之前有些问题。谢谢你的帮助！ :) 你应该更新这一行incoming_flat = tf.reshape(incoming, shape=[-1])，使用shape=[-1, C]，这就是我想要的，然后我会给你分赏。 - gab

@gabriele，很高兴听到这个消息。我更新了我的答案。 :) - thushv89