tf.function으로 성능 높이기

tf.function으로 성능 높이기 (Better performance with tf.function)

파이썬 코드를 그대로 실행하면 일일이 연산을 해석하는 오버헤드가 생겨요. tf.function 은 파이썬 코드를 파이썬 독립적인 데이터플로우 그래프 로 변환하는 도구예요. 그래프로 만들어 실행하면 특히 작은 연산이 많은 프로그램에서 속도가 크게 빨라져요.

출처: TensorFlow 공식 문서 — Better performance with tf.function

기본 사용법

tf.function 을 데코레이터로 붙이거나 직접 호출해 함수를 감싸면, 그 함수는 호출 시 그래프로 실행돼요. 그 안에서 tf.Variable 이나 tf.GradientTape 같은 기능도 그대로 쓸 수 있어요.

v = tf.Variable(1.0)
with tf.GradientTape() as tape:
    result = add(v, 1.0)
tape.gradient(result, v)
# <tf.Tensor: shape=(), dtype=float32, numpy=1.0>

작은 연산이 많은 그래프에서 eager 코드보다 더 빠를 수 있어요. 다만 컨볼루션처럼 비싼 연산이 몇 개뿐인 그래프라면 큰 속도 향상을 기대하기는 어려워요.

트레이싱(tracing)이란

tf.function 은 프로그램을 tf.Graph 로 실행해요. 그 과정에서 두 단계가 있는데, 첫 단계인 tracing 에서 새 tf.Graph 를 만들고, 파이썬 코드는 정상적으로 실행되지만 TensorFlow 연산들은 지연(deferred)되어 그래프에만 캡처되고 실제로 실행되지는 않아요. 트레이싱은 비싼 작업이에요.

리트레이싱 조절

리트레이싱은 입력 유형별로 올바른 그래프를 만드는 데 도움이 되지만 비용이 커요. 입력의 dtype 등에 따라 새 trace가 생긴다는 점을 이해하고, 불필요한 재트레이싱이 일어나지 않도록 입력 형태를 일관되게 유지하는 게 좋아요.

ConcreteFunction 얻기

트레이싱이 끝난 뒤 특정 입력 시그니처에 해당하는 ConcreteFunction 을 얻을 수 있어요. "추적된 함수"를 출력하는 코드를 보면 한 번 trace가 실행됐다가 캐시에서 재사용되는 동작을 확인할 수 있어요.

# Executing traced function
tf.Tensor(b'aa', shape=(), dtype=string)
tf.Tensor(b'bb', shape=(), dtype=string)

AutoGraph 변환

AutoGraph는 tf.function 에서 기본으로 켜져 있는 라이브러리로, 파이썬 eager 코드의 일부를 그래프 호환 TensorFlow 연산으로 변환해요. if, for, while 같은 제어 흐름이 여기에 해당돼요. 덕분에 친숙한 파이썬 문법으로 그래프를 작성할 수 있어요.

제약 사항

파이썬 부작용(side effect) 실행: tf.py_function 은 그래프에 연결돼야 하므로 모든 입출력을 텐서로 캐스팅해요.
tf.function 의 모든 출력은 반환값이어야 해요: 함수 밖에서 정의된 텐서에 접근하려 하면 "out of scope" 오류가 나요.

tf.function으로 성능 높이기