keras 如何在TensorFlow中高效地分配Tensor切片

enyaitl3 于 2022-11-13 发布在其他

关注(0)|答案(2)|浏览(128)

我想在TensorFlow 2.x中为我的一个模型中的输入Tensor切片分配一些值（我使用的是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

(you可以对以下项进行健全性检查：

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

）
但正如你所看到的，写起来非常复杂，而且这只是针对1D（+批量大小）Tensor上的静态更新次数（这里是1）。
我想做的是更投入一点，我认为用tensor_scatter_nd_add写它将是一场噩梦。
很多关于这个主题的问答都涉及到了变量而不是Tensor的情况（例如this或this）。这里提到pytorch确实支持这个，所以我很惊讶最近没有看到任何tf成员对这个主题的回应。这个答案对我没有什么帮助，因为我需要某种掩码生成，这也会很糟糕。
问题是这样的：在没有tensor_scatter_nd_add的情况下，我如何高效地进行切片分配（计算方面、内存方面和代码方面）？关键在于我希望它尽可能地动态，这意味着inputs的形状可以是可变的。
(For如果有人好奇我正在尝试将这段代码翻译成tf）。
此问题最初发布在GitHub问题中。

keras

来源：https://stackoverflow.com/questions/62092147/how-to-efficiently-assign-to-a-slice-of-a-tensor-in-tensorflow

2条答案

按热度按时间

ccgok5k51#

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

"""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，）。在下面的代码中，我指定了输入shape为（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()

我用TF2.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)

P.S.还有一些其他的方法，比如使用tf.numpy_function（），或者生成掩码函数。

赞(0）回复(0）举报 2022-11-13

piah890a2#

它似乎没有产生这样的错误：

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.

赞(0）回复(0）举报 2022-11-13

我来回答

keras 如何在TensorFlow中高效地分配Tensor切片

2条答案

相关问题

热门标签

最新问答