본문 바로가기

AI 프레임워크 (파이토치,텐서플로우)

TF - eager execution , JIT (@tf.function , jax)

텐서플로우는 기본적으로 static graph 를 구축한다.

// tf.session()

TF2.0 에서는 eager execution 모드를 지원하는데, 이를 통해 파이썬 처럼 늘여놓고, 그래프를 생성하지 않아도 텐서를 만들고 값을 계산할 수 있게 해준다.

 

이런 eager execution 은 성능에 영향을 주기 때문에, 일반 연산을 그래프 연산으로 변환할 수 있다. 

// @tf.function 을 붙인다.

@tf.function 은 JIT 컴파일을 지원하여서, 런타임시에 특정코드를 기계어로 번역해준다.  XLA(Accelerated Linear Algebra) 기능도 포함한다.

XLA = 선형대수학용 컴파일러 - 내부에서 그래프를 최적화함. 여러 연산을 하나의 커널로 만들어 풀이하기 때문에, 이 때문에 중간의 값이 남지 않아 값 추적이 어려운 것으로 보인다. -> 매우 빠르지만, tace 가 어렵다.

Example from (텐서플로우 2.0 프로그래밍 p.5)

@tf.function
def add_1_or_10(x,b):
	if b:
    	x += 1.
	else:
    	x += 10.
    return x
    
result = add_1_or_10(tf.constant(1.),True).numpy()

print(result)

 

- 최근 JAX 라이브러리가 굉장히 뜨고 있는데, 이는 TF 의 XLA 기능을 빼 온 기능이다.  

@jit 데코레이터로 이용할 수 있다.