如何在 TensorFlow 中有效地分配给张量的切片 [英] How to efficiently assign to a slice of a tensor in TensorFlow

问题描述

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

I want to assign some values to slices of an input tensor in one of my model in TensorFlow 2.x (I am using 2.2 but ready to accept a solution for 2.1). A non-working template of what I am trying to do is:

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]))) 我得到以下错误:

of course when building this (AddToEven().build(tf.TensorShape([None, None]))) I get the following error:

TypeError: 'Tensor' object does not support item assignment

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

I can achieve this simple example via the following:

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

(您可以通过以下方式进行完整性检查:

(you can sanity-check with:

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

)

但是正如您所见，编写起来非常复杂.这仅适用于 1D(+ 批量大小)张量的静态更新次数(此处为 1).

But as you can see it's very complicated to write. And this is only for a static number of updates (here 1) on a 1D (+ batch size) tensor.

我想做的事情有点复杂，我认为用 tensor_scatter_nd_add 编写它会是一场噩梦.

What I want to do is a bit more involved and I think writing it with tensor_scatter_nd_add is going to be a nightmare.

许多当前关于该主题的 QA 涵盖了变量而不是张量的情况(参见例如 this 或 this).提到这里确实pytorch支持这一点，所以我很惊讶最近没有看到任何 tf 成员对该主题的回应.这个答案并没有真正帮助我，因为我需要一些那种面具生成也会很糟糕.

A lot of the current QAs on the topic cover the case for variables but not tensors (see e.g. this or this). It is mentionned here that indeed pytorch supports this, so I am surprised to see no response from any tf members on that topic recently. This answer doesn't really help me, because I will need some kind of mask generation which is going to be awful as well.

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

The question is thus: how can I do slice assignment efficiently (computation-wise, memory-wise and code-wise) w/o tensor_scatter_nd_add? The trick is that I want this to be as dynamical as possible, meaning that the shape of the inputs could be variable.

(对于任何好奇的人，我正在尝试翻译 这段代码在tf中).

(For anyone curious I am trying to translate this code in tf).

此问题最初发布于在 GitHub 问题中.

推荐答案

这里是另一种基于二进制掩码的解决方案.

Here is another solution based on binary mask.

"""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)是这样的.

Result (with TF 2.1) is like this.

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

--------以下是之前的回答--------

-------- Below is the previous answer --------

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

You need to create tf.Variable in build() method. It also allows dynamic size by shape=(None,). In the code below, I specified the input shape as (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 测试了这段代码

I tested this code with TF 2.1.0 and 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()，或者生成掩码函数.

P.S. There are some other ways, such as using tf.numpy_function(), or generating mask function.

这篇关于如何在 TensorFlow 中有效地分配给张量的切片的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持IT屋！