如何设计神经网络的行为像For-Loop？

Posted 2021-03-19

我想让图像中的每个像素沿着x轴移动一些值。数组 "displacements "包含了每个像素的位移值。

输入： 图像240×320×3 数组 "位移 "240×320。

    def apply_displacements(displacements, img):
        displacements = displacements * -100
        result_img = np.zeros(shape=(img_height, img_width, 3), dtype=float)

        for i in range(0, img_height):
            for j in range(0, img_width):
                disp = int(j + displacements[i][j])
                for k in range(0, 3):
                    if 0 <= disp <= (img_width-1):
                        result_img[i][disp][k] = img[0][i][j][k]
        return result_img

输出：包含原始图像中的像素的图像，但它们已被位移数组中的值转换。

这个函数对于我正在研究的神经网络来说是必要的，但是一个Loop是不可区分的。我读到过，你可以创建一个神经网络来模仿一个函数的行为，因此这个网络是可以区分的。而这正是我正在尝试的。

但是我应该使用什么样的NN模型呢？是否有一些已知的网络架构可以有效地完成这样的任务？卷积NN，因为我正在处理图像，但我不确定有多少层，层的大小应该是多少等等。这将是监督的，因为我有原始函数应该被模拟，NN应该生成与原始函数相同的输出。

我试着在google上搜索，但我都不知道我的搜索查询应该是什么。"Neural Network to simulate simple For-Loop"？

对NN-Architecture有什么建议吗？

答案

我不确定这是否是你想要的，或者你可以使用的东西，但几乎作为一种练习，我把你放在那里的算法的TensorFlow版本放在一起。我是为沿垂直和水平坐标的任意位移图做的，所以你的问题将是它的一个子案例。它和你的算法略有不同，这里如果两个像素被移动到同一个位移坐标上，它们的值就会被聚合，而不是只取最后一个。另外，这个算法使用双线性插值来计算分数位移，这使得它可以进行微分。

import tensorflow as tf

def image_displace_tf(img, disp):
    # Get image shape
    s = tf.shape(img, out_type=tf.int64)
    # Image shape as floats
    sf = tf.dtypes.cast(s, disp.dtype)
    # Make coordinate grid
    bb, ii, jj = tf.meshgrid(tf.range(s[0]), tf.range(s[1]), tf.range(s[2]),
                             indexing='ij')
    # Compute displaced coordinates
    coords = tf.stack([bb, ii, jj], axis=-1)
    # Add a null "batch displacement"
    disp_pad = tf.pad(disp, [[0, 0], [0, 0], [0, 0], [1, 0]])
    coords_disp = tf.dtypes.cast(coords, disp.dtype) + disp_pad
    # Mask displacements out of range
    mask = ((coords_disp[..., 1] >= 0) & (coords_disp[..., 1] < sf[1] - 1) &
            (coords_disp[..., 2] >= 0) & (coords_disp[..., 2] < sf[2] - 1))
    coords = tf.boolean_mask(coords, mask)
    coords_disp = tf.boolean_mask(coords_disp, mask)
    # Compute interpolation alpha values for bilinear interpolation
    alpha_1 = tf.math.floormod(coords_disp, 1.0)
    alpha_1 = tf.expand_dims(alpha_1, axis=-1)
    alpha_0 = 1 - alpha_1
    alpha_00 = alpha_0[..., 1, :] * alpha_0[..., 2, :]
    alpha_01 = alpha_0[..., 1, :] * alpha_1[..., 2, :]
    alpha_10 = alpha_1[..., 1, :] * alpha_0[..., 2, :]
    alpha_11 = alpha_1[..., 1, :] * alpha_1[..., 2, :]
    # Begin and end indices for each dimension
    idx_00 = tf.dtypes.cast(coords_disp, tf.int64)
    idx_01 = idx_00 + [0, 0, 1]
    idx_10 = idx_00 + [0, 1, 0]
    idx_11 = idx_00 + [0, 1, 1]
    # Values at begin and end for each dimension scaled by their alpha values
    img_coords = tf.gather_nd(img, coords)
    value_00 = alpha_00 * img_coords
    value_01 = alpha_01 * img_coords
    value_10 = alpha_10 * img_coords
    value_11 = alpha_11 * img_coords
    # Concatenate all indices and values
    idx_all = tf.concat([idx_00, idx_01, idx_10, idx_11], axis=0)
    value_all = tf.concat([value_00, value_01, value_10, value_11], axis=0)
    # Make aggregated result
    return tf.scatter_nd(idx_all, value_all, s)

我不知道它是否能如你所愿，但你可以同时取图像和位移图的梯度。这是你可以使用它的方法。

import tensorflow as tf
import matplotlib.pyplot as plt

tf.random.set_seed(0)
plt.close('all')
# Make a radial image
x = tf.linspace(-1.0, 1.0, 400)
y = tf.linspace(-1.0, 1.0, 300)
r = tf.math.sqrt(tf.math.square(x) + tf.math.square(tf.expand_dims(y, 1)))
img = tf.math.maximum(0.8 - r, 0.0)
# Give it three channels
img = tf.tile(tf.expand_dims(img, axis=-1), [1, 1, 3])
# Add batch dimension
img = tf.expand_dims(img, axis=0)
# Show image
plt.figure()
plt.imshow(img.numpy()[0])
plt.title('Source image')
plt.show()

# Make a wavy displacement map along the horizontal axis
s = tf.shape(img)
disp_i = tf.zeros(s[:-1], dtype=img.dtype)
_, ii, jj = tf.meshgrid(*(tf.linspace(0.0, 1.0, si) for si in s[:3]), indexing='ij')
disp_j = 5.0 * (tf.sin(40.0 * jj + tf.sin(20.0 * ii)))
disp = tf.stack([disp_i, disp_j], axis=-1)
# Show map
plt.figure()
plt.imshow(0.5 + 0.5 * disp.numpy()[0, ..., 1])
plt.title('Horizontal displacement map')
plt.show()

# Do displacement
with tf.GradientTape() as g:
    g.watch(img)
    g.watch(disp)
    img_disp = image_displace_tf(img, disp)
    # Some minimization goal - e.g. sum of squared pixel values
    goal = tf.math.reduce_sum(tf.square(img_disp))
# Show result
plt.figure()
plt.imshow(img_disp.numpy()[0])
plt.title('Displaced image')
plt.show()
# Show gradients
img_grad, disp_grad = g.gradient(goal, [img, disp])
# Gradient of image
plt.figure()
plt.imshow(img_grad.numpy()[0, ..., 0])
plt.title('Image gradient')
plt.show()
# Gradient of horizontal coorindate of displacement map
plt.figure()
plt.imshow(disp_grad.numpy()[0, ..., 1])
plt.title('Horizontal displacement gradient')
plt.show()

输出图像