Mixed precision 학습

숫자 포맷 두 개, 학습 루프 하나

mixed precision 은 이름 그대로 섞어 써 — FLOP 거의 다 먹는 무거운 연산 (matmul, convolution) 은 float16 (또는 bfloat16) 으로 돌리고, range 가 필요한 부분은 float32 로 둬. 속도는 16-bit 연산에서 나와. 근데 weight 의 master copy 랑 gradient 누적은 32-bit 로 유지해 — 작은 update 가 rounding 에 먹혀 사라지면 안 되니까.

켜는 건 한 줄

set_global_policy('mixed_float16') 를 model 만들기 *전* 에 호출하고, 그 다음은 평소대로 build + train. cast 는 Keras 가 알아서 적절한 경계에 끼워 줘. 디테일 하나는 챙길 값어치 있어 — 마지막 Dense 는 raw logit (softmax 없이) 으로 두고 from_logits=True 랑 짝지어. float16 softmax 로 overflow 나는 대신 numerically stable 하게 loss 를 계산하거든.

공짜 속도의 함정

2-3 배는 *맞는 tensor 유닛이 하드웨어에 있을 때만* 공짜야. Tensor Core 없는 GPU (또는 CPU) 면 메모리는 아끼지만 속도는 거의 안 붙고, cast overhead 만 더해. 그러니까 가정하지 말고 측정해 — measure-it win 이야.

Code

mixed_float16 켜고 logit 은 float32 로 유지·python

# Enable mixed precision globally
keras.mixed_precision.set_global_policy("mixed_float16")

# Build and train normally — Keras handles the casting
model = keras.Sequential([
    keras.Input(shape=(784,)),
    layers.Dense(256, activation="relu"),  # Computes in float16
    layers.Dense(10),                        # Raw logits (no activation)
])

# Use from_logits for numerical stability with mixed precision
model.compile(
    optimizer="adam",
    loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
)

숫자 포맷 두 개, 학습 루프 하나

켜는 건 한 줄

공짜 속도의 함정

Code

External links

Exercise

Progress

댓글 0