Eager / Graph / @tf.function 정신모델

머릿속에 박아둘 가장 중요한 모델 한 가지

TF 2.x는 목적에 따라 두 mode를 왔다 갔다 해. 디버깅은 eager, 속도는 graph. 둘 사이를 뒤집는 데코레이터가 @tf.function이야. 이 정신모델만 제대로 잡으면 나머지 TF가 다 말이 돼.

Eager mode: op 하나하나가 즉시 돌고, concrete tensor 돌려주고, .numpy() 호출 가능. print(x) 작동. pdb 작동. x.shape가 실제 데이터 반영. TF 2.x default이고 진짜 Python처럼 느껴지는 유일한 mode.

Graph mode: 함수에 @tf.function 씌우면 TF가 첫 호출에 tracing해 — symbolic placeholder tensor로 한 번 돌리면서 static computation graph 기록. 두 번째부터는 컴파일된 graph가 runtime에 직접 돌아서 Python overhead 건너뛰고 op fusion + XLA 가능. tight inner loop는 보통 2–10배 빨라져.

함정: @tf.function 안의 print()는 tracing 중에만 돌고, 매 호출 때 안 돌아. print가 갑자기 안 찍힌다면 이게 원인. tf.print()로 바꿔 — 이건 traced graph의 일부로 들어가.

Code

Eager — exactly like NumPy·python

import tensorflow as tf

x = tf.constant([1.0, 2.0, 3.0])
print(x + 1)  # tf.Tensor([2. 3. 4.], shape=(3,), dtype=float32)
print(tf.executing_eagerly())  # True

# Mix freely with NumPy
import numpy as np
print(np.mean(x))   # 2.0
print(x.numpy())    # array([1., 2., 3.], dtype=float32)

@tf.function — same code, compiled to a graph·python

import tensorflow as tf

@tf.function
def train_step(x, y, model, optimizer, loss_fn):
    with tf.GradientTape() as tape:
        logits = model(x, training=True)
        loss = loss_fn(y, logits)
    grads = tape.gradient(loss, model.trainable_weights)
    optimizer.apply_gradients(zip(grads, model.trainable_weights))
    return loss

# Decorate ONCE outside the loop. Inside the loop = retrace per call.
for x_batch, y_batch in train_ds:
    loss = train_step(x_batch, y_batch, model, optimizer, loss_fn)

When print() lies and tf.print() tells the truth·python

@tf.function
def step(x):
    print('python print:', x)  # only on tracing, NOT every call
    tf.print('tf print:', x)   # fires every call, prints to stderr
    return x + 1

for i in range(3):
    step(tf.constant(float(i)))

Eager / Graph / @tf.function 정신모델

머릿속에 박아둘 가장 중요한 모델 한 가지

Code

External links

Exercise

Progress

댓글 0