Python TensorFlow中成批掩蔽任意正方形区域

我在TensorFlow 2.0中有一个大的3D张量，形状为[batch_size，N，N]。批处理中的每个样本都包含一个我想要掩盖的方形ROI（ROI的大小在一个名为inner_dims的单独数组中定义）。我在谷歌上搜索了很多次，但都没能找到一个看似有效的解决方案

示例：

下面的例子显示了我想要实现的目标。在本例中，我手动创建了一个符合我要求的掩码：

batch\u size=5
N=4
#inputs.shape:[5,4,4]
内部直径=[2,3,1,4,2]
面具=
[
[
[1 1 0 0]
[1 1 0 0]
[0 0 0 0]
[0 0 0 0]
]
[
[1 1 1 0]
[1 1 1 0]
[1 1 1 0]
[0 0 0 0]
]
[
[1 0 0 0]
[0 0 0 0]
[0 0 0 0]
[0 0 0 0]
]
[
[1 1 1 1]
[1 1 1 1]
[1 1 1 1]
[1 1 1 1]
]
[
[1 1 0 0]
[1 1 0 0]
[0 0 0 0]
[0 0 0 0]
]
]
输出=tf.乘法（输入，掩码）

我尝试过的

我尝试了将受影响的值赋值为零的天真方法，如下所示：

范围内的i（0，len（内部尺寸））：
输入[i，内部尺寸[i]：，：]=0
输入[i，：内部尺寸[i]，内部尺寸[i]：]=0

这（如预期）给了我一个项目分配错误

TypeError:'tensorflow.python.framework.ops.tensor'对象不支持项分配

我的问题

---

如何以快速且GPU友好的方式实现所需的掩蔽？动态生成掩码是更好的方法，还是直接分配方法可以发挥作用？

我刚刚想出了一个简单的解决方案，使用和广播解决了这个问题

但是，我不确定预先创建这样的遮罩是否是最有效的方法

N = 4
inner_dims = [2, 3, 1, 4, 2]

# This mask is a 5xN tensor.
seq_mask = tf.sequence_mask(inner_dims, N)

# Expand dimensions to 5x1xN, and 5xNx1, respectively.
seq_mask1 = tf.expand_dims(seq_mask, axis=1)
seq_mask2 = tf.expand_dims(seq_mask, axis=2)

# Thanks to broadcasting, the logical_and will yield a 5xNxN tensor.
mask = tf.logical_and(seq_mask1, seq_mask2)
mask = tf.cast(mask, dtype=tf.float32)

outputs=tf.multiply(inputs, mask)