我正在尝试对batch_size为20的一批大小为[128128,3]的图像执行一些转换。我已经在一个图像上完成了部分转换，但我不确定如何将其扩展到整个批次。
x是形状的批处理图像张量[20128128,3]
grads也是形状张量[20128128,3]
我想找到每个图像的grads的最大值，即我将获得20个最大值。让我们将image_i的最大值称为max_grads_i。
add_max_grads是形状张量[20128128,3]
我想要更新add_max_grads的值，使得对于image_I，我需要在grads张量中的max_grads_I的位置更新相应的add_max_grads，并且我需要在3个通道上更新它。
例如，对于图像i，形状为[128128,3]，如果在grads[50][60][1]处找到max_grads_i，则我希望用相同的常数值更新[50][60][0]、[50][6][1]、[50][60][2]处的add_max_grads。
最终，这需要扩展到批次中的所有图像，以创建具有形状的add_max_grads[20128128,3]。
我怎样才能做到这一点？
目前，我正在使用tf.math.reduce_max(grads,(1,2,3))为批处理中的每个图像找到grads张量的最大值，该批处理返回形状为[20]的张量，但我停留在其余部分。
如果我理解正确的话，我想这就是你想要的。
这是一个独立的例子。
import tensorflow as tf
#tshape = (20,128,128,3) # The real shape needed
tshape = (2,4,4,3) # A small shape, to allow print output to be readable
# A sample x vector, set to zeros here so we can see what gets updated easily
x = tf.zeros(shape=tshape, dtype=tf.float32)
# grads, same shape as x, let's use a random matrix here
grads = tf.random.normal(shape=tshape)
print(f"grads: n{grads}")
# ---- Now calculate the indices we will use in tf.tensor_scatter_nd_update
# 4-D mask tensor (m, x, y, c) with 1 in max locations (by m), zeros otherwise
mask4d = tf.cast(grads == tf.reduce_max(grads, axis=(1,2,3), keepdims=True), dtype=tf.int32)
# 3D mask tensor (m, x, y) with 1 in max pixel locations (by m) across channel
mask3d = tf.reduce_max(mask4d, axis=3)
# indices of maximum values, each item gives m, x, y
indices = tf.where(mask3d)
print(f"nindicesn{indices}")
# ---- Now calculate the updates we will use in tf.tensor_scatter_nd_update. This has shape (m, c)
newval = 999
updates = tf.constant(newval, shape=(tshape[0], tshape[3]), dtype=tf.float32)
# ---- Now modify x by scattering newval through it. Each entry in indices indicates a slice x[m, x, y, :]
x_updated = tf.tensor_scatter_nd_update(x, indices, updates)
print(f"nx_updatedn{x_updated}")