(Tensorflow 2) Назначение состояний в режиме графика

#tensorflow #keras #state #lstm

Вопрос:

Я пытаюсь обновить ячейки и скрытые состояния LSTM в режиме графика.

В режиме ожидания он работает так, как я и предполагал, но когда я использую режим графика, он выдает ошибку/ или дает нежелательные результаты (состояние не обновлено)

Ниже приведен мой код, который демонстрирует мою проблему:

 import tensorflow as tf
tf.random.set_seed(4)
import numpy as np
np.random.seed(2)
from tensorflow.keras.models import Model
from tensorflow.keras.layers import Input, LSTM, Dense

lstm_dim = 64
output_dim = 10

inputs = Input(batch_shape=(1, 1, output_dim))
outputs = LSTM(lstm_dim, return_sequences=True, stateful=True)(inputs)
outputs = Dense(output_dim, activation='softmax')(outputs)
model = Model(inputs, outputs)
 

Модель генерирует одномерный вектор на основе заданных входных данных и векторов состояний.

И я определил две функции, одна из которых-функция выполнения в режиме ожидания

 def myfunc_eager(states):
    model.layers[1].reset_states(states=[states[0],states[1]])

    startidx = 0
    zerovec = tf.zeros((1,1,output_dim))
    samplevec = tf.tensor_scatter_nd_update(zerovec, [[0, 0, startidx]], [1])
    idxlist = []

    for i in range(5):
        o = model.predict(samplevec)
        sampleidx = tf.math.argmax(o[0,0])
        samplevec = tf.tensor_scatter_nd_update(zerovec, [[0, 0, sampleidx]], [1])
        idxlist.append(sampleidx.numpy())

    return idxlist
 

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

 @tf.function
def myfunc_graph(states):
    model.layers[1].reset_states()
    
    #Trial 1
    #model.layers[1].reset_states(states=[states[0],states[1]])
    #TypeError: __array__() takes 1 positional argument but 2 were given
    
    #Trial 2
    #model.layers[1].states[0] = states[0]
    #model.layers[1].states[1] = states[1]
    #AttributeError: 'Tensor' object has no attribute 'assign'

    #Trial 3
    #model.layers[1].states[0] = tf.keras.backend.variable(states[0])
    #model.layers[1].states[1] = tf.keras.backend.variable(states[1])
    #ValueError: tf.function-decorated function tried to create variables on non-first call.
    
    #Trial 4
    #do nothing
    #it operates, but of course, states are not updated as intended
    
    startidx = 0
    zerovec = tf.zeros((1,1,output_dim))
    samplevec = tf.tensor_scatter_nd_update(zerovec, [[0, 0, startidx]], [1])
    idxlist = []
    for i in range(5):
        o = model(samplevec)
        sampleidx = tf.math.argmax(o[0,0])
        samplevec = tf.tensor_scatter_nd_update(zerovec, [[0, 0, sampleidx]], [1])
        idxlist.append(sampleidx)

    return idxlist
 

And result:

 states = [10*np.random.rand(1, lstm_dim), 10*np.random.rand(1, lstm_dim)]

print("eager mode execution")
for i in range(10):
    result = myfunc_eager(states)
    print(result)
print("graph mode execution")
for i in range(10):
    result = myfunc_graph(states)
    print([a.numpy() for a in result])
 
 eager mode execution
[2, 2, 2, 7, 4]
[2, 2, 2, 7, 4]
[2, 2, 2, 7, 4]
[2, 2, 2, 7, 4]
[2, 2, 2, 7, 4]
[2, 2, 2, 7, 4]
[2, 2, 2, 7, 4]
[2, 2, 2, 7, 4]
[2, 2, 2, 7, 4]
[2, 2, 2, 7, 4]
graph mode execution
[1, 7, 7, 7, 7]
[1, 7, 7, 7, 7]
[1, 7, 7, 7, 7]
[1, 7, 7, 7, 7]
[1, 7, 7, 7, 7]
[1, 7, 7, 7, 7]
[1, 7, 7, 7, 7]
[1, 7, 7, 7, 7]
[1, 7, 7, 7, 7]
[1, 7, 7, 7, 7]
 

I also tried to make a class that contains myfunc_graph, (to avoid creating variables inside a graph-mode function) but that wasn’t successful either.

Why reset_states() method does not work with states in the graph mode? Is there any solution for this problem? thanks.