Как эффективно назначить срез тензора в TensorFlow - PullRequest
Новая
       147

Как эффективно назначить срез тензора в TensorFlow

3 голосов
/ 29 мая 2020

Я хочу присвоить некоторые значения срезам входного тензора в одной из моих моделей в 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]))

)

Но, как вы можете видеть, это очень сложно написать. И это только для статического c количества обновлений (здесь 1) на тензоре 1D (+ размер пакета).

То, что я хочу сделать, требует более сложного подхода, и я думаю, что напишу его с помощью tensor_scatter_nd_add будет кошмаром.

Многие текущие QA на topi c охватывают случай переменных, но не тензоров (см., Например, this или this ). Здесь упоминается здесь , что pytorch действительно поддерживает это, поэтому я удивлен, что в последнее время не вижу ответа от каких-либо членов tf на этот c topi. Этот ответ мне не очень помогает, потому что мне понадобится какая-то генерация маски, которая тоже будет ужасной.

Вопрос в следующем: как я могу выполнить назначение срезов эффективно (с точки зрения вычислений, памяти и кода) без tensor_scatter_nd_add? Хитрость в том, что я хочу, чтобы это было как можно более динамично, а это означает, что форма inputs может быть переменной.

(Для всех, кому интересно, я пытаюсь перевести этот код in tf).

Этот вопрос изначально был размещен в выпуске GitHub .

Ответы [ 3 ]

1 голос
/ 04 июня 2020

Вот еще одно решение, основанное на бинарной маске. 

"""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)

-------- Ниже приведен предыдущий ответ --------

Вам нужно создать tf.Variable в методе build (). Это также позволяет динамически c размер по 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 () или создание функции маски.

1 голос
/ 29 мая 2020

Кажется, я не выдаю ошибок с этим: 

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 голосов
/ 29 мая 2020

кажется, что это обходной путь, но я предлагаю вам его протестировать 

class AddToEven(Model):
    def call(self, inputs):
        outputs = inputs
        zeros = tf.Variable(lambda : tf.zeros_like(outputs), trainable=False)
        zeros[:, ::2].assign(zeros[:, ::2]+1)
        outputs = outputs + zeros
        return outputs

model = AddToEven()
model(tf.zeros([1, 10]))

вывод 

<tf.Tensor: shape=(1, 10), dtype=float32, numpy=array([[1., 0., 1., 0., 1., 0., 1., 0., 1., 0.]], dtype=float32)>

модифицированный рабочий пример с TF 2.2 

class AddToEvenVar(Layer):

    def __init__(self, inp_dim, **kwargs):
        super(AddToEvenVar, self).__init__(**kwargs)
        self.zeros = tf.Variable(tf.zeros(inp_dim), trainable=False)
        self.zeros[:, ::2].assign(self.zeros[:, ::2]+1)
        self.zeros = tf.expand_dims(self.zeros, 0)

    def call(self, inputs):
        return inputs + self.zeros

inp = Input((1,10))
out = AddToEvenVar((1,10))(inp)
out = Conv1D(1, 3, padding='same')(out)
model = Model(inp, out)
model.compile(loss='mse', optimizer='adam')

xx = tf.zeros([100, 1, 10])
model.fit(xx, xx, epochs=10, verbose = 1)
...