Keras 3 멀티백엔드 아키텍처: JAX·TensorFlow·PyTorch를 한 API로 묶는 법

François Chollet과 Matthew Watson이 공동 집필한 Deep Learning with Python 3판이 deeplearningwithpython.io를 통해 무료로 온라인 공개됐다. 이 책은 Keras 3의 멀티백엔드 지원을 전면 반영하여 TensorFlow·JAX·PyTorch 중 어떤 백엔드로도 예제 코드를 실행할 수 있도록 설계됐으며, 트랜스포머·GPT형 LLM 구축·디퓨전 모델 등 2026년 현재의 생성형 AI 실무 패턴을 포괄한다.

프레임워크 종속에서 벗어난 3판

Deep Learning with Python 1판(2017)은 Keras와 TensorFlow를 처음 대중화한 교재로, 전 세계 수십만 명의 딥러닝 입문자가 참고했다. 2021년 출간된 2판은 TensorFlow 2.x의 Eager Execution과 tf.keras 통합을 다뤘다. 이번 3판은 Keras가 특정 프레임워크 종속성에서 완전히 벗어나 멀티백엔드 프레임워크로 재설계된 Keras 3를 기반으로 하며, 코드 예제 전체가 JAX·TensorFlow·PyTorch 세 백엔드 위에서 동일하게 동작하도록 작성됐다.

무료 온라인 공개는 딥러닝 교육의 접근성 측면에서 중요한 사건이다. Manning Publications와의 협의하에 제공되는 이 웹 버전은 완성된 전체 내용을 포함하며, 코드 노트북은 GitHub(fchollet/deep-learning-with-python-notebooks)에서 함께 공개된다. 실습 환경을 갖춘 누구나 최신 딥러닝 이론과 실무 패턴을 비용 없이 학습할 수 있는 환경이 마련된 셈이다.

환경 변수 하나로 갈아 끼우는 백엔드

Keras 3의 가장 근본적인 혁신은 단일 고수준 API가 여러 수치 연산 백엔드를 플러그인 방식으로 교체할 수 있도록 설계된 아키텍처다. 사용자는 모델 정의·레이어 구성·훈련 루프 코드를 한 번만 작성하고, 실행 시 환경 변수나 설정 파일로 백엔드를 선택한다.

KERAS_BACKEND=jax python train.py
KERAS_BACKEND=torch python train.py
KERAS_BACKEND=tensorflow python train.py

백엔드 선택은 import keras 실행 시점에 고정되며, 이후 런타임 중에는 변경할 수 없다. 이 제약은 각 백엔드의 연산 그래프 초기화 방식이 상이하기 때문에 발생하는 필연적 설계 결정이다.

아래 다이어그램은 Keras 3의 전체 백엔드 추상화 레이어를 나타낸다.

⤢✕

Keras 3 Operations 모듈은 이중 모드(two-phase) 설계를 채택한다. 실제 데이터 없이 모델 구조를 정의할 때, 연산은 형상(shape)과 자료형(dtype) 정보만을 담은 KerasTensor 객체를 반환하는 심볼릭 단계(Symbolic Phase)에서는 모델 전체의 연산 그래프가 추적되며 형상 추론과 디버깅이 가능해진다. 실제 데이터가 입력되면 동일한 Operations 함수가 선택된 백엔드의 구체 구현으로 위임되는 이거 단계(Eager Phase)에서는 JAX라면 jax.numpy, TensorFlow라면 tf, PyTorch라면 torch 함수가 호출된다. 이 설계 덕분에 코드는 백엔드와 무관하게 형상 추론·모델 요약·디버깅 등의 메타 기능을 일관되게 제공받는다.

자동 미분은 각 백엔드가 제공하는 네이티브 구현을 Keras Trainer가 추상화하여 제공한다.

백엔드	자동 미분 구현	특징
TensorFlow	`tf.GradientTape`	이거·그래프 혼합 지원
JAX	`jax.grad` / `jax.value_and_grad`	함수형, 순수 변환
PyTorch	`torch.autograd`	동적 그래프, 디버깅 친화

디바이스 배치는 환경이 인식한 가용 디바이스(GPU/TPU/CPU)에 Keras가 자동으로 텐서를 배치한다. JAX 백엔드를 선택했을 때 특정 모델 아키텍처에서 최대 350%의 속도 향상이 보고된 사례도 있으며, 백엔드 선택 자체가 성능 최적화의 첫 번째 수단이 된다.

Functional API로 쓸 것인가, 서브클래싱으로 쓸 것인가

Keras는 모델을 구성하는 세 가지 방법을 제공한다: Sequential API, Functional API, 서브클래싱(Subclassing). 실무에서는 Functional API와 서브클래싱의 선택이 핵심적이다.

⤢✕

Functional API는 레이어를 노드로, 텐서 흐름을 엣지로 표현하는 방향성 비순환 그래프(DAG) 모델을 직관적으로 표현한다. 다중 입력·다중 출력·잔차 연결·공유 레이어 등의 구조를 깔끔하게 기술할 수 있으며, 모델 구조가 정적으로 정의되므로 형상 추론·시각화·직렬화가 용이하다.

import keras

inputs = keras.Input(shape=(784,))
x = keras.layers.Dense(256, activation="relu")(inputs)
x = keras.layers.Dropout(0.3)(x)
outputs = keras.layers.Dense(10, activation="softmax")(x)

model = keras.Model(inputs=inputs, outputs=outputs)

서브클래싱은 keras.Model을 상속하고 __init__에서 레이어를 정의하며 call 메서드에서 순전파를 구현한다. Python의 동적 제어 흐름(조건문·반복문)을 모델 로직에 직접 활용할 수 있어 트랜스포머의 어텐션 마스킹, 재귀적 모듈, 동적 그래프 구조 등에 적합하다.

class ResidualBlock(keras.Model):
    def __init__(self, filters):
        super().__init__()
        self.conv1 = keras.layers.Conv2D(filters, 3, padding="same", activation="relu")
        self.conv2 = keras.layers.Conv2D(filters, 3, padding="same")
        self.add = keras.layers.Add()

    def call(self, inputs, training=False):
        residual = inputs
        x = self.conv1(inputs)
        x = self.conv2(x)
        return keras.activations.relu(self.add([x, residual]))

혼합 정밀도(Mixed Precision)는 float32 대신 float16 또는 bfloat16으로 연산하여 현대 GPU에서 훈련 속도를 3배 이상, TPU에서 60% 향상시키는 기법이다.

keras.mixed_precision.set_global_policy("mixed_float16")

model = keras.Sequential([
    keras.layers.Dense(4096, activation="relu"),
    keras.layers.Dense(10, activation="softmax", dtype="float32")  # 출력 레이어는 float32 유지
])

출력 레이어는 수치 안정성을 위해 float32로 명시하는 것이 관례다. 커스텀 훈련 루프에서 float16을 사용할 때는 수치 언더플로 방지를 위해 LossScaleOptimizer를 반드시 래핑해야 한다. 커스텀 훈련 루프는 model.fit()이 제공하지 않는 세밀한 제어가 필요할 때 사용한다. 그래디언트 클리핑, 단계별 학습률 조정, 다중 손실 항목 조합, 메타 학습 구현 등에 활용된다.

optimizer = keras.optimizers.Adam(learning_rate=1e-4)
loss_fn = keras.losses.SparseCategoricalCrossentropy()

@tf.function
def train_step(x_batch, y_batch):
    with tf.GradientTape() as tape:
        predictions = model(x_batch, training=True)
        loss = loss_fn(y_batch, predictions)
    gradients = tape.gradient(loss, model.trainable_weights)
    optimizer.apply_gradients(zip(gradients, model.trainable_weights))
    return loss

전이 학습은 대규모 데이터로 사전 훈련된 모델의 특징 추출 능력을 소규모 태스크에 재활용하는 핵심 실무 패턴이다. 3판은 두 단계 파인튜닝 전략을 권장한다. 특징 추출 단계에서는 사전 훈련 모델의 가중치를 동결(trainable=False)하고 새로 추가한 분류 헤드만 훈련한다. 학습률은 상대적으로 높게 설정해도 무방하다. 파인튜닝 단계에서는 분류 헤드가 수렴하면 기반 모델의 상위 레이어 일부를 해동(trainable=True)하고 매우 낮은 학습률(일반적으로 1e-5 이하)로 전체 모델을 공동 훈련한다. 학습률이 높으면 사전 훈련으로 획득한 표현이 파괴되므로 주의가 필요하다.

base_model = keras.applications.EfficientNetV2B0(include_top=False, weights="imagenet")
base_model.trainable = False  # (1) 동결

x = keras.layers.GlobalAveragePooling2D()(base_model.output)
outputs = keras.layers.Dense(num_classes, activation="softmax")(x)
model = keras.Model(base_model.input, outputs)
model.compile(optimizer="adam", loss="sparse_categorical_crossentropy")
model.fit(train_dataset, epochs=10)

# (2) 파인튜닝
base_model.trainable = True
model.compile(optimizer=keras.optimizers.Adam(1e-5), loss="sparse_categorical_crossentropy")
model.fit(train_dataset, epochs=5)

배치 크기가 GPU 메모리를 넘어설 때

GPU 메모리 제약으로 원하는 배치 크기를 한 번에 처리하지 못할 때, 그래디언트 누적(Gradient Accumulation)은 여러 미니배치에 걸쳐 그래디언트를 누적한 뒤 한 번에 파라미터를 업데이트하여 효과적인 배치 크기를 늘린다.

유효 배치 크기 = 미니배치 크기 × 누적 스텝 수
예) 미니배치 32 × 누적 8스텝 = 유효 배치 256

그래디언트 누적의 핵심은 누적 스텝 동안 optimizer.apply_gradients()를 호출하지 않고, 스텝이 끝난 뒤에만 파라미터를 갱신하는 것이다. 이를 통해 메모리 효율과 배치 크기 효과를 동시에 얻는다.

⤢✕

배치 크기를 K배 늘릴 때는 학습률도 K배로 선형 스케일링하는 Linear Scaling Rule을 적용하는 것이 일반적이나, 최근 NeurIPS 2025 연구에서는 소규모 배치에서의 Adam 하이퍼파라미터 조정 전략이 재고돼야 한다는 주장도 제기됐다.

Keras 3는 백엔드별 분산 훈련 전략을 추상화하여 제공한다.

전략	적합한 상황	설명
데이터 병렬화	일반적 대규모 훈련	동일 모델을 여러 GPU에 복제, 배치를 분할 처리 후 그래디언트 동기화
모델 병렬화	단일 GPU에 올릴 수 없는 대형 모델	레이어별로 다른 디바이스에 배치
파이프라인 병렬화	모델·데이터 병렬화 혼합	미니배치를 마이크로배치로 분할, 파이프라인 스테이지별 처리

TensorFlow 백엔드에서는 tf.distribute.MirroredStrategy, JAX 백엔드에서는 jax.experimental.maps와 pjit, PyTorch 백엔드에서는 torch.nn.parallel.DistributedDataParallel이 Keras Trainer 내부에서 활용된다. 사용자는 백엔드 종속 코드를 작성하지 않고도 동일한 Keras 모델로 멀티-GPU 훈련을 실행할 수 있다.

대규모 훈련에서 체크포인트는 단순한 백업이 아니라 훈련 재개·앙상블·조기 종료 판단의 핵심 인프라다. Keras는 keras.callbacks.ModelCheckpoint를 통해 에포크별·배치별 저장, 최고 성능 가중치만 보존, 파일 이름 포매팅 등을 지원한다.

checkpoint_cb = keras.callbacks.ModelCheckpoint(
    filepath="checkpoints/model_epoch{epoch:03d}_val{val_loss:.4f}.keras",
    monitor="val_loss",
    save_best_only=True,
    save_weights_only=False,
    verbose=1
)

early_stop_cb = keras.callbacks.EarlyStopping(
    monitor="val_loss",
    patience=5,
    restore_best_weights=True
)

model.fit(
    train_dataset,
    validation_data=val_dataset,
    epochs=100,
    callbacks=[checkpoint_cb, early_stop_cb]
)

장기 훈련 시에는 체크포인트와 함께 옵티마이저 상태도 저장해야 정확한 재개가 가능하다. save_weights_only=False로 전체 모델(아키텍처·가중치·옵티마이저 상태)을 Keras 네이티브 포맷(.keras)으로 저장하는 것이 권장된다. 활성화 체크포인팅(Gradient Checkpointing)은 순전파 시 중간 활성화를 저장하지 않고 역전파 시 재계산하여 메모리를 절감하는 별도 기법으로, 메모리와 연산량의 트레이드오프를 허용할 때 활용한다.

Deep Learning with Python 3판의 무료 공개는 Keras 3가 실현한 멀티백엔드 추상화의 철학을 가장 잘 보여주는 사례다. JAX·TensorFlow·PyTorch 세 생태계를 단일 고수준 API로 아우르는 Keras 3 아키텍처는, 실무자가 프레임워크 선택의 부담 없이 모델 설계와 훈련 최적화에만 집중할 수 있는 환경을 만든다. Functional API와 서브클래싱의 적절한 혼합, 혼합 정밀도·그래디언트 누적·분산 훈련의 조합, 체계적인 체크포인트 관리가 현대 딥러닝 실무 설계의 기본 패턴이며, 이 책은 그 패턴을 가장 명료하게 코드로 설명하는 레퍼런스로 자리매김했다.

프레임워크 종속에서 벗어난 3판

환경 변수 하나로 갈아 끼우는 백엔드

Functional API로 쓸 것인가, 서브클래싱으로 쓸 것인가

배치 크기가 GPU 메모리를 넘어설 때

Sources

함께 읽기