tf.functon в тензорном потоке с петлей

Question

Я пытаюсь обучить модель, используя TensorFlow. В моем коде есть узкое место, из-за которого оптимизация потока в Tensor достигла sh. Я обнаружил, что это, скорее всего, относится к части в моем коде, где у меня есть oop:

, вот минимальный рабочий пример:

import numpy  as np 
import tensorflow as tf
import scipy.optimize 

kon = 0.01
mu  = 1.5
fi  = 0.5 
kappa = 22
w = (1-theta)

n =100
xs = tf.random.normal(shape=(n,), stddev=0.2)
eps = tf.random.normal(shape=(n,), stddev=0.17)
z = tf.sigmoid(tf.random.normal(shape=(n,), stddev=0.22))




def my_function(z, eps, x0):
    def F(hi):
     return (mu/fi)*np.log(hi) -(1-mu)*kappa*(hi)**(1+(1/fi))-mu*(np.log(w*ei*xs)-np.log(kon))-np.log(ze)

    hvec = np.empty((0,))
    # leisure today

    for ze,ei,xs in zip(z, eps, x0):
        ei=np.exp(ei)
        xs=np.exp(xs)
        htemp = scipy.optimize.newton_krylov(F, 0.5)
        hvec = np.append(hvec, htemp)
    return hvec

, если я использую tf.function чтобы украсить свою функцию я получаю эту ошибку.

@tf.function
def my_function(z, eps, x0):
    def F(hi):
     return (mu/fi)*np.log(hi) -(1-mu)*kappa*(hi)**(1+(1/fi))-mu*(np.log(w*ei*xs)-np.log(kon))-np.log(ze)

    hvec = np.empty((0,))
    # leisure today

    for ze,ei,xs in zip(z, eps, x0):
        ei=np.exp(ei)
        xs=np.exp(xs)
        htemp = scipy.optimize.newton_krylov(F, 0.5)
        hvec = np.append(hvec, htemp)
    return hvec


htest=my_function(z,eps,x0)   

СООБЩЕНИЕ ОБ ОШИБКЕ

OperatorNotAllowedInGraphError: in converted code:

   .....

    OperatorNotAllowedInGraphError: iterating over `tf.Tensor` is not allowed: AutoGraph did not convert this function. Try decorating it directly with @tf.function.

Я пытался следовать этому маршруту:

@tf.function(input_signature=[tf.TensorSpec(None, tf.float32)]) 
def tf_function(input): 
  y = tf.numpy_function(my_function, [input], tf.float32) 
  return y

, но сообщение об ошибке, которое я получаю один раз Я называю tf_function:

htestTF= tf_function(z,eps,x0) 

Ошибка

TypeError: When input_signature is provided, only pass arguments covered by it. Received 3 argument(s).

Может ли кто-нибудь, имеющий опыт, помочь мне отладить это?

Answer 1

Я считаю, что это функция zip, автограф в настоящее время не преобразуется: вам нужно использовать for i in range(len(z)), а затем ze = z[i].

При этом ваша функция использует только NumPy, поэтому он не будет работать должным образом: для функции tf.function требуется код TensorFlow.

Чтобы обучиться работе с TensorFlow, вам необходимо заменить все вызовы np.* на их эквивалент tf.*, заменить hvec на hvec = tf.TensorArray(...) и замените scipy.optimize.newton_krylov реализацией на основе TF.