如何有效地分配给 TensorFlow 中的张量切片

Posted 2023-02-23

【中文标题】如何有效地分配给 TensorFlow 中的张量切片

【英文标题】：How to efficiently assign to a slice of a tensor in TensorFlow

【发布时间】：2020-09-17 09:19:18

【问题描述】：

我想在 TensorFlow 2.x 中的一个模型中为输入张量的切片分配一些值（我正在使用 2.2，但准备接受 2.1 的解决方案）。 我正在尝试做的一个非工作模板是：

import tensorflow as tf
from tensorflow.keras.models import Model

class AddToEven(Model):
    def call(self, inputs):
        outputs = inputs
        outputs[:, ::2] += inputs[:, ::2]
        return outputs

当然，在构建这个 (AddToEven().build(tf.TensorShape([None, None]))) 时，我收到以下错误：

TypeError: 'Tensor' object does not support item assignment

我可以通过以下方式实现这个简单的示例：

class AddToEvenScatter(Model):
    def call(self, inputs):
        batch_size = tf.shape(inputs)[0]
        n = tf.shape(inputs)[-1]
        update_indices = tf.range(0, n, delta=2)[:, None]
        scatter_nd_perm = [1, 0]
        inputs_reshaped = tf.transpose(inputs, scatter_nd_perm)
        outputs = tf.tensor_scatter_nd_add(
            inputs_reshaped,
            indices=update_indices,
            updates=inputs_reshaped[::2],
        )
        outputs = tf.transpose(outputs, scatter_nd_perm)
        return outputs

（您可以通过以下方式进行完整性检查：

model = AddToEvenScatter()
model.build(tf.TensorShape([None, None]))
model(tf.ones([1, 10]))

)

但是正如您所见，编写起来非常复杂。这仅适用于一维（+批量）张量的静态更新次数（此处为 1）。

我想要做的是更多的参与，我认为用tensor_scatter_nd_add 编写它将会是一场噩梦。

当前很多关于该主题的 QA 都涵盖了变量但不包括张量的情况（参见例如 this 或 this）。 提到here 确实 pytorch 支持这一点，所以我很惊讶地看到最近没有任何 tf 成员对此主题的回应。 This answer 并没有真正帮助我，因为我需要生成某种蒙版，这也会很糟糕。

因此，问题是：如何在没有tensor_scatter_nd_add 的情况下有效地进行切片分配（计算方面、内存方面和代码方面）？诀窍是我希望它尽可能动态，这意味着inputs 的形状可以是可变的。

（对于任何好奇的人，我正在尝试将this code 翻译成 tf）。

此问题最初发布于in a GitHub issue。

【问题讨论】：

由于没有更好的解决方案，我使用tensor_scatter_nd_update为此创建了一个模块。从长远来看，希望我不必诉诸于此。但与此同时，如果有人想使用它，您可以查看here。

【参考方案1】：

这似乎没有产生任何错误：

import tensorflow as tf
from tensorflow.keras.models import Model

class AddToEven(Model):
    def call(self, inputs):
        outputs = inputs
        outputs = outputs[:, ::2] + 1
        return outputs

# tf.Tensor.__iadd__ does not seem to exist, but tf.Tensor.__add__ does. 

【讨论】：

当然，但这不是我想做的。在这里，您正在输出输入中添加 1 的偶数位置切片。所以tf.zeros(1, 10) 的输出将是[1, 1, ,1 ,1 1]。我想要的是[1, 0, 1, 0, 1, 0, 1, 0, 1, 0]，即偶数位置加一的输入张量。

我不明白，你写了“我正在尝试做的一个非工作模板是：”，然后显示一个 TypeError。我在不更改原始类提供的任何功能的情况下修复了错误？

不不，功能与我所说的不一样。我希望我的函数从 [0, 0,0,0] 变为 [1,0,1,0]。在 numpy 或 pytorch 中，这是我的函数会做的事情，但在 tensorflow 中，它似乎并不那么简单。

一个健全性检查是输出的形状应该与输入的形状相同。

我并不反对您提供的代码可能不是您想要的代码。我只是指出您提供了代码，说这是我想要工作的代码，但有一个 TypeError，我修复了 TypeError，现在代码完全按照您的要求工作。听起来您不仅希望代码正常工作，还不确定算法本身，也不确定如何获得所需的输出。这在您的帖子中并不清楚，我以为您想知道为什么会收到 TypeError。 :-)

【参考方案2】：

这是另一个基于二进制掩码的解决方案。

"""Solution based on binary mask.
- We just add this mask to inputs, instead of multiplying."""
class AddToEven(tf.keras.Model):
    def __init__(self):
        super(AddToEven, self).__init__()        

    def build(self, inputshape):
        self.built = True # Actually nothing to build with, becuase we don't have any variables or weights here.

    @tf.function
    def call(self, inputs):
        w = inputs.get_shape()[-1]

        # 1-d mask generation for w-axis (activate even indices only)        
        m_w = tf.range(w)  # [0, 1, 2,... w-1]
        m_w = ((m_w%2)==0) # [True, False, True ,...] with dtype=tf.bool

        # Apply 1-d mask to 2-d input
        m_w = tf.expand_dims(m_w, axis=0) # just extend dimension as to be (1, W)
        m_w = tf.cast(m_w, dtype=inputs.dtype) # in advance, we need to convert dtype

        # Here, we just add this (1, W) mask to (H,W) input magically.
        outputs = inputs + m_w # This add operation is allowed in both TF and numpy!
        return tf.reshape(outputs, inputs.get_shape())

在这里进行完整性检查。

# sanity-check as model
model = AddToEven()
model.build(tf.TensorShape([None, None]))
z = model(tf.zeros([2,4]))
print(z)

结果（TF 2.1）是这样的。

tf.Tensor(
[[1. 0. 1. 0.]
 [1. 0. 1. 0.]], shape=(2, 4), dtype=float32)

--------下面是上一个答案--------

您需要在 build() 方法中创建 tf.Variable。 它还允许通过 shape=(None,) 实现动态大小。 在下面的代码中，我将输入形状指定为 (None, None)。

class AddToEven(tf.keras.Model):
    def __init__(self):
        super(AddToEven, self).__init__()

    def build(self, inputshape):
        self.v = tf.Variable(initial_value=tf.zeros((0,0)), shape=(None, None), trainable=False, dtype=tf.float32)

    @tf.function
    def call(self, inputs):
        self.v.assign(inputs)
        self.v[:, ::2].assign(self.v[:, ::2] + 1)
        return self.v.value()

我用 TF 2.1.0 和 TF1.15 测试了这段代码

# test
add_to_even = AddToEven()
z = add_to_even(tf.zeros((2,4)))
print(z)

结果：

tf.Tensor(
[[1. 0. 1. 0.]
 [1. 0. 1. 0.]], shape=(2, 4), dtype=float32)

附：还有一些其他的方法，比如使用 tf.numpy_function()，或者生成掩码函数。

【讨论】：

我接受这个答案，因为它提供了一种解决问题的方法。如果我没记错的话，我认为我看到的主要缺点是，对于模型中的每个切片分配，您需要创建一个不同的缓冲区 tf.Variable。

你将如何在这里生成掩码？

@ZaccharieRamzi 这就是我的解决方案告诉你的

不，因为您的解决方案没有使用可变形状

@Zaccharie Ramzi 我用基于（添加）掩码的新解决方案更新了答案。在这种方法中，模型没有任何变量或权重需要保留。因此，如果您打印 model.summary()，它将显示为空模型。