拨开荷叶行,寻梦已然成。仙女莲花里,翩翩白鹭情。
IMG-LOGO
主页 文章列表 如何将tf.gather_nd用于多维张量

如何将tf.gather_nd用于多维张量

白鹭 - 2022-03-09 1986 0 0

如果我有多维张量,我不完全理解我应该如何使用 tf.gather_nd() 沿某个轴拾取元素。让我们举一个小例子(如果我得到这个简单例子的答案,它也解决了我更复杂的原始问题)。假设我有 rgb 影像,我正在尝试沿通道选择最小的像素值(如果资料顺序为(B、H、W、C),则为最后一个维度)。我知道这可以做到,tf.recude_min(x, axis=-1)但我想知道是否也可以用tf.argmin()and做同样的事情tf.gather_nd()

from skimage import data
import tensorflow as tf
import numpy as np

# Load RGB image from skimage, cast it to float32 and put it in order (B,H,W,C)
image = data.astronaut()
image = tf.cast(image, tf.float32)
image = tf.expand_dims(image, axis=0)

# Take minimum pixel value of each channel in a way number 1
min_along_channels_1 = tf.reduce_min(image, axis=-1)

# Take minimum pixel value of each channel in a way number 2
# The goal is that min_along_channels_1 is equal to min_along_channels_2
idxs = tf.argmin(image, axis=-1)
min_along_channels_2 = tf.gather_nd(image, idxs) # This line gives error :(

uj5u.com热心网友回复:

您将不得不使用tf.meshgrid,这将创建一个由两个一维阵列组成的矩形网格,表示第一维和第二维的张量索引,因为tf.gather_nd需要确切知道从哪里提取跨维度的值。这是一个简化的示例:

import tensorflow as tf

image = tf.random.normal((1, 4, 4, 3))
image = tf.squeeze(image, axis=0)
idx = tf.argmin(image, axis=-1)

ij = tf.stack(tf.meshgrid(
    tf.range(image.shape[0], dtype=tf.int64), 
    tf.range(image.shape[1], dtype=tf.int64),
                              indexing='ij'), axis=-1)

gather_indices = tf.concat([ij, tf.expand_dims(idx, axis=-1)], axis=-1)
result = tf.gather_nd(image, gather_indices)

print('First option -->', tf.reduce_min(image, axis=-1))
print('Second option -->', result)
First option --> tf.Tensor(
[[-0.53245485 -0.29117298 -0.64434254 -0.8209638 ]
 [-0.9386176  -0.5993224  -0.597746   -1.5392851 ]
 [-0.5478666  -1.5280861  -1.0344954  -1.920418  ]
 [-0.5580688  -1.425873   -1.9276617  -1.0668412 ]], shape=(4, 4), dtype=float32)
Second option --> tf.Tensor(
[[-0.53245485 -0.29117298 -0.64434254 -0.8209638 ]
 [-0.9386176  -0.5993224  -0.597746   -1.5392851 ]
 [-0.5478666  -1.5280861  -1.0344954  -1.920418  ]
 [-0.5580688  -1.425873   -1.9276617  -1.0668412 ]], shape=(4, 4), dtype=float32)

或者用你的例子:

from skimage import data
import tensorflow as tf
import numpy as np

image = data.astronaut()
image = tf.cast(image, tf.float32)
image = tf.expand_dims(image, axis=0)

min_along_channels_1 = tf.reduce_min(image, axis=-1)

image = tf.squeeze(image, axis=0)
idx = tf.argmin(image, axis=-1)

ij = tf.stack(tf.meshgrid(
    tf.range(image.shape[0], dtype=tf.int64), 
    tf.range(image.shape[1], dtype=tf.int64),
                              indexing='ij'), axis=-1)

gather_indices = tf.concat([ij, tf.expand_dims(idx, axis=-1)], axis=-1)
min_along_channels_2 = tf.gather_nd(image, gather_indices)

print(tf.equal(min_along_channels_1, min_along_channels_2))
tf.Tensor(
[[[ True  True  True ...  True  True  True]
  [ True  True  True ...  True  True  True]
  [ True  True  True ...  True  True  True]
  ...
  [ True  True  True ...  True  True  True]
  [ True  True  True ...  True  True  True]
  [ True  True  True ...  True  True  True]]], shape=(1, 512, 512), dtype=bool)
标签:

0 评论

发表评论

您的电子邮件地址不会被公开。 必填的字段已做标记 *