[英]How to efficiently assign to a slice of a tensor in TensorFlow
我想為 TensorFlow 2.x 中的 model 之一中的輸入張量切片分配一些值(我正在使用 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]))
)
但是正如你所看到的,寫起來非常復雜。 這僅適用於 static 一維(+ 批量大小)張量上的更新次數(此處為 1)。
我想做的是更多的參與,我認為用tensor_scatter_nd_add
編寫它將是一場噩夢。
當前許多關於該主題的 QA 都涵蓋了變量但不包括張量的情況(例如,參見this或this )。 這里提到確實 pytorch 支持這一點,所以我很驚訝最近沒有任何 tf 成員對此主題的回應。 這個答案並沒有真正幫助我,因為我需要某種面具生成,這也會很糟糕。
因此,問題是:如何在沒有tensor_scatter_nd_add
的情況下有效地進行切片分配(計算方面、內存方面和代碼方面)? 訣竅是我希望它盡可能動態,這意味着inputs
的形狀可以是可變的。
(對於任何好奇的人,我正在嘗試在 tf 中翻譯此代碼)。
這個問題最初發布在 GitHub 問題中。
這是另一個基於二進制掩碼的解決方案。
"""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)
PS還有一些其他的方法,比如使用tf.numpy_function(),或者生成掩碼function。
它似乎沒有產生任何錯誤:
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.
聲明:本站的技術帖子網頁,遵循CC BY-SA 4.0協議,如果您需要轉載,請注明本站網址或者原文地址。任何問題請咨詢:yoyou2525@163.com.