@ tf.function ValueError: Создание переменных при не первом вызове функции, украшенной функцией tf.function, не в состоянии понять поведение

Question

Я хотел бы знать, почему эта функция:

@tf.function
def train(self,TargetNet,epsilon):
    if len(self.experience['s']) < self.min_experiences:
        return 0
    ids=np.random.randint(low=0,high=len(self.replay_buffer['s']),size=self.batch_size)
    states=np.asarray([self.experience['s'][i] for i in ids])
    actions=np.asarray([self.experience['a'][i] for i in ids])
    rewards=np.asarray([self.experience['r'][i] for i in ids])
    next_states=np.asarray([self.experience['s1'][i] for i in ids])
    dones = np.asarray([self.experience['done'][i] for i in ids])
    q_next_actions=self.get_action(next_states,epsilon)
    q_value_next=TargetNet.predict(next_states)
    q_value_next=tf.gather_nd(q_value_next,tf.stack((tf.range(self.batch_size),q_next_actions),axis=1))
    targets=tf.where(dones, rewards, rewards+self.gamma*q_value_next)

    with tf.GradientTape() as tape:
        estimates=tf.math.reduce_sum(self.predict(states)*tf.one_hot(actions,self.num_actions),axis=1)
        loss=tf.math.reduce_sum(tf.square(estimates - targets))
    variables=self.model.trainable_variables
    gradients=tape.gradient(loss,variables)
    self.optimizer.apply_gradients(zip(gradients,variables))

дает ValueError: Создание переменных при не первом вызове функции, украшенной функцией tf.function. Принимая во внимание, что этот код очень похож:

@tf.function
def train(self, TargetNet):
    if len(self.experience['s']) < self.min_experiences:
        return 0
    ids = np.random.randint(low=0, high=len(self.experience['s']), size=self.batch_size)
    states = np.asarray([self.experience['s'][i] for i in ids])
    actions = np.asarray([self.experience['a'][i] for i in ids])
    rewards = np.asarray([self.experience['r'][i] for i in ids])
    states_next = np.asarray([self.experience['s2'][i] for i in ids])
    dones = np.asarray([self.experience['done'][i] for i in ids])
    value_next = np.max(TargetNet.predict(states_next), axis=1)
    actual_values = np.where(dones, rewards, rewards+self.gamma*value_next)

    with tf.GradientTape() as tape:
        selected_action_values = tf.math.reduce_sum(
            self.predict(states) * tf.one_hot(actions, self.num_actions), axis=1)
        loss = tf.math.reduce_sum(tf.square(actual_values - selected_action_values))
    variables = self.model.trainable_variables
    gradients = tape.gradient(loss, variables)
    self.optimizer.apply_gradients(zip(gradients, variables))

Не выдает ошибку. Пожалуйста, помогите мне понять, почему.

EDIT : я удалил параметр epsilon из функция, и она работает. Это потому, что декоратор @ tf.function действителен только для функций с одним аргументом?

Answer 1

Используя функцию tf.function, вы конвертируете содержимое оформленной функции: это означает, что TensorFlow попытается скомпилировать ваш нетерпеливый код в графическое представление.

Переменные, однако, являются специальными объектами. На самом деле, когда вы использовали TensorFlow 1.x (режим графика), вы определяли переменные только один раз, а затем использовали / обновляли их.

В тензорном потоке 2.0, если вы используете чисто жадное выполнение, вы можете объявить и повторно использовать одну и ту же переменную более одного раза, поскольку tf.Variable - в активном режиме - это просто обычный объект Python, который уничтожается, как только функция заканчивается, и переменная, таким образом, выходит из области видимости.

Чтобы TensorFlow мог правильно преобразовывать функцию, которая создает состояние (то есть использует переменные), вы должны разбить область действия функции, объявив переменные вне функции.

Короче , если у вас есть функция, которая работает правильно в активном режиме, например:

def f():
    a = tf.constant([[10,10],[11.,1.]])
    x = tf.constant([[1.,0.],[0.,1.]])
    b = tf.Variable(12.)
    y = tf.matmul(a, x) + b
    return y

Вы должны изменить ее структуру на что-то вроде:

b = None

@tf.function
def f():
    a = tf.constant([[10, 10], [11., 1.]])
    x = tf.constant([[1., 0.], [0., 1.]])
    global b
    if b is None:
        b = tf.Variable(12.)
    y = tf.matmul(a, x) + b
    print("PRINT: ", y)
    tf.print("TF-PRINT: ", y)
    return y

f()

, чтобы она работала правильно с помощью декоратора tf.function.

Я описал этот (и другие) сценарий в нескольких сообщениях в блоге: первая часть анализирует это поведение в разделе Handlin g заявляет, что нарушает область действия функции (однако я предлагаю прочитать ее с самого начала и прочитать также части 2 и 3).