tf.function

tensorflow.org의 [링크] 내용을 정리함.

tf.function

    함수(function)을 호출가능한 Tensorflow graph로 컴파일한다.

tf.function(
   func=None, input_signature=None, autograph=True, experimental_implements=None,
   experimental_autograph_options=None, experimental_relax_shapes=False,
   experimental_compile=None
)

Args

• func
  
  : 컴파일될 함수. 만약 함수가 None이면, tf.function은 단일인자 func로 호출될 수 있는 데코레이터(decorator)를 반환한다. 즉, tf.function(input_signature=...)(func)은 tf.function(func, input_signature=...)과 동일하다. 전자는 데코레이터로 사용될 수 있다.
  
• input_signature
  
  : 함수로 전달될 텐서의 모양과 dtype을 나타내는 tf.TensorSpec 객체의 가능한 중첩 시퀀스(nested sequence). 만약 None이면, 분리한 함수는 각각 추정되는 입력 시그니쳐(signature)에 대해 초기화된다. 만약 input_signature가 지정되면, 함수로의 모든 입력은 반드시 Tensor여야 하고 함수는 **kwargs는 받을 수 없다.
  
• autograph
  
  : 오토그래프가 그래프 추적(tracing)전에 함수에 적용되어야하는지 여부. 데이터 종속(Data-dependent) 제어 흐름은 autograph=True여야만 한다. 좀 더 자세한 내용은 [tf.function and AutoGraph guide]를 참고하라.
  
• experimental_implements
  
  : 만약 제공된다면, 이를 구현한 얄려진(known) 함수의 이름을 포함한다.
  
  예를 들면, "mycompany.my_recurrent_cell"은 추론함수내 속성으로 저장된다. 이는 그려면 직렬화된(serialized) 함수를 처리할 때 탐지(detect)할 수 있게 된다. 자세한 내용은 [standardizing composite ops]을 참고하라. 이 속성을 구성하는 예제에 대해서는 [example]를 참고하라. 이후 나오는 코드는 자동적으로 탐지하고 "embedded_matmul"을 구현한 함수를 대신한다. 그리고 TFLite가 그 자신의 구현을 대신하는 것을 허용한다. 예를 들어 텐서플로 사용자는 이 속성을 @tf.function(expreimental_implements="embedded_matmul")파라미터를 사용하여 지정하는 것으로 사용자의 함수 또한 embedded_matmul을 구현한 것으로 표시하는 속성을 사용할 수 있다.
  
• experimental_autograph_options
  
  : (Optional) [tf.autograph.experimental.Feature] 값의 튜플
  
• experimental_relax_shapes
  
  : True면, tf.function은 더 적게 즉, 입력 모양에서 덜 다루어지는 그래프를 생성한다.
  
• experimental_compile
  : True면, 함수는 항상 [XLA]로 컴파일 된다. XLA는 몇가지 경우(즉, TPU, XLA_GPU, dense tensor 계산)에서 더 효과적일 수 있다.
  

Returns

: func가 None이 아니면, 컴파일된 함수를 실행할(그리고 0 또는 더 많은 tf.Tensor 객체를 반환하는) callable을 반환한다. func가 None이면, 데코레이터를 반환한다. 이는 단일 func 인자로 호출되어질 때, callable을 반환한다.

Raises

: experimental_compile을 사용하지만, XLA 지원이 비활성일 때 ValueError가 발생.

    tf.function은 func에서 텐서플로 연산을 추적-컴파일(trace-compiling)하는 것에 의해 생성된 텐서플로 그래프(tf.Graph)를 실행하는 callable을 만든다.

Example usage:

@tf.function
def f(x, y):
    return x ** 2 + y
x = tf.constant([2, 3])
y = tf.constant([3, -2])
f(x, y)

Result>>

<tf.Tensor: ... numpy=array([7, 7], ...)>

특징
    func은 if, for, while, break, continue, return문장을 포함하여 데이터 의존적(data-dependent) 제어흐름을 사용한다.

@tf.function
def f(x):
    if tf.reduce_sum(x) > 0:
        return x * x
    else:
        return -x // 2
f(tf.constant(-2))

Result>>

<tf.Tensor: ... numpy=1>

    func의 클로져(closure)는 tf.Tensor와 tf.Variable 객체를 포함한다.

@tf.function
def f():
    return x ** 2 + y
x = tf.constant([-2, -3])
y = tf.Variable([3, -2])
f()

Result>>

<tf.Tensor: ... numpy=array([7, 7], ...)>

    func는 또한 tf.print, tf.Variable등과 같은 부가적인 영향으로 연산(ops)를 사용한다.

v = tf.Variable(1)
@tf.function
def f(x):
    for i in tf.range(x):
        v.assign_add(i)
f(3)
v

Result>>

<tf.Tensor: ... numpy=4>

Key Point: func이 추적될 때, 어떤 파이썬 부과효과(리스트에 추가, print로 출력 등)는 단지 한번만 발생한다. tf.function에서 실행된 부과효과를 얻으려면 TF ops로 작성되어야 한다.

>>리스트 추가 사용
l = []
@tf.function
def f(x):
    for i in x:
        l.append(i + 1)    # Caution! Will only happen once when tracing
f(tf.constant([1, 2, 3]))
l

Result>>

[<tf.Tensor ... >]

    대신 tf.TensorArray같은 텐서플로 콜렉션을 사용한다.

>>tf.TensorArray 사용
@tf.function
def f(x):
    ta = tf.TensorArray(dtype=tf.int32, size=0, dynamic_size=True)
    for i in range(len(x)):
        ta = ta.write(i, x[i] + 1)
    return ta.stack()
f(tf.constant([1, 2, 3]))

Result>>

<tf.Tensor ..., numpy=array([2, 3, 4], ...)>

    tf.function은 다형성(polymorphic)이다.

    내부적으로 tf.function은 모양과 dtype에 특화된 더 효과적인 그래프를 만들수 있기 때문에 다른 데이터 타입 또는 모양을 가진 인자를 지원하기 위해 하나 이상의 그래프를 만들 수 있다. tf.function은 또한 opaque 객채로 어떠한 순수 파이썬 값이라도 다룰 수 있고 마주치는 파인썬 인자의 각 셋에 대해 분리된 그래프를 만든다.

    개별적인 그래프를 얻으려면 tf.function으로 생성된 callable의 get_concrete_funtion 메서드를 사용하자. 이 메서느는 func으로 동일한 인자로 호출되고 tf.Graph 객체를 반환한다.

@tf.function
def f(x):
    return x + 1
isinstance(f.get_concrete_function(1).graph, tf.Graph)

Result>>

True

Caution: tf.function에 인자로 파이썬 스칼라 또는 리스트를 전달하는 것은 항상 새로운 그래프를 만들게 한다. 이를 방지하려면, 가능한 Tensor로 숫자형(numeric) 인자를 전달한다.

@tf.function
def f(x):
    return tf.abs(x)
f1 = f.get_concrete_function(1)
f2 = f.get_concrete_function(2)  # Slow - builds new graph
f1 is f2

f1 = f.get_concrete_function(tf.constant(1))
f2 = f.get_concrete_function(tf.constant(2))  # Fast - reuses f1
f1 is f2

Result>>

False
True

    파이썬 숫자형 인자는 오로지 신경망에서 레이어의 숫자같은 하이퍼파라미터같은 구분되는 값을 갖을 때만 사용해야 한다.

Input signatures

    텐서 인자에 대해서, tf.function은 입력 모양과 데이터타입의 모든 유일한 셋에 대해 분리된 그래프를 초기화한다. 아래 예제는 각각 다른 모양에 특화된 두개의 분리된 그래프를 생성한다.

@tf.function
def f(x):
    return x + 1
vector = tf.constant([1.0, 1.0])
matrix = tf.constant([[3.0]])
f.get_concrete_function(vector) is f.get_concrete_function(matrix)

Result>>

False

    "Input signature"는 추적되는 그래프를 제어하기 위해 tf.function에 선택적으로 제공될 수 있다. Input signature는 tf.TensorSpec 객체를 사용하여 함수에 각각의 텐서 인자의 모양과 타입을 나타낸다. 더 일반적인 모양도 사용될 수 있다. 이는 텐서가 동적인 모양을 가질때 여러개의 그래프가 생성되는 것을 피할수 있게 한다. 또한 사용될 수 있는 텐서의 모양과 데이터타입을 제한한다.

@tf.function(
    input_signature=[tf.TensorSpec(shape=None, dtype=tf.float32)])
def f(x):
    return x + 1
vector = tf.constant([1.0, 1.0])
matrix = tf.constant([[3.0]])
f.get_concrete_function(vector) is f.get_concrete_function(matrix)

Result>>

True

    변수는 오로지 한번만 생성된다.

    tf.function은 단지 처음으로 호출될 때에만 새로운 tf.Variable 객체를 생성한다.

class MyModule(tf.Module):
    def __init__(self):
        self.v = None

    @tf.function
    def call(self, x):
        if self.v is None:
            self.v = tf.Variable(tf.ones_like(x))
        return self.v * x

    보통 tf.function외부에서 tf.Variable같은 stateful 객체를 생성하고 이를 인자로 전달한다.