0. Introduction
TF2.x 에선 Eager Execution이 기본 설정이다. Eager Execution(이하 Eager) 은 유연하고 직관적이지만 성능이나 배포측면에선 비효율적이다. 이를 해결하기 위해 Tensorflow에서는 tf.function()을 사용해 Graph 연산으로 변환할 수 있다. (특정 조건에선 Eager가 더 빠를 때도 있다)
1. Basic Usage
Decorator
사용 방법은 함수 위에 데코레이터로 @tf.function()을 달아주면 된다. 이렇게 래핑된 함수는 Tensorflow 연산과 같아진다. (1) Eager 상태로도 사용이 가능하고 (2) Gradient도 추적이 가능해진다. 일반적으로 Eager보다 Graph가 빠르다. 작은 연산을 여러번 할수록 Graph 연산이 유리한데, Convolution처럼 비싼 연산을 적은 횟수로 하는 Ops에선 속도 향상이 미미하거나 오히려 느려질 때도 있다.
import timeit
import tensorflow as tf
conv_layer = tf.keras.layers.Conv2D(100, 3)
@tf.function
def conv_fn(image):
return conv_layer(image)
image = tf.zeros([1, 200, 200, 100])
# Warm up
conv_layer(image); conv_fn(image)
print("Eager conv:", timeit.timeit(lambda: conv_layer(image), number=10))
print("Function conv:", timeit.timeit(lambda: conv_fn(image), number=10))
print("Note how there's not much difference in performance for convolutions")
Python
복사
Tracing ️
️What is Tracing
1.
tf.function()은 프로그램을 Graph에서 동작하게 한다. 그러나 Graph는 Eager상태의 TF 연산을 전부 표현할 수 없다. 예를들어 Python은 연산의 다형성을 지원하지만, Graph는 dtype과 Dimension을 명시해야 한다. 또 (1) Comman Line의 Arg를 읽거나, (2) 복잡한 Python Object를 사용하거나, (3) Error를 발생시키거나 하는 작업 중 어느것도 Graph에서 사용할 수 없다.
2.
tf.function()은 코드를 두 단계로 분리해서 이런 Gap을 메꾼다.
a.
Tracing
•
새로운 Grpah를 만든다. 이때 Python 코드는 정상동작하고, TF Ops는 Graph에 넣을것과 실행하지 않을것으로 나눈뒤 실행을 미룬다.
b.
Graph
•
Tracing 과정에서 실행이 지연된 연산이 실행된다. Tracing Stage보다 훨씬 적은 시간이 걸린다.
3.
tf.function()은 입력에 따라 Tracing을 생략할 때도 있다(이 규칙은 아래에서 자세히 설명한다). Tracing 생략이 자주 일어날 수록 최적화가 이루어 지는것이다. 아래의 코드가 좋은 예시인데, 입력하는 dtype이 일정하면 Graph를 재사용하고,바꾸면 Tracing이 진행된다.
Code
1.
tf.function() 과 Tracing에 대해 다시 정리하면 다음과 같다.
•
Graph는 TF 연산의 Raw하고, 언어에 구애받지 않는 Portable한 표현 방식이다.
•
Grpah는 ConcreteFunction으로 래핑된다.
•
tf.function()은 캐시되어있는 ConcreteFunction을 관리하고, 입력에 맞는걸 가져온다.
•
tf.function()은 Python function을 래핑하고 Function Object를 리턴한다.
•
Tracing은 Graph를 만들고, ConcreteFunction으로 래핑하는 동작을 말한다.
•
Tracing을 줄일수록 TF의 Performance를 최적화 할 수 있다.
Rule of Tracing
1.
일단 함수가 호출이 되면, 모든 Arguments들의 tf.types.experimental.TraceType 을 이용해 각각 대응되는 ConcreteFunction이 있는지 확인한다.
2.
대응되는 게 있으면 재사용하고, 없으면 새로운 ConcreteFunction이 Trace된다. 대응되는게 여러 개면, 가장 Specific한게 선택된다.
3.
Specific함을 비교할때 사용되는 개념이 subtyping이다.
a.
예를들어 TensorShape([1, 2])는 TensorShape([None, None])의 subtype이다.
b.
TensorShape([1, 2])의 입력이 들어오면 TensorShape([None, None])를 입력으로 갖는 ConcreteFunction이 대응될 수 있다.
c.
TensorShape([1, None]) 를 입력으로 갖는 ConcreteFunction이 있다면, 이게 더 Specific 하므로 우선적으로 선택된다.
4.
TraceType은 Input에 따라 정해지는데 자세한 내용은 아래와 같다
•
Tensor : dtype과 shape이 파라미터이다. Fixed Dim이 unknown Dim의 subtype이다.
•
Variable : Tensor와 유사한데 Variable은 Unique한 ID도 갖고있다.
•
Python Primitive Value : 대응되는 Value 자체가 Type이 된다.
예를들어 3은 int가 아니라 LiteralTraceType<3> 이 된다.
•
Python ordered containers : 내부 값의 TraceType + Container TraceType을 가진다.
[2,1] ⇒ ListTraceType<LiteralTraceType<1>, LiteralTraceType<2>>
[1,2] ⇒ ListTraceType<LiteralTraceType<2>, LiteralTraceType<1>>
•
Python Mapping : Key는 그대로 사용되고, Value는 TraceType으로 변환된다.
•
Other Object : Generic TraceType이어서 Matching될 때 까지 기다린다.
•
Object implemented (__tf_tracing_type__) : 해당 method가 반환하는 것
Controlling Retracing
Retracing은 TF가 각 입력에 잘 맞는 Graph를 만들도록 하지만, 비싼 Operation이다. Function이 호출될 때 마다 Retrace를 한다면 tf.function을 쓰지 않을 때 보다 느려진다. 따라서 불필요한 Tracing이 일어나지 않도록 조절해야 한다. 사용할 수 있는 테크닉은 다음과 같다.
tf.function에 고정된 input_signature 전달하기
Code
유연함을 위해 Unknown Dimension을 input_signature로 전달하기
가변 형태의 입력이 들어오면, 그때마다 Retrace가 일어날 수 있다. 그러한 입력에 대해서 None을 이용해 Unknown Dimension을 지정해 두면 Trace가 일어나지 않는다. (해보면 Fixed에 비해 느리긴 하지만 Retrace에 드는 비용이 훨씬 크다)
Code
Python Literal 대신 Tensor를 전달하기
Python Primitive Value는 값 자체가 TraceType이 된다. 즉 값이 바뀌면 Retracing이 일어난다는 뜻인데, 일반적으로 Primitive Value는 HyperParameter를 설정하기 위해 사용된다는 것을 생각해보면 자연스럽다. (num_layers=10, training=True, nonlinearity='relu') 그러나 그렇지 않은 경우에는 Primitive Value로 전달하면 불필요한 Retracing이 발생하므로, Tensor로 바꿔주는게 좋다.
Code
Tracing Protocol 사용하기
Python Type은 가능한 tf.experimental.ExtensionType로 변환되어야 하는것이 권장되는데, 이것의 TraceType은 TypeSpec과 큰 관결이 있다. ExtensionType의 Tracing 과정을 제어하려면 TypeSpec에 대해 override를 하면 된다.
Code
Debugging
일반적으로 Eager가 바로 바로 결과를 볼 수 있으므로 디버깅 하기 쉽다. @tf.funciton 데코레이팅 전에 Eager에서 에러가 없는지 부터 확인을 해야한다. 이를 위해 tf.config.run_functions_eagerly(True)를 사용할 수 있는데, 저 명령어를 걸어주면 tf.function이 비활성화 된다. 디버깅을 위한 팁은 다음과 같다.
1.
print() 는 Tracing 할 때만 작동한다, Trace를 추적할 때 유용하다
2.
tf.print() 는 항상 작동한다. Execution 안에서의 Value를 추적할 때 유용하다.
3.
tf.debugging.enable_check_numerics 은 NaN이나 Inf가 만들어지는 시점을 추적하기 유용하다
4.
pdb() 도 Tracing 할 때 무슨일이 있는지 추적하기 유용한데, AutoGraph 변환이 이루어지면 Drop된다.
2. AutoGraph Transformations
AutoGraph는 tf.function 에 기본적으로 켜져있는 라이브러리이며, Eager Code를 Graph에 맞는 TF Ops로 변환해준다. if, for, while 같은 제어문이 포함된다.
Conditionals
AutoGraph는 if <조건>절이 Tensor일때, 똑같은 동작을 하는 tf.cond절로 바꾼다. Python condiction절은 Tracing 중에 실행 되는데, 분기 중 그때 실행된 조건만 Graph에 추가된다. AutoGraph 없이는, trace된 Graph가 반대 조건을 받을 수 없다.
tf.cond는 trace할 때 각각의 분기를 Graph로 만들고, 실행될 때 Graph를 선택한다.
Loops
Overall
(1) AutoGraph는 for,while절을 똑같은 동작을 하는 tf looping ops로 바꾼다.
•
for x in y : y가 Tensor면 tf.while_loop로 컨버팅된다. y가 tf.data.Dataset이면 Dataset 관련 ops로 변환된다.
•
while <조건> : condition이 Tensor면 tf.while_loop로 컨버팅 된다.
(2) Python Loop는 Tracing과정에서 실행되며, Loop내내 Graph에 연산을 추가한다.
(3) TF Loop는 Body에 대해 Trace하고 이때 반복을 얼마나 실행할지 다이나믹하게 정한다. Body는 Graph에 한번만 등장한다.
Looping over Python data ️
️tf.function을 Python, NumPy Data에 대해 Loop돌리는 경우 매 Iter마다 Graph를 만들어 내는 함정에 빠지게 된다. 만약 Training Loop를 tf.function 으로 래핑하고 싶다면, Data를 tf.data.Dataset으로 만들어 주면 된다. 그러면 AutoGraph가 알아서 Unroll해준다.
Python, Numpy Data를 Dataset으로 래핑할때 (1)from_generator (2) from_tensors 두 선택지가 있는데 전자는 데이터를 Python에서 보관하고 tf-py function으로 연산을 변환해 성능 관련 문제가 생길 수 있고, 후자는 데이터를 거대한 Graph에 tf.constant Node로 보관해 메모리 관련 문제가 생길 수 있다.
TFRecordDataset, CsvDataset을 사용하면 이런 문제 없이 Python을 사용하지 않고 파일로 부터 데이터를 다룰 수 수 있기 때문에 가장 효율적이다.
Accumulating values in a loop
Loop를 돌며 값을 쌓는 패턴은 아주 흔한데, 일반적으로 dict, list에 값을 넣는 식으로 수행한다. 그러나 이런식의 연산은 TF에선 예기치 않은 Side Effect가 생길 수 있다. 따라서 dict, list 대신 tf.TensorArray 를 이용해 값을 저장하는 것이 권장된다.
batch_size = 2
seq_len = 3
feature_size = 4
def rnn_step(inp, state):
return inp + state
@tf.function
def dynamic_rnn(rnn_step, input_data, initial_state):
# [batch, time, features] -> [time, batch, features]
input_data = tf.transpose(input_data, [1, 0, 2])
max_seq_len = input_data.shape[0]
states = tf.TensorArray(tf.float32, size=max_seq_len)
state = initial_state
for i in tf.range(max_seq_len):
state = rnn_step(input_data[i], state)
states = states.write(i, state)
return tf.transpose(states.stack(), [1, 0, 2])
dynamic_rnn(rnn_step,
tf.random.uniform([batch_size, seq_len, feature_size]),
tf.zeros([batch_size, feature_size]))
Python
복사
3. Limitations
Executing Python Side Effects
print, append, global variable 같은 연산은 Function 안에서 2번 실행되거나, 아예 실행이 안되는 등 예상치 못한 방식으로 동작할 수도 있다. 이런 Side Effect는 처음 Input을 넣어서 호출할 때만 발생한다. 이후엔 Python Code 실행 없이 Graph를 재활용 하기 때문이다.